src/main.py

from datetime import datetime
import logging,os
import fastapi
from fastapi import Body, Depends
import uvicorn
from fastapi import HTTPException , status
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI as Response
from sse_starlette.sse import EventSourceResponse
from starlette.responses import StreamingResponse
from starlette.requests import Request
from pydantic import BaseModel, Extra
from enum import Enum
from typing import List, Dict, Any, Generator, Optional, cast, Callable
from chromaIntf import ChromaIntf
import baseInfra.dropbox_handler as dbh
import traceback

chromaIntf=ChromaIntf()

class PathRequest(BaseModel):
    dir: str = "/"

class MetaD(BaseModel):
    timestamp: Optional[str]= datetime.now().isoformat()
    class Config:
        allow_population_by_field_name = True
        extra = Extra.allow

class DocWithMeta(BaseModel):
    text: str = ""
    metadata: Optional[MetaD] = MetaD()

async def catch_exceptions_middleware(
    request: Request, call_next: Callable[[Request], Any]
) -> Response:
    try:
        #print("In exception cater middleware")
        #print(request.headers)
        #print(await request.body())
        return await call_next(request)
    except Exception as e:
        print(repr(e))
        print("from the catch exception")
        traceback.print_exc()
        return JSONResponse(content={"error": repr(e)}, status_code=500)



app = fastapi.FastAPI(title="Maya Persistence")
app.middleware("http")(catch_exceptions_middleware)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
api_base="/api/v1"

@app.post(api_base+"/getMatchingDocs")
async def get_matching_docs(inStr: str ) -> Any:
    """
    Gets the query embeddings and uses metadata appropriately and gets the matching docs for query
    TODO: Add parameter for type of query and number of docs to return
    TODO: Add parameter to return the source information as well
    """
    return await chromaIntf.getRelevantDocs(inStr)

@app.post(api_base+"/addTextDocument")
async def add_text_document(inDoc: DocWithMeta ) -> Any:
    """
     Add text and metadata to the db
    """
    print("Received request for")
    print(inDoc) 
    return await chromaIntf.addText(inDoc.text,inDoc.metadata)

@app.get(api_base+"/persist")
async def persist_db():
    return await chromaIntf.persist()

@app.get(api_base+"/reset")
async def reset_db():
    return await dbh.restoreFolder("db")

@app.post(api_base+"/walk")
def walk(path: PathRequest):
    print("Received walk request")
    dirs=[]
    fileList=[]
    try:
        for root, items, files in os.walk(path.dir,topdown=True):
            for item in items:
                dirs.append(item)
            for filea in files:
                fileList.append(filea)
    except Exception:
        print("got exception",Exception)
    response= JSONResponse(content= {"dirs":dirs,"files":fileList})
    return response

@app.get(api_base+"/list")
async def list_docs():
    return await chromaIntf.listDocs()

print(__name__)

if __name__ == '__main__' or __name__ == "src.main":
    uvicorn.run(app, host="0.0.0.0", port=8000)