File: extensions/discord/__init__.py




File: extensions/discord/bot.py

import os
import re

import discord
import requests
from discord.ext import commands
import dotenv

dotenv.load_dotenv()

# Replace 'YOUR_BOT_TOKEN' with your bot's token
TOKEN = os.getenv("DISCORD_TOKEN")
PREFIX = '@DocsGPT'
BASE_API_URL = 'http://localhost:7091'

intents = discord.Intents.default()
intents.message_content = True

bot = commands.Bot(command_prefix=PREFIX, intents=intents)


def split_string(input_str):
    pattern = r'^<@!?{0}>\s*'.format(bot.user.id)
    match = re.match(pattern, input_str)
    if match:
        content = input_str[match.end():].strip()
        return str(bot.user.id), content
    return None, input_str


@bot.event
async def on_ready():
    print(f'{bot.user.name} has connected to Discord!')


async def fetch_answer(question):
    data = {
        'sender': 'discord',
        'question': question,
        'history': ''
    }
    headers = {"Content-Type": "application/json",
               "Accept": "application/json"}
    response = requests.post(BASE_API_URL + '/api/answer', json=data, headers=headers)
    if response.status_code == 200:
        return response.json()['answer']
    return 'Sorry, I could not fetch the answer.'


@bot.event
async def on_message(message):
    if message.author == bot.user:
        return

    content = message.content.strip()
    prefix, content = split_string(content)
    if prefix is None:
        return

    part_prefix = str(bot.user.id)
    if part_prefix == prefix:
        answer = await fetch_answer(content)
        await message.channel.send(answer)

    await bot.process_commands(message)


bot.run(TOKEN)



File: extensions/chatwoot/__init__.py




File: extensions/chatwoot/app.py

import os
import pprint

import dotenv
import requests
from flask import Flask, request

dotenv.load_dotenv()
docsgpt_url = os.getenv("docsgpt_url")
chatwoot_url = os.getenv("chatwoot_url")
docsgpt_key = os.getenv("docsgpt_key")
chatwoot_token = os.getenv("chatwoot_token")
# account_id = os.getenv("account_id")
# assignee_id = os.getenv("assignee_id")
label_stop = "human-requested"


def send_to_bot(sender, message):
    data = {
        'sender': sender,
        'question': message,
        'api_key': docsgpt_key,
        'embeddings_key': docsgpt_key,
        'history': ''
    }
    headers = {"Content-Type": "application/json",
               "Accept": "application/json"}

    r = requests.post(f'{docsgpt_url}/api/answer',
                      json=data, headers=headers)
    return r.json()['answer']


def send_to_chatwoot(account, conversation, message):
    data = {
        'content': message
    }
    url = f"{chatwoot_url}/api/v1/accounts/{account}/conversations/{conversation}/messages"
    headers = {"Content-Type": "application/json",
               "Accept": "application/json",
               "api_access_token": f"{chatwoot_token}"}

    r = requests.post(url,
                      json=data, headers=headers)
    return r.json()


app = Flask(__name__)


@app.route('/docsgpt', methods=['POST'])
def docsgpt():
    data = request.get_json()
    pp = pprint.PrettyPrinter(indent=4)
    pp.pprint(data)
    try:
        message_type = data['message_type']
    except KeyError:
        return "Not a message"
    message = data['content']
    conversation = data['conversation']['id']
    contact = data['sender']['id']
    account = data['account']['id']
    assignee = data['conversation']['meta']['assignee']['id']
    print(account)
    print(label_stop)
    print(data['conversation']['labels'])
    print(assignee)

    if label_stop in data['conversation']['labels']:
        return "Label stop"
    # elif str(account) != str(account_id):
    #     return "Not the right account"

    # elif str(assignee) != str(assignee_id):
    #     return "Not the right assignee"

    if (message_type == "incoming"):
        bot_response = send_to_bot(contact, message)
        create_message = send_to_chatwoot(
            account, conversation, bot_response)
    else:
        return "Not an incoming message"

    return create_message


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=80)



File: scripts/ingest.py

import os
import sys
from collections import defaultdict
from typing import List, Optional

import dotenv
import nltk
import typer

from parser.file.bulk import SimpleDirectoryReader
from parser.java2doc import extract_functions_and_classes as extract_java
from parser.js2doc import extract_functions_and_classes as extract_js
from parser.open_ai_func import call_openai_api, get_user_permission
from parser.py2doc import extract_functions_and_classes as extract_py
from parser.py2doc import transform_to_docs
from parser.schema.base import Document
from parser.token_func import group_split

dotenv.load_dotenv()

app = typer.Typer(add_completion=False)

nltk.download('punkt', quiet=True)
nltk.download('averaged_perceptron_tagger', quiet=True)


def metadata_from_filename(title):
    return {'title': title}

# Splits all files in specified folder to documents
@app.command()
def ingest(yes: bool = typer.Option(False, "-y", "--yes", prompt=False,
                                    help="Whether to skip price confirmation"),
           dir: Optional[List[str]] = typer.Option(["inputs"],
                                                   help="""List of paths to directory for index creation.
                                                        E.g. --dir inputs --dir inputs2"""),
           file: Optional[List[str]] = typer.Option(None,
                                                    help="""File paths to use (Optional; overrides dir).
                                                        E.g. --file inputs/1.md --file inputs/2.md"""),
           recursive: Optional[bool] = typer.Option(True, help="Whether to recursively search in subdirectories."),
           limit: Optional[int] = typer.Option(None, help="Maximum number of files to read."),
           formats: Optional[List[str]] = typer.Option([".rst", ".md"],
                                                       help="""List of required extensions (list with .)
                                                        Currently supported: 
                                                        .rst, .md, .pdf, .docx, .csv, .epub, .html, .mdx"""),
           exclude: Optional[bool] = typer.Option(True, help="Whether to exclude hidden files (dotfiles)."),
           sample: Optional[bool] = typer.Option(False,
                                                 help="Whether to output sample of the first 5 split documents."),
           token_check: Optional[bool] = typer.Option(True, help="Whether to group small documents and split large."),
           min_tokens: Optional[int] = typer.Option(150, help="Minimum number of tokens to not group."),
           max_tokens: Optional[int] = typer.Option(2000, help="Maximum number of tokens to not split."),
           ):
    """
        Creates index from specified location or files.
        By default /inputs folder is used, .rst and .md are parsed.
    """

    def process_one_docs(directory, folder_name):
        raw_docs = SimpleDirectoryReader(input_dir=directory, input_files=file, recursive=recursive,
                                         required_exts=formats, num_files_limit=limit,
                                         exclude_hidden=exclude, file_metadata=metadata_from_filename).load_data()

        # Here we split the documents, as needed, into smaller chunks.
        # We do this due to the context limits of the LLMs.
        raw_docs = group_split(documents=raw_docs, min_tokens=min_tokens, max_tokens=max_tokens,
                               token_check=token_check)
        # Old method
        # text_splitter = RecursiveCharacterTextSplitter()
        # docs = text_splitter.split_documents(raw_docs)

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

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

        # Here we check for command line arguments for bot calls.
        # If no argument exists or the yes is not True, then the
        # user permission is requested to call the API.
        if len(sys.argv) > 1 and yes:
            call_openai_api(docs, folder_name)
        else:
            get_user_permission(docs, folder_name)


    folder_counts = defaultdict(int)
    folder_names = []
    for dir_path in dir:
        folder_name = os.path.basename(os.path.normpath(dir_path))
        folder_counts[folder_name] += 1
        if folder_counts[folder_name] > 1:
            folder_name = f"{folder_name}_{folder_counts[folder_name]}"
        folder_names.append(folder_name)

    for directory, folder_name in zip(dir, folder_names):
        process_one_docs(directory, folder_name)


@app.command()
def convert(dir: Optional[str] = typer.Option("inputs",
                                              help="""Path to directory to make documentation for.
                                                        E.g. --dir inputs """),
            formats: Optional[str] = typer.Option("py",
                                                  help="""Required language. 
                                                        py, js, java supported for now""")):
    """
            Creates documentation linked to original functions from specified location.
            By default /inputs folder is used, .py is parsed.
    """
    # Using a dictionary to map between the formats and their respective extraction functions
    # makes the code more scalable. When adding more formats in the future, 
    # you only need to update the extraction_functions dictionary.
    extraction_functions = {
    'py': extract_py,
    'js': extract_js,
    'java': extract_java
    }

    if formats in extraction_functions:
        functions_dict, classes_dict = extraction_functions[formats](dir)
    else:
        raise Exception("Sorry, language not supported yet")                                   
    transform_to_docs(functions_dict, classes_dict, formats, dir)


if __name__ == "__main__":
    app()



File: scripts/code_docs_gen.py

import ast
import json
from pathlib import Path

import dotenv
from langchain_community.llms import OpenAI
from langchain.prompts import PromptTemplate

dotenv.load_dotenv()

ps = list(Path("inputs").glob("**/*.py"))
data = []
sources = []
for p in ps:
    with open(p) as f:
        data.append(f.read())
    sources.append(p)


def get_functions_in_class(node):
    functions = []
    functions_code = []
    for child in node.body:
        if isinstance(child, ast.FunctionDef):
            functions.append(child.name)
            functions_code.append(ast.unparse(child))

    return functions, functions_code


def get_classes_and_functions(source_code):
    tree = ast.parse(source_code)
    classes = {}
    for node in tree.body:
        if isinstance(node, ast.ClassDef):
            class_name = node.name
            function_name, function = get_functions_in_class(node)
            # join function name and function code
            functions = dict(zip(function_name, function))
            classes[class_name] = functions
    return classes


structure_dict = {}
c1 = 0
for code in data:
    classes = get_classes_and_functions(ast.parse(code))
    source = str(sources[c1])
    structure_dict[source] = classes
    c1 += 1

# save the structure dict as json
with open('structure_dict.json', 'w') as f:
    json.dump(structure_dict, f)

if not Path("outputs").exists():
    Path("outputs").mkdir()

c1 = len(structure_dict)
c2 = 0
for source, classes in structure_dict.items():
    c2 += 1
    print(f"Processing file {c2}/{c1}")
    f1 = len(classes)
    f2 = 0
    for class_name, functions in classes.items():
        f2 += 1
        print(f"Processing class {f2}/{f1}")
        source_w = source.replace("inputs/", "")
        source_w = source_w.replace(".py", ".txt")
        if not Path(f"outputs/{source_w}").exists():
            with open(f"outputs/{source_w}", "w") as f:
                f.write(f"Class: {class_name}")
        else:
            with open(f"outputs/{source_w}", "a") as f:
                f.write(f"\n\nClass: {class_name}")
        # append class name to the front
        for function in functions:
            b1 = len(functions)
            b2 = 0
            print(f"Processing function {b2}/{b1}")
            b2 += 1
            prompt = PromptTemplate(
                input_variables=["code"],
                template="Code: \n{code}, \nDocumentation: ",
            )
            llm = OpenAI(temperature=0)
            response = llm(prompt.format(code=functions[function]))

            if not Path(f"outputs/{source_w}").exists():
                with open(f"outputs/{source_w}", "w") as f:
                    f.write(f"Function: {functions[function]}, \nDocumentation: {response}")
            else:
                with open(f"outputs/{source_w}", "a") as f:
                    f.write(f"\n\nFunction: {functions[function]}, \nDocumentation: {response}")



File: scripts/__init__.py




File: scripts/migrate_to_v1_vectorstore.py

import pymongo
import os

def migrate_to_v1_vectorstore_mongo():
    client = pymongo.MongoClient("mongodb://localhost:27017/")
    db = client["docsgpt"]
    vectors_collection = db["vectors"]
    sources_collection = db["sources"]

    for vector in vectors_collection.find():
        if "location" in vector:
            del vector["location"]
        if "retriever" not in vector:
            vector["retriever"] = "classic"
            vector["remote_data"] = None
        vectors_collection.update_one({"_id": vector["_id"]}, {"$set": vector})

    # move data from vectors_collection to sources_collection
    for vector in vectors_collection.find():
        sources_collection.insert_one(vector)

    vectors_collection.drop()

    client.close()

def migrate_faiss_to_v1_vectorstore():
    client = pymongo.MongoClient("mongodb://localhost:27017/")
    db = client["docsgpt"]
    vectors_collection = db["vectors"]

    for vector in vectors_collection.find():
        old_path = f"./application/indexes/{vector['user']}/{vector['name']}"
        new_path = f"./application/indexes/{vector['_id']}"
        try:
            os.rename(old_path, new_path)
        except OSError as e:
            print(f"Error moving {old_path} to {new_path}: {e}")

    client.close()

def migrate_mongo_atlas_vector_to_v1_vectorstore():
    client = pymongo.MongoClient("mongodb+srv://<username>:<password>@<cluster>/<dbname>?retryWrites=true&w=majority")
    db = client["docsgpt"]
    vectors_collection = db["vectors"]

    # mongodb atlas collection
    documents_collection = db["documents"]

    for vector in vectors_collection.find():
        documents_collection.update_many({"store": vector["user"] + "/" + vector["name"]}, {"$set": {"source_id": str(vector["_id"])}})

    client.close()

migrate_faiss_to_v1_vectorstore()
migrate_to_v1_vectorstore_mongo()


File: scripts/old/ingest_rst.py

import pickle
import sys
from argparse import ArgumentParser
from pathlib import Path

import dotenv
import faiss
import tiktoken
from langchain_openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS


def num_tokens_from_string(string: str, encoding_name: str) -> int:
    # Function to convert string to tokens and estimate user cost.
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    total_price = ((num_tokens / 1000) * 0.0004)
    return num_tokens, total_price


def call_openai_api():
    # Function to create a vector store from the documents and save it to disk.
    store = FAISS.from_texts(docs, OpenAIEmbeddings(), metadatas=metadatas)
    faiss.write_index(store.index, "docs.index")
    store.index = None
    with open("faiss_store.pkl", "wb") as f:
        pickle.dump(store, f)


def get_user_permission():
    # Function to ask user permission to call the OpenAI api and spend their OpenAI funds.
    # Here we convert the docs list to a string and calculate the number of OpenAI tokens the string represents.
    docs_content = (" ".join(docs))
    tokens, total_price = num_tokens_from_string(string=docs_content, encoding_name="cl100k_base")
    # Here we print the number of tokens and the approx user cost with some visually appealing formatting.
    print(f"Number of Tokens = {format(tokens, ',d')}")
    print(f"Approx Cost = ${format(total_price, ',.2f')}")
    # Here we check for user permission before calling the API.
    user_input = input("Price Okay? (Y/N) \n").lower()
    if user_input == "y":
        call_openai_api()
    elif user_input == "":
        call_openai_api()
    else:
        print("The API was not called. No money was spent.")


# Load .env file
dotenv.load_dotenv()

ap = ArgumentParser("Script for training DocsGPT on .rst documentation files.")
ap.add_argument("-i", "--inputs",
                type=str,
                default="inputs",
                help="Directory containing documentation files")
args = ap.parse_args()

# Here we load in the data in the format that Notion exports it in.
ps = list(Path(args.inputs).glob("**/*.rst"))

# parse all child directories
data = []
sources = []
for p in ps:
    with open(p) as f:
        data.append(f.read())
    sources.append(p)

# Here we split the documents, as needed, into smaller chunks.
# We do this due to the context limits of the LLMs.
text_splitter = CharacterTextSplitter(chunk_size=1500, separator="\n")
docs = []
metadatas = []
for i, d in enumerate(data):
    splits = text_splitter.split_text(d)
    docs.extend(splits)
    metadatas.extend([{"source": sources[i]}] * len(splits))

# Here we check for command line arguments for bot calls.
# If no argument exists or the permission_bypass_flag argument is not '-y',
# user permission is requested to call the API.
if len(sys.argv) > 1:
    permission_bypass_flag = sys.argv[1]
    if permission_bypass_flag == '-y':
        call_openai_api()
    else:
        get_user_permission()
else:
    get_user_permission()



File: scripts/old/__init__.py




File: scripts/old/ingest_rst_sphinx.py

import os
import pickle
import shutil
import sys
from argparse import ArgumentParser
from pathlib import Path

import dotenv
import faiss
import tiktoken
from langchain_openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS
from sphinx.cmd.build import main as sphinx_main


def convert_rst_to_txt(src_dir, dst_dir):
    # Check if the source directory exists
    if not os.path.exists(src_dir):
        raise Exception("Source directory does not exist")
    # Walk through the source directory
    for root, dirs, files in os.walk(src_dir):
        for file in files:
            # Check if the file has .rst extension
            if file.endswith(".rst"):
                # Construct the full path of the file
                src_file = os.path.join(root, file.replace(".rst", ""))
                # Convert the .rst file to .txt file using sphinx-build
                args = f". -b text -D extensions=sphinx.ext.autodoc " \
                       f"-D master_doc={src_file} " \
                       f"-D source_suffix=.rst " \
                       f"-C {dst_dir} "
                sphinx_main(args.split())
            elif file.endswith(".md"):
                # Rename the .md file to .rst file
                src_file = os.path.join(root, file)
                dst_file = os.path.join(root, file.replace(".md", ".rst"))
                os.rename(src_file, dst_file)
                # Convert the .rst file to .txt file using sphinx-build
                args = f". -b text -D extensions=sphinx.ext.autodoc " \
                       f"-D master_doc={dst_file} " \
                       f"-D source_suffix=.rst " \
                       f"-C {dst_dir} "
                sphinx_main(args.split())


def num_tokens_from_string(string: str, encoding_name: str) -> int:
    # Function to convert string to tokens and estimate user cost.
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    total_price = ((num_tokens / 1000) * 0.0004)
    return num_tokens, total_price


def call_openai_api():
    # Function to create a vector store from the documents and save it to disk.
    store = FAISS.from_texts(docs, OpenAIEmbeddings(), metadatas=metadatas)
    faiss.write_index(store.index, "docs.index")
    store.index = None
    with open("faiss_store.pkl", "wb") as f:
        pickle.dump(store, f)


def get_user_permission():
    # Function to ask user permission to call the OpenAI api and spend their OpenAI funds.
    # Here we convert the docs list to a string and calculate the number of OpenAI tokens the string represents.
    docs_content = (" ".join(docs))
    tokens, total_price = num_tokens_from_string(string=docs_content, encoding_name="cl100k_base")
    # Here we print the number of tokens and the approx user cost with some visually appealing formatting.
    print(f"Number of Tokens = {format(tokens, ',d')}")
    print(f"Approx Cost = ${format(total_price, ',.2f')}")
    # Here we check for user permission before calling the API.
    user_input = input("Price Okay? (Y/N) \n").lower()
    if user_input == "y":
        call_openai_api()
    elif user_input == "":
        call_openai_api()
    else:
        print("The API was not called. No money was spent.")


ap = ArgumentParser("Script for training DocsGPT on Sphinx documentation")
ap.add_argument("-i", "--inputs",
                type=str,
                default="inputs",
                help="Directory containing documentation files")
args = ap.parse_args()

# Load .env file
dotenv.load_dotenv()

# Directory to vector
src_dir = args.inputs
dst_dir = "tmp"

convert_rst_to_txt(src_dir, dst_dir)

# Here we load in the data in the format that Notion exports it in.
ps = list(Path("tmp/" + src_dir).glob("**/*.txt"))

# parse all child directories
data = []
sources = []
for p in ps:
    with open(p) as f:
        data.append(f.read())
    sources.append(p)

# Here we split the documents, as needed, into smaller chunks.
# We do this due to the context limits of the LLMs.
text_splitter = CharacterTextSplitter(chunk_size=1500, separator="\n")
docs = []
metadatas = []
for i, d in enumerate(data):
    splits = text_splitter.split_text(d)
    docs.extend(splits)
    metadatas.extend([{"source": sources[i]}] * len(splits))

# Here we check for command line arguments for bot calls.
# If no argument exists or the permission_bypass_flag argument is not '-y',
# user permission is requested to call the API.
if len(sys.argv) > 1:
    permission_bypass_flag = sys.argv[1]
    if permission_bypass_flag == '-y':
        call_openai_api()
    else:
        get_user_permission()
else:
    get_user_permission()

# Delete tmp folder
# Commented out for now
shutil.rmtree(dst_dir)



File: scripts/parser/py2doc.py

import ast
import os
from pathlib import Path

import tiktoken
from langchain_community.llms import OpenAI
from langchain.prompts import PromptTemplate


def find_files(directory):
    files_list = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.py'):
                files_list.append(os.path.join(root, file))
    return files_list


