from fastapi import APIRouter, HTTPException, Query from functools import lru_cache from pydantic import BaseModel from app.config.constants import API_PREFIX from app.config.settings import get_settings from app.graph.neo4j_client import Neo4jClient from app.modeling.graph_schema import GraphSchema from app.query.query_router import QueryRoutingError from app.services.chat_logger import SupabaseChatLogger from app.services.query_service import QueryService from app.utils.logger import get_logger logger = get_logger("api.routes_query") router = APIRouter(prefix=f"{API_PREFIX}/query", tags=["query"]) # --------------------------------------------------------------------------- # Request / Response models # --------------------------------------------------------------------------- class AskRequest(BaseModel): question: str schema_json: dict | None = None session_id: str | None = None class AskResponse(BaseModel): answer: str explanation: str query_used: str nodes: list[dict] edges: list[dict] data: list[dict] metadata: dict class ChatLoggerStatusResponse(BaseModel): enabled: bool connected: bool table: str message: str class ChatHistoryItem(BaseModel): created_at: str session_id: str | None = None question: str answer: str explanation: str query_used: str metadata: dict node_count: int edge_count: int record_count: int status: str error: str | None = None class ChatHistoryResponse(BaseModel): session_id: str items: list[ChatHistoryItem] class ChatClearResponse(BaseModel): session_id: str deleted_count: int # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _build_query_service(schema: GraphSchema) -> tuple[QueryService, Neo4jClient]: """Wire up QueryService from settings. Returns (service, client) so caller can close client.""" settings = get_settings() client = Neo4jClient(settings.neo4j_uri, settings.neo4j_user, settings.neo4j_password) service = QueryService.create(settings, client, schema) return service, client @lru_cache(maxsize=1) def _get_chat_logger() -> SupabaseChatLogger: return SupabaseChatLogger(get_settings()) def _load_schema(schema_json: dict | None) -> GraphSchema: """Load schema from request body or from the saved config file.""" if schema_json: return GraphSchema.model_validate(schema_json) # Try loading from saved config files import json from pathlib import Path for filename in ("graph_schema_curated.json", "graph_schema.json"): path = Path(filename) if path.exists(): with open(path, "r", encoding="utf-8") as f: return GraphSchema.model_validate(json.load(f)) raise HTTPException( status_code=400, detail="No schema provided and no schema file found. " "Run /graph/preview + /graph/build first, or pass schema_json in the request.", ) def _is_rate_limited_error(exc: Exception) -> bool: text = str(exc).lower() return ( "rate limit" in text or "rate_limit" in text or "too many requests" in text or "429" in text ) # --------------------------------------------------------------------------- # Endpoints # --------------------------------------------------------------------------- @router.post("/ask", response_model=AskResponse) def ask_question(request: AskRequest): """Natural language query endpoint. Accepts a question, routes it through the full pipeline (guardrails -> intent extraction -> template -> Cypher -> execution -> formatting), and returns a structured response. """ schema = _load_schema(request.schema_json) service, client = _build_query_service(schema) chat_logger = _get_chat_logger() try: response = service.ask(request.question) response_metadata = dict(response.metadata) if request.session_id: response_metadata["session_id"] = request.session_id chat_logger.log( session_id=request.session_id, question=request.question, answer=response.answer, explanation=response.explanation, query_used=response.query_used, metadata=response_metadata, node_count=len(response.nodes), edge_count=len(response.edges), record_count=len(response.data), status="success", ) return AskResponse( answer=response.answer, explanation=response.explanation, query_used=response.query_used, nodes=[n.model_dump() for n in response.nodes], edges=[e.model_dump() for e in response.edges], data=response.data, metadata=response_metadata, ) except QueryRoutingError as e: logger.warning("Query routing failed at stage '%s': %s", e.stage, e) chat_logger.log( session_id=request.session_id, question=request.question, answer="", explanation="", query_used="", metadata={"stage": e.stage}, node_count=0, edge_count=0, record_count=0, status="routing_error", error=str(e), ) raise HTTPException(status_code=422, detail={ "message": str(e), "stage": e.stage, }) except Exception as e: logger.error("Query failed: %s", e) if _is_rate_limited_error(e): chat_logger.log( session_id=request.session_id, question=request.question, answer="", explanation="", query_used="", metadata={}, node_count=0, edge_count=0, record_count=0, status="rate_limited", error=str(e), ) raise HTTPException( status_code=429, detail=( "LLM rate limit exceeded. Please retry in a few minutes, " "or switch to a higher-quota model/provider." ), ) chat_logger.log( session_id=request.session_id, question=request.question, answer="", explanation="", query_used="", metadata={}, node_count=0, edge_count=0, record_count=0, status="error", error=str(e), ) raise HTTPException(status_code=500, detail=str(e)) finally: client.close() @router.get("/chat/status", response_model=ChatLoggerStatusResponse) def chat_status(): logger_state = _get_chat_logger().health() return ChatLoggerStatusResponse(**logger_state) @router.get("/chat/history", response_model=ChatHistoryResponse) def chat_history(session_id: str = Query(...), limit: int = Query(200, ge=1, le=500)): items = _get_chat_logger().list_session(session_id=session_id, limit=limit) return ChatHistoryResponse(session_id=session_id, items=[ChatHistoryItem(**item) for item in items]) @router.delete("/chat/history", response_model=ChatClearResponse) def clear_chat_history(session_id: str = Query(...)): deleted_count = _get_chat_logger().clear_session(session_id=session_id) return ChatClearResponse(session_id=session_id, deleted_count=deleted_count)