app.py

import fastapi as api
from typing import Annotated
from fastapi.security import OAuth2PasswordBearer, OAuth2AuthorizationCodeBearer, OAuth2PasswordRequestForm
from model.document import Document, PlainTextDocument, JsonDocument
import sys
from model.user import User
from fastapi import FastAPI, File, UploadFile
from di import initialize_di_for_app
import gradio as gr
import os
import json
SETTINGS, STORAGE, EMBEDDING, INDEX = initialize_di_for_app()
user_json_str = STORAGE.load('user.json')
USER = User.parse_raw(user_json_str)

oauth2_scheme  = OAuth2PasswordBearer(tokenUrl="/api/v1/auth/token")
app = api.FastAPI()
app.openapi_version = "3.0.0"
users = [USER]
async def get_current_user(token: str = api.Depends(oauth2_scheme)):
    '''
    Get current user
    '''
    for user in users:
        if user.user_name == token:
            return user
    
    raise api.HTTPException(status_code=401, detail="Invalid authentication credentials")


@app.post("/api/v1/auth/token")
async def login(form_data: Annotated[OAuth2PasswordRequestForm, api.Depends()]):
    '''
    Login to get a token
    '''
    return {"access_token": form_data.username}

@app.post("/api/v1/uploadfile/", include_in_schema=False)
def create_upload_file(file: UploadFile = api.File(...)) -> Document:
    '''
    Upload a file
    '''
    fileUrl = f'{USER.user_name}-{file.filename}'
    STORAGE.save(fileUrl, file.read())

    # create plainTextDocument if the file is a text file
    if file.filename.endswith('.txt'):
        return PlainTextDocument(
            name=file.filename,
            status='uploading',
            url=fileUrl,
            embedding=EMBEDDING,
            storage=STORAGE,
        )
    else:
        raise api.HTTPException(status_code=400, detail="File type not supported")
        

### /api/v1/.well-known
#### Get /openapi.json
# Get the openapi json file
@app.get("/api/v1/.well-known/openapi.json")
async def get_openapi():
    '''
    otherwise return 401
    '''

    # get a list of document names + description
    document_list = [[doc.name, doc.description] for doc in USER.documents]

    # get openapi json from api
    openapi = app.openapi().copy()

    openapi['info']['title'] = 'DocumentSearch'
    description = f'''Search documents with a query.
    ## Documents
    {document_list}
    '''

    openapi['info']['description'] = description

    # update description in /api/v1/search
    openapi['paths']['/api/v1/search']['get']['description'] += f'''
Available documents:
{document_list}
'''

    # filter out unnecessary endpoints
    openapi['paths'] = {
        '/api/v1/search': openapi['paths']['/api/v1/search'],
    }

    # remove components
    openapi['components'] = {}

    # return the openapi json
    return openapi



### /api/v1/document
#### Get /list
# Get the list of documents
@app.get("/api/v1/document/list")
# async def get_document_list(user: Annotated[User, api.Depends(get_current_user)]) -> list[Document]:
async def get_document_list() -> list[Document]:
    '''
    Get the list of documents
    '''
    return USER.documents

#### Post /upload
# Upload a document
@app.post("/api/v1/document/upload")
# def upload_document(user: Annotated[User, api.Depends(get_current_user)], document: Annotated[Document, api.Depends(create_upload_file)]):
def upload_document(document: Annotated[Document, api.Depends(create_upload_file)]):
    '''
    Upload a document
    '''
    document.status = 'processing'
    INDEX.load_or_update_document(user, document, progress)
    document.status = 'done'
    USER.documents.append(document)

#### Get /delete
# Delete a document
@app.get("/api/v1/document/delete")
# async def delete_document(user: Annotated[User, api.Depends(get_current_user)], document_name: str):
async def delete_document(document_name: str):
    '''
    Delete a document
    '''
    for doc in USER.documents:
        if doc.name == document_name:
            STORAGE.delete(doc.url)
            INDEX.remove_document(USER, doc)
            USER.documents.remove(doc)
            return
    
    raise api.HTTPException(status_code=404, detail="Document not found")

# Query the index
@app.get("/api/v1/search", operation_id=None,)
def search(
    # user: Annotated[User, api.Depends(get_current_user)],
    query: str,
    document_name: str = None,
    top_k: int = 10,
    threshold: float = 0.5):
    '''
    Search documents with a query. It will return [top_k] results with a score higher than [threshold].
    query: the query string, required
    document_name: the document name, optional. You can provide this parameter to search in a specific document.
    top_k: the number of results to return, optional. Default to 10.
    threshold: the threshold of the results, optional. Default to 0.5.
    '''
    if document_name:
        for doc in USER.documents:
            if doc.name == document_name:
                return INDEX.query_document(USER, doc, query, top_k, threshold)

        raise api.HTTPException(status_code=404, detail="Document not found")
    else:
        return INDEX.query_index(USER, query, top_k, threshold)