def extract_functions(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        functions = {}
        tree = ast.parse(source_code)
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                func_name = node.name
                func_def = ast.get_source_segment(source_code, node)
                functions[func_name] = func_def
    return functions


def extract_classes(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        classes = {}
        tree = ast.parse(source_code)
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                class_name = node.name
                function_names = []
                for subnode in ast.walk(node):
                    if isinstance(subnode, ast.FunctionDef):
                        function_names.append(subnode.name)
                classes[class_name] = ", ".join(function_names)
    return classes


def extract_functions_and_classes(directory):
    files = find_files(directory)
    functions_dict = {}
    classes_dict = {}
    for file in files:
        functions = extract_functions(file)
        if functions:
            functions_dict[file] = functions
        classes = extract_classes(file)
        if classes:
            classes_dict[file] = classes
    return functions_dict, classes_dict


def parse_functions(functions_dict, formats, dir):
    c1 = len(functions_dict)
    for i, (source, functions) in enumerate(functions_dict.items(), start=1):
        print(f"Processing file {i}/{c1}")
        source_w = source.replace(dir + "/", "").replace("." + formats, ".md")
        subfolders = "/".join(source_w.split("/")[:-1])
        Path(f"outputs/{subfolders}").mkdir(parents=True, exist_ok=True)
        for j, (name, function) in enumerate(functions.items(), start=1):
            print(f"Processing function {j}/{len(functions)}")
            prompt = PromptTemplate(
                input_variables=["code"],
                template="Code: \n{code}, \nDocumentation: ",
            )
            llm = OpenAI(temperature=0)
            response = llm(prompt.format(code=function))
            mode = "a" if Path(f"outputs/{source_w}").exists() else "w"
            with open(f"outputs/{source_w}", mode) as f:
                f.write(
                    f"\n\n# Function name: {name} \n\nFunction: \n```\n{function}\n```, \nDocumentation: \n{response}")


def parse_classes(classes_dict, formats, dir):
    c1 = len(classes_dict)
    for i, (source, classes) in enumerate(classes_dict.items()):
        print(f"Processing file {i + 1}/{c1}")
        source_w = source.replace(dir + "/", "").replace("." + formats, ".md")
        subfolders = "/".join(source_w.split("/")[:-1])
        Path(f"outputs/{subfolders}").mkdir(parents=True, exist_ok=True)
        for name, function_names in classes.items():
            print(f"Processing Class {i + 1}/{c1}")
            prompt = PromptTemplate(
                input_variables=["class_name", "functions_names"],
                template="Class name: {class_name} \nFunctions: {functions_names}, \nDocumentation: ",
            )
            llm = OpenAI(temperature=0)
            response = llm(prompt.format(class_name=name, functions_names=function_names))

            with open(f"outputs/{source_w}", "a" if Path(f"outputs/{source_w}").exists() else "w") as f:
                f.write(f"\n\n# Class name: {name} \n\nFunctions: \n{function_names}, \nDocumentation: \n{response}")


def transform_to_docs(functions_dict, classes_dict, formats, dir):
    docs_content = ''.join([str(key) + str(value) for key, value in functions_dict.items()])
    docs_content += ''.join([str(key) + str(value) for key, value in classes_dict.items()])

    num_tokens = len(tiktoken.get_encoding("cl100k_base").encode(docs_content))
    total_price = ((num_tokens / 1000) * 0.02)

    print(f"Number of Tokens = {num_tokens:,d}")
    print(f"Approx Cost = ${total_price:,.2f}")

    user_input = input("Price Okay? (Y/N)\n").lower()
    if user_input == "y" or user_input == "":
        if not Path("outputs").exists():
            Path("outputs").mkdir()
        parse_functions(functions_dict, formats, dir)
        parse_classes(classes_dict, formats, dir)
        print("All done!")
    else:
        print("The API was not called. No money was spent.")



File: scripts/parser/token_func.py

import re
from math import ceil
from typing import List

import tiktoken
from parser.schema.base import Document

def separate_header_and_body(text):
    header_pattern = r"^(.*?\n){3}"
    match = re.match(header_pattern, text)
    header = match.group(0)
    body = text[len(header):]
    return header, body


def group_documents(documents: List[Document], min_tokens: int, max_tokens: int) -> List[Document]:
    docs = []
    current_group = None

    for doc in documents:
        doc_len = len(tiktoken.get_encoding("cl100k_base").encode(doc.text))

        if current_group is None:
            current_group = Document(text=doc.text, doc_id=doc.doc_id, embedding=doc.embedding,
                                     extra_info=doc.extra_info)
        elif len(tiktoken.get_encoding("cl100k_base").encode(
                current_group.text)) + doc_len < max_tokens and doc_len < min_tokens:
            current_group.text += " " + doc.text
        else:
            docs.append(current_group)
            current_group = Document(text=doc.text, doc_id=doc.doc_id, embedding=doc.embedding,
                                     extra_info=doc.extra_info)

    if current_group is not None:
        docs.append(current_group)

    return docs


def split_documents(documents: List[Document], max_tokens: int) -> List[Document]:
    docs = []
    for doc in documents:
        token_length = len(tiktoken.get_encoding("cl100k_base").encode(doc.text))
        if token_length <= max_tokens:
            docs.append(doc)
        else:
            header, body = separate_header_and_body(doc.text)
            if len(tiktoken.get_encoding("cl100k_base").encode(header)) > max_tokens:
                body = doc.text
                header = ""
            num_body_parts = ceil(token_length / max_tokens)
            part_length = ceil(len(body) / num_body_parts)
            body_parts = [body[i:i + part_length] for i in range(0, len(body), part_length)]
            for i, body_part in enumerate(body_parts):
                new_doc = Document(text=header + body_part.strip(),
                                   doc_id=f"{doc.doc_id}-{i}",
                                   embedding=doc.embedding,
                                   extra_info=doc.extra_info)
                docs.append(new_doc)
    return docs


def group_split(documents: List[Document], max_tokens: int = 2000, min_tokens: int = 150, token_check: bool = True):
    if not token_check:
        return documents
    print("Grouping small documents")
    try:
        documents = group_documents(documents=documents, min_tokens=min_tokens, max_tokens=max_tokens)
    except Exception:
        print("Grouping failed, try running without token_check")
    print("Separating large documents")
    try:
        documents = split_documents(documents=documents, max_tokens=max_tokens)
    except Exception:
        print("Grouping failed, try running without token_check")
    return documents



File: scripts/parser/__init__.py





File: scripts/parser/java2doc.py

import os

import javalang


def find_files(directory):
    files_list = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.java'):
                files_list.append(os.path.join(root, file))
    return files_list


def extract_functions(file_path):
    with open(file_path, "r") as file:
        java_code = file.read()
        methods = {}
        tree = javalang.parse.parse(java_code)
        for _, node in tree.filter(javalang.tree.MethodDeclaration):
            method_name = node.name
            start_line = node.position.line - 1
            end_line = start_line
            brace_count = 0
            for line in java_code.splitlines()[start_line:]:
                end_line += 1
                brace_count += line.count("{") - line.count("}")
                if brace_count == 0:
                    break
            method_source_code = "\n".join(java_code.splitlines()[start_line:end_line])
            methods[method_name] = method_source_code
    return methods


def extract_classes(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        classes = {}
        tree = javalang.parse.parse(source_code)
        for class_decl in tree.types:
            class_name = class_decl.name
            declarations = []
            methods = []
            for field_decl in class_decl.fields:
                field_name = field_decl.declarators[0].name
                field_type = field_decl.type.name
                declarations.append(f"{field_type} {field_name}")
            for method_decl in class_decl.methods:
                methods.append(method_decl.name)
            class_string = "Declarations: " + ", ".join(declarations) + "\n  Method name: " + ", ".join(methods)
            classes[class_name] = class_string
    return classes


def extract_functions_and_classes(directory):
    files = find_files(directory)
    functions_dict = {}
    classes_dict = {}
    for file in files:
        functions = extract_functions(file)
        if functions:
            functions_dict[file] = functions
        classes = extract_classes(file)
        if classes:
            classes_dict[file] = classes
    return functions_dict, classes_dict



File: scripts/parser/js2doc.py

import os

import escodegen
import esprima


def find_files(directory):
    files_list = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.js'):
                files_list.append(os.path.join(root, file))
    return files_list


def extract_functions(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        functions = {}
        tree = esprima.parseScript(source_code)
        for node in tree.body:
            if node.type == 'FunctionDeclaration':
                func_name = node.id.name if node.id else '<anonymous>'
                functions[func_name] = escodegen.generate(node)
            elif node.type == 'VariableDeclaration':
                for declaration in node.declarations:
                    if declaration.init and declaration.init.type == 'FunctionExpression':
                        func_name = declaration.id.name if declaration.id else '<anonymous>'
                        functions[func_name] = escodegen.generate(declaration.init)
            elif node.type == 'ClassDeclaration':
                for subnode in node.body.body:
                    if subnode.type == 'MethodDefinition':
                        func_name = subnode.key.name
                        functions[func_name] = escodegen.generate(subnode.value)
                    elif subnode.type == 'VariableDeclaration':
                        for declaration in subnode.declarations:
                            if declaration.init and declaration.init.type == 'FunctionExpression':
                                func_name = declaration.id.name if declaration.id else '<anonymous>'
                                functions[func_name] = escodegen.generate(declaration.init)
        return functions


def extract_classes(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        classes = {}
        tree = esprima.parseScript(source_code)
        for node in tree.body:
            if node.type == 'ClassDeclaration':
                class_name = node.id.name
                function_names = []
                for subnode in node.body.body:
                    if subnode.type == 'MethodDefinition':
                        function_names.append(subnode.key.name)
                classes[class_name] = ", ".join(function_names)
    return classes


def extract_functions_and_classes(directory):
    files = find_files(directory)
    functions_dict = {}
    classes_dict = {}
    for file in files:
        functions = extract_functions(file)
        if functions:
            functions_dict[file] = functions
        classes = extract_classes(file)
        if classes:
            classes_dict[file] = classes
    return functions_dict, classes_dict



File: scripts/parser/open_ai_func.py

import os

import tiktoken
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from retry import retry


# from langchain.embeddings import HuggingFaceEmbeddings
# from langchain.embeddings import HuggingFaceInstructEmbeddings
# from langchain.embeddings import CohereEmbeddings


def num_tokens_from_string(string: str, encoding_name: str) -> tuple[int, float]:
    # Function to convert string to tokens and estimate user cost.
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    total_price = (num_tokens / 1000) * 0.0004
    return num_tokens, total_price


@retry(tries=10, delay=60)
def store_add_texts_with_retry(store, i):
    store.add_texts([i.page_content], metadatas=[i.metadata])
    # store_pine.add_texts([i.page_content], metadatas=[i.metadata])


def call_openai_api(docs, folder_name):
    # Function to create a vector store from the documents and save it to disk.

    # create output folder if it doesn't exist
    if not os.path.exists(f"outputs/{folder_name}"):
        os.makedirs(f"outputs/{folder_name}")

    from tqdm import tqdm

    docs_test = [docs[0]]
    # remove the first element from docs
    docs.pop(0)
    # cut first n docs if you want to restart
    # docs = docs[:n]
    c1 = 0
    # pinecone.init(
    #     api_key="",  # find at app.pinecone.io
    #     environment="us-east1-gcp"  # next to api key in console
    # )
    # index_name = "pandas"
    if (  # azure
        os.environ.get("OPENAI_API_BASE")
        and os.environ.get("OPENAI_API_VERSION")
        and os.environ.get("AZURE_DEPLOYMENT_NAME")
        and os.environ.get("AZURE_EMBEDDINGS_DEPLOYMENT_NAME")
    ):
        os.environ["OPENAI_API_TYPE"] = "azure"
        openai_embeddings = OpenAIEmbeddings(model=os.environ.get("AZURE_EMBEDDINGS_DEPLOYMENT_NAME"))
    else:
        openai_embeddings = OpenAIEmbeddings()
    store = FAISS.from_documents(docs_test, openai_embeddings)
    # store_pine = Pinecone.from_documents(docs_test, OpenAIEmbeddings(), index_name=index_name)

    # Uncomment for MPNet embeddings
    # model_name = "sentence-transformers/all-mpnet-base-v2"
    # hf = HuggingFaceEmbeddings(model_name=model_name)
    # store = FAISS.from_documents(docs_test, hf)
    for i in tqdm(
        docs, desc="Embedding 🦖", unit="docs", total=len(docs), bar_format="{l_bar}{bar}| Time Left: {remaining}"
    ):
        try:
            store_add_texts_with_retry(store, i)
        except Exception as e:
            print(e)
            print("Error on ", i)
            print("Saving progress")
            print(f"stopped at {c1} out of {len(docs)}")
            store.save_local(f"outputs/{folder_name}")
            break
        c1 += 1
    store.save_local(f"outputs/{folder_name}")


def get_user_permission(docs, folder_name):
    # Function to ask user permission to call the OpenAI api and spend their OpenAI funds.
    # Here we convert the docs list to a string and calculate the number of OpenAI tokens the string represents.
    # docs_content = (" ".join(docs))
    docs_content = ""
    for doc in docs:
        docs_content += doc.page_content

    tokens, total_price = num_tokens_from_string(string=docs_content, encoding_name="cl100k_base")
    # Here we print the number of tokens and the approx user cost with some visually appealing formatting.
    print(f"Number of Tokens = {format(tokens, ',d')}")
    print(f"Approx Cost = ${format(total_price, ',.2f')}")
    # Here we check for user permission before calling the API.
    user_input = input("Price Okay? (Y/N) \n").lower()
    if user_input == "y":
        call_openai_api(docs, folder_name)
    elif user_input == "":
        call_openai_api(docs, folder_name)
    else:
        print("The API was not called. No money was spent.")



File: scripts/parser/file/markdown_parser.py

"""Markdown parser.

Contains parser for md files.

"""
import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union, cast

import tiktoken
from parser.file.base_parser import BaseParser


class MarkdownParser(BaseParser):
    """Markdown parser.

    Extract text from markdown files.
    Returns dictionary with keys as headers and values as the text between headers.

    """

    def __init__(
            self,
            *args: Any,
            remove_hyperlinks: bool = True,
            remove_images: bool = True,
            max_tokens: int = 2048,
            # remove_tables: bool = True,
            **kwargs: Any,
    ) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._remove_hyperlinks = remove_hyperlinks
        self._remove_images = remove_images
        self._max_tokens = max_tokens
        # self._remove_tables = remove_tables

    def tups_chunk_append(self, tups: List[Tuple[Optional[str], str]], current_header: Optional[str],
                          current_text: str):
        """Append to tups chunk."""
        num_tokens = len(tiktoken.get_encoding("cl100k_base").encode(current_text))
        if num_tokens > self._max_tokens:
            chunks = [current_text[i:i + self._max_tokens] for i in range(0, len(current_text), self._max_tokens)]
            for chunk in chunks:
                tups.append((current_header, chunk))
        else:
            tups.append((current_header, current_text))
        return tups

    def markdown_to_tups(self, markdown_text: str) -> List[Tuple[Optional[str], str]]:
        """Convert a markdown file to a dictionary.

        The keys are the headers and the values are the text under each header.

        """
        markdown_tups: List[Tuple[Optional[str], str]] = []
        lines = markdown_text.split("\n")

        current_header = None
        current_text = ""

        for line in lines:
            header_match = re.match(r"^#+\s", line)
            if header_match:
                if current_header is not None:
                    if current_text == "" or None:
                        continue
                    markdown_tups = self.tups_chunk_append(markdown_tups, current_header, current_text)

                current_header = line
                current_text = ""
            else:
                current_text += line + "\n"
        markdown_tups = self.tups_chunk_append(markdown_tups, current_header, current_text)

        if current_header is not None:
            # pass linting, assert keys are defined
            markdown_tups = [
                (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
                for key, value in markdown_tups
            ]
        else:
            markdown_tups = [
                (key, re.sub("\n", "", value)) for key, value in markdown_tups
            ]

        return markdown_tups

    def remove_images(self, content: str) -> str:
        """Get a dictionary of a markdown file from its path."""
        pattern = r"!{1}\[\[(.*)\]\]"
        content = re.sub(pattern, "", content)
        return content

    # def remove_tables(self, content: str) -> List[List[str]]:
    #     """Convert markdown tables to nested lists."""
    #     table_rows_pattern = r"((\r?\n){2}|^)([^\r\n]*\|[^\r\n]*(\r?\n)?)+(?=(\r?\n){2}|$)"
    #     table_cells_pattern = r"([^\|\r\n]*)\|"
    #
    #     table_rows = re.findall(table_rows_pattern, content, re.MULTILINE)
    #     table_lists = []
    #     for row in table_rows:
    #         cells = re.findall(table_cells_pattern, row[2])
    #         cells = [cell.strip() for cell in cells if cell.strip()]
    #         table_lists.append(cells)
    #     return str(table_lists)

    def remove_hyperlinks(self, content: str) -> str:
        """Get a dictionary of a markdown file from its path."""
        pattern = r"\[(.*?)\]\((.*?)\)"
        content = re.sub(pattern, r"\1", content)
        return content

    def _init_parser(self) -> Dict:
        """Initialize the parser with the config."""
        return {}

    def parse_tups(
            self, filepath: Path, errors: str = "ignore"
    ) -> List[Tuple[Optional[str], str]]:
        """Parse file into tuples."""
        with open(filepath, "r", encoding='utf8') as f:
            try:
                content = f.read()
            except (Exception,) as e:
                print(f'Error a file: "{filepath}"')
                raise e
        if self._remove_hyperlinks:
            content = self.remove_hyperlinks(content)
        if self._remove_images:
            content = self.remove_images(content)
        # if self._remove_tables:
        #     content = self.remove_tables(content)
        markdown_tups = self.markdown_to_tups(content)
        return markdown_tups

    def parse_file(
            self, filepath: Path, errors: str = "ignore"
    ) -> Union[str, List[str]]:
        """Parse file into string."""
        tups = self.parse_tups(filepath, errors=errors)
        results = []
        # TODO: don't include headers right now
        for header, value in tups:
            if header is None:
                results.append(value)
            else:
                results.append(f"\n\n{header}\n{value}")
        return results



File: scripts/parser/file/html_parser.py

"""HTML parser.

Contains parser for html files.

"""
import re
from pathlib import Path
from typing import Dict, Union

from parser.file.base_parser import BaseParser


class HTMLParser(BaseParser):
    """HTML parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, list[str]]:
        """Parse file.

            Returns:
            Union[str, List[str]]: a string or a List of strings.
        """
        try:
            from unstructured.partition.html import partition_html
            from unstructured.staging.base import convert_to_isd
            from unstructured.cleaners.core import clean
        except ImportError:
            raise ValueError("unstructured package is required to parse HTML files.")

        # Using the unstructured library to convert the html to isd format
        # isd sample : isd = [
        #   {"text": "My Title", "type": "Title"},
        #   {"text": "My Narrative", "type": "NarrativeText"}
        # ]
        with open(file, "r", encoding="utf-8") as fp:
            elements = partition_html(file=fp)
            isd = convert_to_isd(elements)

        # Removing non ascii charactwers from isd_el['text']
        for isd_el in isd:
            isd_el['text'] = isd_el['text'].encode("ascii", "ignore").decode()

        # Removing all the \n characters from isd_el['text'] using regex and replace with single space
        # Removing all the extra spaces  from isd_el['text'] using regex and replace with single space
        for isd_el in isd:
            isd_el['text'] = re.sub(r'\n', ' ', isd_el['text'], flags=re.MULTILINE | re.DOTALL)
            isd_el['text'] = re.sub(r"\s{2,}", " ", isd_el['text'], flags=re.MULTILINE | re.DOTALL)

        # more cleaning: extra_whitespaces, dashes, bullets, trailing_punctuation
        for isd_el in isd:
            clean(isd_el['text'], extra_whitespace=True, dashes=True, bullets=True, trailing_punctuation=True)

        # Creating a list of all the indexes of isd_el['type'] = 'Title'
        title_indexes = [i for i, isd_el in enumerate(isd) if isd_el['type'] == 'Title']

        # Creating 'Chunks' - List of lists of strings 
        # each list starting with isd_el['type'] = 'Title' and all the data till the next 'Title'
        # Each Chunk can be thought of as an individual set of data, which can be sent to the model
        # Where Each Title is grouped together with the data under it

        Chunks = [[]]
        final_chunks = list(list())

        for i, isd_el in enumerate(isd):
            if i in title_indexes:
                Chunks.append([])
            Chunks[-1].append(isd_el['text'])

        # Removing all the chunks with sum of length of all the strings in the chunk < 25
        # TODO: This value can be a user defined variable
        for chunk in Chunks:
            # sum of length of all the strings in the chunk
            sum = 0
            sum += len(str(chunk))
            if sum < 25:
                Chunks.remove(chunk)
            else:
                # appending all the approved chunks to final_chunks as a single string       
                final_chunks.append(" ".join([str(item) for item in chunk]))
        return final_chunks



File: scripts/parser/file/rst_parser.py

"""reStructuredText parser.

Contains parser for md files.

"""
import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from parser.file.base_parser import BaseParser


class RstParser(BaseParser):
    """reStructuredText parser.

    Extract text from .rst files.
    Returns dictionary with keys as headers and values as the text between headers.

    """

    def __init__(
            self,
            *args: Any,
            remove_hyperlinks: bool = True,
            remove_images: bool = True,
            remove_table_excess: bool = True,
            remove_interpreters: bool = True,
            remove_directives: bool = True,
            remove_whitespaces_excess: bool = True,
            # Be careful with remove_characters_excess, might cause data loss
            remove_characters_excess: bool = True,
            **kwargs: Any,
    ) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._remove_hyperlinks = remove_hyperlinks
        self._remove_images = remove_images
        self._remove_table_excess = remove_table_excess
        self._remove_interpreters = remove_interpreters
        self._remove_directives = remove_directives
        self._remove_whitespaces_excess = remove_whitespaces_excess
        self._remove_characters_excess = remove_characters_excess

    def rst_to_tups(self, rst_text: str) -> List[Tuple[Optional[str], str]]:
        """Convert a reStructuredText file to a dictionary.

        The keys are the headers and the values are the text under each header.

        """
        rst_tups: List[Tuple[Optional[str], str]] = []
        lines = rst_text.split("\n")

        current_header = None
        current_text = ""

        for i, line in enumerate(lines):
            header_match = re.match(r"^[^\S\n]*[-=]+[^\S\n]*$", line)
            if header_match and i > 0 and (
                    len(lines[i - 1].strip()) == len(header_match.group().strip()) or lines[i - 2] == lines[i - 2]):
                if current_header is not None:
                    if current_text == "" or None:
                        continue
                    # removes the next heading from current Document
                    if current_text.endswith(lines[i - 1] + "\n"):
                        current_text = current_text[:len(current_text) - len(lines[i - 1] + "\n")]
                    rst_tups.append((current_header, current_text))

                current_header = lines[i - 1]
                current_text = ""
            else:
                current_text += line + "\n"

        rst_tups.append((current_header, current_text))

        # TODO: Format for rst
        #
        # if current_header is not None:
        #     # pass linting, assert keys are defined
        #     rst_tups = [
        #         (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
        #         for key, value in rst_tups
        #     ]
        # else:
        #     rst_tups = [
        #         (key, re.sub("\n", "", value)) for key, value in rst_tups
        #     ]

        if current_header is None:
            rst_tups = [
                (key, re.sub("\n", "", value)) for key, value in rst_tups
            ]
        return rst_tups

    def remove_images(self, content: str) -> str:
        pattern = r"\.\. image:: (.*)"
        content = re.sub(pattern, "", content)
        return content

    def remove_hyperlinks(self, content: str) -> str:
        pattern = r"`(.*?) <(.*?)>`_"
        content = re.sub(pattern, r"\1", content)
        return content

    def remove_directives(self, content: str) -> str:
        """Removes reStructuredText Directives"""
        pattern = r"`\.\.([^:]+)::"
        content = re.sub(pattern, "", content)
        return content

    def remove_interpreters(self, content: str) -> str:
        """Removes reStructuredText Interpreted Text Roles"""
        pattern = r":(\w+):"
        content = re.sub(pattern, "", content)
        return content

    def remove_table_excess(self, content: str) -> str:
        """Pattern to remove grid table separators"""
        pattern = r"^\+[-]+\+[-]+\+$"
        content = re.sub(pattern, "", content, flags=re.MULTILINE)
        return content

    def remove_whitespaces_excess(self, content: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
        """Pattern to match 2 or more consecutive whitespaces"""
        pattern = r"\s{2,}"
        content = [(key, re.sub(pattern, "  ", value)) for key, value in content]
        return content

    def remove_characters_excess(self, content: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
        """Pattern to match 2 or more consecutive characters"""
        pattern = r"(\S)\1{2,}"
        content = [(key, re.sub(pattern, r"\1\1\1", value, flags=re.MULTILINE)) for key, value in content]
        return content

    def _init_parser(self) -> Dict:
        """Initialize the parser with the config."""
        return {}

    def parse_tups(
            self, filepath: Path, errors: str = "ignore"
    ) -> List[Tuple[Optional[str], str]]:
        """Parse file into tuples."""
        with open(filepath, "r") as f:
            content = f.read()
        if self._remove_hyperlinks:
            content = self.remove_hyperlinks(content)
        if self._remove_images:
            content = self.remove_images(content)
        if self._remove_table_excess:
            content = self.remove_table_excess(content)
        if self._remove_directives:
            content = self.remove_directives(content)
        if self._remove_interpreters:
            content = self.remove_interpreters(content)
        rst_tups = self.rst_to_tups(content)
        if self._remove_whitespaces_excess:
            rst_tups = self.remove_whitespaces_excess(rst_tups)
        if self._remove_characters_excess:
            rst_tups = self.remove_characters_excess(rst_tups)
        return rst_tups

    def parse_file(
            self, filepath: Path, errors: str = "ignore"
    ) -> Union[str, List[str]]:
        """Parse file into string."""
        tups = self.parse_tups(filepath, errors=errors)
        results = []
        # TODO: don't include headers right now
        for header, value in tups:
            if header is None:
                results.append(value)
            else:
                results.append(f"\n\n{header}\n{value}")
        return results



File: scripts/parser/file/bulk.py

"""Simple reader that reads files of different formats from a directory."""
import logging
from parser.file.base import BaseReader
from parser.file.base_parser import BaseParser
from parser.file.docs_parser import DocxParser, PDFParser
from parser.file.epub_parser import EpubParser
from parser.file.html_parser import HTMLParser
from parser.file.markdown_parser import MarkdownParser
from parser.file.rst_parser import RstParser
from parser.file.tabular_parser import PandasCSVParser
from parser.schema.base import Document
from pathlib import Path
from typing import Callable, Dict, List, Optional, Union

DEFAULT_FILE_EXTRACTOR: Dict[str, BaseParser] = {
    ".pdf": PDFParser(),
    ".docx": DocxParser(),
    ".csv": PandasCSVParser(),
    ".epub": EpubParser(),
    ".md": MarkdownParser(),
    ".rst": RstParser(),
    ".html": HTMLParser(),
    ".mdx": MarkdownParser(),
}


class SimpleDirectoryReader(BaseReader):
    """Simple directory reader.

    Can read files into separate documents, or concatenates
    files into one document text.

    Args:
        input_dir (str): Path to the directory.
        input_files (List): List of file paths to read (Optional; overrides input_dir)
        exclude_hidden (bool): Whether to exclude hidden files (dotfiles).
        errors (str): how encoding and decoding errors are to be handled,
              see https://docs.python.org/3/library/functions.html#open
        recursive (bool): Whether to recursively search in subdirectories.
            False by default.
        required_exts (Optional[List[str]]): List of required extensions.
            Default is None.
        file_extractor (Optional[Dict[str, BaseParser]]): A mapping of file
            extension to a BaseParser class that specifies how to convert that file
            to text. See DEFAULT_FILE_EXTRACTOR.
        num_files_limit (Optional[int]): Maximum number of files to read.
            Default is None.
        file_metadata (Optional[Callable[str, Dict]]): A function that takes
            in a filename and returns a Dict of metadata for the Document.
            Default is None.
    """

    def __init__(
            self,
            input_dir: Optional[str] = None,
            input_files: Optional[List] = None,
            exclude_hidden: bool = True,
            errors: str = "ignore",
            recursive: bool = True,
            required_exts: Optional[List[str]] = None,
            file_extractor: Optional[Dict[str, BaseParser]] = None,
            num_files_limit: Optional[int] = None,
            file_metadata: Optional[Callable[[str], Dict]] = None,
    ) -> None:
        """Initialize with parameters."""
        super().__init__()

        if not input_dir and not input_files:
            raise ValueError("Must provide either `input_dir` or `input_files`.")

        self.errors = errors

        self.recursive = recursive
        self.exclude_hidden = exclude_hidden
        self.required_exts = required_exts
        self.num_files_limit = num_files_limit
        print("input_files")
        print(input_files)

        if input_files:
            self.input_files = []
            for path in input_files:
                input_file = Path(path)
                self.input_files.append(input_file)
        elif input_dir:
            self.input_dir = Path(input_dir)
            self.input_files = self._add_files(self.input_dir)

        self.file_extractor = file_extractor or DEFAULT_FILE_EXTRACTOR
        self.file_metadata = file_metadata

    def _add_files(self, input_dir: Path) -> List[Path]:
        """Add files."""
        input_files = sorted(input_dir.iterdir())
        new_input_files = []
        dirs_to_explore = []
        for input_file in input_files:
            if input_file.is_dir():
                if self.recursive:
                    dirs_to_explore.append(input_file)
            elif self.exclude_hidden and input_file.name.startswith("."):
                continue
            elif (
                    self.required_exts is not None
                    and input_file.suffix not in self.required_exts
            ):
                continue
            else:
                new_input_files.append(input_file)

        for dir_to_explore in dirs_to_explore:
            sub_input_files = self._add_files(dir_to_explore)
            new_input_files.extend(sub_input_files)

        if self.num_files_limit is not None and self.num_files_limit > 0:
            new_input_files = new_input_files[0: self.num_files_limit]

        # print total number of files added
        logging.debug(
            f"> [SimpleDirectoryReader] Total files added: {len(new_input_files)}"
        )

        return new_input_files

    def load_data(self, concatenate: bool = False) -> List[Document]:
        """Load data from the input directory.

        Args:
            concatenate (bool): whether to concatenate all files into one document.
                If set to True, file metadata is ignored.
                False by default.

        Returns:
            List[Document]: A list of documents.

        """
        data: Union[str, List[str]] = ""
        data_list: List[str] = []
        metadata_list = []
        for input_file in self.input_files:
            if input_file.suffix in self.file_extractor:
                parser = self.file_extractor[input_file.suffix]
                if not parser.parser_config_set:
                    parser.init_parser()
                data = parser.parse_file(input_file, errors=self.errors)
            else:
                # do standard read
                with open(input_file, "r", errors=self.errors) as f:
                    data = f.read()
            if isinstance(data, List):
                data_list.extend(data)
                if self.file_metadata is not None:
                    for _ in range(len(data)):
                        metadata_list.append(self.file_metadata(str(input_file)))
            else:
                data_list.append(str(data))
                if self.file_metadata is not None:
                    metadata_list.append(self.file_metadata(str(input_file)))

            

        if concatenate:
            return [Document("\n".join(data_list))]
        elif self.file_metadata is not None:
            return [Document(d, extra_info=m) for d, m in zip(data_list, metadata_list)]
        else:
            return [Document(d) for d in data_list]



File: scripts/parser/file/__init__.py




File: scripts/parser/file/docs_parser.py

"""Docs parser.

Contains parsers for docx, pdf files.

"""
from pathlib import Path
from typing import Dict

from parser.file.base_parser import BaseParser


class PDFParser(BaseParser):
    """PDF parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> str:
        """Parse file."""
        try:
            import PyPDF2
        except ImportError:
            raise ValueError("PyPDF2 is required to read PDF files.")
        text_list = []
        with open(file, "rb") as fp:
            # Create a PDF object
            pdf = PyPDF2.PdfReader(fp)

            # Get the number of pages in the PDF document
            num_pages = len(pdf.pages)

            # Iterate over every page
            for page in range(num_pages):
                # Extract the text from the page
                page_text = pdf.pages[page].extract_text()
                text_list.append(page_text)
        text = "\n".join(text_list)

        return text


class DocxParser(BaseParser):
    """Docx parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> str:
        """Parse file."""
        try:
            import docx2txt
        except ImportError:
            raise ValueError("docx2txt is required to read Microsoft Word files.")

        text = docx2txt.process(file)

        return text



File: scripts/parser/file/base_parser.py

"""Base parser and config class."""

from abc import abstractmethod
from pathlib import Path
from typing import Dict, List, Optional, Union


class BaseParser:
    """Base class for all parsers."""

    def __init__(self, parser_config: Optional[Dict] = None):
        """Init params."""
        self._parser_config = parser_config

    def init_parser(self) -> None:
        """Init parser and store it."""
        parser_config = self._init_parser()
        self._parser_config = parser_config

    @property
    def parser_config_set(self) -> bool:
        """Check if parser config is set."""
        return self._parser_config is not None

    @property
    def parser_config(self) -> Dict:
        """Check if parser config is set."""
        if self._parser_config is None:
            raise ValueError("Parser config not set.")
        return self._parser_config

    @abstractmethod
    def _init_parser(self) -> Dict:
        """Initialize the parser with the config."""

    @abstractmethod
    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
        """Parse file."""



File: scripts/parser/file/tabular_parser.py

"""Tabular parser.

Contains parsers for tabular data files.

"""
from pathlib import Path
from typing import Any, Dict, List, Union

from parser.file.base_parser import BaseParser


class CSVParser(BaseParser):
    """CSV parser.

    Args:
        concat_rows (bool): whether to concatenate all rows into one document.
            If set to False, a Document will be created for each row.
            True by default.

    """

    def __init__(self, *args: Any, concat_rows: bool = True, **kwargs: Any) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._concat_rows = concat_rows

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
        """Parse file.

        Returns:
            Union[str, List[str]]: a string or a List of strings.

        """
        try:
            import csv
        except ImportError:
            raise ValueError("csv module is required to read CSV files.")
        text_list = []
        with open(file, "r") as fp:
            csv_reader = csv.reader(fp)
            for row in csv_reader:
                text_list.append(", ".join(row))
        if self._concat_rows:
            return "\n".join(text_list)
        else:
            return text_list


class PandasCSVParser(BaseParser):
    r"""Pandas-based CSV parser.

    Parses CSVs using the separator detection from Pandas `read_csv`function.
    If special parameters are required, use the `pandas_config` dict.

    Args:
        concat_rows (bool): whether to concatenate all rows into one document.
            If set to False, a Document will be created for each row.
            True by default.

        col_joiner (str): Separator to use for joining cols per row.
            Set to ", " by default.

        row_joiner (str): Separator to use for joining each row.
            Only used when `concat_rows=True`.
            Set to "\n" by default.

        pandas_config (dict): Options for the `pandas.read_csv` function call.
            Refer to https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html
            for more information.
            Set to empty dict by default, this means pandas will try to figure
            out the separators, table head, etc. on its own.

    """

    def __init__(
            self,
            *args: Any,
            concat_rows: bool = True,
            col_joiner: str = ", ",
            row_joiner: str = "\n",
            pandas_config: dict = {},
            **kwargs: Any
    ) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._concat_rows = concat_rows
        self._col_joiner = col_joiner
        self._row_joiner = row_joiner
        self._pandas_config = pandas_config

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
        """Parse file."""
        try:
            import pandas as pd
        except ImportError:
            raise ValueError("pandas module is required to read CSV files.")

        df = pd.read_csv(file, **self._pandas_config)

        text_list = df.apply(
            lambda row: (self._col_joiner).join(row.astype(str).tolist()), axis=1
        ).tolist()

        if self._concat_rows:
            return (self._row_joiner).join(text_list)
        else:
            return text_list



File: scripts/parser/file/openapi3_parser.py

from urllib.parse import urlparse

from openapi_parser import parse

try:
    from scripts.parser.file.base_parser import BaseParser
except ModuleNotFoundError:
    from base_parser import BaseParser


class OpenAPI3Parser(BaseParser):
    def init_parser(self) -> None:
        return super().init_parser()

    def get_base_urls(self, urls):
        base_urls = []
        for i in urls:
            parsed_url = urlparse(i)
            base_url = parsed_url.scheme + "://" + parsed_url.netloc
            if base_url not in base_urls:
                base_urls.append(base_url)
        return base_urls

    def get_info_from_paths(self, path):
        info = ""
        if path.operations:
            for operation in path.operations:
                info += (
                    f"\n{operation.method.value}="
                    f"{operation.responses[0].description}"
                )
        return info

    def parse_file(self, file_path):
        data = parse(file_path)
        results = ""
        base_urls = self.get_base_urls(link.url for link in data.servers)
        base_urls = ",".join([base_url for base_url in base_urls])
        results += f"Base URL:{base_urls}\n"
        i = 1
        for path in data.paths:
            info = self.get_info_from_paths(path)
            results += (
                f"Path{i}: {path.url}\n"
                f"description: {path.description}\n"
                f"parameters: {path.parameters}\nmethods: {info}\n"
            )
            i += 1
        with open("results.txt", "w") as f:
            f.write(results)
        return results



File: scripts/parser/file/epub_parser.py

"""Epub parser.

Contains parsers for epub files.
"""

from pathlib import Path
from typing import Dict

from parser.file.base_parser import BaseParser


class EpubParser(BaseParser):
    """Epub Parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> str:
        """Parse file."""
        try:
            import ebooklib
            from ebooklib import epub
        except ImportError:
            raise ValueError("`EbookLib` is required to read Epub files.")
        try:
            import html2text
        except ImportError:
            raise ValueError("`html2text` is required to parse Epub files.")

        text_list = []
        book = epub.read_epub(file, options={"ignore_ncx": True})

        # Iterate through all chapters.
        for item in book.get_items():
            # Chapters are typically located in epub documents items.
            if item.get_type() == ebooklib.ITEM_DOCUMENT:
                text_list.append(
                    html2text.html2text(item.get_content().decode("utf-8"))
                )

        text = "\n".join(text_list)
        return text



File: scripts/parser/file/base.py

"""Base reader class."""
from abc import abstractmethod
from typing import Any, List

from langchain.docstore.document import Document as LCDocument
from parser.schema.base import Document


class BaseReader:
    """Utilities for loading data from a directory."""

    @abstractmethod
    def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]:
        """Load data from the input directory."""

    def load_langchain_documents(self, **load_kwargs: Any) -> List[LCDocument]:
        """Load data in LangChain document format."""
        docs = self.load_data(**load_kwargs)
        return [d.to_langchain_format() for d in docs]



File: scripts/parser/schema/__init__.py




File: scripts/parser/schema/base.py

"""Base schema for readers."""
from dataclasses import dataclass

from langchain.docstore.document import Document as LCDocument
from parser.schema.schema import BaseDocument


@dataclass
class Document(BaseDocument):
    """Generic interface for a data document.

    This document connects to data sources.

    """

    def __post_init__(self) -> None:
        """Post init."""
        if self.text is None:
            raise ValueError("text field not set.")

    @classmethod
    def get_type(cls) -> str:
        """Get Document type."""
        return "Document"

    def to_langchain_format(self) -> LCDocument:
        """Convert struct to LangChain document format."""
        metadata = self.extra_info or {}
        return LCDocument(page_content=self.text, metadata=metadata)

    @classmethod
    def from_langchain_format(cls, doc: LCDocument) -> "Document":
        """Convert struct from LangChain document format."""
        return cls(text=doc.page_content, extra_info=doc.metadata)



File: scripts/parser/schema/schema.py

"""Base schema for data structures."""
from abc import abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from dataclasses_json import DataClassJsonMixin


@dataclass
class BaseDocument(DataClassJsonMixin):
    """Base document.

    Generic abstract interfaces that captures both index structs
    as well as documents.

    """

    # TODO: consolidate fields from Document/IndexStruct into base class
    text: Optional[str] = None
    doc_id: Optional[str] = None
    embedding: Optional[List[float]] = None

    # extra fields
    extra_info: Optional[Dict[str, Any]] = None

    @classmethod
    @abstractmethod
    def get_type(cls) -> str:
        """Get Document type."""

    def get_text(self) -> str:
        """Get text."""
        if self.text is None:
            raise ValueError("text field not set.")
        return self.text

    def get_doc_id(self) -> str:
        """Get doc_id."""
        if self.doc_id is None:
            raise ValueError("doc_id not set.")
        return self.doc_id

    @property
    def is_doc_id_none(self) -> bool:
        """Check if doc_id is None."""
        return self.doc_id is None

    def get_embedding(self) -> List[float]:
        """Get embedding.

        Errors if embedding is None.

        """
        if self.embedding is None:
            raise ValueError("embedding not set.")
        return self.embedding

    @property
    def extra_info_str(self) -> Optional[str]:
        """Extra info string."""
        if self.extra_info is None:
            return None

        return "\n".join([f"{k}: {str(v)}" for k, v in self.extra_info.items()])



File: application/worker.py

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

import requests
from bson.objectid import ObjectId

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
from application.utils import count_tokens_docs



# Define a function to extract metadata from a given filename.
def metadata_from_filename(title):
    return {"title": title}


# 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__))))


def extract_zip_recursive(zip_path, extract_to, current_depth=0, max_depth=5):
    """
    Recursively extract zip files with a limit on recursion depth.

    Args:
        zip_path (str): Path to the zip file to be extracted.
        extract_to (str): Destination path for extracted files.
        current_depth (int): Current depth of recursion.
        max_depth (int): Maximum allowed depth of recursion to prevent infinite loops.
    """
    if current_depth > max_depth:
        logging.warning(f"Reached maximum recursion depth of {max_depth}")
        return

    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(extract_to)
    os.remove(zip_path)  # Remove the zip file after extracting

    # Check for nested zip files and extract them
    for root, dirs, files in os.walk(extract_to):
        for file in files:
            if file.endswith(".zip"):
                # If a nested zip file is found, extract it recursively
                file_path = os.path.join(root, file)
                extract_zip_recursive(file_path, root, current_depth + 1, max_depth)


# Define the main function for ingesting and processing documents.
def ingest_worker(self, directory, formats, name_job, filename, user, retriever="classic"):
    """
    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.
        retriever (str): Type of retriever to use for processing the documents.

    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
    recursion_depth = 2
    full_path = os.path.join(directory, user, name_job)

    logging.info(f"Ingest file: {full_path}", extra={"user": user, "job": name_job})
    # 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
    )
    file = response.content

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

    # check if file is .zip and extract it
    if filename.endswith(".zip"):
        extract_zip_recursive(os.path.join(full_path, filename), full_path, 0, recursion_depth)

    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]
    id = ObjectId()

    call_openai_api(docs, full_path, id, self)
    tokens = count_tokens_docs(docs)
    self.update_state(state="PROGRESS", meta={"current": 100})

    if sample:
        for i in range(min(5, len(raw_docs))):
            logging.info(f"Sample document {i}: {raw_docs[i]}")

    # 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, "tokens": tokens, "retriever": retriever, "id": str(id), 'type': 'local'}
    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)
    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, loader, directory="temp", retriever="classic"):
    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})
    logging.info(f"Remote job: {full_path}", extra={"user": user, "job": name_job, source_data: source_data})

    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]
    tokens = count_tokens_docs(docs)
    id = ObjectId()
    call_openai_api(docs, full_path, id, 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, "tokens": tokens, "retriever": retriever, 
                 "id": str(id), 'type': loader, 'remote_data': source_data}
    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)
    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}



File: application/error.py

from flask import jsonify
from werkzeug.http import HTTP_STATUS_CODES


def response_error(code_status, message=None):
    payload = {'error': HTTP_STATUS_CODES.get(code_status, "something went wrong")}
    if message:
        payload['message'] = message
    response = jsonify(payload)
    response.status_code = code_status
    return response


def bad_request(status_code=400, message=''):
    return response_error(code_status=status_code, message=message)



File: application/celery_init.py

from celery import Celery
from application.core.settings import settings
from celery.signals import setup_logging

def make_celery(app_name=__name__):
    celery = Celery(app_name, broker=settings.CELERY_BROKER_URL, backend=settings.CELERY_RESULT_BACKEND)
    celery.conf.update(settings)
    return celery

@setup_logging.connect
def config_loggers(*args, **kwargs):
    from application.core.logging_config import setup_logging
    setup_logging()

celery = make_celery()



File: application/__init__.py




File: application/celeryconfig.py

import os

broker_url = os.getenv("CELERY_BROKER_URL")
result_backend = os.getenv("CELERY_RESULT_BACKEND")

task_serializer = 'json'
result_serializer = 'json'
accept_content = ['json']



File: application/utils.py

import tiktoken

_encoding = None

def get_encoding():
    global _encoding
    if _encoding is None:
        _encoding = tiktoken.get_encoding("cl100k_base")
    return _encoding

def num_tokens_from_string(string: str) -> int:
    encoding = get_encoding()
    num_tokens = len(encoding.encode(string))
    return num_tokens

def count_tokens_docs(docs):
    docs_content = ""
    for doc in docs:
        docs_content += doc.page_content

    tokens = num_tokens_from_string(docs_content)
    return tokens


File: application/app.py

import platform
import dotenv
from application.celery_init import celery
from flask import Flask, request, redirect
from application.core.settings import settings
from application.api.user.routes import user
from application.api.answer.routes import answer
from application.api.internal.routes import internal
from application.core.logging_config import setup_logging

if platform.system() == "Windows":
    import pathlib
    pathlib.PosixPath = pathlib.WindowsPath

dotenv.load_dotenv()
setup_logging()

app = Flask(__name__)
app.register_blueprint(user)
app.register_blueprint(answer)
app.register_blueprint(internal)
app.config.update(
    UPLOAD_FOLDER="inputs",
    CELERY_BROKER_URL=settings.CELERY_BROKER_URL,
    CELERY_RESULT_BACKEND=settings.CELERY_RESULT_BACKEND,
    MONGO_URI=settings.MONGO_URI
)
celery.config_from_object("application.celeryconfig")

@app.route("/")
def home():
    if request.remote_addr in ('0.0.0.0', '127.0.0.1', 'localhost', '172.18.0.1'):
        return redirect('http://localhost:5173')
    else:
        return 'Welcome to DocsGPT Backend!'

@app.after_request
def after_request(response):
    response.headers.add("Access-Control-Allow-Origin", "*")
    response.headers.add("Access-Control-Allow-Headers", "Content-Type,Authorization")
    response.headers.add("Access-Control-Allow-Methods", "GET,PUT,POST,DELETE,OPTIONS")
    return response

if __name__ == "__main__":
    app.run(debug=settings.FLASK_DEBUG_MODE, port=7091)




File: application/usage.py

import sys
from pymongo import MongoClient
from datetime import datetime
from application.core.settings import settings
from application.utils import num_tokens_from_string

mongo = MongoClient(settings.MONGO_URI)
db = mongo["docsgpt"]
usage_collection = db["token_usage"]


def update_token_usage(user_api_key, token_usage):
    if "pytest" in sys.modules:
        return
    usage_data = {
        "api_key": user_api_key,
        "prompt_tokens": token_usage["prompt_tokens"],
        "generated_tokens": token_usage["generated_tokens"],
        "timestamp": datetime.now(),
    }
    usage_collection.insert_one(usage_data)


def gen_token_usage(func):
    def wrapper(self, model, messages, stream, **kwargs):
        for message in messages:
            self.token_usage["prompt_tokens"] += num_tokens_from_string(message["content"])
        result = func(self, model, messages, stream, **kwargs)
        self.token_usage["generated_tokens"] += num_tokens_from_string(result)
        update_token_usage(self.user_api_key, self.token_usage)
        return result

    return wrapper


def stream_token_usage(func):
    def wrapper(self, model, messages, stream, **kwargs):
        for message in messages:
            self.token_usage["prompt_tokens"] += num_tokens_from_string(message["content"])
        batch = []
        result = func(self, model, messages, stream, **kwargs)
        for r in result:
            batch.append(r)
            yield r
        for line in batch:
            self.token_usage["generated_tokens"] += num_tokens_from_string(line)
        update_token_usage(self.user_api_key, self.token_usage)

    return wrapper



File: application/wsgi.py

from application.app import app
from application.core.settings import settings

if __name__ == "__main__":
    app.run(debug=settings.FLASK_DEBUG_MODE, port=7091)



File: application/llm/premai.py

from application.llm.base import BaseLLM
from application.core.settings import settings


class PremAILLM(BaseLLM):

    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
        from premai import Prem

        super().__init__(*args, **kwargs)
        self.client = Prem(api_key=api_key)
        self.api_key = api_key
        self.user_api_key = user_api_key
        self.project_id = settings.PREMAI_PROJECT_ID

    def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
        response = self.client.chat.completions.create(
            model=model,
            project_id=self.project_id,
            messages=messages,
            stream=stream,
            **kwargs
        )

        return response.choices[0].message["content"]

    def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):
        response = self.client.chat.completions.create(
            model=model,
            project_id=self.project_id,
            messages=messages,
            stream=stream,
            **kwargs
        )

        for line in response:
            if line.choices[0].delta["content"] is not None:
                yield line.choices[0].delta["content"]



File: application/llm/llm_creator.py

from application.llm.openai import OpenAILLM, AzureOpenAILLM
from application.llm.sagemaker import SagemakerAPILLM
from application.llm.huggingface import HuggingFaceLLM
from application.llm.llama_cpp import LlamaCpp
from application.llm.anthropic import AnthropicLLM
from application.llm.docsgpt_provider import DocsGPTAPILLM
from application.llm.premai import PremAILLM


class LLMCreator:
    llms = {
        "openai": OpenAILLM,
        "azure_openai": AzureOpenAILLM,
        "sagemaker": SagemakerAPILLM,
        "huggingface": HuggingFaceLLM,
        "llama.cpp": LlamaCpp,
        "anthropic": AnthropicLLM,
        "docsgpt": DocsGPTAPILLM,
        "premai": PremAILLM,
    }

    @classmethod
    def create_llm(cls, type, api_key, user_api_key, *args, **kwargs):
        llm_class = cls.llms.get(type.lower())
        if not llm_class:
            raise ValueError(f"No LLM class found for type {type}")
        return llm_class(api_key, user_api_key, *args, **kwargs)



File: application/llm/__init__.py




File: application/llm/llama_cpp.py

from application.llm.base import BaseLLM
from application.core.settings import settings
import threading

class LlamaSingleton:
    _instances = {}
    _lock = threading.Lock()  # Add a lock for thread synchronization

    @classmethod
    def get_instance(cls, llm_name):
        if llm_name not in cls._instances:
            try:
                from llama_cpp import Llama
            except ImportError:
                raise ImportError(
                    "Please install llama_cpp using pip install llama-cpp-python"
                )
            cls._instances[llm_name] = Llama(model_path=llm_name, n_ctx=2048)
        return cls._instances[llm_name]

    @classmethod
    def query_model(cls, llm, prompt, **kwargs):
        with cls._lock:
            return llm(prompt, **kwargs)


class LlamaCpp(BaseLLM):
    def __init__(
        self,
        api_key=None,
        user_api_key=None,
        llm_name=settings.MODEL_PATH,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.api_key = api_key
        self.user_api_key = user_api_key
        self.llama = LlamaSingleton.get_instance(llm_name)

    def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
        result = LlamaSingleton.query_model(self.llama, prompt, max_tokens=150, echo=False)
        return result["choices"][0]["text"].split("### Answer \n")[-1]

    def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
        result = LlamaSingleton.query_model(self.llama, prompt, max_tokens=150, echo=False, stream=stream)
        for item in result:
            for choice in item["choices"]:
                yield choice["text"]


File: application/llm/docsgpt_provider.py

from application.llm.base import BaseLLM
import json
import requests


class DocsGPTAPILLM(BaseLLM):

    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.api_key = api_key
        self.user_api_key = user_api_key
        self.endpoint = "https://llm.docsgpt.co.uk"

    def _raw_gen(self, baseself, model, messages, stream=False, *args, **kwargs):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"

        response = requests.post(
            f"{self.endpoint}/answer", json={"prompt": prompt, "max_new_tokens": 30}
        )
        response_clean = response.json()["a"].replace("###", "")

        return response_clean

    def _raw_gen_stream(self, baseself, model, messages, stream=True, *args, **kwargs):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"

        # send prompt to endpoint /stream
        response = requests.post(
            f"{self.endpoint}/stream",
            json={"prompt": prompt, "max_new_tokens": 256},
            stream=True,
        )

        for line in response.iter_lines():
            if line:
                # data = json.loads(line)
                data_str = line.decode("utf-8")
                if data_str.startswith("data: "):
                    data = json.loads(data_str[6:])
                    yield data["a"]



File: application/llm/openai.py

from application.llm.base import BaseLLM
from application.core.settings import settings



class OpenAILLM(BaseLLM):

    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
        from openai import OpenAI

        super().__init__(*args, **kwargs)
        if settings.OPENAI_BASE_URL:
            self.client = OpenAI(
                api_key=api_key,
                base_url=settings.OPENAI_BASE_URL
            )
        else:
            self.client = OpenAI(api_key=api_key)
        self.api_key = api_key
        self.user_api_key = user_api_key

    def _raw_gen(
        self,
        baseself,
        model,
        messages,
        stream=False,
        engine=settings.AZURE_DEPLOYMENT_NAME,
        **kwargs
    ):  
        response = self.client.chat.completions.create(
            model=model, messages=messages, stream=stream, **kwargs
        )

        return response.choices[0].message.content

    def _raw_gen_stream(
        self,
        baseself,
        model,
        messages,
        stream=True,
        engine=settings.AZURE_DEPLOYMENT_NAME,
        **kwargs
    ):  
        response = self.client.chat.completions.create(
            model=model, messages=messages, stream=stream, **kwargs
        )

        for line in response:
            # import sys
            # print(line.choices[0].delta.content, file=sys.stderr)
            if line.choices[0].delta.content is not None:
                yield line.choices[0].delta.content


class AzureOpenAILLM(OpenAILLM):

    def __init__(
        self, openai_api_key, openai_api_base, openai_api_version, deployment_name
    ):
        super().__init__(openai_api_key)
        self.api_base = (settings.OPENAI_API_BASE,)
        self.api_version = (settings.OPENAI_API_VERSION,)
        self.deployment_name = (settings.AZURE_DEPLOYMENT_NAME,)
        from openai import AzureOpenAI

        self.client = AzureOpenAI(
            api_key=openai_api_key,
            api_version=settings.OPENAI_API_VERSION,
            api_base=settings.OPENAI_API_BASE,
            deployment_name=settings.AZURE_DEPLOYMENT_NAME,
        )



File: application/llm/sagemaker.py

from application.llm.base import BaseLLM
from application.core.settings import settings
import json
import io


class LineIterator:
    """
    A helper class for parsing the byte stream input.

    The output of the model will be in the following format:
    ```
    b'{"outputs": [" a"]}\n'
    b'{"outputs": [" challenging"]}\n'
    b'{"outputs": [" problem"]}\n'
    ...
    ```

    While usually each PayloadPart event from the event stream will contain a byte array
    with a full json, this is not guaranteed and some of the json objects may be split across
    PayloadPart events. For example:
    ```
    {'PayloadPart': {'Bytes': b'{"outputs": '}}
    {'PayloadPart': {'Bytes': b'[" problem"]}\n'}}
    ```

    This class accounts for this by concatenating bytes written via the 'write' function
    and then exposing a method which will return lines (ending with a '\n' character) within
    the buffer via the 'scan_lines' function. It maintains the position of the last read
    position to ensure that previous bytes are not exposed again.
    """

    def __init__(self, stream):
        self.byte_iterator = iter(stream)
        self.buffer = io.BytesIO()
        self.read_pos = 0

    def __iter__(self):
        return self

    def __next__(self):
        while True:
            self.buffer.seek(self.read_pos)
            line = self.buffer.readline()
            if line and line[-1] == ord("\n"):
                self.read_pos += len(line)
                return line[:-1]
            try:
                chunk = next(self.byte_iterator)
            except StopIteration:
                if self.read_pos < self.buffer.getbuffer().nbytes:
                    continue
                raise
            if "PayloadPart" not in chunk:
                print("Unknown event type:" + chunk)
                continue
            self.buffer.seek(0, io.SEEK_END)
            self.buffer.write(chunk["PayloadPart"]["Bytes"])


class SagemakerAPILLM(BaseLLM):

    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
        import boto3

        runtime = boto3.client(
            "runtime.sagemaker",
            aws_access_key_id="xxx",
            aws_secret_access_key="xxx",
            region_name="us-west-2",
        )

        super().__init__(*args, **kwargs)
        self.api_key = api_key
        self.user_api_key = user_api_key
        self.endpoint = settings.SAGEMAKER_ENDPOINT
        self.runtime = runtime

    def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"

        # Construct payload for endpoint
        payload = {
            "inputs": prompt,
            "stream": False,
            "parameters": {
                "do_sample": True,
                "temperature": 0.1,
                "max_new_tokens": 30,
                "repetition_penalty": 1.03,
                "stop": ["</s>", "###"],
            },
        }
        body_bytes = json.dumps(payload).encode("utf-8")

        # Invoke the endpoint
        response = self.runtime.invoke_endpoint(
            EndpointName=self.endpoint, ContentType="application/json", Body=body_bytes
        )
        result = json.loads(response["Body"].read().decode())
        import sys

        print(result[0]["generated_text"], file=sys.stderr)
        return result[0]["generated_text"][len(prompt) :]

    def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"

        # Construct payload for endpoint
        payload = {
            "inputs": prompt,
            "stream": True,
            "parameters": {
                "do_sample": True,
                "temperature": 0.1,
                "max_new_tokens": 512,
                "repetition_penalty": 1.03,
                "stop": ["</s>", "###"],
            },
        }
        body_bytes = json.dumps(payload).encode("utf-8")

        # Invoke the endpoint
        response = self.runtime.invoke_endpoint_with_response_stream(
            EndpointName=self.endpoint, ContentType="application/json", Body=body_bytes
        )
        # result = json.loads(response['Body'].read().decode())
        event_stream = response["Body"]
        start_json = b"{"
        for line in LineIterator(event_stream):
            if line != b"" and start_json in line:
                # print(line)
                data = json.loads(line[line.find(start_json) :].decode("utf-8"))
                if data["token"]["text"] not in ["</s>", "###"]:
                    print(data["token"]["text"], end="")
                    yield data["token"]["text"]



File: application/llm/huggingface.py

from application.llm.base import BaseLLM


class HuggingFaceLLM(BaseLLM):

    def __init__(
        self,
        api_key=None,
        user_api_key=None,
        llm_name="Arc53/DocsGPT-7B",
        q=False,
        *args,
        **kwargs,
    ):
        global hf

        from langchain.llms import HuggingFacePipeline

        if q:
            import torch
            from transformers import (
                AutoModelForCausalLM,
                AutoTokenizer,
                pipeline,
                BitsAndBytesConfig,
            )

            tokenizer = AutoTokenizer.from_pretrained(llm_name)
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16,
            )
            model = AutoModelForCausalLM.from_pretrained(
                llm_name, quantization_config=bnb_config
            )
        else:
            from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

            tokenizer = AutoTokenizer.from_pretrained(llm_name)
            model = AutoModelForCausalLM.from_pretrained(llm_name)

        super().__init__(*args, **kwargs)
        self.api_key = api_key
        self.user_api_key = user_api_key
        pipe = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer,
            max_new_tokens=2000,
            device_map="auto",
            eos_token_id=tokenizer.eos_token_id,
        )
        hf = HuggingFacePipeline(pipeline=pipe)

    def _raw_gen(self, baseself, model, messages, stream=False, **kwargs):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"

        result = hf(prompt)

        return result.content

    def _raw_gen_stream(self, baseself, model, messages, stream=True, **kwargs):

        raise NotImplementedError("HuggingFaceLLM Streaming is not implemented yet.")



File: application/llm/anthropic.py

from application.llm.base import BaseLLM
from application.core.settings import settings


class AnthropicLLM(BaseLLM):

    def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
        from anthropic import Anthropic, HUMAN_PROMPT, AI_PROMPT

        super().__init__(*args, **kwargs)
        self.api_key = (
            api_key or settings.ANTHROPIC_API_KEY
        )  # If not provided, use a default from settings
        self.user_api_key = user_api_key
        self.anthropic = Anthropic(api_key=self.api_key)
        self.HUMAN_PROMPT = HUMAN_PROMPT
        self.AI_PROMPT = AI_PROMPT

    def _raw_gen(
        self, baseself, model, messages, stream=False, max_tokens=300, **kwargs
    ):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Context \n {context} \n ### Question \n {user_question}"
        if stream:
            return self.gen_stream(model, prompt, stream, max_tokens, **kwargs)

        completion = self.anthropic.completions.create(
            model=model,
            max_tokens_to_sample=max_tokens,
            stream=stream,
            prompt=f"{self.HUMAN_PROMPT} {prompt}{self.AI_PROMPT}",
        )
        return completion.completion

    def _raw_gen_stream(
        self, baseself, model, messages, stream=True, max_tokens=300, **kwargs
    ):
        context = messages[0]["content"]
        user_question = messages[-1]["content"]
        prompt = f"### Context \n {context} \n ### Question \n {user_question}"
        stream_response = self.anthropic.completions.create(
            model=model,
            prompt=f"{self.HUMAN_PROMPT} {prompt}{self.AI_PROMPT}",
            max_tokens_to_sample=max_tokens,
            stream=True,
        )

        for completion in stream_response:
            yield completion.completion



File: application/llm/base.py

from abc import ABC, abstractmethod
from application.usage import gen_token_usage, stream_token_usage


class BaseLLM(ABC):
    def __init__(self):
        self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}

    def _apply_decorator(self, method, decorator, *args, **kwargs):
        return decorator(method, *args, **kwargs)

    @abstractmethod
    def _raw_gen(self, model, messages, stream, *args, **kwargs):
        pass

    def gen(self, model, messages, stream=False, *args, **kwargs):
        return self._apply_decorator(self._raw_gen, gen_token_usage)(
            self, model=model, messages=messages, stream=stream, *args, **kwargs
        )

    @abstractmethod
    def _raw_gen_stream(self, model, messages, stream, *args, **kwargs):
        pass

    def gen_stream(self, model, messages, stream=True, *args, **kwargs):
        return self._apply_decorator(self._raw_gen_stream, stream_token_usage)(
            self, model=model, messages=messages, stream=stream, *args, **kwargs
        )



File: application/core/logging_config.py

from logging.config import dictConfig

def setup_logging():
    dictConfig({
        'version': 1,
        'formatters': {
            'default': {
                'format': '[%(asctime)s] %(levelname)s in %(module)s: %(message)s',
            }
        },
        "handlers": {
            "console": {
                "class": "logging.StreamHandler",
                "stream": "ext://sys.stdout",
                "formatter": "default",
            }
        },
        'root': {
            'level': 'INFO',
            'handlers': ['console'],
        },
    })


File: application/core/__init__.py




File: application/core/settings.py

from pathlib import Path
from typing import Optional
import os

from pydantic_settings import BaseSettings

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


class Settings(BaseSettings):
    LLM_NAME: str = "docsgpt"
    MODEL_NAME: Optional[str] = None # if LLM_NAME is openai, MODEL_NAME can be gpt-4 or gpt-3.5-turbo
    EMBEDDINGS_NAME: str = "huggingface_sentence-transformers/all-mpnet-base-v2"
    CELERY_BROKER_URL: str = "redis://localhost:6379/0"
    CELERY_RESULT_BACKEND: str = "redis://localhost:6379/1"
    MONGO_URI: str = "mongodb://localhost:27017/docsgpt"
    MODEL_PATH: str = os.path.join(current_dir, "models/docsgpt-7b-f16.gguf")
    DEFAULT_MAX_HISTORY: int = 150
    MODEL_TOKEN_LIMITS: dict = {"gpt-3.5-turbo": 4096, "claude-2": 1e5}
    UPLOAD_FOLDER: str = "inputs"
    VECTOR_STORE: str = "faiss"  # "faiss" or "elasticsearch" or "qdrant" or "milvus"
    RETRIEVERS_ENABLED: list = ["classic_rag", "duckduck_search"] # also brave_search

    API_URL: str = "http://localhost:7091"  # backend url for celery worker

    API_KEY: Optional[str] = None  # LLM api key
    EMBEDDINGS_KEY: Optional[str] = None  # api key for embeddings (if using openai, just copy API_KEY)
    OPENAI_API_BASE: Optional[str] = None  # azure openai api base url
    OPENAI_API_VERSION: Optional[str] = None  # azure openai api version
    AZURE_DEPLOYMENT_NAME: Optional[str] = None  # azure deployment name for answering
    AZURE_EMBEDDINGS_DEPLOYMENT_NAME: Optional[str] = None  # azure deployment name for embeddings
    OPENAI_BASE_URL: Optional[str] = None # openai base url for open ai compatable models

    # elasticsearch
    ELASTIC_CLOUD_ID: Optional[str] = None  # cloud id for elasticsearch
    ELASTIC_USERNAME: Optional[str] = None  # username for elasticsearch
    ELASTIC_PASSWORD: Optional[str] = None  # password for elasticsearch
    ELASTIC_URL: Optional[str] = None  # url for elasticsearch
    ELASTIC_INDEX: Optional[str] = "docsgpt"  # index name for elasticsearch

    # SageMaker config
    SAGEMAKER_ENDPOINT: Optional[str] = None  # SageMaker endpoint name
    SAGEMAKER_REGION: Optional[str] = None  # SageMaker region name
    SAGEMAKER_ACCESS_KEY: Optional[str] = None  # SageMaker access key
    SAGEMAKER_SECRET_KEY: Optional[str] = None  # SageMaker secret key

    # prem ai project id
    PREMAI_PROJECT_ID: Optional[str] = None

    # Qdrant vectorstore config
    QDRANT_COLLECTION_NAME: Optional[str] = "docsgpt"
    QDRANT_LOCATION: Optional[str] = None
    QDRANT_URL: Optional[str] = None
    QDRANT_PORT: Optional[int] = 6333
    QDRANT_GRPC_PORT: int = 6334
    QDRANT_PREFER_GRPC: bool = False
    QDRANT_HTTPS: Optional[bool] = None
    QDRANT_API_KEY: Optional[str] = None
    QDRANT_PREFIX: Optional[str] = None
    QDRANT_TIMEOUT: Optional[float] = None
    QDRANT_HOST: Optional[str] = None
    QDRANT_PATH: Optional[str] = None
    QDRANT_DISTANCE_FUNC: str = "Cosine"

    # Milvus vectorstore config
    MILVUS_COLLECTION_NAME: Optional[str] = "docsgpt"
    MILVUS_URI: Optional[str] = "./milvus_local.db"   # milvus lite version as default
    MILVUS_TOKEN: Optional[str] = ""

    BRAVE_SEARCH_API_KEY: Optional[str] = None

    FLASK_DEBUG_MODE: bool = False


path = Path(__file__).parent.parent.absolute()
settings = Settings(_env_file=path.joinpath(".env"), _env_file_encoding="utf-8")



File: application/parser/py2doc.py

import ast
import os
from pathlib import Path

import tiktoken
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate


def find_files(directory):
    files_list = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.py'):
                files_list.append(os.path.join(root, file))
    return files_list


def extract_functions(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        functions = {}
        tree = ast.parse(source_code)
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                func_name = node.name
                func_def = ast.get_source_segment(source_code, node)
                functions[func_name] = func_def
    return functions


def extract_classes(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        classes = {}
        tree = ast.parse(source_code)
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                class_name = node.name
                function_names = []
                for subnode in ast.walk(node):
                    if isinstance(subnode, ast.FunctionDef):
                        function_names.append(subnode.name)
                classes[class_name] = ", ".join(function_names)
    return classes


def extract_functions_and_classes(directory):
    files = find_files(directory)
    functions_dict = {}
    classes_dict = {}
    for file in files:
        functions = extract_functions(file)
        if functions:
            functions_dict[file] = functions
        classes = extract_classes(file)
        if classes:
            classes_dict[file] = classes
    return functions_dict, classes_dict


def parse_functions(functions_dict, formats, dir):
    c1 = len(functions_dict)
    for i, (source, functions) in enumerate(functions_dict.items(), start=1):
        print(f"Processing file {i}/{c1}")
        source_w = source.replace(dir + "/", "").replace("." + formats, ".md")
        subfolders = "/".join(source_w.split("/")[:-1])
        Path(f"outputs/{subfolders}").mkdir(parents=True, exist_ok=True)
        for j, (name, function) in enumerate(functions.items(), start=1):
            print(f"Processing function {j}/{len(functions)}")
            prompt = PromptTemplate(
                input_variables=["code"],
                template="Code: \n{code}, \nDocumentation: ",
            )
            llm = OpenAI(temperature=0)
            response = llm(prompt.format(code=function))
            mode = "a" if Path(f"outputs/{source_w}").exists() else "w"
            with open(f"outputs/{source_w}", mode) as f:
                f.write(
                    f"\n\n# Function name: {name} \n\nFunction: \n```\n{function}\n```, \nDocumentation: \n{response}")


def parse_classes(classes_dict, formats, dir):
    c1 = len(classes_dict)
    for i, (source, classes) in enumerate(classes_dict.items()):
        print(f"Processing file {i + 1}/{c1}")
        source_w = source.replace(dir + "/", "").replace("." + formats, ".md")
        subfolders = "/".join(source_w.split("/")[:-1])
        Path(f"outputs/{subfolders}").mkdir(parents=True, exist_ok=True)
        for name, function_names in classes.items():
            print(f"Processing Class {i + 1}/{c1}")
            prompt = PromptTemplate(
                input_variables=["class_name", "functions_names"],
                template="Class name: {class_name} \nFunctions: {functions_names}, \nDocumentation: ",
            )
            llm = OpenAI(temperature=0)
            response = llm(prompt.format(class_name=name, functions_names=function_names))

            with open(f"outputs/{source_w}", "a" if Path(f"outputs/{source_w}").exists() else "w") as f:
                f.write(f"\n\n# Class name: {name} \n\nFunctions: \n{function_names}, \nDocumentation: \n{response}")


def transform_to_docs(functions_dict, classes_dict, formats, dir):
    docs_content = ''.join([str(key) + str(value) for key, value in functions_dict.items()])
    docs_content += ''.join([str(key) + str(value) for key, value in classes_dict.items()])

    num_tokens = len(tiktoken.get_encoding("cl100k_base").encode(docs_content))
    total_price = ((num_tokens / 1000) * 0.02)

    print(f"Number of Tokens = {num_tokens:,d}")
    print(f"Approx Cost = ${total_price:,.2f}")

    user_input = input("Price Okay? (Y/N)\n").lower()
    if user_input == "y" or user_input == "":
        if not Path("outputs").exists():
            Path("outputs").mkdir()
        parse_functions(functions_dict, formats, dir)
        parse_classes(classes_dict, formats, dir)
        print("All done!")
    else:
        print("The API was not called. No money was spent.")



File: application/parser/token_func.py

import re
from math import ceil
from typing import List

import tiktoken
from application.parser.schema.base import Document


def separate_header_and_body(text):
    header_pattern = r"^(.*?\n){3}"
    match = re.match(header_pattern, text)
    header = match.group(0)
    body = text[len(header):]
    return header, body


def group_documents(documents: List[Document], min_tokens: int, max_tokens: int) -> List[Document]:
    docs = []
    current_group = None

    for doc in documents:
        doc_len = len(tiktoken.get_encoding("cl100k_base").encode(doc.text))

        # Check if current group is empty or if the document can be added based on token count and matching metadata
        if (current_group is None or 
            (len(tiktoken.get_encoding("cl100k_base").encode(current_group.text)) + doc_len < max_tokens and 
             doc_len < min_tokens and 
             current_group.extra_info == doc.extra_info)):
            if current_group is None:
                current_group = doc  # Use the document directly to retain its metadata
            else:
                current_group.text += " " + doc.text  # Append text to the current group
        else:
            docs.append(current_group)
            current_group = doc  # Start a new group with the current document

    if current_group is not None:
        docs.append(current_group)

    return docs


def split_documents(documents: List[Document], max_tokens: int) -> List[Document]:
    docs = []
    for doc in documents:
        token_length = len(tiktoken.get_encoding("cl100k_base").encode(doc.text))
        if token_length <= max_tokens:
            docs.append(doc)
        else:
            header, body = separate_header_and_body(doc.text)
            if len(tiktoken.get_encoding("cl100k_base").encode(header)) > max_tokens:
                body = doc.text
                header = ""
            num_body_parts = ceil(token_length / max_tokens)
            part_length = ceil(len(body) / num_body_parts)
            body_parts = [body[i:i + part_length] for i in range(0, len(body), part_length)]
            for i, body_part in enumerate(body_parts):
                new_doc = Document(text=header + body_part.strip(),
                                   doc_id=f"{doc.doc_id}-{i}",
                                   embedding=doc.embedding,
                                   extra_info=doc.extra_info)
                docs.append(new_doc)
    return docs


def group_split(documents: List[Document], max_tokens: int = 2000, min_tokens: int = 150, token_check: bool = True):
    if not token_check:
        return documents
    print("Grouping small documents")
    try:
        documents = group_documents(documents=documents, min_tokens=min_tokens, max_tokens=max_tokens)
    except Exception:
        print("Grouping failed, try running without token_check")
    print("Separating large documents")
    try:
        documents = split_documents(documents=documents, max_tokens=max_tokens)
    except Exception:
        print("Grouping failed, try running without token_check")
    return documents



File: application/parser/__init__.py





File: application/parser/java2doc.py

import os

import javalang


def find_files(directory):
    files_list = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.java'):
                files_list.append(os.path.join(root, file))
    return files_list


def extract_functions(file_path):
    with open(file_path, "r") as file:
        java_code = file.read()
        methods = {}
        tree = javalang.parse.parse(java_code)
        for _, node in tree.filter(javalang.tree.MethodDeclaration):
            method_name = node.name
            start_line = node.position.line - 1
            end_line = start_line
            brace_count = 0
            for line in java_code.splitlines()[start_line:]:
                end_line += 1
                brace_count += line.count("{") - line.count("}")
                if brace_count == 0:
                    break
            method_source_code = "\n".join(java_code.splitlines()[start_line:end_line])
            methods[method_name] = method_source_code
    return methods


def extract_classes(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        classes = {}
        tree = javalang.parse.parse(source_code)
        for class_decl in tree.types:
            class_name = class_decl.name
            declarations = []
            methods = []
            for field_decl in class_decl.fields:
                field_name = field_decl.declarators[0].name
                field_type = field_decl.type.name
                declarations.append(f"{field_type} {field_name}")
            for method_decl in class_decl.methods:
                methods.append(method_decl.name)
            class_string = "Declarations: " + ", ".join(declarations) + "\n  Method name: " + ", ".join(methods)
            classes[class_name] = class_string
    return classes


def extract_functions_and_classes(directory):
    files = find_files(directory)
    functions_dict = {}
    classes_dict = {}
    for file in files:
        functions = extract_functions(file)
        if functions:
            functions_dict[file] = functions
        classes = extract_classes(file)
        if classes:
            classes_dict[file] = classes
    return functions_dict, classes_dict



File: application/parser/js2doc.py

import os

import escodegen
import esprima


def find_files(directory):
    files_list = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.js'):
                files_list.append(os.path.join(root, file))
    return files_list


def extract_functions(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        functions = {}
        tree = esprima.parseScript(source_code)
        for node in tree.body:
            if node.type == 'FunctionDeclaration':
                func_name = node.id.name if node.id else '<anonymous>'
                functions[func_name] = escodegen.generate(node)
            elif node.type == 'VariableDeclaration':
                for declaration in node.declarations:
                    if declaration.init and declaration.init.type == 'FunctionExpression':
                        func_name = declaration.id.name if declaration.id else '<anonymous>'
                        functions[func_name] = escodegen.generate(declaration.init)
            elif node.type == 'ClassDeclaration':
                for subnode in node.body.body:
                    if subnode.type == 'MethodDefinition':
                        func_name = subnode.key.name
                        functions[func_name] = escodegen.generate(subnode.value)
                    elif subnode.type == 'VariableDeclaration':
                        for declaration in subnode.declarations:
                            if declaration.init and declaration.init.type == 'FunctionExpression':
                                func_name = declaration.id.name if declaration.id else '<anonymous>'
                                functions[func_name] = escodegen.generate(declaration.init)
        return functions


def extract_classes(file_path):
    with open(file_path, 'r') as file:
        source_code = file.read()
        classes = {}
        tree = esprima.parseScript(source_code)
        for node in tree.body:
            if node.type == 'ClassDeclaration':
                class_name = node.id.name
                function_names = []
                for subnode in node.body.body:
                    if subnode.type == 'MethodDefinition':
                        function_names.append(subnode.key.name)
                classes[class_name] = ", ".join(function_names)
    return classes


def extract_functions_and_classes(directory):
    files = find_files(directory)
    functions_dict = {}
    classes_dict = {}
    for file in files:
        functions = extract_functions(file)
        if functions:
            functions_dict[file] = functions
        classes = extract_classes(file)
        if classes:
            classes_dict[file] = classes
    return functions_dict, classes_dict



File: application/parser/open_ai_func.py

import os

from application.vectorstore.vector_creator import VectorCreator
from application.core.settings import settings
from retry import retry


# from langchain_community.embeddings import HuggingFaceEmbeddings
# from langchain_community.embeddings import HuggingFaceInstructEmbeddings
# from langchain_community.embeddings import CohereEmbeddings


@retry(tries=10, delay=60)
def store_add_texts_with_retry(store, i, id):
    # add source_id to the metadata 
    i.metadata["source_id"] = str(id)
    store.add_texts([i.page_content], metadatas=[i.metadata])
    # store_pine.add_texts([i.page_content], metadatas=[i.metadata])


def call_openai_api(docs, folder_name, id, task_status):
    # Function to create a vector store from the documents and save it to disk

    if not os.path.exists(f"{folder_name}"):
        os.makedirs(f"{folder_name}")

    from tqdm import tqdm

    c1 = 0
    if settings.VECTOR_STORE == "faiss":
        docs_init = [docs[0]]
        docs.pop(0)

        store = VectorCreator.create_vectorstore(
            settings.VECTOR_STORE,
            docs_init=docs_init,
            source_id=f"{folder_name}",
            embeddings_key=os.getenv("EMBEDDINGS_KEY"),
        )
    else:
        store = VectorCreator.create_vectorstore(
            settings.VECTOR_STORE,
            source_id=str(id),
            embeddings_key=os.getenv("EMBEDDINGS_KEY"),
        )
    # Uncomment for MPNet embeddings
    # model_name = "sentence-transformers/all-mpnet-base-v2"
    # hf = HuggingFaceEmbeddings(model_name=model_name)
    # store = FAISS.from_documents(docs_test, hf)
    s1 = len(docs)
    for i in tqdm(
        docs,
        desc="Embedding 🦖",
        unit="docs",
        total=len(docs),
        bar_format="{l_bar}{bar}| Time Left: {remaining}",
    ):
        try:
            task_status.update_state(
                state="PROGRESS", meta={"current": int((c1 / s1) * 100)}
            )
            store_add_texts_with_retry(store, i, id)
        except Exception as e:
            print(e)
            print("Error on ", i)
            print("Saving progress")
            print(f"stopped at {c1} out of {len(docs)}")
            store.save_local(f"{folder_name}")
            break
        c1 += 1
    if settings.VECTOR_STORE == "faiss":
        store.save_local(f"{folder_name}")





File: application/parser/file/markdown_parser.py

"""Markdown parser.

Contains parser for md files.

"""
import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union, cast

import tiktoken
from application.parser.file.base_parser import BaseParser


class MarkdownParser(BaseParser):
    """Markdown parser.

    Extract text from markdown files.
    Returns dictionary with keys as headers and values as the text between headers.

    """

    def __init__(
            self,
            *args: Any,
            remove_hyperlinks: bool = True,
            remove_images: bool = True,
            max_tokens: int = 2048,
            # remove_tables: bool = True,
            **kwargs: Any,
    ) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._remove_hyperlinks = remove_hyperlinks
        self._remove_images = remove_images
        self._max_tokens = max_tokens
        # self._remove_tables = remove_tables

    def tups_chunk_append(self, tups: List[Tuple[Optional[str], str]], current_header: Optional[str],
                          current_text: str):
        """Append to tups chunk."""
        num_tokens = len(tiktoken.get_encoding("cl100k_base").encode(current_text))
        if num_tokens > self._max_tokens:
            chunks = [current_text[i:i + self._max_tokens] for i in range(0, len(current_text), self._max_tokens)]
            for chunk in chunks:
                tups.append((current_header, chunk))
        else:
            tups.append((current_header, current_text))
        return tups

    def markdown_to_tups(self, markdown_text: str) -> List[Tuple[Optional[str], str]]:
        """Convert a markdown file to a dictionary.

        The keys are the headers and the values are the text under each header.

        """
        markdown_tups: List[Tuple[Optional[str], str]] = []
        lines = markdown_text.split("\n")

        current_header = None
        current_text = ""

        for line in lines:
            header_match = re.match(r"^#+\s", line)
            if header_match:
                if current_header is not None:
                    if current_text == "" or None:
                        continue
                    markdown_tups = self.tups_chunk_append(markdown_tups, current_header, current_text)

                current_header = line
                current_text = ""
            else:
                current_text += line + "\n"
        markdown_tups = self.tups_chunk_append(markdown_tups, current_header, current_text)

        if current_header is not None:
            # pass linting, assert keys are defined
            markdown_tups = [
                (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
                for key, value in markdown_tups
            ]
        else:
            markdown_tups = [
                (key, re.sub("\n", "", value)) for key, value in markdown_tups
            ]

        return markdown_tups

    def remove_images(self, content: str) -> str:
        """Get a dictionary of a markdown file from its path."""
        pattern = r"!{1}\[\[(.*)\]\]"
        content = re.sub(pattern, "", content)
        return content

    # def remove_tables(self, content: str) -> List[List[str]]:
    #     """Convert markdown tables to nested lists."""
    #     table_rows_pattern = r"((\r?\n){2}|^)([^\r\n]*\|[^\r\n]*(\r?\n)?)+(?=(\r?\n){2}|$)"
    #     table_cells_pattern = r"([^\|\r\n]*)\|"
    #
    #     table_rows = re.findall(table_rows_pattern, content, re.MULTILINE)
    #     table_lists = []
    #     for row in table_rows:
    #         cells = re.findall(table_cells_pattern, row[2])
    #         cells = [cell.strip() for cell in cells if cell.strip()]
    #         table_lists.append(cells)
    #     return str(table_lists)

    def remove_hyperlinks(self, content: str) -> str:
        """Get a dictionary of a markdown file from its path."""
        pattern = r"\[(.*?)\]\((.*?)\)"
        content = re.sub(pattern, r"\1", content)
        return content

    def _init_parser(self) -> Dict:
        """Initialize the parser with the config."""
        return {}

    def parse_tups(
            self, filepath: Path, errors: str = "ignore"
    ) -> List[Tuple[Optional[str], str]]:
        """Parse file into tuples."""
        with open(filepath, "r") as f:
            content = f.read()
        if self._remove_hyperlinks:
            content = self.remove_hyperlinks(content)
        if self._remove_images:
            content = self.remove_images(content)
        # if self._remove_tables:
        #     content = self.remove_tables(content)
        markdown_tups = self.markdown_to_tups(content)
        return markdown_tups

    def parse_file(
            self, filepath: Path, errors: str = "ignore"
    ) -> Union[str, List[str]]:
        """Parse file into string."""
        tups = self.parse_tups(filepath, errors=errors)
        results = []
        # TODO: don't include headers right now
        for header, value in tups:
            if header is None:
                results.append(value)
            else:
                results.append(f"\n\n{header}\n{value}")
        return results



File: application/parser/file/html_parser.py

"""HTML parser.

Contains parser for html files.

"""
from pathlib import Path
from typing import Dict, Union

from application.parser.file.base_parser import BaseParser


class HTMLParser(BaseParser):
    """HTML parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, list[str]]:
        from langchain_community.document_loaders import BSHTMLLoader

        loader = BSHTMLLoader(file)
        data = loader.load()        
        return data



File: application/parser/file/rst_parser.py

"""reStructuredText parser.

Contains parser for md files.

"""
import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from application.parser.file.base_parser import BaseParser


class RstParser(BaseParser):
    """reStructuredText parser.

    Extract text from .rst files.
    Returns dictionary with keys as headers and values as the text between headers.

    """

    def __init__(
            self,
            *args: Any,
            remove_hyperlinks: bool = True,
            remove_images: bool = True,
            remove_table_excess: bool = True,
            remove_interpreters: bool = True,
            remove_directives: bool = True,
            remove_whitespaces_excess: bool = True,
            # Be careful with remove_characters_excess, might cause data loss
            remove_characters_excess: bool = True,
            **kwargs: Any,
    ) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._remove_hyperlinks = remove_hyperlinks
        self._remove_images = remove_images
        self._remove_table_excess = remove_table_excess
        self._remove_interpreters = remove_interpreters
        self._remove_directives = remove_directives
        self._remove_whitespaces_excess = remove_whitespaces_excess
        self._remove_characters_excess = remove_characters_excess

    def rst_to_tups(self, rst_text: str) -> List[Tuple[Optional[str], str]]:
        """Convert a reStructuredText file to a dictionary.

        The keys are the headers and the values are the text under each header.

        """
        rst_tups: List[Tuple[Optional[str], str]] = []
        lines = rst_text.split("\n")

        current_header = None
        current_text = ""

        for i, line in enumerate(lines):
            header_match = re.match(r"^[^\S\n]*[-=]+[^\S\n]*$", line)
            if header_match and i > 0 and (
                    len(lines[i - 1].strip()) == len(header_match.group().strip()) or lines[i - 2] == lines[i - 2]):
                if current_header is not None:
                    if current_text == "" or None:
                        continue
                    # removes the next heading from current Document
                    if current_text.endswith(lines[i - 1] + "\n"):
                        current_text = current_text[:len(current_text) - len(lines[i - 1] + "\n")]
                    rst_tups.append((current_header, current_text))

                current_header = lines[i - 1]
                current_text = ""
            else:
                current_text += line + "\n"

        rst_tups.append((current_header, current_text))

        # TODO: Format for rst
        #
        # if current_header is not None:
        #     # pass linting, assert keys are defined
        #     rst_tups = [
        #         (re.sub(r"#", "", cast(str, key)).strip(), re.sub(r"<.*?>", "", value))
        #         for key, value in rst_tups
        #     ]
        # else:
        #     rst_tups = [
        #         (key, re.sub("\n", "", value)) for key, value in rst_tups
        #     ]

        if current_header is None:
            rst_tups = [
                (key, re.sub("\n", "", value)) for key, value in rst_tups
            ]
        return rst_tups

    def remove_images(self, content: str) -> str:
        pattern = r"\.\. image:: (.*)"
        content = re.sub(pattern, "", content)
        return content

    def remove_hyperlinks(self, content: str) -> str:
        pattern = r"`(.*?) <(.*?)>`_"
        content = re.sub(pattern, r"\1", content)
        return content

    def remove_directives(self, content: str) -> str:
        """Removes reStructuredText Directives"""
        pattern = r"`\.\.([^:]+)::"
        content = re.sub(pattern, "", content)
        return content

    def remove_interpreters(self, content: str) -> str:
        """Removes reStructuredText Interpreted Text Roles"""
        pattern = r":(\w+):"
        content = re.sub(pattern, "", content)
        return content

    def remove_table_excess(self, content: str) -> str:
        """Pattern to remove grid table separators"""
        pattern = r"^\+[-]+\+[-]+\+$"
        content = re.sub(pattern, "", content, flags=re.MULTILINE)
        return content

    def remove_whitespaces_excess(self, content: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
        """Pattern to match 2 or more consecutive whitespaces"""
        pattern = r"\s{2,}"
        content = [(key, re.sub(pattern, "  ", value)) for key, value in content]
        return content

    def remove_characters_excess(self, content: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
        """Pattern to match 2 or more consecutive characters"""
        pattern = r"(\S)\1{2,}"
        content = [(key, re.sub(pattern, r"\1\1\1", value, flags=re.MULTILINE)) for key, value in content]
        return content

    def _init_parser(self) -> Dict:
        """Initialize the parser with the config."""
        return {}

    def parse_tups(
            self, filepath: Path, errors: str = "ignore"
    ) -> List[Tuple[Optional[str], str]]:
        """Parse file into tuples."""
        with open(filepath, "r") as f:
            content = f.read()
        if self._remove_hyperlinks:
            content = self.remove_hyperlinks(content)
        if self._remove_images:
            content = self.remove_images(content)
        if self._remove_table_excess:
            content = self.remove_table_excess(content)
        if self._remove_directives:
            content = self.remove_directives(content)
        if self._remove_interpreters:
            content = self.remove_interpreters(content)
        rst_tups = self.rst_to_tups(content)
        if self._remove_whitespaces_excess:
            rst_tups = self.remove_whitespaces_excess(rst_tups)
        if self._remove_characters_excess:
            rst_tups = self.remove_characters_excess(rst_tups)
        return rst_tups

    def parse_file(
            self, filepath: Path, errors: str = "ignore"
    ) -> Union[str, List[str]]:
        """Parse file into string."""
        tups = self.parse_tups(filepath, errors=errors)
        results = []
        # TODO: don't include headers right now
        for header, value in tups:
            if header is None:
                results.append(value)
            else:
                results.append(f"\n\n{header}\n{value}")
        return results



File: application/parser/file/bulk.py

"""Simple reader that reads files of different formats from a directory."""
import logging
from pathlib import Path
from typing import Callable, Dict, List, Optional, Union

from application.parser.file.base import BaseReader
from application.parser.file.base_parser import BaseParser
from application.parser.file.docs_parser import DocxParser, PDFParser
from application.parser.file.epub_parser import EpubParser
from application.parser.file.html_parser import HTMLParser
from application.parser.file.markdown_parser import MarkdownParser
from application.parser.file.rst_parser import RstParser
from application.parser.file.tabular_parser import PandasCSVParser
from application.parser.schema.base import Document

DEFAULT_FILE_EXTRACTOR: Dict[str, BaseParser] = {
    ".pdf": PDFParser(),
    ".docx": DocxParser(),
    ".csv": PandasCSVParser(),
    ".epub": EpubParser(),
    ".md": MarkdownParser(),
    ".rst": RstParser(),
    ".html": HTMLParser(),
    ".mdx": MarkdownParser(),
}


class SimpleDirectoryReader(BaseReader):
    """Simple directory reader.

    Can read files into separate documents, or concatenates
    files into one document text.

    Args:
        input_dir (str): Path to the directory.
        input_files (List): List of file paths to read (Optional; overrides input_dir)
        exclude_hidden (bool): Whether to exclude hidden files (dotfiles).
        errors (str): how encoding and decoding errors are to be handled,
              see https://docs.python.org/3/library/functions.html#open
        recursive (bool): Whether to recursively search in subdirectories.
            False by default.
        required_exts (Optional[List[str]]): List of required extensions.
            Default is None.
        file_extractor (Optional[Dict[str, BaseParser]]): A mapping of file
            extension to a BaseParser class that specifies how to convert that file
            to text. See DEFAULT_FILE_EXTRACTOR.
        num_files_limit (Optional[int]): Maximum number of files to read.
            Default is None.
        file_metadata (Optional[Callable[str, Dict]]): A function that takes
            in a filename and returns a Dict of metadata for the Document.
            Default is None.
    """

    def __init__(
            self,
            input_dir: Optional[str] = None,
            input_files: Optional[List] = None,
            exclude_hidden: bool = True,
            errors: str = "ignore",
            recursive: bool = True,
            required_exts: Optional[List[str]] = None,
            file_extractor: Optional[Dict[str, BaseParser]] = None,
            num_files_limit: Optional[int] = None,
            file_metadata: Optional[Callable[[str], Dict]] = None,
    ) -> None:
        """Initialize with parameters."""
        super().__init__()

        if not input_dir and not input_files:
            raise ValueError("Must provide either `input_dir` or `input_files`.")

        self.errors = errors

        self.recursive = recursive
        self.exclude_hidden = exclude_hidden
        self.required_exts = required_exts
        self.num_files_limit = num_files_limit

        if input_files:
            self.input_files = []
            for path in input_files:
                print(path)
                input_file = Path(path)
                self.input_files.append(input_file)
        elif input_dir:
            self.input_dir = Path(input_dir)
            self.input_files = self._add_files(self.input_dir)

        self.file_extractor = file_extractor or DEFAULT_FILE_EXTRACTOR
        self.file_metadata = file_metadata

    def _add_files(self, input_dir: Path) -> List[Path]:
        """Add files."""
        input_files = sorted(input_dir.iterdir())
        new_input_files = []
        dirs_to_explore = []
        for input_file in input_files:
            if input_file.is_dir():
                if self.recursive:
                    dirs_to_explore.append(input_file)
            elif self.exclude_hidden and input_file.name.startswith("."):
                continue
            elif (
                    self.required_exts is not None
                    and input_file.suffix not in self.required_exts
            ):
                continue
            else:
                new_input_files.append(input_file)

        for dir_to_explore in dirs_to_explore:
            sub_input_files = self._add_files(dir_to_explore)
            new_input_files.extend(sub_input_files)

        if self.num_files_limit is not None and self.num_files_limit > 0:
            new_input_files = new_input_files[0: self.num_files_limit]

        # print total number of files added
        logging.debug(
            f"> [SimpleDirectoryReader] Total files added: {len(new_input_files)}"
        )

        return new_input_files

    def load_data(self, concatenate: bool = False) -> List[Document]:
        """Load data from the input directory.

        Args:
            concatenate (bool): whether to concatenate all files into one document.
                If set to True, file metadata is ignored.
                False by default.

        Returns:
            List[Document]: A list of documents.

        """
        data: Union[str, List[str]] = ""
        data_list: List[str] = []
        metadata_list = []
        for input_file in self.input_files:
            if input_file.suffix in self.file_extractor:
                parser = self.file_extractor[input_file.suffix]
                if not parser.parser_config_set:
                    parser.init_parser()
                data = parser.parse_file(input_file, errors=self.errors)
            else:
                # do standard read
                with open(input_file, "r", errors=self.errors) as f:
                    data = f.read()
            # Prepare metadata for this file
            if self.file_metadata is not None:
                file_metadata = self.file_metadata(str(input_file))
            else:
                # Provide a default empty metadata
                file_metadata = {'title': '', 'store': ''}
                # TODO: Find a case with no metadata and check if breaks anything 

            if isinstance(data, List):
                # Extend data_list with each item in the data list
                data_list.extend([str(d) for d in data])
                # For each item in the data list, add the file's metadata to metadata_list
                metadata_list.extend([file_metadata for _ in data])
            else:
                # Add the single piece of data to data_list
                data_list.append(str(data))
                # Add the file's metadata to metadata_list
                metadata_list.append(file_metadata)

        if concatenate:
            return [Document("\n".join(data_list))]
        elif self.file_metadata is not None:
            return [Document(d, extra_info=m) for d, m in zip(data_list, metadata_list)]
        else:
            return [Document(d) for d in data_list]



File: application/parser/file/__init__.py





File: application/parser/file/docs_parser.py

"""Docs parser.

Contains parsers for docx, pdf files.

"""
from pathlib import Path
from typing import Dict

from application.parser.file.base_parser import BaseParser


class PDFParser(BaseParser):
    """PDF parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> str:
        """Parse file."""
        try:
            import PyPDF2
        except ImportError:
            raise ValueError("PyPDF2 is required to read PDF files.")
        text_list = []
        with open(file, "rb") as fp:
            # Create a PDF object
            pdf = PyPDF2.PdfReader(fp)

            # Get the number of pages in the PDF document
            num_pages = len(pdf.pages)

            # Iterate over every page
            for page in range(num_pages):
                # Extract the text from the page
                page_text = pdf.pages[page].extract_text()
                text_list.append(page_text)
        text = "\n".join(text_list)

        return text


class DocxParser(BaseParser):
    """Docx parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> str:
        """Parse file."""
        try:
            import docx2txt
        except ImportError:
            raise ValueError("docx2txt is required to read Microsoft Word files.")

        text = docx2txt.process(file)

        return text



File: application/parser/file/base_parser.py

"""Base parser and config class."""

from abc import abstractmethod
from pathlib import Path
from typing import Dict, List, Optional, Union


class BaseParser:
    """Base class for all parsers."""

    def __init__(self, parser_config: Optional[Dict] = None):
        """Init params."""
        self._parser_config = parser_config

    def init_parser(self) -> None:
        """Init parser and store it."""
        parser_config = self._init_parser()
        self._parser_config = parser_config

    @property
    def parser_config_set(self) -> bool:
        """Check if parser config is set."""
        return self._parser_config is not None

    @property
    def parser_config(self) -> Dict:
        """Check if parser config is set."""
        if self._parser_config is None:
            raise ValueError("Parser config not set.")
        return self._parser_config

    @abstractmethod
    def _init_parser(self) -> Dict:
        """Initialize the parser with the config."""

    @abstractmethod
    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
        """Parse file."""



File: application/parser/file/tabular_parser.py

"""Tabular parser.

Contains parsers for tabular data files.

"""
from pathlib import Path
from typing import Any, Dict, List, Union

from application.parser.file.base_parser import BaseParser


class CSVParser(BaseParser):
    """CSV parser.

    Args:
        concat_rows (bool): whether to concatenate all rows into one document.
            If set to False, a Document will be created for each row.
            True by default.

    """

    def __init__(self, *args: Any, concat_rows: bool = True, **kwargs: Any) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._concat_rows = concat_rows

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
        """Parse file.

        Returns:
            Union[str, List[str]]: a string or a List of strings.

        """
        try:
            import csv
        except ImportError:
            raise ValueError("csv module is required to read CSV files.")
        text_list = []
        with open(file, "r") as fp:
            csv_reader = csv.reader(fp)
            for row in csv_reader:
                text_list.append(", ".join(row))
        if self._concat_rows:
            return "\n".join(text_list)
        else:
            return text_list


class PandasCSVParser(BaseParser):
    r"""Pandas-based CSV parser.

    Parses CSVs using the separator detection from Pandas `read_csv`function.
    If special parameters are required, use the `pandas_config` dict.

    Args:
        concat_rows (bool): whether to concatenate all rows into one document.
            If set to False, a Document will be created for each row.
            True by default.

        col_joiner (str): Separator to use for joining cols per row.
            Set to ", " by default.

        row_joiner (str): Separator to use for joining each row.
            Only used when `concat_rows=True`.
            Set to "\n" by default.

        pandas_config (dict): Options for the `pandas.read_csv` function call.
            Refer to https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html
            for more information.
            Set to empty dict by default, this means pandas will try to figure
            out the separators, table head, etc. on its own.

    """

    def __init__(
            self,
            *args: Any,
            concat_rows: bool = True,
            col_joiner: str = ", ",
            row_joiner: str = "\n",
            pandas_config: dict = {},
            **kwargs: Any
    ) -> None:
        """Init params."""
        super().__init__(*args, **kwargs)
        self._concat_rows = concat_rows
        self._col_joiner = col_joiner
        self._row_joiner = row_joiner
        self._pandas_config = pandas_config

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
        """Parse file."""
        try:
            import pandas as pd
        except ImportError:
            raise ValueError("pandas module is required to read CSV files.")

        df = pd.read_csv(file, **self._pandas_config)

        text_list = df.apply(
            lambda row: (self._col_joiner).join(row.astype(str).tolist()), axis=1
        ).tolist()

        if self._concat_rows:
            return (self._row_joiner).join(text_list)
        else:
            return text_list



File: application/parser/file/openapi3_parser.py

from urllib.parse import urlparse

from openapi_parser import parse

try:
    from application.parser.file.base_parser import BaseParser
except ModuleNotFoundError:
    from base_parser import BaseParser


class OpenAPI3Parser(BaseParser):
    def init_parser(self) -> None:
        return super().init_parser()

    def get_base_urls(self, urls):
        base_urls = []
        for i in urls:
            parsed_url = urlparse(i)
            base_url = parsed_url.scheme + "://" + parsed_url.netloc
            if base_url not in base_urls:
                base_urls.append(base_url)
        return base_urls

    def get_info_from_paths(self, path):
        info = ""
        if path.operations:
            for operation in path.operations:
                info += (
                    f"\n{operation.method.value}="
                    f"{operation.responses[0].description}"
                )
        return info

    def parse_file(self, file_path):
        data = parse(file_path)
        results = ""
        base_urls = self.get_base_urls(link.url for link in data.servers)
        base_urls = ",".join([base_url for base_url in base_urls])
        results += f"Base URL:{base_urls}\n"
        i = 1
        for path in data.paths:
            info = self.get_info_from_paths(path)
            results += (
                f"Path{i}: {path.url}\n"
                f"description: {path.description}\n"
                f"parameters: {path.parameters}\nmethods: {info}\n"
            )
            i += 1
        with open("results.txt", "w") as f:
            f.write(results)
        return results



File: application/parser/file/epub_parser.py

"""Epub parser.

Contains parsers for epub files.
"""

from pathlib import Path
from typing import Dict

from application.parser.file.base_parser import BaseParser


class EpubParser(BaseParser):
    """Epub Parser."""

    def _init_parser(self) -> Dict:
        """Init parser."""
        return {}

    def parse_file(self, file: Path, errors: str = "ignore") -> str:
        """Parse file."""
        try:
            import ebooklib
            from ebooklib import epub
        except ImportError:
            raise ValueError("`EbookLib` is required to read Epub files.")
        try:
            import html2text
        except ImportError:
            raise ValueError("`html2text` is required to parse Epub files.")

        text_list = []
        book = epub.read_epub(file, options={"ignore_ncx": True})

        # Iterate through all chapters.
        for item in book.get_items():
            # Chapters are typically located in epub documents items.
            if item.get_type() == ebooklib.ITEM_DOCUMENT:
                text_list.append(
                    html2text.html2text(item.get_content().decode("utf-8"))
                )

        text = "\n".join(text_list)
        return text



File: application/parser/file/base.py

"""Base reader class."""
from abc import abstractmethod
from typing import Any, List

from langchain.docstore.document import Document as LCDocument
from application.parser.schema.base import Document


class BaseReader:
    """Utilities for loading data from a directory."""

    @abstractmethod
    def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]:
        """Load data from the input directory."""

    def load_langchain_documents(self, **load_kwargs: Any) -> List[LCDocument]:
        """Load data in LangChain document format."""
        docs = self.load_data(**load_kwargs)
        return [d.to_langchain_format() for d in docs]



File: application/parser/schema/__init__.py





File: application/parser/schema/base.py

"""Base schema for readers."""
from dataclasses import dataclass

from langchain.docstore.document import Document as LCDocument
from application.parser.schema.schema import BaseDocument


@dataclass
class Document(BaseDocument):
    """Generic interface for a data document.

    This document connects to data sources.

    """

    def __post_init__(self) -> None:
        """Post init."""
        if self.text is None:
            raise ValueError("text field not set.")

    @classmethod
    def get_type(cls) -> str:
        """Get Document type."""
        return "Document"

    def to_langchain_format(self) -> LCDocument:
        """Convert struct to LangChain document format."""
        metadata = self.extra_info or {}
        return LCDocument(page_content=self.text, metadata=metadata)

    @classmethod
    def from_langchain_format(cls, doc: LCDocument) -> "Document":
        """Convert struct from LangChain document format."""
        return cls(text=doc.page_content, extra_info=doc.metadata)



File: application/parser/schema/schema.py

"""Base schema for data structures."""
from abc import abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from dataclasses_json import DataClassJsonMixin


@dataclass
class BaseDocument(DataClassJsonMixin):
    """Base document.

    Generic abstract interfaces that captures both index structs
    as well as documents.

    """

    # TODO: consolidate fields from Document/IndexStruct into base class
    text: Optional[str] = None
    doc_id: Optional[str] = None
    embedding: Optional[List[float]] = None

    # extra fields
    extra_info: Optional[Dict[str, Any]] = None

    @classmethod
    @abstractmethod
    def get_type(cls) -> str:
        """Get Document type."""

    def get_text(self) -> str:
        """Get text."""
        if self.text is None:
            raise ValueError("text field not set.")
        return self.text

    def get_doc_id(self) -> str:
        """Get doc_id."""
        if self.doc_id is None:
            raise ValueError("doc_id not set.")
        return self.doc_id

    @property
    def is_doc_id_none(self) -> bool:
        """Check if doc_id is None."""
        return self.doc_id is None

    def get_embedding(self) -> List[float]:
        """Get embedding.

        Errors if embedding is None.

        """
        if self.embedding is None:
            raise ValueError("embedding not set.")
        return self.embedding

    @property
    def extra_info_str(self) -> Optional[str]:
        """Extra info string."""
        if self.extra_info is None:
            return None

        return "\n".join([f"{k}: {str(v)}" for k, v in self.extra_info.items()])



File: application/parser/remote/sitemap_loader.py

import requests
import re  # Import regular expression library
import xml.etree.ElementTree as ET
from application.parser.remote.base import BaseRemote

class SitemapLoader(BaseRemote):
    def __init__(self, limit=20):
        from langchain_community.document_loaders import WebBaseLoader
        self.loader = WebBaseLoader
        self.limit = limit  # Adding limit to control the number of URLs to process

    def load_data(self, inputs):
        sitemap_url= inputs
        # Check if the input is a list and if it is, use the first element
        if isinstance(sitemap_url, list) and sitemap_url:
            url = sitemap_url[0]

        urls = self._extract_urls(sitemap_url)
        if not urls:
            print(f"No URLs found in the sitemap: {sitemap_url}")
            return []

        # Load content of extracted URLs
        documents = []
        processed_urls = 0  # Counter for processed URLs
        for url in urls:
            if self.limit is not None and processed_urls >= self.limit:
                break  # Stop processing if the limit is reached

            try:
                loader = self.loader([url])
                documents.extend(loader.load())
                processed_urls += 1  # Increment the counter after processing each URL
            except Exception as e:
                print(f"Error processing URL {url}: {e}")
                continue

        return documents

    def _extract_urls(self, sitemap_url):
        try:
            response = requests.get(sitemap_url)
            response.raise_for_status()  # Raise an exception for HTTP errors
        except (requests.exceptions.HTTPError, requests.exceptions.ConnectionError) as e:
            print(f"Failed to fetch sitemap: {sitemap_url}. Error: {e}")
            return []

        # Determine if this is a sitemap or a URL
        if self._is_sitemap(response):
            # It's a sitemap, so parse it and extract URLs
            return self._parse_sitemap(response.content)
        else:
            # It's not a sitemap, return the URL itself
            return [sitemap_url]

    def _is_sitemap(self, response):
        content_type = response.headers.get('Content-Type', '')
        if 'xml' in content_type or response.url.endswith('.xml'):
            return True

        if '<sitemapindex' in response.text or '<urlset' in response.text:
            return True

        return False

    def _parse_sitemap(self, sitemap_content):
        # Remove namespaces
        sitemap_content = re.sub(' xmlns="[^"]+"', '', sitemap_content.decode('utf-8'), count=1)

        root = ET.fromstring(sitemap_content)

        urls = []
        for loc in root.findall('.//url/loc'):
            urls.append(loc.text)

        # Check for nested sitemaps
        for sitemap in root.findall('.//sitemap/loc'):
            nested_sitemap_url = sitemap.text
            urls.extend(self._extract_urls(nested_sitemap_url))

        return urls



File: application/parser/remote/github_loader.py




File: application/parser/remote/telegram.py

from langchain.document_loader import TelegramChatApiLoader
from application.parser.remote.base import BaseRemote

class TelegramChatApiRemote(BaseRemote):
    def _init_parser(self, *args, **load_kwargs):
        self.loader = TelegramChatApiLoader(**load_kwargs)
        return {}

    def parse_file(self, *args, **load_kwargs):

        return 


File: application/parser/remote/reddit_loader.py

from application.parser.remote.base import BaseRemote
from langchain_community.document_loaders import RedditPostsLoader


class RedditPostsLoaderRemote(BaseRemote):
    def load_data(self, inputs):
        data = eval(inputs)
        client_id = data.get("client_id")
        client_secret = data.get("client_secret")
        user_agent = data.get("user_agent")
        categories = data.get("categories", ["new", "hot"])
        mode = data.get("mode", "subreddit")
        search_queries = data.get("search_queries")
        number_posts = data.get("number_posts", 10)
        self.loader = RedditPostsLoader(
            client_id=client_id,
            client_secret=client_secret,
            user_agent=user_agent,
            categories=categories,
            mode=mode,
            search_queries=search_queries,
            number_posts=number_posts,
        )
        documents = self.loader.load()
        print(f"Loaded {len(documents)} documents from Reddit")
        return documents



File: application/parser/remote/web_loader.py

from application.parser.remote.base import BaseRemote
from langchain_community.document_loaders import WebBaseLoader

headers = {
    "User-Agent": "Mozilla/5.0",
    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*"
    ";q=0.8",
    "Accept-Language": "en-US,en;q=0.5",
    "Referer": "https://www.google.com/",
    "DNT": "1",
    "Connection": "keep-alive",
    "Upgrade-Insecure-Requests": "1",
}


class WebLoader(BaseRemote):
    def __init__(self):
        self.loader = WebBaseLoader

    def load_data(self, inputs):
        urls = inputs
        if isinstance(urls, str):
            urls = [urls]
        documents = []
        for url in urls:
            try:
                loader = self.loader([url], header_template=headers)
                documents.extend(loader.load())
            except Exception as e:
                print(f"Error processing URL {url}: {e}")
                continue
        return documents



File: application/parser/remote/crawler_loader.py

import requests
from urllib.parse import urlparse, urljoin
from bs4 import BeautifulSoup
from application.parser.remote.base import BaseRemote

class CrawlerLoader(BaseRemote):
    def __init__(self, limit=10):
        from langchain_community.document_loaders import WebBaseLoader
        self.loader = WebBaseLoader  # Initialize the document loader
        self.limit = limit  # Set the limit for the number of pages to scrape

    def load_data(self, inputs):
        url = inputs
        # Check if the input is a list and if it is, use the first element
        if isinstance(url, list) and url:
            url = url[0]

        # Check if the URL scheme is provided, if not, assume http
        if not urlparse(url).scheme:
            url = "http://" + url

        visited_urls = set()  # Keep track of URLs that have been visited
        base_url = urlparse(url).scheme + "://" + urlparse(url).hostname  # Extract the base URL
        urls_to_visit = [url]  # List of URLs to be visited, starting with the initial URL
        loaded_content = []  # Store the loaded content from each URL

        # Continue crawling until there are no more URLs to visit
        while urls_to_visit:
            current_url = urls_to_visit.pop(0)  # Get the next URL to visit
            visited_urls.add(current_url)  # Mark the URL as visited

            # Try to load and process the content from the current URL
            try:
                response = requests.get(current_url)  # Fetch the content of the current URL
                response.raise_for_status()  # Raise an exception for HTTP errors
                loader = self.loader([current_url])  # Initialize the document loader for the current URL
                loaded_content.extend(loader.load())  # Load the content and add it to the loaded_content list
            except Exception as e:
                # Print an error message if loading or processing fails and continue with the next URL
                print(f"Error processing URL {current_url}: {e}")
                continue

            # Parse the HTML content to extract all links
            soup = BeautifulSoup(response.text, 'html.parser')
            all_links = [
                urljoin(current_url, a['href'])
                for a in soup.find_all('a', href=True)
                if base_url in urljoin(current_url, a['href'])  # Ensure links are from the same domain
            ]

            # Add new links to the list of URLs to visit if they haven't been visited yet
            urls_to_visit.extend([link for link in all_links if link not in visited_urls])
            urls_to_visit = list(set(urls_to_visit))  # Remove duplicate URLs

            # Stop crawling if the limit of pages to scrape is reached
            if self.limit is not None and len(visited_urls) >= self.limit:
                break

        return loaded_content  # Return the loaded content from all visited URLs



File: application/parser/remote/base.py

"""Base reader class."""
from abc import abstractmethod
from typing import Any, List

from langchain.docstore.document import Document as LCDocument
from application.parser.schema.base import Document


class BaseRemote:
    """Utilities for loading data from a directory."""

    @abstractmethod
    def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]:
        """Load data from the input directory."""

    def load_langchain_documents(self, **load_kwargs: Any) -> List[LCDocument]:
        """Load data in LangChain document format."""
        docs = self.load_data(**load_kwargs)
        return [d.to_langchain_format() for d in docs]



File: application/parser/remote/remote_creator.py

from application.parser.remote.sitemap_loader import SitemapLoader
from application.parser.remote.crawler_loader import CrawlerLoader
from application.parser.remote.web_loader import WebLoader
from application.parser.remote.reddit_loader import RedditPostsLoaderRemote


class RemoteCreator:
    loaders = {
        "url": WebLoader,
        "sitemap": SitemapLoader,
        "crawler": CrawlerLoader,
        "reddit": RedditPostsLoaderRemote,
    }

    @classmethod
    def create_loader(cls, type, *args, **kwargs):
        loader_class = cls.loaders.get(type.lower())
        if not loader_class:
            raise ValueError(f"No LLM class found for type {type}")
        return loader_class(*args, **kwargs)



File: application/retriever/brave_search.py

import json
from application.retriever.base import BaseRetriever
from application.core.settings import settings
from application.llm.llm_creator import LLMCreator
from application.utils import num_tokens_from_string
from langchain_community.tools import BraveSearch


class BraveRetSearch(BaseRetriever):

    def __init__(
        self,
        question,
        source,
        chat_history,
        prompt,
        chunks=2,
        token_limit=150,
        gpt_model="docsgpt",
        user_api_key=None,
    ):
        self.question = question
        self.source = source
        self.chat_history = chat_history
        self.prompt = prompt
        self.chunks = chunks
        self.gpt_model = gpt_model
        self.token_limit = (
            token_limit
            if token_limit
            < settings.MODEL_TOKEN_LIMITS.get(
                self.gpt_model, settings.DEFAULT_MAX_HISTORY
            )
            else settings.MODEL_TOKEN_LIMITS.get(
                self.gpt_model, settings.DEFAULT_MAX_HISTORY
            )
        )
        self.user_api_key = user_api_key

    def _get_data(self):
        if self.chunks == 0:
            docs = []
        else:
            search = BraveSearch.from_api_key(
                api_key=settings.BRAVE_SEARCH_API_KEY,
                search_kwargs={"count": int(self.chunks)},
            )
            results = search.run(self.question)
            results = json.loads(results)

            docs = []
            for i in results:
                try:
                    title = i["title"]
                    link = i["link"]
                    snippet = i["snippet"]
                    docs.append({"text": snippet, "title": title, "link": link})
                except IndexError:
                    pass
        if settings.LLM_NAME == "llama.cpp":
            docs = [docs[0]]

        return docs

    def gen(self):
        docs = self._get_data()

        # join all page_content together with a newline
        docs_together = "\n".join([doc["text"] for doc in docs])
        p_chat_combine = self.prompt.replace("{summaries}", docs_together)
        messages_combine = [{"role": "system", "content": p_chat_combine}]
        for doc in docs:
            yield {"source": doc}

        if len(self.chat_history) > 1:
            tokens_current_history = 0
            # count tokens in history
            self.chat_history.reverse()
            for i in self.chat_history:
                if "prompt" in i and "response" in i:
                    tokens_batch = num_tokens_from_string(i["prompt"]) + num_tokens_from_string(
                        i["response"]
                    )
                    if tokens_current_history + tokens_batch < self.token_limit:
                        tokens_current_history += tokens_batch
                        messages_combine.append(
                            {"role": "user", "content": i["prompt"]}
                        )
                        messages_combine.append(
                            {"role": "system", "content": i["response"]}
                        )
        messages_combine.append({"role": "user", "content": self.question})

        llm = LLMCreator.create_llm(
            settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=self.user_api_key
        )

        completion = llm.gen_stream(model=self.gpt_model, messages=messages_combine)
        for line in completion:
            yield {"answer": str(line)}

    def search(self):
        return self._get_data()

    def get_params(self):
        return {
            "question": self.question,
            "source": self.source,
            "chat_history": self.chat_history,
            "prompt": self.prompt,
            "chunks": self.chunks,
            "token_limit": self.token_limit,
            "gpt_model": self.gpt_model,
            "user_api_key": self.user_api_key
        }



File: application/retriever/__init__.py




File: application/retriever/classic_rag.py

from application.retriever.base import BaseRetriever
from application.core.settings import settings
from application.vectorstore.vector_creator import VectorCreator
from application.llm.llm_creator import LLMCreator

from application.utils import num_tokens_from_string


class ClassicRAG(BaseRetriever):

    def __init__(
        self,
        question,
        source,
        chat_history,
        prompt,
        chunks=2,
        token_limit=150,
        gpt_model="docsgpt",
        user_api_key=None,
    ):
        self.question = question
        self.vectorstore = source['active_docs'] if 'active_docs' in source else None
        self.chat_history = chat_history
        self.prompt = prompt
        self.chunks = chunks
        self.gpt_model = gpt_model
        self.token_limit = (
            token_limit
            if token_limit
            < settings.MODEL_TOKEN_LIMITS.get(
                self.gpt_model, settings.DEFAULT_MAX_HISTORY
            )
            else settings.MODEL_TOKEN_LIMITS.get(
                self.gpt_model, settings.DEFAULT_MAX_HISTORY
            )
        )
        self.user_api_key = user_api_key

    def _get_data(self):
        if self.chunks == 0:
            docs = []
        else:
            docsearch = VectorCreator.create_vectorstore(
                settings.VECTOR_STORE, self.vectorstore, settings.EMBEDDINGS_KEY
            )
            docs_temp = docsearch.search(self.question, k=self.chunks)
            print(docs_temp)
            docs = [
                {
                    "title": i.metadata.get(
                        "title", i.metadata.get("post_title", i.page_content)
                    ).split("/")[-1],
                    "text": i.page_content,
                    "source": (
                        i.metadata.get("source")
                        if i.metadata.get("source")
                        else "local"
                    ),
                }
                for i in docs_temp
            ]
        if settings.LLM_NAME == "llama.cpp":
            docs = [docs[0]]

        return docs

    def gen(self):
        docs = self._get_data()

        # join all page_content together with a newline
        docs_together = "\n".join([doc["text"] for doc in docs])
        p_chat_combine = self.prompt.replace("{summaries}", docs_together)
        messages_combine = [{"role": "system", "content": p_chat_combine}]
        for doc in docs:
            yield {"source": doc}

        if len(self.chat_history) > 1:
            tokens_current_history = 0
            # count tokens in history
            self.chat_history.reverse()
            for i in self.chat_history:
                if "prompt" in i and "response" in i:
                    tokens_batch = num_tokens_from_string(i["prompt"]) + num_tokens_from_string(
                        i["response"]
                    )
                    if tokens_current_history + tokens_batch < self.token_limit:
                        tokens_current_history += tokens_batch
                        messages_combine.append(
                            {"role": "user", "content": i["prompt"]}
                        )
                        messages_combine.append(
                            {"role": "system", "content": i["response"]}
                        )
        messages_combine.append({"role": "user", "content": self.question})

        llm = LLMCreator.create_llm(
            settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=self.user_api_key
        )

        completion = llm.gen_stream(model=self.gpt_model, messages=messages_combine)
        for line in completion:
            yield {"answer": str(line)}

    def search(self):
        return self._get_data()
    
    def get_params(self):
        return {
            "question": self.question,
            "source": self.vectorstore,
            "chat_history": self.chat_history,
            "prompt": self.prompt,
            "chunks": self.chunks,
            "token_limit": self.token_limit,
            "gpt_model": self.gpt_model,
            "user_api_key": self.user_api_key
        }



File: application/retriever/retriever_creator.py

from application.retriever.classic_rag import ClassicRAG
from application.retriever.duckduck_search import DuckDuckSearch
from application.retriever.brave_search import BraveRetSearch



class RetrieverCreator:
    retrievers = {
        'classic': ClassicRAG,
        'duckduck_search': DuckDuckSearch,
        'brave_search': BraveRetSearch,
        'default': ClassicRAG
    }

    @classmethod
    def create_retriever(cls, type, *args, **kwargs):
        retiever_class = cls.retrievers.get(type.lower())
        if not retiever_class:
            raise ValueError(f"No retievers class found for type {type}")
        return retiever_class(*args, **kwargs)


File: application/retriever/duckduck_search.py

from application.retriever.base import BaseRetriever
from application.core.settings import settings
from application.llm.llm_creator import LLMCreator
from application.utils import num_tokens_from_string
from langchain_community.tools import DuckDuckGoSearchResults
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper


class DuckDuckSearch(BaseRetriever):

    def __init__(
        self,
        question,
        source,
        chat_history,
        prompt,
        chunks=2,
        token_limit=150,
        gpt_model="docsgpt",
        user_api_key=None,
    ):
        self.question = question
        self.source = source
        self.chat_history = chat_history
        self.prompt = prompt
        self.chunks = chunks
        self.gpt_model = gpt_model
        self.token_limit = (
            token_limit
            if token_limit
            < settings.MODEL_TOKEN_LIMITS.get(
                self.gpt_model, settings.DEFAULT_MAX_HISTORY
            )
            else settings.MODEL_TOKEN_LIMITS.get(
                self.gpt_model, settings.DEFAULT_MAX_HISTORY
            )
        )
        self.user_api_key = user_api_key

    def _parse_lang_string(self, input_string):
        result = []
        current_item = ""
        inside_brackets = False
        for char in input_string:
            if char == "[":
                inside_brackets = True
            elif char == "]":
                inside_brackets = False
                result.append(current_item)
                current_item = ""
            elif inside_brackets:
                current_item += char

        if inside_brackets:
            result.append(current_item)

        return result

    def _get_data(self):
        if self.chunks == 0:
            docs = []
        else:
            wrapper = DuckDuckGoSearchAPIWrapper(max_results=self.chunks)
            search = DuckDuckGoSearchResults(api_wrapper=wrapper)
            results = search.run(self.question)
            results = self._parse_lang_string(results)

            docs = []
            for i in results:
                try:
                    text = i.split("title:")[0]
                    title = i.split("title:")[1].split("link:")[0]
                    link = i.split("link:")[1]
                    docs.append({"text": text, "title": title, "link": link})
                except IndexError:
                    pass
        if settings.LLM_NAME == "llama.cpp":
            docs = [docs[0]]

        return docs

    def gen(self):
        docs = self._get_data()

        # join all page_content together with a newline
        docs_together = "\n".join([doc["text"] for doc in docs])
        p_chat_combine = self.prompt.replace("{summaries}", docs_together)
        messages_combine = [{"role": "system", "content": p_chat_combine}]
        for doc in docs:
            yield {"source": doc}

        if len(self.chat_history) > 1:
            tokens_current_history = 0
            # count tokens in history
            self.chat_history.reverse()
            for i in self.chat_history:
                if "prompt" in i and "response" in i:
                    tokens_batch = num_tokens_from_string(i["prompt"]) + num_tokens_from_string(
                        i["response"]
                    )
                    if tokens_current_history + tokens_batch < self.token_limit:
                        tokens_current_history += tokens_batch
                        messages_combine.append(
                            {"role": "user", "content": i["prompt"]}
                        )
                        messages_combine.append(
                            {"role": "system", "content": i["response"]}
                        )
        messages_combine.append({"role": "user", "content": self.question})

        llm = LLMCreator.create_llm(
            settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=self.user_api_key
        )

        completion = llm.gen_stream(model=self.gpt_model, messages=messages_combine)
        for line in completion:
            yield {"answer": str(line)}

    def search(self):
        return self._get_data()
    
    def get_params(self):
        return {
            "question": self.question,
            "source": self.source,
            "chat_history": self.chat_history,
            "prompt": self.prompt,
            "chunks": self.chunks,
            "token_limit": self.token_limit,
            "gpt_model": self.gpt_model,
            "user_api_key": self.user_api_key
        }



File: application/retriever/base.py

from abc import ABC, abstractmethod


class BaseRetriever(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def gen(self, *args, **kwargs):
        pass

    @abstractmethod
    def search(self, *args, **kwargs):
        pass

    @abstractmethod
    def get_params(self):
        pass



File: application/api/__init__.py




File: application/api/answer/__init__.py




File: application/api/answer/routes.py

import asyncio
import os
import sys
from flask import Blueprint, request, Response, current_app
import json
import datetime
import logging
import traceback

from pymongo import MongoClient
from bson.objectid import ObjectId
from bson.dbref import DBRef

from application.core.settings import settings
from application.llm.llm_creator import LLMCreator
from application.retriever.retriever_creator import RetrieverCreator
from application.error import bad_request

logger = logging.getLogger(__name__)

mongo = MongoClient(settings.MONGO_URI)
db = mongo["docsgpt"]
conversations_collection = db["conversations"]
sources_collection = db["sources"]
prompts_collection = db["prompts"]
api_key_collection = db["api_keys"]
user_logs_collection = db["user_logs"]
answer = Blueprint("answer", __name__)

gpt_model = ""
# to have some kind of default behaviour
if settings.LLM_NAME == "openai":
    gpt_model = "gpt-3.5-turbo"
elif settings.LLM_NAME == "anthropic":
    gpt_model = "claude-2"

if settings.MODEL_NAME:  # in case there is particular model name configured
    gpt_model = settings.MODEL_NAME

# load the prompts
current_dir = os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
with open(os.path.join(current_dir, "prompts", "chat_combine_default.txt"), "r") as f:
    chat_combine_template = f.read()

with open(os.path.join(current_dir, "prompts", "chat_reduce_prompt.txt"), "r") as f:
    chat_reduce_template = f.read()

with open(os.path.join(current_dir, "prompts", "chat_combine_creative.txt"), "r") as f:
    chat_combine_creative = f.read()

with open(os.path.join(current_dir, "prompts", "chat_combine_strict.txt"), "r") as f:
    chat_combine_strict = f.read()

api_key_set = settings.API_KEY is not None
embeddings_key_set = settings.EMBEDDINGS_KEY is not None


async def async_generate(chain, question, chat_history):
    result = await chain.arun({"question": question, "chat_history": chat_history})
    return result


def run_async_chain(chain, question, chat_history):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    result = {}
    try:
        answer = loop.run_until_complete(async_generate(chain, question, chat_history))
    finally:
        loop.close()
    result["answer"] = answer
    return result


def get_data_from_api_key(api_key):
    data = api_key_collection.find_one({"key": api_key})
    # # Raise custom exception if the API key is not found
    if data is None:
        raise Exception("Invalid API Key, please generate new key", 401)

    if "retriever" not in data:
        data["retriever"] = None

    if "source" in data and isinstance(data["source"], DBRef):
        source_doc = db.dereference(data["source"])
        data["source"] = str(source_doc["_id"])
        if "retriever" in source_doc:
            data["retriever"] = source_doc["retriever"]
    else:
        data["source"] = {}
    return data


def get_retriever(source_id: str):
    doc = sources_collection.find_one({"_id": ObjectId(source_id)})
    if doc is None:
        raise Exception("Source document does not exist", 404)
    retriever_name = None if "retriever" not in doc else doc["retriever"]
    return retriever_name


def is_azure_configured():
    return (
        settings.OPENAI_API_BASE
        and settings.OPENAI_API_VERSION
        and settings.AZURE_DEPLOYMENT_NAME
    )


def save_conversation(conversation_id, question, response, source_log_docs, llm):
    if conversation_id is not None and conversation_id != "None":
        conversations_collection.update_one(
            {"_id": ObjectId(conversation_id)},
            {
                "$push": {
                    "queries": {
                        "prompt": question,
                        "response": response,
                        "sources": source_log_docs,
                    }
                }
            },
        )

    else:
        # create new conversation
        # generate summary
        messages_summary = [
            {
                "role": "assistant",
                "content": "Summarise following conversation in no more than 3 "
                "words, respond ONLY with the summary, use the same "
                "language as the system \n\nUser: "
                + question
                + "\n\n"
                + "AI: "
                + response,
            },
            {
                "role": "user",
                "content": "Summarise following conversation in no more than 3 words, "
                "respond ONLY with the summary, use the same language as the "
                "system",
            },
        ]

        completion = llm.gen(model=gpt_model, messages=messages_summary, max_tokens=30)
        conversation_id = conversations_collection.insert_one(
            {
                "user": "local",
                "date": datetime.datetime.utcnow(),
                "name": completion,
                "queries": [
                    {
                        "prompt": question,
                        "response": response,
                        "sources": source_log_docs,
                    }
                ],
            }
        ).inserted_id
    return conversation_id


def get_prompt(prompt_id):
    if prompt_id == "default":
        prompt = chat_combine_template
    elif prompt_id == "creative":
        prompt = chat_combine_creative
    elif prompt_id == "strict":
        prompt = chat_combine_strict
    else:
        prompt = prompts_collection.find_one({"_id": ObjectId(prompt_id)})["content"]
    return prompt


def complete_stream(
    question, retriever, conversation_id, user_api_key, isNoneDoc=False
):

    try:
        response_full = ""
        source_log_docs = []
        answer = retriever.gen()
        for line in answer:
            if "answer" in line:
                response_full += str(line["answer"])
                data = json.dumps(line)
                yield f"data: {data}\n\n"
            elif "source" in line:
                source_log_docs.append(line["source"])

        if isNoneDoc:
            for doc in source_log_docs:
                doc["source"] = "None"

        llm = LLMCreator.create_llm(
            settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=user_api_key
        )
        if user_api_key is None:
            conversation_id = save_conversation(
                conversation_id, question, response_full, source_log_docs, llm
            )
            # send data.type = "end" to indicate that the stream has ended as json
            data = json.dumps({"type": "id", "id": str(conversation_id)})
            yield f"data: {data}\n\n"

        retriever_params = retriever.get_params()
        user_logs_collection.insert_one(
            {
                "action": "stream_answer",
                "level": "info",
                "user": "local",
                "api_key": user_api_key,
                "question": question,
                "response": response_full,
                "sources": source_log_docs,
                "retriever_params": retriever_params,
                "timestamp": datetime.datetime.now(datetime.timezone.utc),
            }
        )

        data = json.dumps({"type": "end"})
        yield f"data: {data}\n\n"
    except Exception as e:
        print("\033[91merr", str(e), file=sys.stderr)
        data = json.dumps(
            {
                "type": "error",
                "error": "Please try again later. We apologize for any inconvenience.",
                "error_exception": str(e),
            }
        )
        yield f"data: {data}\n\n"
        return


@answer.route("/stream", methods=["POST"])
def stream():
    try:
        data = request.get_json()
        question = data["question"]
        if "history" not in data:
            history = []
        else:
            history = data["history"]
            history = json.loads(history)
        if "conversation_id" not in data:
            conversation_id = None
        else:
            conversation_id = data["conversation_id"]
        if "prompt_id" in data:
            prompt_id = data["prompt_id"]
        else:
            prompt_id = "default"
        if "selectedDocs" in data and data["selectedDocs"] is None:
            chunks = 0
        elif "chunks" in data:
            chunks = int(data["chunks"])
        else:
            chunks = 2
        if "token_limit" in data:
            token_limit = data["token_limit"]
        else:
            token_limit = settings.DEFAULT_MAX_HISTORY

        ## retriever can be "brave_search, duckduck_search or classic"
        retriever_name = data["retriever"] if "retriever" in data else "classic"

        # check if active_docs or api_key is set
        if "api_key" in data:
            data_key = get_data_from_api_key(data["api_key"])
            chunks = int(data_key["chunks"])
            prompt_id = data_key["prompt_id"]
            source = {"active_docs": data_key["source"]}
            retriever_name = data_key["retriever"] or retriever_name
            user_api_key = data["api_key"]

        elif "active_docs" in data:
            source = {"active_docs": data["active_docs"]}
            retriever_name = get_retriever(data["active_docs"]) or retriever_name
            user_api_key = None

        else:
            source = {}
            user_api_key = None

        current_app.logger.info(
            f"/stream - request_data: {data}, source: {source}",
            extra={"data": json.dumps({"request_data": data, "source": source})},
        )

        prompt = get_prompt(prompt_id)

        retriever = RetrieverCreator.create_retriever(
            retriever_name,
            question=question,
            source=source,
            chat_history=history,
            prompt=prompt,
            chunks=chunks,
            token_limit=token_limit,
            gpt_model=gpt_model,
            user_api_key=user_api_key,
        )

        return Response(
            complete_stream(
                question=question,
                retriever=retriever,
                conversation_id=conversation_id,
                user_api_key=user_api_key,
                isNoneDoc=data.get("isNoneDoc"),
            ),
            mimetype="text/event-stream",
        )

    except ValueError:
        message = "Malformed request body"
        print("\033[91merr", str(message), file=sys.stderr)
        return Response(
            error_stream_generate(message),
            status=400,
            mimetype="text/event-stream",
        )
    except Exception as e:
        current_app.logger.error(
            f"/stream - error: {str(e)} - traceback: {traceback.format_exc()}",
            extra={"error": str(e), "traceback": traceback.format_exc()},
        )
        message = e.args[0]
        status_code = 400
        # # Custom exceptions with two arguments, index 1 as status code
        if len(e.args) >= 2:
            status_code = e.args[1]
        return Response(
            error_stream_generate(message),
            status=status_code,
            mimetype="text/event-stream",
        )


def error_stream_generate(err_response):
    data = json.dumps({"type": "error", "error": err_response})
    yield f"data: {data}\n\n"


@answer.route("/api/answer", methods=["POST"])
def api_answer():
    data = request.get_json()
    question = data["question"]
    if "history" not in data:
        history = []
    else:
        history = data["history"]
    if "conversation_id" not in data:
        conversation_id = None
    else:
        conversation_id = data["conversation_id"]
    print("-" * 5)
    if "prompt_id" in data:
        prompt_id = data["prompt_id"]
    else:
        prompt_id = "default"
    if "chunks" in data:
        chunks = int(data["chunks"])
    else:
        chunks = 2
    if "token_limit" in data:
        token_limit = data["token_limit"]
    else:
        token_limit = settings.DEFAULT_MAX_HISTORY

    ## retriever can be brave_search, duckduck_search or classic
    retriever_name = data["retriever"] if "retriever" in data else "classic"

    # use try and except  to check for exception
    try:
        # check if the vectorstore is set
        if "api_key" in data:
            data_key = get_data_from_api_key(data["api_key"])
            chunks = int(data_key["chunks"])
            prompt_id = data_key["prompt_id"]
            source = {"active_docs": data_key["source"]}
            retriever_name = data_key["retriever"] or retriever_name
            user_api_key = data["api_key"]
        elif "active_docs" in data:
            source = {"active_docs": data["active_docs"]}
            retriever_name = get_retriever(data["active_docs"]) or retriever_name
            user_api_key = None
        else:
            source = {}
            user_api_key = None

        prompt = get_prompt(prompt_id)

        current_app.logger.info(
            f"/api/answer - request_data: {data}, source: {source}",
            extra={"data": json.dumps({"request_data": data, "source": source})},
        )

        retriever = RetrieverCreator.create_retriever(
            retriever_name,
            question=question,
            source=source,
            chat_history=history,
            prompt=prompt,
            chunks=chunks,
            token_limit=token_limit,
            gpt_model=gpt_model,
            user_api_key=user_api_key,
        )
        source_log_docs = []
        response_full = ""
        for line in retriever.gen():
            if "source" in line:
                source_log_docs.append(line["source"])
            elif "answer" in line:
                response_full += line["answer"]

        if data.get("isNoneDoc"):
            for doc in source_log_docs:
                doc["source"] = "None"

        llm = LLMCreator.create_llm(
            settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=user_api_key
        )

        result = {"answer": response_full, "sources": source_log_docs}
        result["conversation_id"] = str(
            save_conversation(
                conversation_id, question, response_full, source_log_docs, llm
            )
        )
        retriever_params = retriever.get_params()
        user_logs_collection.insert_one(
            {
                "action": "api_answer",
                "level": "info",
                "user": "local",
                "api_key": user_api_key,
                "question": question,
                "response": response_full,
                "sources": source_log_docs,
                "retriever_params": retriever_params,
                "timestamp": datetime.datetime.now(datetime.timezone.utc),
            }
        )

        return result
    except Exception as e:
        current_app.logger.error(
            f"/api/answer - error: {str(e)} - traceback: {traceback.format_exc()}",
            extra={"error": str(e), "traceback": traceback.format_exc()},
        )
        return bad_request(500, str(e))


@answer.route("/api/search", methods=["POST"])
def api_search():
    data = request.get_json()
    question = data["question"]
    if "chunks" in data:
        chunks = int(data["chunks"])
    else:
        chunks = 2
    if "api_key" in data:
        data_key = get_data_from_api_key(data["api_key"])
        chunks = int(data_key["chunks"])
        source = {"active_docs":data_key["source"]}
        user_api_key = data["api_key"]
    elif "active_docs" in data:
        source = {"active_docs": data["active_docs"]}
        user_api_key = None
    else:
        source = {}
        user_api_key = None

    if "retriever" in data:
        retriever_name = data["retriever"]
    else:
        retriever_name = "classic"
    if "token_limit" in data:
        token_limit = data["token_limit"]
    else:
        token_limit = settings.DEFAULT_MAX_HISTORY

    current_app.logger.info(
        f"/api/answer - request_data: {data}, source: {source}",
        extra={"data": json.dumps({"request_data": data, "source": source})},
    )

    retriever = RetrieverCreator.create_retriever(
        retriever_name,
        question=question,
        source=source,
        chat_history=[],
        prompt="default",
        chunks=chunks,
        token_limit=token_limit,
        gpt_model=gpt_model,
        user_api_key=user_api_key,
    )
    docs = retriever.search()

    retriever_params = retriever.get_params()
    user_logs_collection.insert_one(
        {
            "action": "api_search",
            "level": "info",
            "user": "local",
            "api_key": user_api_key,
            "question": question,
            "sources": docs,
            "retriever_params": retriever_params,
            "timestamp": datetime.datetime.now(datetime.timezone.utc),
        }
    )

    if data.get("isNoneDoc"):
        for doc in docs:
            doc["source"] = "None"

    return docs



File: application/api/internal/__init__.py




File: application/api/internal/routes.py

import os
import datetime
from flask import Blueprint, request, send_from_directory
from pymongo import MongoClient
from werkzeug.utils import secure_filename
from bson.objectid import ObjectId

from application.core.settings import settings
mongo = MongoClient(settings.MONGO_URI)
db = mongo["docsgpt"]
conversations_collection = db["conversations"]
sources_collection = db["sources"]

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


internal = Blueprint('internal', __name__)
@internal.route("/api/download", methods=["get"])
def download_file():
    user = secure_filename(request.args.get("user"))
    job_name = secure_filename(request.args.get("name"))
    filename = secure_filename(request.args.get("file"))
    save_dir = os.path.join(current_dir, settings.UPLOAD_FOLDER, user, job_name)
    return send_from_directory(save_dir, filename, as_attachment=True)



@internal.route("/api/upload_index", methods=["POST"])
def upload_index_files():
    """Upload two files(index.faiss, index.pkl) to the user's folder."""
    if "user" not in request.form:
        return {"status": "no user"}
    user = secure_filename(request.form["user"])
    if "name" not in request.form:
        return {"status": "no name"}
    job_name = secure_filename(request.form["name"])
    tokens = secure_filename(request.form["tokens"])
    retriever = secure_filename(request.form["retriever"])
    id = secure_filename(request.form["id"])
    type = secure_filename(request.form["type"])
    remote_data = secure_filename(request.form["remote_data"]) if "remote_data" in  request.form else None

    save_dir = os.path.join(current_dir, "indexes", str(id))
    if settings.VECTOR_STORE == "faiss":
        if "file_faiss" not in request.files:
            print("No file part")
            return {"status": "no file"}
        file_faiss = request.files["file_faiss"]
        if file_faiss.filename == "":
            return {"status": "no file name"}
        if "file_pkl" not in request.files:
            print("No file part")
            return {"status": "no file"}
        file_pkl = request.files["file_pkl"]
        if file_pkl.filename == "":
            return {"status": "no file name"}
        # saves index files
        
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        file_faiss.save(os.path.join(save_dir, "index.faiss"))
        file_pkl.save(os.path.join(save_dir, "index.pkl"))
    # create entry in sources_collection
    sources_collection.insert_one(
        {
            "_id": ObjectId(id),
            "user": user,
            "name": job_name,
            "language": job_name,
            "date": datetime.datetime.now().strftime("%d/%m/%Y %H:%M:%S"),
            "model": settings.EMBEDDINGS_NAME,
            "type": type,
            "tokens": tokens,
            "retriever": retriever,
            "remote_data": remote_data
        }
    )
    return {"status": "ok"}


File: application/api/user/tasks.py

from application.worker import ingest_worker, remote_worker
from application.celery_init import celery

@celery.task(bind=True)
def ingest(self, directory, formats, name_job, filename, user):
    resp = ingest_worker(self, directory, formats, name_job, filename, user)
    return resp

@celery.task(bind=True)
def ingest_remote(self, source_data, job_name, user, loader):
    resp = remote_worker(self, source_data, job_name, user, loader)
    return resp



File: application/api/user/__init__.py




File: application/api/user/routes.py

import datetime
import os
import shutil
import uuid

from bson.binary import Binary, UuidRepresentation
from bson.dbref import DBRef
from bson.objectid import ObjectId
from flask import Blueprint, jsonify, request
from pymongo import MongoClient
from werkzeug.utils import secure_filename

from application.api.user.tasks import ingest, ingest_remote

from application.core.settings import settings
from application.vectorstore.vector_creator import VectorCreator

mongo = MongoClient(settings.MONGO_URI)
db = mongo["docsgpt"]
conversations_collection = db["conversations"]
sources_collection = db["sources"]
prompts_collection = db["prompts"]
feedback_collection = db["feedback"]
api_key_collection = db["api_keys"]
token_usage_collection = db["token_usage"]
shared_conversations_collections = db["shared_conversations"]
user_logs_collection = db["user_logs"]

user = Blueprint("user", __name__)

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


def generate_minute_range(start_date, end_date):
    return {
        (start_date + datetime.timedelta(minutes=i)).strftime("%Y-%m-%d %H:%M:00"): 0
        for i in range(int((end_date - start_date).total_seconds() // 60) + 1)
    }


def generate_hourly_range(start_date, end_date):
    return {
        (start_date + datetime.timedelta(hours=i)).strftime("%Y-%m-%d %H:00"): 0
        for i in range(int((end_date - start_date).total_seconds() // 3600) + 1)
    }


def generate_date_range(start_date, end_date):
    return {
        (start_date + datetime.timedelta(days=i)).strftime("%Y-%m-%d"): 0
        for i in range((end_date - start_date).days + 1)
    }


@user.route("/api/delete_conversation", methods=["POST"])
def delete_conversation():
    # deletes a conversation from the database
    conversation_id = request.args.get("id")
    # write to mongodb
    conversations_collection.delete_one(
        {
            "_id": ObjectId(conversation_id),
        }
    )

    return {"status": "ok"}


@user.route("/api/delete_all_conversations", methods=["GET"])
def delete_all_conversations():
    user_id = "local"
    conversations_collection.delete_many({"user": user_id})
    return {"status": "ok"}


@user.route("/api/get_conversations", methods=["get"])
def get_conversations():
    # provides a list of conversations
    conversations = conversations_collection.find().sort("date", -1).limit(30)
    list_conversations = []
    for conversation in conversations:
        list_conversations.append(
            {"id": str(conversation["_id"]), "name": conversation["name"]}
        )

    # list_conversations = [{"id": "default", "name": "default"}, {"id": "jeff", "name": "jeff"}]

    return jsonify(list_conversations)


@user.route("/api/get_single_conversation", methods=["get"])
def get_single_conversation():
    # provides data for a conversation
    conversation_id = request.args.get("id")
    conversation = conversations_collection.find_one({"_id": ObjectId(conversation_id)})
    return jsonify(conversation["queries"])


@user.route("/api/update_conversation_name", methods=["POST"])
def update_conversation_name():
    # update data for a conversation
    data = request.get_json()
    id = data["id"]
    name = data["name"]
    conversations_collection.update_one({"_id": ObjectId(id)}, {"$set": {"name": name}})
    return {"status": "ok"}


@user.route("/api/feedback", methods=["POST"])
def api_feedback():
    data = request.get_json()
    question = data["question"]
    answer = data["answer"]
    feedback = data["feedback"]
    new_doc = {
        "question": question,
        "answer": answer,
        "feedback": feedback,
        "timestamp": datetime.datetime.now(datetime.timezone.utc),
    }
    if "api_key" in data:
        new_doc["api_key"] = data["api_key"]
    feedback_collection.insert_one(new_doc)
    return {"status": "ok"}


@user.route("/api/delete_by_ids", methods=["get"])
def delete_by_ids():
    """Delete by ID. These are the IDs in the vectorstore"""

    ids = request.args.get("path")
    if not ids:
        return {"status": "error"}

    if settings.VECTOR_STORE == "faiss":
        result = sources_collection.delete_index(ids=ids)
        if result:
            return {"status": "ok"}
    return {"status": "error"}


@user.route("/api/delete_old", methods=["get"])
def delete_old():
    """Delete old indexes."""
    import shutil

    source_id = request.args.get("source_id")
    doc = sources_collection.find_one(
        {
            "_id": ObjectId(source_id),
            "user": "local",
        }
    )
    if doc is None:
        return {"status": "not found"}, 404
    if settings.VECTOR_STORE == "faiss":
        try:
            shutil.rmtree(os.path.join(current_dir, str(doc["_id"])))
        except FileNotFoundError:
            pass
    else:
        vetorstore = VectorCreator.create_vectorstore(
            settings.VECTOR_STORE, source_id=str(doc["_id"])
        )
        vetorstore.delete_index()
    sources_collection.delete_one(
        {
            "_id": ObjectId(source_id),
        }
    )

    return {"status": "ok"}


@user.route("/api/upload", methods=["POST"])
def upload_file():
    """Upload a file to get vectorized and indexed."""
    if "user" not in request.form:
        return {"status": "no user"}
    user = secure_filename(request.form["user"])
    if "name" not in request.form:
        return {"status": "no name"}
    job_name = secure_filename(request.form["name"])
    # check if the post request has the file part
    files = request.files.getlist("file")

    if not files or all(file.filename == "" for file in files):
        return {"status": "no file name"}

    # Directory where files will be saved
    save_dir = os.path.join(current_dir, settings.UPLOAD_FOLDER, user, job_name)
    os.makedirs(save_dir, exist_ok=True)

    if len(files) > 1:
        # Multiple files; prepare them for zip
        temp_dir = os.path.join(save_dir, "temp")
        os.makedirs(temp_dir, exist_ok=True)

        for file in files:
            filename = secure_filename(file.filename)
            file.save(os.path.join(temp_dir, filename))

        # Use shutil.make_archive to zip the temp directory
        zip_path = shutil.make_archive(
            base_name=os.path.join(save_dir, job_name), format="zip", root_dir=temp_dir
        )
        final_filename = os.path.basename(zip_path)

        # Clean up the temporary directory after zipping
        shutil.rmtree(temp_dir)
    else:
        # Single file
        file = files[0]
        final_filename = secure_filename(file.filename)
        file_path = os.path.join(save_dir, final_filename)
        file.save(file_path)

    # Call ingest with the single file or zipped file
    task = ingest.delay(
        settings.UPLOAD_FOLDER,
        [".rst", ".md", ".pdf", ".txt", ".docx", ".csv", ".epub", ".html", ".mdx"],
        job_name,
        final_filename,
        user,
    )

    return {"status": "ok", "task_id": task.id}


@user.route("/api/remote", methods=["POST"])
def upload_remote():
    """Upload a remote source to get vectorized and indexed."""
    if "user" not in request.form:
        return {"status": "no user"}
    user = secure_filename(request.form["user"])
    if "source" not in request.form:
        return {"status": "no source"}
    source = secure_filename(request.form["source"])
    if "name" not in request.form:
        return {"status": "no name"}
    job_name = secure_filename(request.form["name"])
    if "data" not in request.form:
        print("No data")
        return {"status": "no data"}
    source_data = request.form["data"]

    if source_data:
        task = ingest_remote.delay(
            source_data=source_data, job_name=job_name, user=user, loader=source
        )
        task_id = task.id
        return {"status": "ok", "task_id": task_id}
    else:
        return {"status": "error"}


@user.route("/api/task_status", methods=["GET"])
def task_status():
    """Get celery job status."""
    task_id = request.args.get("task_id")
    from application.celery_init import celery

    task = celery.AsyncResult(task_id)
    task_meta = task.info
    return {"status": task.status, "result": task_meta}


@user.route("/api/combine", methods=["GET"])
def combined_json():
    user = "local"
    """Provide json file with combined available indexes."""
    # get json from https://d3dg1063dc54p9.cloudfront.net/combined.json

    data = [
        {
            "name": "default",
            "date": "default",
            "model": settings.EMBEDDINGS_NAME,
            "location": "remote",
            "tokens": "",
            "retriever": "classic",
        }
    ]
    # structure: name, language, version, description, fullName, date, docLink
    # append data from sources_collection in sorted order in descending order of date
    for index in sources_collection.find({"user": user}).sort("date", -1):
        data.append(
            {
                "id": str(index["_id"]),
                "name": index["name"],
                "date": index["date"],
                "model": settings.EMBEDDINGS_NAME,
                "location": "local",
                "tokens": index["tokens"] if ("tokens" in index.keys()) else "",
                "retriever": (
                    index["retriever"] if ("retriever" in index.keys()) else "classic"
                ),
            }
        )
    if "duckduck_search" in settings.RETRIEVERS_ENABLED:
        data.append(
            {
                "name": "DuckDuckGo Search",
                "date": "duckduck_search",
                "model": settings.EMBEDDINGS_NAME,
                "location": "custom",
                "tokens": "",
                "retriever": "duckduck_search",
            }
        )
    if "brave_search" in settings.RETRIEVERS_ENABLED:
        data.append(
            {
                "name": "Brave Search",
                "language": "en",
                "date": "brave_search",
                "model": settings.EMBEDDINGS_NAME,
                "location": "custom",
                "tokens": "",
                "retriever": "brave_search",
            }
        )

    return jsonify(data)


@user.route("/api/docs_check", methods=["POST"])
def check_docs():
    data = request.get_json()

    vectorstore = "vectors/" + secure_filename(data["docs"])
    if os.path.exists(vectorstore) or data["docs"] == "default":
        return {"status": "exists"}
    else:
        return {"status": "not found"}


@user.route("/api/create_prompt", methods=["POST"])
def create_prompt():
    data = request.get_json()
    content = data["content"]
    name = data["name"]
    if name == "":
        return {"status": "error"}
    user = "local"
    resp = prompts_collection.insert_one(
        {
            "name": name,
            "content": content,
            "user": user,
        }
    )
    new_id = str(resp.inserted_id)
    return {"id": new_id}


@user.route("/api/get_prompts", methods=["GET"])
def get_prompts():
    user = "local"
    prompts = prompts_collection.find({"user": user})
    list_prompts = []
    list_prompts.append({"id": "default", "name": "default", "type": "public"})
    list_prompts.append({"id": "creative", "name": "creative", "type": "public"})
    list_prompts.append({"id": "strict", "name": "strict", "type": "public"})
    for prompt in prompts:
        list_prompts.append(
            {"id": str(prompt["_id"]), "name": prompt["name"], "type": "private"}
        )

    return jsonify(list_prompts)


@user.route("/api/get_single_prompt", methods=["GET"])
def get_single_prompt():
    prompt_id = request.args.get("id")
    if prompt_id == "default":
        with open(
            os.path.join(current_dir, "prompts", "chat_combine_default.txt"), "r"
        ) as f:
            chat_combine_template = f.read()
        return jsonify({"content": chat_combine_template})
    elif prompt_id == "creative":
        with open(
            os.path.join(current_dir, "prompts", "chat_combine_creative.txt"), "r"
        ) as f:
            chat_reduce_creative = f.read()
        return jsonify({"content": chat_reduce_creative})
    elif prompt_id == "strict":
        with open(
            os.path.join(current_dir, "prompts", "chat_combine_strict.txt"), "r"
        ) as f:
            chat_reduce_strict = f.read()
        return jsonify({"content": chat_reduce_strict})

    prompt = prompts_collection.find_one({"_id": ObjectId(prompt_id)})
    return jsonify({"content": prompt["content"]})


@user.route("/api/delete_prompt", methods=["POST"])
def delete_prompt():
    data = request.get_json()
    id = data["id"]
    prompts_collection.delete_one(
        {
            "_id": ObjectId(id),
        }
    )
    return {"status": "ok"}


@user.route("/api/update_prompt", methods=["POST"])
def update_prompt_name():
    data = request.get_json()
    id = data["id"]
    name = data["name"]
    content = data["content"]
    # check if name is null
    if name == "":
        return {"status": "error"}
    prompts_collection.update_one(
        {"_id": ObjectId(id)}, {"$set": {"name": name, "content": content}}
    )
    return {"status": "ok"}


@user.route("/api/get_api_keys", methods=["GET"])
def get_api_keys():
    user = "local"
    keys = api_key_collection.find({"user": user})
    list_keys = []
    for key in keys:
        if "source" in key and isinstance(key["source"], DBRef):
            source = db.dereference(key["source"])
            if source is None:
                continue
            else:
                source_name = source["name"]
        elif "retriever" in key:
            source_name = key["retriever"]
        else:
            continue

        list_keys.append(
            {
                "id": str(key["_id"]),
                "name": key["name"],
                "key": key["key"][:4] + "..." + key["key"][-4:],
                "source": source_name,
                "prompt_id": key["prompt_id"],
                "chunks": key["chunks"],
            }
        )
    return jsonify(list_keys)


@user.route("/api/create_api_key", methods=["POST"])
def create_api_key():
    data = request.get_json()
    name = data["name"]
    prompt_id = data["prompt_id"]
    chunks = data["chunks"]
    key = str(uuid.uuid4())
    user = "local"
    new_api_key = {
        "name": name,
        "key": key,
        "user": user,
        "prompt_id": prompt_id,
        "chunks": chunks,
    }
    if "source" in data and ObjectId.is_valid(data["source"]):
        new_api_key["source"] = DBRef("sources", ObjectId(data["source"]))
    if "retriever" in data:
        new_api_key["retriever"] = data["retriever"]
    resp = api_key_collection.insert_one(new_api_key)
    new_id = str(resp.inserted_id)
    return {"id": new_id, "key": key}


@user.route("/api/delete_api_key", methods=["POST"])
def delete_api_key():
    data = request.get_json()
    id = data["id"]
    api_key_collection.delete_one(
        {
            "_id": ObjectId(id),
        }
    )
    return {"status": "ok"}


# route to share conversation
##isPromptable should be passed through queries
@user.route("/api/share", methods=["POST"])
def share_conversation():
    try:
        data = request.get_json()
        user = "local" if "user" not in data else data["user"]
        conversation_id = data["conversation_id"]
        isPromptable = request.args.get("isPromptable").lower() == "true"

        conversation = conversations_collection.find_one(
            {"_id": ObjectId(conversation_id)}
        )
        if conversation is None:
            raise Exception("Conversation does not exist")
        current_n_queries = len(conversation["queries"])

        ##generate binary representation of uuid
        explicit_binary = Binary.from_uuid(uuid.uuid4(), UuidRepresentation.STANDARD)

        if isPromptable:
            prompt_id = "default" if "prompt_id" not in data else data["prompt_id"]
            chunks = "2" if "chunks" not in data else data["chunks"]

            name = conversation["name"] + "(shared)"
            new_api_key_data = {
                "prompt_id": prompt_id,
                "chunks": chunks,
                "user": user,
            }
            if "source" in data and ObjectId.is_valid(data["source"]):
                new_api_key_data["source"] = DBRef("sources", ObjectId(data["source"]))
            elif "retriever" in data:
                new_api_key_data["retriever"] = data["retriever"]

            pre_existing_api_document = api_key_collection.find_one(new_api_key_data)
            if pre_existing_api_document:
                api_uuid = pre_existing_api_document["key"]
                pre_existing = shared_conversations_collections.find_one(
                    {
                        "conversation_id": DBRef(
                            "conversations", ObjectId(conversation_id)
                        ),
                        "isPromptable": isPromptable,
                        "first_n_queries": current_n_queries,
                        "user": user,
                        "api_key": api_uuid,
                    }
                )
                if pre_existing is not None:
                    return (
                        jsonify(
                            {
                                "success": True,
                                "identifier": str(pre_existing["uuid"].as_uuid()),
                            }
                        ),
                        200,
                    )
                else:
                    shared_conversations_collections.insert_one(
                        {
                            "uuid": explicit_binary,
                            "conversation_id": {
                                "$ref": "conversations",
                                "$id": ObjectId(conversation_id),
                            },
                            "isPromptable": isPromptable,
                            "first_n_queries": current_n_queries,
                            "user": user,
                            "api_key": api_uuid,
                        }
                    )
                    return jsonify(
                        {"success": True, "identifier": str(explicit_binary.as_uuid())}
                    )
            else:

                api_uuid = str(uuid.uuid4())
                new_api_key_data["key"] = api_uuid
                new_api_key_data["name"] = name
                if "source" in data and ObjectId.is_valid(data["source"]):
                    new_api_key_data["source"] = DBRef(
                        "sources", ObjectId(data["source"])
                    )
                if "retriever" in data:
                    new_api_key_data["retriever"] = data["retriever"]
                api_key_collection.insert_one(new_api_key_data)
                shared_conversations_collections.insert_one(
                    {
                        "uuid": explicit_binary,
                        "conversation_id": {
                            "$ref": "conversations",
                            "$id": ObjectId(conversation_id),
                        },
                        "isPromptable": isPromptable,
                        "first_n_queries": current_n_queries,
                        "user": user,
                        "api_key": api_uuid,
                    }
                )
            ## Identifier as route parameter in frontend
            return (
                jsonify(
                    {"success": True, "identifier": str(explicit_binary.as_uuid())}
                ),
                201,
            )

        ##isPromptable = False
        pre_existing = shared_conversations_collections.find_one(
            {
                "conversation_id": DBRef("conversations", ObjectId(conversation_id)),
                "isPromptable": isPromptable,
                "first_n_queries": current_n_queries,
                "user": user,
            }
        )
        if pre_existing is not None:
            return (
                jsonify(
                    {"success": True, "identifier": str(pre_existing["uuid"].as_uuid())}
                ),
                200,
            )
        else:
            shared_conversations_collections.insert_one(
                {
                    "uuid": explicit_binary,
                    "conversation_id": {
                        "$ref": "conversations",
                        "$id": ObjectId(conversation_id),
                    },
                    "isPromptable": isPromptable,
                    "first_n_queries": current_n_queries,
                    "user": user,
                }
            )
            ## Identifier as route parameter in frontend
            return (
                jsonify(
                    {"success": True, "identifier": str(explicit_binary.as_uuid())}
                ),
                201,
            )
    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400


# route to get publicly shared conversations
@user.route("/api/shared_conversation/<string:identifier>", methods=["GET"])
def get_publicly_shared_conversations(identifier: str):
    try:
        query_uuid = Binary.from_uuid(
            uuid.UUID(identifier), UuidRepresentation.STANDARD
        )
        shared = shared_conversations_collections.find_one({"uuid": query_uuid})
        conversation_queries = []
        if (
            shared
            and "conversation_id" in shared
            and isinstance(shared["conversation_id"], DBRef)
        ):
            # Resolve the DBRef
            conversation_ref = shared["conversation_id"]
            conversation = db.dereference(conversation_ref)
            if conversation is None:
                return (
                    jsonify(
                        {
                            "sucess": False,
                            "error": "might have broken url or the conversation does not exist",
                        }
                    ),
                    404,
                )
            conversation_queries = conversation["queries"][
                : (shared["first_n_queries"])
            ]
        else:
            return (
                jsonify(
                    {
                        "sucess": False,
                        "error": "might have broken url or the conversation does not exist",
                    }
                ),
                404,
            )
        date = conversation["_id"].generation_time.isoformat()
        res = {
            "success": True,
            "queries": conversation_queries,
            "title": conversation["name"],
            "timestamp": date,
        }
        if shared["isPromptable"] and "api_key" in shared:
            res["api_key"] = shared["api_key"]
        return jsonify(res), 200
    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400


@user.route("/api/get_message_analytics", methods=["POST"])
def get_message_analytics():
    data = request.get_json()
    api_key_id = data.get("api_key_id")
    filter_option = data.get("filter_option", "last_30_days")

    try:
        api_key = (
            api_key_collection.find_one({"_id": ObjectId(api_key_id)})["key"]
            if api_key_id
            else None
        )
    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400
    end_date = datetime.datetime.now(datetime.timezone.utc)

    if filter_option == "last_hour":
        start_date = end_date - datetime.timedelta(hours=1)
        group_format = "%Y-%m-%d %H:%M:00"
        group_stage = {
            "$group": {
                "_id": {
                    "minute": {
                        "$dateToString": {"format": group_format, "date": "$date"}
                    }
                },
                "total_messages": {"$sum": 1},
            }
        }

    elif filter_option == "last_24_hour":
        start_date = end_date - datetime.timedelta(hours=24)
        group_format = "%Y-%m-%d %H:00"
        group_stage = {
            "$group": {
                "_id": {
                    "hour": {"$dateToString": {"format": group_format, "date": "$date"}}
                },
                "total_messages": {"$sum": 1},
            }
        }

    else:
        if filter_option in ["last_7_days", "last_15_days", "last_30_days"]:
            filter_days = (
                6
                if filter_option == "last_7_days"
                else (14 if filter_option == "last_15_days" else 29)
            )
        else:
            return jsonify({"success": False, "error": "Invalid option"}), 400
        start_date = end_date - datetime.timedelta(days=filter_days)
        start_date = start_date.replace(hour=0, minute=0, second=0, microsecond=0)
        end_date = end_date.replace(hour=23, minute=59, second=59, microsecond=999999)
        group_format = "%Y-%m-%d"
        group_stage = {
            "$group": {
                "_id": {
                    "day": {"$dateToString": {"format": group_format, "date": "$date"}}
                },
                "total_messages": {"$sum": 1},
            }
        }

    try:
        match_stage = {
            "$match": {
                "date": {"$gte": start_date, "$lte": end_date},
            }
        }
        if api_key:
            match_stage["$match"]["api_key"] = api_key
        message_data = conversations_collection.aggregate(
            [
                match_stage,
                group_stage,
                {"$sort": {"_id": 1}},
            ]
        )

        if filter_option == "last_hour":
            intervals = generate_minute_range(start_date, end_date)
        elif filter_option == "last_24_hour":
            intervals = generate_hourly_range(start_date, end_date)
        else:
            intervals = generate_date_range(start_date, end_date)

        daily_messages = {interval: 0 for interval in intervals}

        for entry in message_data:
            if filter_option == "last_hour":
                daily_messages[entry["_id"]["minute"]] = entry["total_messages"]
            elif filter_option == "last_24_hour":
                daily_messages[entry["_id"]["hour"]] = entry["total_messages"]
            else:
                daily_messages[entry["_id"]["day"]] = entry["total_messages"]

    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400

    return jsonify({"success": True, "messages": daily_messages}), 200


@user.route("/api/get_token_analytics", methods=["POST"])
def get_token_analytics():
    data = request.get_json()
    api_key_id = data.get("api_key_id")
    filter_option = data.get("filter_option", "last_30_days")

    try:
        api_key = (
            api_key_collection.find_one({"_id": ObjectId(api_key_id)})["key"]
            if api_key_id
            else None
        )
    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400
    end_date = datetime.datetime.now(datetime.timezone.utc)

    if filter_option == "last_hour":
        start_date = end_date - datetime.timedelta(hours=1)
        group_format = "%Y-%m-%d %H:%M:00"
        group_stage = {
            "$group": {
                "_id": {
                    "minute": {
                        "$dateToString": {"format": group_format, "date": "$timestamp"}
                    }
                },
                "total_tokens": {
                    "$sum": {"$add": ["$prompt_tokens", "$generated_tokens"]}
                },
            }
        }

    elif filter_option == "last_24_hour":
        start_date = end_date - datetime.timedelta(hours=24)
        group_format = "%Y-%m-%d %H:00"
        group_stage = {
            "$group": {
                "_id": {
                    "hour": {
                        "$dateToString": {"format": group_format, "date": "$timestamp"}
                    }
                },
                "total_tokens": {
                    "$sum": {"$add": ["$prompt_tokens", "$generated_tokens"]}
                },
            }
        }

    else:
        if filter_option in ["last_7_days", "last_15_days", "last_30_days"]:
            filter_days = (
                6
                if filter_option == "last_7_days"
                else (14 if filter_option == "last_15_days" else 29)
            )
        else:
            return jsonify({"success": False, "error": "Invalid option"}), 400
        start_date = end_date - datetime.timedelta(days=filter_days)
        start_date = start_date.replace(hour=0, minute=0, second=0, microsecond=0)
        end_date = end_date.replace(hour=23, minute=59, second=59, microsecond=999999)
        group_format = "%Y-%m-%d"
        group_stage = {
            "$group": {
                "_id": {
                    "day": {
                        "$dateToString": {"format": group_format, "date": "$timestamp"}
                    }
                },
                "total_tokens": {
                    "$sum": {"$add": ["$prompt_tokens", "$generated_tokens"]}
                },
            }
        }

    try:
        match_stage = {
            "$match": {
                "timestamp": {"$gte": start_date, "$lte": end_date},
            }
        }
        if api_key:
            match_stage["$match"]["api_key"] = api_key

        token_usage_data = token_usage_collection.aggregate(
            [
                match_stage,
                group_stage,
                {"$sort": {"_id": 1}},
            ]
        )

        if filter_option == "last_hour":
            intervals = generate_minute_range(start_date, end_date)
        elif filter_option == "last_24_hour":
            intervals = generate_hourly_range(start_date, end_date)
        else:
            intervals = generate_date_range(start_date, end_date)

        daily_token_usage = {interval: 0 for interval in intervals}

        for entry in token_usage_data:
            if filter_option == "last_hour":
                daily_token_usage[entry["_id"]["minute"]] = entry["total_tokens"]
            elif filter_option == "last_24_hour":
                daily_token_usage[entry["_id"]["hour"]] = entry["total_tokens"]
            else:
                daily_token_usage[entry["_id"]["day"]] = entry["total_tokens"]

    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400

    return jsonify({"success": True, "token_usage": daily_token_usage}), 200


@user.route("/api/get_feedback_analytics", methods=["POST"])
def get_feedback_analytics():
    data = request.get_json()
    api_key_id = data.get("api_key_id")
    filter_option = data.get("filter_option", "last_30_days")

    try:
        api_key = (
            api_key_collection.find_one({"_id": ObjectId(api_key_id)})["key"]
            if api_key_id
            else None
        )
    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400
    end_date = datetime.datetime.now(datetime.timezone.utc)

    if filter_option == "last_hour":
        start_date = end_date - datetime.timedelta(hours=1)
        group_format = "%Y-%m-%d %H:%M:00"
        group_stage_1 = {
            "$group": {
                "_id": {
                    "minute": {
                        "$dateToString": {"format": group_format, "date": "$timestamp"}
                    },
                    "feedback": "$feedback",
                },
                "count": {"$sum": 1},
            }
        }
        group_stage_2 = {
            "$group": {
                "_id": "$_id.minute",
                "likes": {
                    "$sum": {
                        "$cond": [
                            {"$eq": ["$_id.feedback", "LIKE"]},
                            "$count",
                            0,
                        ]
                    }
                },
                "dislikes": {
                    "$sum": {
                        "$cond": [
                            {"$eq": ["$_id.feedback", "DISLIKE"]},
                            "$count",
                            0,
                        ]
                    }
                },
            }
        }

    elif filter_option == "last_24_hour":
        start_date = end_date - datetime.timedelta(hours=24)
        group_format = "%Y-%m-%d %H:00"
        group_stage_1 = {
            "$group": {
                "_id": {
                    "hour": {
                        "$dateToString": {"format": group_format, "date": "$timestamp"}
                    },
                    "feedback": "$feedback",
                },
                "count": {"$sum": 1},
            }
        }
        group_stage_2 = {
            "$group": {
                "_id": "$_id.hour",
                "likes": {
                    "$sum": {
                        "$cond": [
                            {"$eq": ["$_id.feedback", "LIKE"]},
                            "$count",
                            0,
                        ]
                    }
                },
                "dislikes": {
                    "$sum": {
                        "$cond": [
                            {"$eq": ["$_id.feedback", "DISLIKE"]},
                            "$count",
                            0,
                        ]
                    }
                },
            }
        }

    else:
        if filter_option in ["last_7_days", "last_15_days", "last_30_days"]:
            filter_days = (
                6
                if filter_option == "last_7_days"
                else (14 if filter_option == "last_15_days" else 29)
            )
        else:
            return jsonify({"success": False, "error": "Invalid option"}), 400
        start_date = end_date - datetime.timedelta(days=filter_days)
        start_date = start_date.replace(hour=0, minute=0, second=0, microsecond=0)
        end_date = end_date.replace(hour=23, minute=59, second=59, microsecond=999999)
        group_format = "%Y-%m-%d"
        group_stage_1 = {
            "$group": {
                "_id": {
                    "day": {
                        "$dateToString": {"format": group_format, "date": "$timestamp"}
                    },
                    "feedback": "$feedback",
                },
                "count": {"$sum": 1},
            }
        }
        group_stage_2 = {
            "$group": {
                "_id": "$_id.day",
                "likes": {
                    "$sum": {
                        "$cond": [
                            {"$eq": ["$_id.feedback", "LIKE"]},
                            "$count",
                            0,
                        ]
                    }
                },
                "dislikes": {
                    "$sum": {
                        "$cond": [
                            {"$eq": ["$_id.feedback", "DISLIKE"]},
                            "$count",
                            0,
                        ]
                    }
                },
            }
        }

    try:
        match_stage = {
            "$match": {
                "timestamp": {"$gte": start_date, "$lte": end_date},
            }
        }
        if api_key:
            match_stage["$match"]["api_key"] = api_key

        feedback_data = feedback_collection.aggregate(
            [
                match_stage,
                group_stage_1,
                group_stage_2,
                {"$sort": {"_id": 1}},
            ]
        )

        if filter_option == "last_hour":
            intervals = generate_minute_range(start_date, end_date)
        elif filter_option == "last_24_hour":
            intervals = generate_hourly_range(start_date, end_date)
        else:
            intervals = generate_date_range(start_date, end_date)

        daily_feedback = {
            interval: {"positive": 0, "negative": 0} for interval in intervals
        }

        for entry in feedback_data:
            daily_feedback[entry["_id"]] = {
                "positive": entry["likes"],
                "negative": entry["dislikes"],
            }

    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400

    return jsonify({"success": True, "feedback": daily_feedback}), 200


@user.route("/api/get_user_logs", methods=["POST"])
def get_user_logs():
    data = request.get_json()
    page = int(data.get("page", 1))
    api_key_id = data.get("api_key_id")
    page_size = int(data.get("page_size", 10))
    skip = (page - 1) * page_size

    try:
        api_key = (
            api_key_collection.find_one({"_id": ObjectId(api_key_id)})["key"]
            if api_key_id
            else None
        )
    except Exception as err:
        print(err)
        return jsonify({"success": False, "error": str(err)}), 400

    query = {}
    if api_key:
        query = {"api_key": api_key}
    items_cursor = (
        user_logs_collection.find(query)
        .sort("timestamp", -1)
        .skip(skip)
        .limit(page_size + 1)
    )
    items = list(items_cursor)

    results = []
    for item in items[:page_size]:
        results.append(
            {
                "id": str(item.get("_id")),
                "action": item.get("action"),
                "level": item.get("level"),
                "user": item.get("user"),
                "question": item.get("question"),
                "sources": item.get("sources"),
                "retriever_params": item.get("retriever_params"),
                "timestamp": item.get("timestamp"),
            }
        )
    has_more = len(items) > page_size

    return (
        jsonify(
            {
                "success": True,
                "logs": results,
                "page": page,
                "page_size": page_size,
                "has_more": has_more,
            }
        ),
        200,
    )



File: application/vectorstore/elasticsearch.py

from application.vectorstore.base import BaseVectorStore
from application.core.settings import settings
from application.vectorstore.document_class import Document
import elasticsearch




class ElasticsearchStore(BaseVectorStore):
    _es_connection = None  # Class attribute to hold the Elasticsearch connection

    def __init__(self, source_id, embeddings_key, index_name=settings.ELASTIC_INDEX):
        super().__init__()
        self.source_id = source_id.replace("application/indexes/", "").rstrip("/")
        self.embeddings_key = embeddings_key
        self.index_name = index_name
        
        if ElasticsearchStore._es_connection is None:
            connection_params = {}
            if settings.ELASTIC_URL:
                connection_params["hosts"] = [settings.ELASTIC_URL]
                connection_params["http_auth"] = (settings.ELASTIC_USERNAME, settings.ELASTIC_PASSWORD)
            elif settings.ELASTIC_CLOUD_ID:
                connection_params["cloud_id"] = settings.ELASTIC_CLOUD_ID
                connection_params["basic_auth"] = (settings.ELASTIC_USERNAME, settings.ELASTIC_PASSWORD)
            else:
                raise ValueError("Please provide either elasticsearch_url or cloud_id.")

            

            ElasticsearchStore._es_connection = elasticsearch.Elasticsearch(**connection_params)
            
        self.docsearch = ElasticsearchStore._es_connection

    def connect_to_elasticsearch(
        *,
        es_url = None,
        cloud_id = None,
        api_key = None,
        username = None,
        password = None,
    ):
        try:
            import elasticsearch
        except ImportError:
            raise ImportError(
                "Could not import elasticsearch python package. "
                "Please install it with `pip install elasticsearch`."
            )

        if es_url and cloud_id:
            raise ValueError(
                "Both es_url and cloud_id are defined. Please provide only one."
            )

        connection_params = {}

        if es_url:
            connection_params["hosts"] = [es_url]
        elif cloud_id:
            connection_params["cloud_id"] = cloud_id
        else:
            raise ValueError("Please provide either elasticsearch_url or cloud_id.")

        if api_key:
            connection_params["api_key"] = api_key
        elif username and password:
            connection_params["basic_auth"] = (username, password)

        es_client = elasticsearch.Elasticsearch(
            **connection_params,
        )
        try:
            es_client.info()
        except Exception as e:
            raise e

        return es_client

    def search(self, question, k=2, index_name=settings.ELASTIC_INDEX, *args, **kwargs):
        embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)
        vector = embeddings.embed_query(question)
        knn = {
            "filter": [{"match": {"metadata.source_id.keyword": self.source_id}}],
            "field": "vector",
            "k": k,
            "num_candidates": 100,
            "query_vector": vector,
        }
        full_query = {
            "knn": knn,
            "query": {
                "bool": {
                    "must": [
                        {
                            "match": {
                                "text": {
                                    "query": question,
                                }
                            }
                        }
                    ],
                    "filter": [{"match": {"metadata.source_id.keyword": self.source_id}}],
                }
            },
            "rank": {"rrf": {}},
        }
        resp = self.docsearch.search(index=self.index_name, query=full_query['query'], size=k, knn=full_query['knn'])
        # create Documents objects from the results page_content ['_source']['text'], metadata ['_source']['metadata']
        doc_list = []
        for hit in resp['hits']['hits']:
            
            doc_list.append(Document(page_content = hit['_source']['text'], metadata = hit['_source']['metadata']))
        return doc_list

    def _create_index_if_not_exists(
            self, index_name, dims_length
        ):

        if self._es_connection.indices.exists(index=index_name):
            print(f"Index {index_name} already exists.")

        else:

            indexSettings = self.index(
                dims_length=dims_length,
            )
            self._es_connection.indices.create(index=index_name, **indexSettings)

    def index(
            self,
            dims_length,
        ):
        return {
            "mappings": {
                "properties": {
                    "vector": {
                        "type": "dense_vector",
                        "dims": dims_length,
                        "index": True,
                        "similarity": "cosine",
                    },
                }
            }
        }

    def add_texts(
        self,
        texts,
        metadatas = None,
        ids = None,
        refresh_indices = True,
        create_index_if_not_exists = True,
        bulk_kwargs = None,
        **kwargs,
        ):
        
        from elasticsearch.helpers import BulkIndexError, bulk

        bulk_kwargs = bulk_kwargs or {}
        import uuid
        embeddings = []
        ids = ids or [str(uuid.uuid4()) for _ in texts]
        requests = []
        embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)

        vectors = embeddings.embed_documents(list(texts))

        dims_length = len(vectors[0])

        if create_index_if_not_exists:
            self._create_index_if_not_exists(
                index_name=self.index_name, dims_length=dims_length
            )

        for i, (text, vector) in enumerate(zip(texts, vectors)):
            metadata = metadatas[i] if metadatas else {}

            requests.append(
                {
                    "_op_type": "index",
                    "_index": self.index_name,
                    "text": text,
                    "vector": vector,
                    "metadata": metadata,
                    "_id": ids[i],
                }
            )


        if len(requests) > 0:
            try:
                success, failed = bulk(
                    self._es_connection,
                    requests,
                    stats_only=True,
                    refresh=refresh_indices,
                    **bulk_kwargs,
                )
                return ids
            except BulkIndexError as e:
                print(f"Error adding texts: {e}")
                firstError = e.errors[0].get("index", {}).get("error", {})
                print(f"First error reason: {firstError.get('reason')}")
                raise e

        else:
            return []

    def delete_index(self):
        self._es_connection.delete_by_query(index=self.index_name, query={"match": {
                                      "metadata.source_id.keyword": self.source_id}},)



File: application/vectorstore/milvus.py

from typing import List, Optional
from uuid import uuid4


from application.core.settings import settings
from application.vectorstore.base import BaseVectorStore


class MilvusStore(BaseVectorStore):
    def __init__(self, path: str = "", embeddings_key: str = "embeddings"):
        super().__init__()
        from langchain_milvus import Milvus

        connection_args = {
            "uri": settings.MILVUS_URI,
            "token": settings.MILVUS_TOKEN,
        }
        self._docsearch = Milvus(
            embedding_function=self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key),
            collection_name=settings.MILVUS_COLLECTION_NAME,
            connection_args=connection_args,
        )
        self._path = path

    def search(self, question, k=2, *args, **kwargs):
        return self._docsearch.similarity_search(query=question, k=k, filter={"path": self._path} *args, **kwargs)

    def add_texts(self, texts: List[str], metadatas: Optional[List[dict]], *args, **kwargs):
        ids = [str(uuid4()) for _ in range(len(texts))]

        return self._docsearch.add_texts(texts=texts, metadatas=metadatas, ids=ids, *args, **kwargs)

    def save_local(self, *args, **kwargs):
        pass

    def delete_index(self, *args, **kwargs):
        pass



File: application/vectorstore/document_class.py

class Document(str):
    """Class for storing a piece of text and associated metadata."""

    def __new__(cls, page_content: str, metadata: dict):
        instance = super().__new__(cls, page_content)
        instance.page_content = page_content
        instance.metadata = metadata
        return instance



File: application/vectorstore/qdrant.py

from langchain_community.vectorstores.qdrant import Qdrant
from application.vectorstore.base import BaseVectorStore
from application.core.settings import settings
from qdrant_client import models


class QdrantStore(BaseVectorStore):
    def __init__(self, source_id: str = "", embeddings_key: str = "embeddings"):
        self._filter = models.Filter(
            must=[
                models.FieldCondition(
                    key="metadata.source_id",
                    match=models.MatchValue(value=source_id.replace("application/indexes/", "").rstrip("/")),
                )
            ]
        )

        self._docsearch = Qdrant.construct_instance(
            ["TEXT_TO_OBTAIN_EMBEDDINGS_DIMENSION"],
            embedding=self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key),
            collection_name=settings.QDRANT_COLLECTION_NAME,
            location=settings.QDRANT_LOCATION,
            url=settings.QDRANT_URL,
            port=settings.QDRANT_PORT,
            grpc_port=settings.QDRANT_GRPC_PORT,
            https=settings.QDRANT_HTTPS,
            prefer_grpc=settings.QDRANT_PREFER_GRPC,
            api_key=settings.QDRANT_API_KEY,
            prefix=settings.QDRANT_PREFIX,
            timeout=settings.QDRANT_TIMEOUT,
            path=settings.QDRANT_PATH,
            distance_func=settings.QDRANT_DISTANCE_FUNC,
        )

    def search(self, *args, **kwargs):
        return self._docsearch.similarity_search(filter=self._filter, *args, **kwargs)

    def add_texts(self, *args, **kwargs):
        return self._docsearch.add_texts(*args, **kwargs)

    def save_local(self, *args, **kwargs):
        pass

    def delete_index(self, *args, **kwargs):
        return self._docsearch.client.delete(
            collection_name=settings.QDRANT_COLLECTION_NAME, points_selector=self._filter
        )



File: application/vectorstore/mongodb.py

from application.vectorstore.base import BaseVectorStore
from application.core.settings import settings
from application.vectorstore.document_class import Document

class MongoDBVectorStore(BaseVectorStore):
    def __init__(
        self,
        source_id: str = "",
        embeddings_key: str = "embeddings",
        collection: str = "documents",
        index_name: str = "vector_search_index",
        text_key: str = "text",
        embedding_key: str = "embedding",
        database: str = "docsgpt",
    ):
        self._index_name = index_name
        self._text_key = text_key
        self._embedding_key = embedding_key
        self._embeddings_key = embeddings_key
        self._mongo_uri = settings.MONGO_URI
        self._source_id = source_id.replace("application/indexes/", "").rstrip("/")
        self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)

        try:
            import pymongo
        except ImportError:
            raise ImportError(
                "Could not import pymongo python package. "
                "Please install it with `pip install pymongo`."
            )

        self._client = pymongo.MongoClient(self._mongo_uri)
        self._database = self._client[database]
        self._collection = self._database[collection]

        
    def search(self, question, k=2, *args, **kwargs):
        query_vector = self._embedding.embed_query(question)

        pipeline = [
            {
                "$vectorSearch": {
                    "queryVector": query_vector, 
                    "path": self._embedding_key,
                    "limit": k, 
                    "numCandidates": k * 10, 
                    "index": self._index_name,
                    "filter": {
                        "source_id": {"$eq": self._source_id}
                    }
                }
            }
        ]

        cursor = self._collection.aggregate(pipeline)
        
        results = []
        for doc in cursor:
            text = doc[self._text_key]
            doc.pop("_id")
            doc.pop(self._text_key)
            doc.pop(self._embedding_key)
            metadata = doc
            results.append(Document(text, metadata))
        return results
    
    def _insert_texts(self, texts, metadatas):
        if not texts:
            return []
        embeddings = self._embedding.embed_documents(texts)
        to_insert = [
            {self._text_key: t, self._embedding_key: embedding, **m}
            for t, m, embedding in zip(texts, metadatas, embeddings)
        ]
        # insert the documents in MongoDB Atlas
        insert_result = self._collection.insert_many(to_insert)
        return insert_result.inserted_ids
    
    def add_texts(self,
        texts,
        metadatas = None,
        ids = None,
        refresh_indices = True,
        create_index_if_not_exists = True,
        bulk_kwargs = None,
        **kwargs,):


        #dims = self._embedding.client[1].word_embedding_dimension
        # # check if index exists
        # if create_index_if_not_exists:
        #     # check if index exists
        #     info = self._collection.index_information()
        #     if self._index_name not in info:
        #         index_mongo = {
        #         "fields": [{
        #             "type": "vector",
        #             "path": self._embedding_key,
        #             "numDimensions": dims,
        #             "similarity": "cosine",
        #         },
        #         {
        #             "type": "filter",
        #             "path": "store"
        #         }]
        #         }
        #         self._collection.create_index(self._index_name, index_mongo)

        batch_size = 100
        _metadatas = metadatas or ({} for _ in texts)
        texts_batch = []
        metadatas_batch = []
        result_ids = []
        for i, (text, metadata) in enumerate(zip(texts, _metadatas)):
            texts_batch.append(text)
            metadatas_batch.append(metadata)
            if (i + 1) % batch_size == 0:
                result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))
                texts_batch = []
                metadatas_batch = []
        if texts_batch:
            result_ids.extend(self._insert_texts(texts_batch, metadatas_batch))
        return result_ids
    
    def delete_index(self, *args, **kwargs):
        self._collection.delete_many({"source_id": self._source_id})


File: application/vectorstore/__init__.py




File: application/vectorstore/faiss.py

from langchain_community.vectorstores import FAISS
from application.vectorstore.base import BaseVectorStore
from application.core.settings import settings
import os

def get_vectorstore(path):
    if path:
        vectorstore = "indexes/"+path
        vectorstore = os.path.join("application", vectorstore)
    else:
        vectorstore = os.path.join("application")

    return vectorstore

class FaissStore(BaseVectorStore):

    def __init__(self, source_id, embeddings_key, docs_init=None):
        super().__init__()
        self.path = get_vectorstore(source_id)
        embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)
        if docs_init:
            self.docsearch = FAISS.from_documents(
                docs_init, embeddings
            )
        else:
            self.docsearch = FAISS.load_local(
                self.path, embeddings, 
                allow_dangerous_deserialization=True
            )
        self.assert_embedding_dimensions(embeddings)

    def search(self, *args, **kwargs):
        return self.docsearch.similarity_search(*args, **kwargs)

    def add_texts(self, *args, **kwargs):
        return self.docsearch.add_texts(*args, **kwargs)

    def save_local(self, *args, **kwargs):
        return self.docsearch.save_local(*args, **kwargs)

    def delete_index(self, *args, **kwargs):
        return self.docsearch.delete(*args, **kwargs)

    def assert_embedding_dimensions(self, embeddings):
        """
        Check that the word embedding dimension of the docsearch index matches
        the dimension of the word embeddings used 
        """
        if settings.EMBEDDINGS_NAME == "huggingface_sentence-transformers/all-mpnet-base-v2":
            try:
                word_embedding_dimension = embeddings.dimension
            except AttributeError as e:
                raise AttributeError("'dimension' attribute not found in embeddings instance. Make sure the embeddings object is properly initialized.") from e
            docsearch_index_dimension = self.docsearch.index.d
            if word_embedding_dimension != docsearch_index_dimension:
                raise ValueError(f"Embedding dimension mismatch: embeddings.dimension ({word_embedding_dimension}) " +
                                 f"!= docsearch index dimension ({docsearch_index_dimension})")


File: application/vectorstore/vector_creator.py

from application.vectorstore.faiss import FaissStore
from application.vectorstore.elasticsearch import ElasticsearchStore
from application.vectorstore.milvus import MilvusStore
from application.vectorstore.mongodb import MongoDBVectorStore
from application.vectorstore.qdrant import QdrantStore


class VectorCreator:
    vectorstores = {
        "faiss": FaissStore,
        "elasticsearch": ElasticsearchStore,
        "mongodb": MongoDBVectorStore,
        "qdrant": QdrantStore,
        "milvus": MilvusStore,
    }

    @classmethod
    def create_vectorstore(cls, type, *args, **kwargs):
        vectorstore_class = cls.vectorstores.get(type.lower())
        if not vectorstore_class:
            raise ValueError(f"No vectorstore class found for type {type}")
        return vectorstore_class(*args, **kwargs)



File: application/vectorstore/base.py

from abc import ABC, abstractmethod
import os
from sentence_transformers import SentenceTransformer
from langchain_openai import OpenAIEmbeddings
from application.core.settings import settings

class EmbeddingsWrapper:
    def __init__(self, model_name, *args, **kwargs):
        self.model = SentenceTransformer(model_name, config_kwargs={'allow_dangerous_deserialization': True}, *args, **kwargs)
        self.dimension = self.model.get_sentence_embedding_dimension()

    def embed_query(self, query: str):
        return self.model.encode(query).tolist()
    
    def embed_documents(self, documents: list):
        return self.model.encode(documents).tolist()
    
    def __call__(self, text):
        if isinstance(text, str):
            return self.embed_query(text)
        elif isinstance(text, list):
            return self.embed_documents(text)
        else:
            raise ValueError("Input must be a string or a list of strings")



class EmbeddingsSingleton:
    _instances = {}

    @staticmethod
    def get_instance(embeddings_name, *args, **kwargs):
        if embeddings_name not in EmbeddingsSingleton._instances:
            EmbeddingsSingleton._instances[embeddings_name] = EmbeddingsSingleton._create_instance(
                embeddings_name, *args, **kwargs
            )
        return EmbeddingsSingleton._instances[embeddings_name]

    @staticmethod
    def _create_instance(embeddings_name, *args, **kwargs):
        embeddings_factory = {
            "openai_text-embedding-ada-002": OpenAIEmbeddings,
            "huggingface_sentence-transformers/all-mpnet-base-v2": lambda: EmbeddingsWrapper("sentence-transformers/all-mpnet-base-v2"),
            "huggingface_sentence-transformers-all-mpnet-base-v2": lambda: EmbeddingsWrapper("sentence-transformers/all-mpnet-base-v2"),
            "huggingface_hkunlp/instructor-large": lambda: EmbeddingsWrapper("hkunlp/instructor-large"),
        }

        if embeddings_name in embeddings_factory:
            return embeddings_factory[embeddings_name](*args, **kwargs)
        else:
            return EmbeddingsWrapper(embeddings_name, *args, **kwargs)

class BaseVectorStore(ABC):
    def __init__(self):
        pass

    @abstractmethod
    def search(self, *args, **kwargs):
        pass

    def is_azure_configured(self):
        return settings.OPENAI_API_BASE and settings.OPENAI_API_VERSION and settings.AZURE_DEPLOYMENT_NAME

    def _get_embeddings(self, embeddings_name, embeddings_key=None):
        if embeddings_name == "openai_text-embedding-ada-002":
            if self.is_azure_configured():
                os.environ["OPENAI_API_TYPE"] = "azure"
                embedding_instance = EmbeddingsSingleton.get_instance(
                    embeddings_name,
                    model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME
                )
            else:
                embedding_instance = EmbeddingsSingleton.get_instance(
                    embeddings_name,
                    openai_api_key=embeddings_key
                )
        elif embeddings_name == "huggingface_sentence-transformers/all-mpnet-base-v2":
            if os.path.exists("./model/all-mpnet-base-v2"):
                embedding_instance = EmbeddingsSingleton.get_instance(
                    embeddings_name="./model/all-mpnet-base-v2",
                )
            else:
                embedding_instance = EmbeddingsSingleton.get_instance(
                    embeddings_name,
                )
        else:
            embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)

        return embedding_instance


