query.py

from fastapi import FastAPI, Request
from fastapi.responses import RedirectResponse
import requests
from pydantic import BaseModel
from tools_config import select_tools, run_tools, get_fundamentals, get_performance, get_recent_dividends_rate
from rag_config import vectorstore_config, weaviate_client_config, close_weaviate_client, retrievers
from contextlib import asynccontextmanager
from chain import final_chain

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.weaviate_client = weaviate_client_config()
    app.state.retrievers = retrievers(app.state.weaviate_client)
    yield
    close_weaviate_client(app.state.weaviate_client)

app = FastAPI(title="Agent_LLM_API", lifespan=lifespan)

class Query(BaseModel):
    question: str
    ticker: str

@app.get("/")
def root():
    return RedirectResponse(url="/docs")

@app.post("/query")
def query(query: Query, request: Request):
    response = final_chain((request.app.state.retrievers)).invoke({
        "question": query.question,
        "ticker": query.ticker
    })
    return response