def receive_signal(signalNumber, frame):
    print('Received:', signalNumber)
    sys.exit()


@app.on_event("startup")
async def startup_event():
    import signal
    signal.signal(signal.SIGINT, receive_signal)
    # startup tasks

@app.on_event("shutdown")
def exit_event():
    # save USER
    STORAGE.save('user.json', USER.model_dump_json())
    print('exit')

user = USER

def gradio_upload_document(file: File):
    file_temp_path = file.name
    # load file
    file_name = os.path.basename(file_temp_path)
    fileUrl = f'{USER.user_name}-{file_name}'
    with open(file_temp_path, 'r', encoding='utf-8') as f:
        STORAGE.save(fileUrl, f.read())

    # create plainTextDocument if the file is a text file
    doc = None
    if file_name.endswith('.txt'):
        doc = PlainTextDocument(
            name=file_name,
            status='uploading',
            url=fileUrl,
            embedding=EMBEDDING,
            storage=STORAGE,
        )
    elif file_name.endswith('.json'):
        doc = JsonDocument(
            name=file_name,
            status='uploading',
            url=fileUrl,
            embedding=EMBEDDING,
            storage=STORAGE,
        )
    else:
        raise api.HTTPException(status_code=400, detail="File type not supported")
    doc.status = 'processing'
    INDEX.load_or_update_document(user, doc)
    doc.status = 'done'
    USER.documents.append(doc)

    return f'uploaded {file_name}'

def gradio_query(query: str, document_name: str = None, top_k: int = 10, threshold: float = 0.5):
    res_or_exception = search(query, document_name, top_k, threshold)
    if isinstance(res_or_exception, api.HTTPException):
        raise res_or_exception
    
    # convert to json string

    records = [record.model_dump(mode='json') for record in res_or_exception]

    return json.dumps(records, indent=4)

with gr.Blocks() as ui:
    gr.Markdown("#llm-memory")

    with gr.Column():
        gr.Markdown(
        """
        ## LLM Memory
        """)        
        with gr.Row():
            user_name = gr.Label(label="User name", value=USER.user_name)

            # url to .well-known/openapi.json
            gr.Label(label=".wellknown/openapi.json", value=f"/api/v1/.well-known/openapi.json")

        # with gr.Tab("avaiable documents"):
        #     available_documents = gr.Label(label="avaiable documents", value="avaiable documents")
        #     refresh_btn = gr.Button(label="refresh", type="button")
        #     refresh_btn.click(lambda: '\r\n'.join([doc.name for doc in USER.documents]), None, available_documents)
        #     documents = USER.documents
        #     for document in documents:
        #         gr.Label(label=document.name, value=document.name)
        # with gr.Tab("upload document"):
        #     with gr.Tab("upload .txt document"):
        #         file = gr.File(label="upload document", type="file", file_types=[".txt"])
        #         output = gr.Label(label="output", value="output")
        #         upload_btn = gr.Button("upload document", type="button")
        #         upload_btn.click(gradio_upload_document, file, output)
        #     with gr.Tab("upload .json document"):
        #         gr.Markdown(
        #         """
        #         The json document should be a list of objects, each object should have a `content` field. If you want to add more fields, you can add them in the `meta_data` field.
        #         For example:
        #         ```json
        #         [
        #             {
        #                 "content": "hello world",
        #                 "meta_data": {
        #                     "title": "hello world",
        #                     "author": "llm-memory"
        #                 }
        #             },
        #             {
        #                 "content": "hello world"
        #                 "meta_data": {
        #                     "title": "hello world",
        #                     "author": "llm-memory"
        #                 }
        #             }
        #         ]
        #         ```

        #         ## Note
        #         - The `meta_data` should be a dict which both keys and values are strings.
        #         """)
        #         file = gr.File(label="upload document", type="file", file_types=[".json"])
        #         output = gr.Label(label="output", value="output")
        #         upload_btn = gr.Button("upload document", type="button")
        #         upload_btn.click(gradio_upload_document, file, output)
        with gr.Tab("search"):
            query = gr.Textbox(label="search", placeholder="Query")
            document = gr.Dropdown(label="document", choices=[None] + [doc.name for doc in USER.documents], placeholder="document, optional")
            top_k = gr.Number(label="top_k", placeholder="top_k, optional", value=10)
            threshold = gr.Number(label="threshold", placeholder="threshold, optional", value=0.5)
            output = gr.Code(label="output", language="json", value="output")
            query_btn = gr.Button("Query")
            query_btn.click(gradio_query, [query, document, top_k, threshold], output, api_name="search")

            
gradio_app = gr.routes.App.create_app(ui)
app.mount("/", gradio_app)

ui.launch()