<h1 align="center">
  DocsGPT  🦖
</h1>

<p align="center">
  <strong>Open-Source Documentation Assistant</strong>
</p>

<p align="left">
  <strong><a href="https://www.docsgpt.cloud/">DocsGPT</a></strong> is a cutting-edge open-source solution that streamlines the process of finding information in the project documentation. With its integration of the powerful <strong>GPT</strong> models, developers can easily ask questions about a project and receive accurate answers.
  
Say goodbye to time-consuming manual searches, and let <strong><a href="https://www.docsgpt.cloud/">DocsGPT</a></strong> help you quickly find the information you need. Try it out and see how it revolutionizes your project documentation experience. Contribute to its development and be a part of the future of AI-powered assistance.
</p>

<div align="center">
  
  <a href="https://github.com/arc53/DocsGPT">![link to main GitHub showing Stars number](https://img.shields.io/github/stars/arc53/docsgpt?style=social)</a>
  <a href="https://github.com/arc53/DocsGPT">![link to main GitHub showing Forks number](https://img.shields.io/github/forks/arc53/docsgpt?style=social)</a>
  <a href="https://github.com/arc53/DocsGPT/blob/main/LICENSE">![link to license file](https://img.shields.io/github/license/arc53/docsgpt)</a>
  <a href="https://discord.gg/n5BX8dh8rU">![link to discord](https://img.shields.io/discord/1070046503302877216)</a>
  <a href="https://twitter.com/docsgptai">![X (formerly Twitter) URL](https://img.shields.io/twitter/follow/docsgptai)</a>

 
</div>

### Production Support / Help for Companies:

We're eager to provide personalized assistance when deploying your DocsGPT to a live environment.

- [Get Enterprise / teams Demo :wave:](https://www.docsgpt.cloud/contact)
- [Send Email :email:](mailto:contact@arc53.com?subject=DocsGPT%20support%2Fsolutions)

![video-example-of-docs-gpt](https://d3dg1063dc54p9.cloudfront.net/videos/demov3.gif)

## Roadmap

You can find our roadmap [here](https://github.com/orgs/arc53/projects/2). Please don't hesitate to contribute or create issues, it helps us improve DocsGPT!

## Our Open-Source Models Optimized for DocsGPT:

| Name                                                                  | Base Model  | Requirements (or similar) |
| --------------------------------------------------------------------- | ----------- | ------------------------- |
| [Docsgpt-7b-mistral](https://huggingface.co/Arc53/docsgpt-7b-mistral)   | Mistral-7b   | 1xA10G gpu                |
| [Docsgpt-14b](https://huggingface.co/Arc53/docsgpt-14b)               | llama-2-14b | 2xA10 gpu's               |
| [Docsgpt-40b-falcon](https://huggingface.co/Arc53/docsgpt-40b-falcon) | falcon-40b  | 8xA10G gpu's              |

If you don't have enough resources to run it, you can use bitsnbytes to quantize.

## End to End AI Framework for Information Retrieval

![Architecture chart](https://github.com/user-attachments/assets/fc6a7841-ddfc-45e6-b5a0-d05fe648cbe2)

## Useful Links

- :mag: :fire: [Cloud Version](https://app.docsgpt.cloud/)

- :speech_balloon: :tada: [Join our Discord](https://discord.gg/n5BX8dh8rU)

- :books: :sunglasses: [Guides](https://docs.docsgpt.cloud/)

- :couple: [Interested in contributing?](https://github.com/arc53/DocsGPT/blob/main/CONTRIBUTING.md)

- :file_folder: :rocket: [How to use any other documentation](https://docs.docsgpt.cloud/Guides/How-to-train-on-other-documentation)

- :house: :closed_lock_with_key: [How to host it locally (so all data will stay on-premises)](https://docs.docsgpt.cloud/Guides/How-to-use-different-LLM)

## Project Structure

- Application - Flask app (main application).

- Extensions - Chrome extension.

- Scripts - Script that creates similarity search index for other libraries.

- Frontend - Frontend uses <a href="https://vitejs.dev/">Vite</a> and <a href="https://react.dev/">React</a>.

## QuickStart

> [!Note]
> Make sure you have [Docker](https://docs.docker.com/engine/install/) installed

On Mac OS or Linux, write:

`./setup.sh`

It will install all the dependencies and allow you to download the local model, use OpenAI or use our LLM API.

Otherwise, refer to this Guide for Windows:

1. Download and open this repository with `git clone https://github.com/arc53/DocsGPT.git`
2. Create a `.env` file in your root directory and set the env variables and `VITE_API_STREAMING` to true or false, depending on whether you want streaming answers or not.
   It should look like this inside:

   ```
   LLM_NAME=[docsgpt or openai or others] 
   VITE_API_STREAMING=true
   API_KEY=[if LLM_NAME is openai]
   ```

   See optional environment variables in the [/.env-template](https://github.com/arc53/DocsGPT/blob/main/.env-template) and [/application/.env_sample](https://github.com/arc53/DocsGPT/blob/main/application/.env_sample) files.

3. Run [./run-with-docker-compose.sh](https://github.com/arc53/DocsGPT/blob/main/run-with-docker-compose.sh).
4. Navigate to http://localhost:5173/.

To stop, just run `Ctrl + C`.

## Development Environments

### Spin up Mongo and Redis

For development, only two containers are used from [docker-compose.yaml](https://github.com/arc53/DocsGPT/blob/main/docker-compose.yaml) (by deleting all services except for Redis and Mongo).
See file [docker-compose-dev.yaml](./docker-compose-dev.yaml).

Run

```
docker compose -f docker-compose-dev.yaml build
docker compose -f docker-compose-dev.yaml up -d
```

### Run the Backend

> [!Note]
> Make sure you have Python 3.10 or 3.11 installed.

1. Export required environment variables or prepare a `.env` file in the project folder:
   - Copy [.env_sample](https://github.com/arc53/DocsGPT/blob/main/application/.env_sample) and create `.env`.

(check out [`application/core/settings.py`](application/core/settings.py) if you want to see more config options.)

2. (optional) Create a Python virtual environment:
   You can follow the [Python official documentation](https://docs.python.org/3/tutorial/venv.html) for virtual environments.

a) On Mac OS and Linux

```commandline
python -m venv venv
. venv/bin/activate
```

b) On Windows

```commandline
python -m venv venv
 venv/Scripts/activate
```

3. Download embedding model and save it in the `model/` folder:
You can use the script below, or download it manually from [here](https://d3dg1063dc54p9.cloudfront.net/models/embeddings/mpnet-base-v2.zip), unzip it and save it in the `model/` folder.

```commandline
wget https://d3dg1063dc54p9.cloudfront.net/models/embeddings/mpnet-base-v2.zip
unzip mpnet-base-v2.zip -d model
rm mpnet-base-v2.zip
```

4. Install dependencies for the backend:

```commandline
pip install -r application/requirements.txt
```

5. Run the app using `flask --app application/app.py run --host=0.0.0.0 --port=7091`.
6. Start worker with `celery -A application.app.celery worker -l INFO`.

### Start Frontend

> [!Note]
> Make sure you have Node version 16 or higher.

1. Navigate to the [/frontend](https://github.com/arc53/DocsGPT/tree/main/frontend) folder.
2. Install the required packages `husky` and `vite` (ignore if already installed).

```commandline
npm install husky -g
npm install vite -g
```

3. Install dependencies by running `npm install --include=dev`.
4. Run the app using `npm run dev`.

## Contributing

Please refer to the [CONTRIBUTING.md](CONTRIBUTING.md) file for information about how to get involved. We welcome issues, questions, and pull requests.

## Code Of Conduct

We as members, contributors, and leaders, pledge to make participation in our community a harassment-free experience for everyone, regardless of age, body size, visible or invisible disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, religion, or sexual identity and orientation. Please refer to the [CODE_OF_CONDUCT.md](CODE_OF_CONDUCT.md) file for more information about contributing.

## Many Thanks To Our Contributors⚡

<a href="https://github.com/arc53/DocsGPT/graphs/contributors" alt="View Contributors">
  <img src="https://contrib.rocks/image?repo=arc53/DocsGPT" alt="Contributors" />
</a>

## License

The source code license is [MIT](https://opensource.org/license/mit/), as described in the [LICENSE](LICENSE) file.

Built with [:bird: :link: LangChain](https://github.com/hwchase17/langchain)
