anubhav77's picture
v0.8.2
5d4719e
raw
history blame
2.66 kB
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
from enum import Enum
from typing import List, Dict, Any, Generator, Optional, cast, Callable
from chroma_intf import *
import baseInfra.dropbox_handler as dbh
class PathRequest(BaseModel):
dir: str = "/"
class DocWithMeta(BaseModel):
text: str = ""
metadata: Optional[dict] = {}
async def catch_exceptions_middleware(
request: Request, call_next: Callable[[Request], Any]
) -> Response:
try:
return await call_next(request)
except Exception as e:
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 getRelevantDocs(inStr)
@app.post(api_base+"/addTextDocument")
async def add_text_document(inDoc: DocWithMeta ) -> Any:
"""
Add text and metadata to the db
"""
return addText(inDoc.text,inDoc.metadata)
@app.get(api_base+"/persist")
async def persist_db():
persist()
return await dbh.backupFolder("db")
@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
print(__name__)
if __name__ == '__main__' or __name__ == "src.main":
uvicorn.run(app, host="0.0.0.0", port=8000)