TetrisGGBOBot
8e986a7bb4
Use /search/stream as the single entry for both fast and deep modes via request params, suppress COT events in fast mode while keeping deep COT, and return 410 for deprecated /search. Made-with: Cursor
102 lines
3.7 KiB
Python
102 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from .agent import AISearchAgent
|
|
from .audit import AuditLogger
|
|
from .config import Settings
|
|
from .models import SearchRequest, SearchResponse
|
|
|
|
|
|
def create_app(enable_streaming: bool = True) -> FastAPI:
|
|
settings = Settings.from_env()
|
|
agent = AISearchAgent(settings)
|
|
audit_logger = AuditLogger(settings.audit_log_path)
|
|
app = FastAPI(title="AI Search Agent", version="1.0.0")
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, str]:
|
|
return {"status": "ok"}
|
|
|
|
@app.post("/search")
|
|
async def search(request: SearchRequest) -> JSONResponse:
|
|
return JSONResponse(
|
|
status_code=410,
|
|
content={
|
|
"error": "SEARCH_ENDPOINT_REMOVED",
|
|
"message": "Use /search/stream with search_mode=fast|deep.",
|
|
"query": request.query,
|
|
"search_mode": agent.resolve_search_mode(request),
|
|
},
|
|
)
|
|
|
|
@app.post("/search/stream")
|
|
async def search_stream(request: SearchRequest):
|
|
if not enable_streaming:
|
|
return JSONResponse(
|
|
status_code=501,
|
|
content={
|
|
"error": "STREAMING_NOT_AVAILABLE_IN_FUNCTIONS",
|
|
"message": "Use the always-on stream service for /search/stream.",
|
|
"stream_endpoint": None,
|
|
"query": request.query,
|
|
"search_mode": agent.resolve_search_mode(request),
|
|
},
|
|
)
|
|
|
|
async def event_generator():
|
|
# Send an immediate event so clients/proxies flush the SSE channel early.
|
|
yield "event: connected\ndata: {}\n\n"
|
|
response_payload = None
|
|
mode = agent.resolve_search_mode(request)
|
|
timeout_seconds = agent.get_time_budget_seconds(request)
|
|
try:
|
|
async with asyncio.timeout(timeout_seconds):
|
|
async for event in agent.run_stream(request):
|
|
event_type = event.get("type", "message")
|
|
payload = event.get("data", {})
|
|
if event_type == "result":
|
|
response_payload = payload
|
|
|
|
yield (
|
|
f"event: {event_type}\n"
|
|
f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
|
)
|
|
except asyncio.TimeoutError:
|
|
timeout_payload = {
|
|
"query": request.query,
|
|
"search_mode": mode,
|
|
"error": f"TIMEOUT_{mode.upper()}_SEARCH",
|
|
"message": f"{mode} search exceeded {timeout_seconds}s time budget",
|
|
}
|
|
yield (
|
|
"event: timeout\n"
|
|
f"data: {json.dumps(timeout_payload, ensure_ascii=False)}\n\n"
|
|
)
|
|
return
|
|
|
|
if response_payload is not None:
|
|
response_obj = SearchResponse.model_validate(response_payload)
|
|
audit_id = audit_logger.write(request, response_obj)
|
|
yield (
|
|
"event: audit\n"
|
|
f"data: {json.dumps({'audit_id': audit_id}, ensure_ascii=False)}\n\n"
|
|
)
|
|
|
|
return StreamingResponse(
|
|
event_generator(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|
|
|
|
return app
|