|
|
from langchain_core.runnables import Runnable |
|
|
from langchain_core.callbacks import BaseCallbackHandler |
|
|
from fastapi import FastAPI, Request, Depends |
|
|
from sse_starlette.sse import EventSourceResponse |
|
|
from sqlalchemy.orm import Session |
|
|
from langserve.serialization import WellKnownLCSerializer |
|
|
from typing import Any, List |
|
|
import crud |
|
|
import models |
|
|
import schemas |
|
|
from database import SessionLocal, engine |
|
|
from chains import simple_chain, formatted_chain, history_chain, rag_chain |
|
|
from prompts import format_chat_history |
|
|
from callbacks import LogResponseCallback |
|
|
import orjson |
|
|
|
|
|
models.Base.metadata.create_all(bind=engine) |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
def get_db(): |
|
|
db = SessionLocal() |
|
|
try: |
|
|
yield db |
|
|
finally: |
|
|
db.close() |
|
|
|
|
|
|
|
|
async def generate_stream(input_data: schemas.BaseModel, callback): |
|
|
stream = rag_chain.stream( |
|
|
input_data.dict(), |
|
|
config={"callbacks": [callback]} |
|
|
) |
|
|
for output in stream: |
|
|
data = orjson.dumps(output).decode("utf-8") |
|
|
yield {'data': data, "event": "data"} |
|
|
yield {"event": "end"} |
|
|
|
|
|
|
|
|
@app.get("/") |
|
|
def greet_json(): |
|
|
return {"Hello": "World!"} |
|
|
|
|
|
|
|
|
@app.post("/simple/stream") |
|
|
async def simple_stream(request: Request): |
|
|
data = await request.json() |
|
|
user_question = schemas.UserQuestion(**data['input']) |
|
|
return EventSourceResponse(generate_stream(user_question, simple_chain)) |
|
|
|
|
|
|
|
|
@app.post("/formatted/stream") |
|
|
async def formatted_stream(request: Request): |
|
|
data = await request.json() |
|
|
user_question = schemas.UserQuestion(**data['input']) |
|
|
return EventSourceResponse(generate_stream(user_question, formatted_chain)) |
|
|
|
|
|
|
|
|
@app.post("/history/stream") |
|
|
async def history_stream(request: Request, db: Session = Depends(get_db)): |
|
|
data = await request.json() |
|
|
user_request = schemas.UserRequest(**data['input']) |
|
|
chat_history = crud.get_user_chat_history(db=db, username=user_request.username) |
|
|
message = schemas.MessageBase(message=user_request.question, type='User') |
|
|
crud.add_message(db, message, user_request.username) |
|
|
|
|
|
history_input = schemas.HistoryInput( |
|
|
question=user_request.question, |
|
|
chat_history=format_chat_history(chat_history) |
|
|
) |
|
|
|
|
|
return EventSourceResponse(generate_stream( |
|
|
history_input, |
|
|
history_chain, |
|
|
[LogResponseCallback(user_request, db)] |
|
|
)) |
|
|
|
|
|
|
|
|
@app.post("/rag/stream") |
|
|
async def rag_stream(request: Request, db: Session = Depends(get_db)): |
|
|
data = await request.json() |
|
|
user_request = schemas.UserRequest(**data['input']) |
|
|
chat_history = crud.get_user_chat_history(db=db, username=user_request.username) |
|
|
message = schemas.MessageBase(message=user_request.question, type='User') |
|
|
crud.add_message(db, message, user_request.username) |
|
|
|
|
|
rag_input = schemas.RagInput( |
|
|
question=user_request.question, |
|
|
chat_history=format_chat_history(chat_history), |
|
|
) |
|
|
|
|
|
callback = LogResponseCallback(user_request, db) |
|
|
return EventSourceResponse(generate_stream(rag_input, callback)) |
|
|
|
|
|
@app.post("/filtered_rag/stream") |
|
|
async def filtered_rag_stream(request: Request, db: Session = Depends(get_db)): |
|
|
data = await request.json() |
|
|
print(data) |
|
|
user_request = schemas.UserRequest(**data['input']) |
|
|
chat_history = crud.get_user_chat_history(db=db, username=user_request.username) |
|
|
message = schemas.MessageBase(message=user_request.question, type='User') |
|
|
crud.add_message(db, message, user_request.username) |
|
|
|
|
|
rag_input = schemas.RagInput( |
|
|
question=user_request.question, |
|
|
chat_history=format_chat_history(chat_history), |
|
|
hybrid_search=True |
|
|
) |
|
|
|
|
|
return EventSourceResponse(generate_stream( |
|
|
rag_input, |
|
|
rag_chain, |
|
|
[LogResponseCallback(user_request, db)] |
|
|
)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
import uvicorn |
|
|
uvicorn.run("main:app", host="localhost", reload=True, port=8002) |
