src/agentcore/models.py

"""Pydantic request/response models shared across the pipeline (Stage 0: I/O contracts)."""
import asyncio
import html
import json
import logging
import os
import re
import time
import uuid
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional, Tuple
from urllib.parse import quote

import httpx
from pydantic import BaseModel, ConfigDict, Field

from src.config import get_settings, LIBBEE_VERSION

logger = logging.getLogger(__name__)

class ChatMessage(BaseModel):
    role: str
    content: str


class SearchContextPayload(BaseModel):
    context_id: str
    topic: str
    display_topic: Optional[str] = None
    intent: str = "search_academic"
    source: str = "primo"
    resource_type: str = "articles"   # "articles" | "books" | "both"
    peer_reviewed: bool = False
    open_access: bool = False
    year_from: Optional[str] = None
    year_to: Optional[str] = None
    ai_tool_query: Optional[str] = None
    primo_boolean_query: Optional[str] = None
    created_at: float = Field(default_factory=time.time)
    # v3.8 — ambiguity gate
    clarification_needed: bool = False
    clarification_message: Optional[str] = None


class ClientStatePayload(BaseModel):
    session_id: Optional[str] = None
    active_search_context_id: Optional[str] = None
    recent_search_contexts: List[SearchContextPayload] = Field(default_factory=list)
    follow_up_action: Optional[str] = None
    follow_up_context_id: Optional[str] = None


class AgentRequest(BaseModel):
    question: str
    # FIX: validated literal — malformed model values caught at Pydantic layer
    model: Literal["gpt", "claude"] = "gpt"
    history: List[ChatMessage] = Field(default_factory=list)
    client_state: Optional[ClientStatePayload] = None


class AgentResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")

    answer: str
    intent: str = "general"
    tools_used: List[str] = Field(default_factory=list)
    search_results: List[dict] = Field(default_factory=list)
    sources: List[dict] = Field(default_factory=list)
    model_used: str = "gpt"
    response_time: float = 0
    corrected_query: Optional[str] = None
    natural_query: Optional[str] = None
    database_query: Optional[str] = None
    ask_librarian: Optional[bool] = None
    is_follow_up: bool = False
    follow_up_question: Optional[str] = None
    follow_up_suggestions: List[dict] = Field(default_factory=list)
    search_context: Optional[dict] = None
    ai_tool_query: Optional[str] = None
    primo_boolean_query: Optional[str] = None
    primo_search_url: Optional[str] = None
    primo_ai_url: Optional[str] = None
    pubmed_search_url: Optional[str] = None