Agentic_Rag/app/schemas.py

from __future__ import annotations

from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field


class RouteName(StrEnum):
    DIRECT_ANSWER = "direct_answer"
    MILVUS_SEARCH = "milvus_search"
    SQL_QUERY = "sql_query"
    WEB_SEARCH = "web_search"
    MULTI_SOURCE = "multi_source"
    CLARIFY = "clarify"
    REFUSE = "refuse"


class RouteDecision(BaseModel):
    needs_retrieval: bool
    intent: str
    routes: list[RouteName] = Field(default_factory=list)
    requires_decomposition: bool = False
    confidence: str = "medium"
    reason_code: str
    filters: dict[str, Any] = Field(default_factory=dict)
    max_rounds: int = 2
    fallback: str = "clarify"


class PlanStep(BaseModel):
    id: str
    tool: RouteName
    query: str
    arguments: dict[str, Any] = Field(default_factory=dict)
    depends_on: list[str] = Field(default_factory=list)


class RetrievalPlan(BaseModel):
    goal: str
    steps: list[PlanStep]


class Evidence(BaseModel):
    source_type: str
    source: str
    content: str
    score: float | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class ToolResult(BaseModel):
    status: str
    tool: str
    data: list[dict[str, Any]] = Field(default_factory=list)
    evidence: list[Evidence] = Field(default_factory=list)
    latency_ms: int = 0
    retryable: bool = False
    error_code: str | None = None
    error_message: str | None = None


class QualityGrade(BaseModel):
    relevant: bool
    sufficient: bool
    missing_aspects: list[str] = Field(default_factory=list)
    recommended_action: str
    reason: str


class QueryRequest(BaseModel):
    query: str = Field(min_length=1, max_length=4000)
    session_id: str = "default"
    debug: bool = False


class QueryResponse(BaseModel):
    answer: str
    citations: list[Evidence]
    route: RouteDecision
    executed_queries: list[dict[str, Any]] = Field(default_factory=list)
    trace: list[dict[str, Any]] = Field(default_factory=list)
    termination_reason: str