DocsGPT/application/worker.py

import os
import shutil
import string
import zipfile
from urllib.parse import urljoin

import nltk
import requests

from application.core.settings import settings
from application.parser.file.bulk import SimpleDirectoryReader
from application.parser.remote.remote_creator import RemoteCreator
from application.parser.open_ai_func import call_openai_api
from application.parser.schema.base import Document
from application.parser.token_func import group_split

try:
    nltk.download('punkt', quiet=True)
    nltk.download('averaged_perceptron_tagger', quiet=True)
except FileExistsError:
    pass


# Define a function to extract metadata from a given filename.
def metadata_from_filename(title):
    store = '/'.join(title.split('/')[1:3])
    return {'title': title, 'store': store}


# Define a function to generate a random string of a given length.
def generate_random_string(length):
    return ''.join([string.ascii_letters[i % 52] for i in range(length)])

current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Define the main function for ingesting and processing documents.
def ingest_worker(self, directory, formats, name_job, filename, user):
    """
    Ingest and process documents.

    Args:
        self: Reference to the instance of the task.
        directory (str): Specifies the directory for ingesting ('inputs' or 'temp').
        formats (list of str): List of file extensions to consider for ingestion (e.g., [".rst", ".md"]).
        name_job (str): Name of the job for this ingestion task.
        filename (str): Name of the file to be ingested.
        user (str): Identifier for the user initiating the ingestion.

    Returns:
        dict: Information about the completed ingestion task, including input parameters and a "limited" flag.
    """
    # directory = 'inputs' or 'temp'
    # formats = [".rst", ".md"]
    input_files = None
    recursive = True
    limit = None
    exclude = True
    # name_job = 'job1'
    # filename = 'install.rst'
    # user = 'local'
    sample = False
    token_check = True
    min_tokens = 150
    max_tokens = 1250
    full_path = directory + '/' + user + '/' + name_job
    import sys
    print(full_path, file=sys.stderr)
    # check if API_URL env variable is set
    file_data = {'name': name_job, 'file': filename, 'user': user}
    response = requests.get(urljoin(settings.API_URL, "/api/download"), params=file_data)
    # check if file is in the response
    print(response, file=sys.stderr)
    file = response.content

    if not os.path.exists(full_path):
        os.makedirs(full_path)
    with open(full_path + '/' + filename, 'wb') as f:
        f.write(file)

    # check if file is .zip and extract it
    if filename.endswith('.zip'):
        with zipfile.ZipFile(full_path + '/' + filename, 'r') as zip_ref:
            zip_ref.extractall(full_path)
        os.remove(full_path + '/' + filename)

    self.update_state(state='PROGRESS', meta={'current': 1})

    raw_docs = SimpleDirectoryReader(input_dir=full_path, input_files=input_files, recursive=recursive,
                                     required_exts=formats, num_files_limit=limit,
                                     exclude_hidden=exclude, file_metadata=metadata_from_filename).load_data()
    raw_docs = group_split(documents=raw_docs, min_tokens=min_tokens, max_tokens=max_tokens, token_check=token_check)

    docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs]

    call_openai_api(docs, full_path, self)
    self.update_state(state='PROGRESS', meta={'current': 100})

    if sample:
        for i in range(min(5, len(raw_docs))):
            print(raw_docs[i].text)

    # get files from outputs/inputs/index.faiss and outputs/inputs/index.pkl
    # and send them to the server (provide user and name in form)
    file_data = {'name': name_job, 'user': user}
    if settings.VECTOR_STORE == "faiss":
        files = {'file_faiss': open(full_path + '/index.faiss', 'rb'),
                'file_pkl': open(full_path + '/index.pkl', 'rb')}
        response = requests.post(urljoin(settings.API_URL, "/api/upload_index"), files=files, data=file_data)
        response = requests.get(urljoin(settings.API_URL, "/api/delete_old?path=" + full_path))
    else:
        response = requests.post(urljoin(settings.API_URL, "/api/upload_index"), data=file_data)

    
    # delete local
    shutil.rmtree(full_path)

    return {
        'directory': directory,
        'formats': formats,
        'name_job': name_job,
        'filename': filename,
        'user': user,
        'limited': False
    }

def remote_worker(self, source_data, name_job, user, directory = 'temp', loader = 'url'):
    # sample = False
    token_check = True
    min_tokens = 150
    max_tokens = 1250
    full_path = directory + '/' + user + '/' + name_job

    if not os.path.exists(full_path):
        os.makedirs(full_path)

    self.update_state(state='PROGRESS', meta={'current': 1})
 
    # source_data {"data": [url]} for url type task just urls
 
    # Use RemoteCreator to load data from URL
    remote_loader = RemoteCreator.create_loader(loader)
    raw_docs = remote_loader.load_data(source_data)

    docs = group_split(documents=raw_docs, min_tokens=min_tokens, max_tokens=max_tokens, token_check=token_check)

    #docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs]

    call_openai_api(docs, full_path, self)
    self.update_state(state='PROGRESS', meta={'current': 100})
    


    # Proceed with uploading and cleaning as in the original function
    file_data = {'name': name_job, 'user': user}
    if settings.VECTOR_STORE == "faiss":
        files = {'file_faiss': open(full_path + '/index.faiss', 'rb'),
                 'file_pkl': open(full_path + '/index.pkl', 'rb')}
        requests.post(urljoin(settings.API_URL, "/api/upload_index"), files=files, data=file_data)
        requests.get(urljoin(settings.API_URL, "/api/delete_old?path=" + full_path))
    else:
        requests.post(urljoin(settings.API_URL, "/api/upload_index"), data=file_data)

    shutil.rmtree(full_path)

    return {
        'urls': source_data,
        'name_job': name_job,
        'user': user,
        'limited': False
    }