Axon/axon/models.py

from typing import Literal
from pydantic import BaseModel, Field


class SearchRequest(BaseModel):
    query: str
    max_results: int = Field(default=10, ge=1, le=20)
    search_depth: Literal["basic", "advanced"] = "basic"
    include_answer: bool = False
    include_raw_content: bool = False
    domains_whitelist: list[str] = []
    domains_blacklist: list[str] = []
    topic: Literal["general", "news"] = "general"
    time_range: Literal["day", "week", "month", "year"] | None = None
    include_private: bool = True


class SearchResult(BaseModel):
    title: str
    url: str
    content: str
    score: float = 0.5
    published_date: str | None = None
    favicon: str | None = None
    source_quality: float | None = None


class SearchResponse(BaseModel):
    query: str
    answer: str | None = None
    results: list[SearchResult]
    response_time: float
    credits_used: int = 1


class ExtractRequest(BaseModel):
    urls: list[str]
    include_images: bool = False


class ExtractResult(BaseModel):
    url: str
    title: str | None = None
    content: str = ""
    success: bool
    error: str | None = None


class ExtractResponse(BaseModel):
    results: list[ExtractResult]


class CrawlRequest(BaseModel):
    url: str
    max_depth: int = Field(default=1, ge=0, le=3)
    max_pages: int = Field(default=20, ge=1, le=50)
    include_raw_content: bool = False
    instructions: str | None = None


class CrawlResult(BaseModel):
    url: str
    title: str | None = None
    content: str = ""
    links: list[str] = []
    depth: int = 0
    success: bool
    error: str | None = None


class CrawlResponse(BaseModel):
    start_url: str
    results: list[CrawlResult]
    response_time: float


class MapRequest(BaseModel):
    url: str
    max_depth: int = Field(default=1, ge=0, le=3)
    max_pages: int = Field(default=50, ge=1, le=200)
    instructions: str | None = None


class MapResponse(BaseModel):
    start_url: str
    urls: list[str]
    response_time: float


class ResearchRequest(BaseModel):
    query: str
    depth: Literal["basic", "deep"] = "basic"
    output_format: Literal["markdown", "json"] = "markdown"
    search_depth: Literal["basic", "advanced"] = "advanced"
    max_results_per_query: int = Field(default=5, ge=1, le=10)


class ResearchResponse(BaseModel):
    query: str
    report: str
    sources: list[SearchResult]
    sub_queries: list[str]
    response_time: float


class MetricsResponse(BaseModel):
    requests_total: int
    requests_by_path: dict[str, int]
    errors_total: int
    average_latency_ms: float


class CitationValidationRequest(BaseModel):
    report: str
    sources: list[SearchResult]


class CitationValidationResponse(BaseModel):
    citation_count: int
    missing_citations: list[int]
    valid: bool


class IngestDocument(BaseModel):
    id: str | None = None
    title: str
    content: str
    url: str | None = None
    metadata: dict[str, str] = {}


class IngestRequest(BaseModel):
    documents: list[IngestDocument]


class IngestResponse(BaseModel):
    indexed: int


class IndexSearchRequest(BaseModel):
    query: str
    max_results: int = Field(default=10, ge=1, le=50)


class IndexDocument(BaseModel):
    id: str
    title: str
    content: str
    url: str | None = None
    score: float


class IndexSearchResponse(BaseModel):
    query: str
    results: list[IndexDocument]


class JobCreateResponse(BaseModel):
    job_id: str
    status: Literal["queued", "running", "succeeded", "failed"]


class JobResponse(BaseModel):
    job_id: str
    kind: str
    status: Literal["queued", "running", "succeeded", "failed"]
    result: dict | None = None
    error: str | None = None


class TavilySearchRequest(BaseModel):
    query: str
    search_depth: Literal["basic", "advanced"] = "basic"
    max_results: int = Field(default=10, ge=1, le=20)
    include_answer: bool | str = False
    include_raw_content: bool = False
    include_domains: list[str] = []
    exclude_domains: list[str] = []
    topic: Literal["general", "news"] = "general"
    time_range: Literal["day", "week", "month", "year"] | None = None


class TavilyExtractRequest(BaseModel):
    urls: str | list[str]
    include_images: bool = False