71 lines
1.8 KiB
Python
71 lines
1.8 KiB
Python
from typing import Optional
|
||
|
||
import uvicorn
|
||
from fastapi import FastAPI, HTTPException
|
||
from pydantic import BaseModel
|
||
|
||
from gbase_agent import init_agent_service
|
||
|
||
app = FastAPI(title="Database Assistant API", version="1.0.0")
|
||
|
||
# Initialize agent globally at startup
|
||
bot = init_agent_service()
|
||
|
||
|
||
class QueryRequest(BaseModel):
|
||
question: str
|
||
file_url: Optional[str] = None
|
||
|
||
|
||
class QueryResponse(BaseModel):
|
||
answer: str
|
||
|
||
|
||
@app.post("/query", response_model=QueryResponse)
|
||
async def query_database(request: QueryRequest):
|
||
"""
|
||
Process a database query using the assistant agent.
|
||
|
||
Args:
|
||
request: QueryRequest containing the query and optional file URL
|
||
|
||
Returns:
|
||
QueryResponse containing the assistant's response
|
||
"""
|
||
try:
|
||
messages = []
|
||
|
||
if request.file_url:
|
||
messages.append(
|
||
{
|
||
"role": "user",
|
||
"content": [{"text":"使用sqlite数据库，用日语回答下面问题："+request.question}, {"file": request.file_url}],
|
||
}
|
||
)
|
||
else:
|
||
messages.append({"role": "user", "content": request.question})
|
||
|
||
responses = []
|
||
for response in bot.run(messages):
|
||
responses.append(response)
|
||
|
||
if responses:
|
||
final_response = responses[-1][-1]
|
||
return QueryResponse(answer=final_response["content"])
|
||
else:
|
||
raise HTTPException(status_code=500, detail="No response from agent")
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
|
||
|
||
|
||
@app.get("/")
|
||
async def root():
|
||
"""Health check endpoint"""
|
||
return {"message": "Database Assistant API is running"}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||
|