""" Pydantic models for Grant Analyst. Strict schemas for grants, search requests, responses, and QA interactions. All core entities use validated Pydantic models instead of loose dicts. """ from __future__ import annotations from typing import List, Optional, Dict, Any from datetime import date, datetime from pydantic import BaseModel, Field, HttpUrl, field_validator from enum import Enum class GrantStatus(str, Enum): """Grant/competition status.""" OPEN = "open" UPCOMING = "upcoming" CLOSED = "closed" UNKNOWN = "unknown" class FundingRates(BaseModel): """Funding rates by organization size.""" micro_small: Optional[int] = Field(None, description="Funding % for micro/small businesses") medium: Optional[int] = Field(None, description="Funding % for medium businesses") large: Optional[int] = Field(None, description="Funding % for large businesses") class FundingInfo(BaseModel): """Grant funding information.""" min: Optional[float] = Field(None, description="Minimum funding amount") max: Optional[float] = Field(None, description="Maximum funding amount") total_pot: Optional[float] = Field(None, description="Total competition pot") rates: Optional[FundingRates] = Field(None, description="Funding rates by org size") class GrantUrls(BaseModel): """URLs associated with a grant.""" overview: Optional[str] = Field(None, description="Overview page URL") apply: Optional[str] = Field(None, description="Application page URL") guidance: Optional[str] = Field(None, description="Guidance page URL") other: List[str] = Field(default_factory=list, description="Other related URLs") class Grant(BaseModel): """ Validated grant/competition model. Core fields extracted from Innovate UK competition snapshots. """ # Identifiers id: str = Field(..., description="Internal grant/competition ID") source: str = Field(default="innovate_uk", description="Data source (e.g., 'innovate_uk')") # Core content title: str = Field(..., description="Grant title") programme: Optional[str] = Field(None, description="Programme/scheme name") # Descriptive content (raw HTML/text from snapshots) summary: Optional[str] = Field(None, description="Summary/description") scope: Optional[str] = Field(None, description="Scope and eligibility criteria") eligibility: Optional[str] = Field(None, description="Detailed eligibility requirements") # Dates open_date: Optional[date] = Field(None, description="Competition opening date") close_date: Optional[datetime] = Field(None, description="Application deadline (may include time)") notify_date: Optional[date] = Field(None, description="Notification date") project_start_from: Optional[date] = Field(None, description="Earliest project start date") # Funding funding: Optional[FundingInfo] = Field(None, description="Funding information") # Status (derived from dates) status: GrantStatus = Field(default=GrantStatus.UNKNOWN, description="Current status") # URLs url: Optional[str] = Field(None, description="Primary URL") urls: Optional[GrantUrls] = Field(None, description="Structured URLs") # Metadata created_at: Optional[datetime] = Field(None, description="Record creation timestamp") updated_at: Optional[datetime] = Field(None, description="Last update timestamp") # Raw sections (for full-text search and citation) sections: Optional[Dict[str, Any]] = Field(None, description="Raw HTML sections from source") @field_validator('status', mode='before') @classmethod def derive_status(cls, v: Any, info) -> GrantStatus: """Derive status from dates if not explicitly set.""" if isinstance(v, GrantStatus): return v if isinstance(v, str): try: return GrantStatus(v.lower()) except ValueError: pass # Derive from dates if available data = info.data close_date = data.get('close_date') open_date = data.get('open_date') if close_date: now = datetime.now().date() close_dt = close_date.date() if isinstance(close_date, datetime) else close_date if isinstance(close_dt, date) and close_dt < now: return GrantStatus.CLOSED if open_date: now = datetime.now().date() if isinstance(open_date, date) and open_date > now: return GrantStatus.UPCOMING if open_date and close_date: return GrantStatus.OPEN return GrantStatus.UNKNOWN class Config: use_enum_values = True class SearchFilters(BaseModel): """Filters for grant search queries.""" sources: Optional[List[str]] = Field(None, description="Filter by data sources") status: Optional[List[GrantStatus]] = Field(None, description="Filter by status (open/upcoming/closed)") min_funding: Optional[float] = Field(None, description="Minimum funding amount") max_funding: Optional[float] = Field(None, description="Maximum funding amount") location: Optional[List[str]] = Field(None, description="Filter by location/region") org_type: Optional[List[str]] = Field(None, description="Filter by organization type") theme_keywords: Optional[List[str]] = Field(None, description="Filter by theme keywords") class SearchHit(BaseModel): """A single search result with score.""" grant: Grant = Field(..., description="The matched grant") score: float = Field(..., description="Relevance score (0-1)") class Config: arbitrary_types_allowed = True class QARequest(BaseModel): """Request model for QA queries.""" query: str = Field(..., description="User's question about grants", max_length=4000) filters: Optional[SearchFilters] = Field(None, description="Optional search filters") org_profile: Optional[Dict[str, Any]] = Field(None, description="Optional organization profile for context") session_id: Optional[str] = Field(None, description="Optional session ID for tracking") use_llm_routing: bool = Field(default=True, description="Use LLM-based query routing") @field_validator('query') @classmethod def validate_query_length(cls, v: str) -> str: """Ensure query is not empty and within limits.""" v = v.strip() if not v: raise ValueError("Query cannot be empty") # Additional check beyond Pydantic's max_length for settings integration from analyzer.config import get_settings settings = get_settings() if len(v) > settings.MAX_QUERY_CHARS: raise ValueError(f"Query exceeds maximum length of {settings.MAX_QUERY_CHARS} characters") return v class ChunkType(str, Enum): """Types of chunks in streaming responses.""" METADATA = "metadata" INTENT = "intent" TOKEN = "token" CITATIONS = "citations" DONE = "done" ERROR = "error" class QAChunk(BaseModel): """ A chunk in the streaming QA response. Supports Server-Sent Events (SSE) and NDJSON streaming. """ type: ChunkType = Field(..., description="Chunk type") # Type-specific fields (only one should be populated based on type) content: Optional[str] = Field(None, description="Text content (for TOKEN chunks)") token: Optional[str] = Field(None, description="Alias for content (for TOKEN chunks)") citations: Optional[List[Dict[str, Any]]] = Field(None, description="Citation list (for CITATIONS chunks)") error: Optional[str] = Field(None, description="Error message (for ERROR chunks)") # Metadata fields session_id: Optional[str] = Field(None, description="Session ID (for METADATA chunks)") query: Optional[str] = Field(None, description="Original query (for METADATA chunks)") intent: Optional[str] = Field(None, description="Detected intent (for INTENT chunks)") latency_ms: Optional[int] = Field(None, description="Total latency (for DONE chunks)") class Config: use_enum_values = True class CitationInfo(BaseModel): """Citation information in response.""" grant_id: str = Field(..., description="Grant/competition ID") title: str = Field(..., description="Grant title") url: Optional[str] = Field(None, description="Grant URL") confidence: Optional[float] = Field(None, description="Citation confidence score (0-1)") class QAResponse(BaseModel): """Complete QA response (non-streaming).""" session_id: str = Field(..., description="Session ID") query: str = Field(..., description="Original query") answer: str = Field(..., description="Generated answer") citations: List[CitationInfo] = Field(default_factory=list, description="Citations used") latency_ms: int = Field(..., description="Total latency in milliseconds") success: bool = Field(default=True, description="Whether request succeeded") error: Optional[str] = Field(None, description="Error message if failed") intent: Optional[str] = Field(None, description="Detected query intent") # Legacy support: simple dict-to-Grant conversion def dict_to_grant(data: Dict[str, Any]) -> Grant: """ Convert a dictionary (from data loader) to a validated Grant model. Handles various field name variations and missing fields gracefully. """ # Extract funding info if present funding_data = data.get("funding", {}) funding = None if funding_data and isinstance(funding_data, dict): rates_data = funding_data.get("rates", {}) rates = FundingRates(**rates_data) if rates_data else None funding = FundingInfo( min=funding_data.get("min"), max=funding_data.get("max"), total_pot=funding_data.get("total_pot"), rates=rates ) # Extract URLs url = data.get("url") urls = None if url: urls = GrantUrls(overview=url, apply=None, guidance=None, other=[]) # Parse dates def parse_date(val: Any) -> Optional[date]: if isinstance(val, date): return val if isinstance(val, datetime): return val.date() if isinstance(val, str): try: return datetime.fromisoformat(val.replace('Z', '+00:00')).date() except: return None return None def parse_datetime(val: Any) -> Optional[datetime]: if isinstance(val, datetime): return val if isinstance(val, str): try: return datetime.fromisoformat(val.replace('Z', '+00:00')) except: # Try parsing as date only dt = parse_date(val) return datetime.combine(dt, datetime.min.time()) if dt else None return None # Extract summary from various sources summary = ( data.get("summary") or data.get("sections", {}).get("summary_raw") or data.get("public_description") or "" ) # Extract scope and eligibility sections = data.get("sections", {}) scope = sections.get("scope_raw") eligibility = sections.get("eligibility_raw") return Grant( id=data.get("id", "unknown"), source=data.get("source", "innovate_uk"), title=data.get("title", "Untitled"), programme=data.get("programme"), summary=summary, scope=scope, eligibility=eligibility, open_date=parse_date(data.get("open_date")), close_date=parse_datetime(data.get("close_date")), notify_date=parse_date(data.get("notify_date")), project_start_from=parse_date(data.get("project_start_from")), funding=funding, url=url, urls=urls, created_at=parse_datetime(data.get("created_at")), updated_at=parse_datetime(data.get("updated_at")), sections=sections if sections else None )