Spaces:
Sleeping
Sleeping
File size: 1,978 Bytes
ced6b34 e5ccef8 ced6b34 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from orator import Session, SQLDatabase, DocumentDatabase
from langchain.chat_models import init_chat_model
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import asyncio
# Initialize FastAPI app
app = FastAPI(title="Orator Chat API")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Use a specific origin in production
allow_credentials=True,
allow_methods=["*"], # Allow all methods (GET, POST, etc.)
allow_headers=["*"], # Allow all headers
)
# Initialize LLM and databases
llm = init_chat_model("o3-mini", model_provider="openai")
chinook_db = SQLDatabase.from_uri("sqlite:///data/Chinook.db")
pricegram_db = DocumentDatabase("data/pricegram.json", top_k=10)
# Initialize session
session = Session(llm=llm, datasources=[chinook_db, pricegram_db])
# Pydantic model for request
class QueryRequest(BaseModel):
query: str
source: int
@app.post("/query/")
async def get_response(request: QueryRequest):
"""Process a query and return the response."""
try:
print("Got Request:", request)
response, logs = session.invoke(request.query, datasource=request.source)
response = {"response": response}
print("Sending Respose:", response)
return response
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query/stream/")
async def stream_response(request: QueryRequest):
"""Stream responses for a given query."""
async def event_generator():
try:
events = session.stream(request.query)
for event in events:
for person, quote in event.items():
yield f"{person}: {quote['messages'][-1].text}\n"
await asyncio.sleep(0.1) # Simulate streaming delay
except Exception as e:
yield f"Error: {str(e)}"
return StreamingResponse(event_generator(), media_type="text/plain")
@app.get("/")
async def root():
return {"message": "Welcome to the Orator Chat API"}
|