response_models.py

"""
Core response models for Hindsight memory system.

These models define the structure of data returned by the core MemoryEngine class.
API response models should be kept separate and convert from these core models to maintain
API stability even if internal models change.
"""

from typing import Any

from pydantic import BaseModel, ConfigDict, Field, field_validator

VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "observation"])


class LLMToolCall(BaseModel):
    """A tool call requested by the LLM."""

    id: str = Field(description="Unique identifier for this tool call")
    name: str = Field(description="Name of the tool to call")
    arguments: dict[str, Any] = Field(description="Arguments to pass to the tool")
    thought_signature: str | None = Field(
        default=None,
        description="Opaque token required by Gemini 3.1+ thinking models to preserve thought context across turns",
    )


class LLMToolCallResult(BaseModel):
    """Result from an LLM call that may include tool calls."""

    content: str | None = Field(default=None, description="Text content if any")
    tool_calls: list[LLMToolCall] = Field(default_factory=list, description="Tool calls requested by the LLM")
    finish_reason: str | None = Field(default=None, description="Reason the LLM stopped: 'stop', 'tool_calls', etc.")
    input_tokens: int = Field(default=0, description="Input tokens used in this call")
    output_tokens: int = Field(default=0, description="Output tokens used in this call")


class ToolCallTrace(BaseModel):
    """A single tool call made during reflect."""

    tool: str = Field(description="Tool name: lookup, recall, learn, expand")
    reason: str | None = Field(default=None, description="Agent's reasoning for making this tool call")
    input: dict = Field(description="Tool input parameters")
    output: dict = Field(description="Tool output/result")
    duration_ms: int = Field(description="Execution time in milliseconds")
    iteration: int = Field(default=0, description="Iteration number (1-based) when this tool was called")


class LLMCallTrace(BaseModel):
    """A single LLM call made during reflect."""

    scope: str = Field(description="Call scope: agent_1, agent_2, final, etc.")
    duration_ms: int = Field(description="Execution time in milliseconds")


class ObservationRef(BaseModel):
    """Reference to an observation accessed during reflect."""

    id: str = Field(description="Observation ID")
    name: str = Field(description="Observation name")
    type: str = Field(description="Observation type: entity, concept, event")
    subtype: str = Field(description="Observation subtype: structural, emergent, learned")
    description: str = Field(description="Brief description")
    summary: str | None = Field(default=None, description="Full summary (when looked up in detail)")


class DirectiveRef(BaseModel):
    """Reference to a directive that was applied during reflect."""

    id: str = Field(description="Directive mental model ID")
    name: str = Field(description="Directive name")
    content: str = Field(description="Directive content")


class TokenUsage(BaseModel):
    """
    Token usage metrics for LLM calls.

    Tracks input/output tokens for a single request to enable
    per-request cost tracking and monitoring.
    """

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "input_tokens": 1500,
                "output_tokens": 500,
                "total_tokens": 2000,
            }
        }
    )

    input_tokens: int = Field(default=0, description="Number of input/prompt tokens consumed")
    output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
    total_tokens: int = Field(default=0, description="Total tokens (input + output)")
    cached_tokens: int = Field(default=0, description="Cached/cache-read prompt tokens, when reported by the provider")

    def __add__(self, other: "TokenUsage") -> "TokenUsage":
        """Allow aggregating token usage from multiple calls."""
        return TokenUsage(
            input_tokens=self.input_tokens + other.input_tokens,
            output_tokens=self.output_tokens + other.output_tokens,
            total_tokens=self.total_tokens + other.total_tokens,
            cached_tokens=self.cached_tokens + other.cached_tokens,
        )


class ExtractedFact(BaseModel):
    """A single candidate fact produced by dry-run extraction (no resolution/links/persistence).

    A deliberate subset of the persisted memory-unit shape — only the fields a fresh extraction
    yields. Storage/consolidation/curation fields (id, document_id, chunk_id, proof_count, state, …)
    are omitted because nothing is stored. Entities are raw, unresolved names.
    """

    text: str = Field(description="The extracted fact text.")
    fact_type: str = Field(description="Perspective classification: 'world' or 'experience'.")
    occurred_start: str | None = Field(default=None, description="ISO timestamp the fact's event started, if dated.")
    occurred_end: str | None = Field(default=None, description="ISO timestamp the fact's event ended, if dated.")
    entities: list[str] = Field(
        default_factory=list, description="Raw (unresolved) entity names mentioned in the fact."
    )


class DryRunExtractionResult(BaseModel):
    """Result of dry-run fact extraction: candidate facts plus aggregated LLM token usage."""

    facts: list[ExtractedFact] = Field(
        default_factory=list, description="Candidate facts the retain step would extract."
    )
    usage: TokenUsage = Field(
        default_factory=TokenUsage, description="Aggregated token usage across the extraction LLM calls."
    )


class DispositionTraits(BaseModel):
    """
    Disposition traits for a memory bank.

    All traits are scored 1-5 where:
    - skepticism: 1=trusting, 5=skeptical (how much to doubt or question information)
    - literalism: 1=flexible interpretation, 5=literal interpretation (how strictly to interpret information)
    - empathy: 1=detached, 5=empathetic (how much to consider emotional context)
    """

    skepticism: int = Field(ge=1, le=5, description="How skeptical vs trusting (1=trusting, 5=skeptical)")
    literalism: int = Field(ge=1, le=5, description="How literally to interpret information (1=flexible, 5=literal)")
    empathy: int = Field(ge=1, le=5, description="How much to consider emotional context (1=detached, 5=empathetic)")

    model_config = ConfigDict(json_schema_extra={"example": {"skepticism": 3, "literalism": 3, "empathy": 3}})


class MemoryFact(BaseModel):
    """
    A single memory fact returned by search or think operations.

    This represents a unit of information stored in the memory system,
    including both the content and metadata.
    """

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "id": "123e4567-e89b-12d3-a456-426614174000",
                "text": "Alice works at Google on the AI team",
                "fact_type": "world",
                "entities": ["Alice", "Google"],
                "context": "work info",
                "occurred_start": "2024-01-15T10:30:00Z",
                "occurred_end": "2024-01-15T10:30:00Z",
                "mentioned_at": "2024-01-15T10:30:00Z",
                "document_id": "session_abc123",
                "metadata": {"source": "slack"},
                "chunk_id": "bank123_session_abc123_0",
                "activation": 0.95,
                "tags": ["user_a", "session_123"],
            }
        }
    )

    id: str = Field(description="Unique identifier for the memory fact")
    text: str = Field(description="The actual text content of the memory")
    fact_type: str = Field(description="Type of fact: 'world', 'experience', 'opinion', or 'observation'")
    entities: list[str] | None = Field(None, description="Entity names mentioned in this fact")
    context: str | None = Field(None, description="Additional context for the memory")
    occurred_start: str | None = Field(None, description="ISO format date when the event started occurring")
    occurred_end: str | None = Field(None, description="ISO format date when the event ended occurring")
    mentioned_at: str | None = Field(None, description="ISO format date when the fact was mentioned/learned")
    document_id: str | None = Field(None, description="ID of the document this memory belongs to")
    metadata: dict[str, str] | None = Field(None, description="User-defined metadata")

    @field_validator("metadata", mode="before")
    @classmethod
    def parse_metadata(cls, v: Any) -> dict[str, str] | None:
        """Parse metadata from JSON string if needed (asyncpg may return JSONB as str)."""
        if v is None:
            return None
        if isinstance(v, str):
            import json

            return json.loads(v)
        return v

    chunk_id: str | None = Field(
        None, description="ID of the chunk this fact was extracted from (format: bank_id_document_id_chunk_index)"
    )
    tags: list[str] | None = Field(None, description="Visibility scope tags associated with this fact")
    source_fact_ids: list[str] | None = Field(
        None,
        description="IDs of source facts this observation was derived from (observation type only, when source_facts is enabled)",
    )


class ChunkInfo(BaseModel):
    """Information about a chunk."""

    chunk_text: str = Field(description="The raw chunk text")
    chunk_index: int = Field(description="Index of the chunk within the document")
    truncated: bool = Field(default=False, description="Whether the chunk was truncated due to token limits")


class ObservationResult(BaseModel):
    """An observation result from recall (consolidated knowledge synthesized from facts)."""

    id: str = Field(description="Unique observation ID")
    text: str = Field(description="The observation text")
    proof_count: int = Field(description="Number of facts supporting this observation")
    relevance: float = Field(default=0.0, description="Relevance score to the query")
    tags: list[str] | None = Field(default=None, description="Tags for visibility scoping")
    source_memory_ids: list[str] = Field(
        default_factory=list, description="IDs of facts that contribute to this observation"
    )


class MentalModelResult(BaseModel):
    """A mental model result from recall (stored reflect response)."""

    id: str = Field(description="Unique mental model ID")
    name: str = Field(description="Human-readable name")
    content: str = Field(description="The synthesized content")
    relevance: float = Field(default=0.0, description="Relevance score to the query")


class RecallResult(BaseModel):
    """
    Result from a recall operation.

    Contains a list of matching memory facts and optional trace information
    for debugging and transparency.
    """

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "results": [
                    {
                        "id": "123e4567-e89b-12d3-a456-426614174000",
                        "text": "Alice works at Google on the AI team",
                        "fact_type": "world",
                        "context": "work info",
                        "occurred_start": "2024-01-15T10:30:00Z",
                        "occurred_end": "2024-01-15T10:30:00Z",
                        "activation": 0.95,
                    }
                ],
                "trace": {"query": "What did Alice say about machine learning?", "num_results": 1},
            }
        }
    )

    results: list[MemoryFact] = Field(description="List of memory facts matching the query")
    trace: dict[str, Any] | None = Field(None, description="Trace information for debugging")
    entities: dict[str, "EntityState"] | None = Field(
        None, description="Entity states for entities mentioned in results (keyed by canonical name)"
    )
    chunks: dict[str, ChunkInfo] | None = Field(
        None, description="Chunks for facts, keyed by '{document_id}_{chunk_index}'"
    )
    source_facts: dict[str, MemoryFact] | None = Field(
        None, description="Source facts for observation-type results, keyed by fact ID"
    )


class ReflectResult(BaseModel):
    """
    Result from a reflect operation.

    Contains the formulated answer, the facts it was based on (organized by type),
    any new opinions that were formed during the reflection process, and optionally
    structured output if a response schema was provided.
    """

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "text": "Based on my knowledge, machine learning is being actively used in healthcare...",
                "based_on": {
                    "world": [
                        {
                            "id": "123e4567-e89b-12d3-a456-426614174000",
                            "text": "Machine learning is used in medical diagnosis",
                            "fact_type": "world",
                            "context": "healthcare",
                            "occurred_start": "2024-01-15T10:30:00Z",
                            "occurred_end": "2024-01-15T10:30:00Z",
                        }
                    ],
                    "experience": [],
                    "opinion": [],
                    "observation": [],
                    "mental-models": [],
                    "directives": [
                        {
                            "id": "directive-123",
                            "name": "Response Style",
                            "rules": ["Always be concise"],
                        }
                    ],
                },
                "structured_output": {"summary": "ML in healthcare", "confidence": 0.9},
                "usage": {"input_tokens": 1500, "output_tokens": 500, "total_tokens": 2000},
            }
        }
    )

    text: str = Field(description="The formulated answer text")
    based_on: dict[str, Any] = Field(
        description="Facts used to formulate the answer, organized by type (world, experience, observation, mental-models, directives)"
    )
    structured_output: dict[str, Any] | None = Field(
        default=None,
        description="Structured output parsed according to the provided response schema. Only present when response_schema was provided.",
    )
    usage: TokenUsage | None = Field(
        default=None,
        description="Token usage metrics for the LLM calls made during this reflect operation.",
    )
    tool_trace: list[ToolCallTrace] = Field(
        default_factory=list,
        description="Trace of tool calls made during reflection. Only present when include.tool_calls is enabled.",
    )
    llm_trace: list[LLMCallTrace] = Field(
        default_factory=list,
        description="Trace of LLM calls made during reflection. Only present when include.tool_calls is enabled.",
    )
    directives_applied: list[DirectiveRef] = Field(
        default_factory=list,
        description="Directive mental models that were applied during this reflection.",
    )


class EntityObservation(BaseModel):
    """
    An observation about an entity.

    Observations are objective facts synthesized from multiple memory facts
    about an entity, without personality influence.
    """

    model_config = ConfigDict(
        json_schema_extra={
            "example": {"text": "John is detail-oriented and works at Google", "mentioned_at": "2024-01-15T10:30:00Z"}
        }
    )

    text: str = Field(description="The observation text")
    mentioned_at: str | None = Field(None, description="ISO format date when this observation was created")


class EntityState(BaseModel):
    """
    Current mental model of an entity.

    Contains observations synthesized from facts about the entity.
    """

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "entity_id": "123e4567-e89b-12d3-a456-426614174000",
                "canonical_name": "John",
                "observations": [
                    {"text": "John is detail-oriented", "mentioned_at": "2024-01-15T10:30:00Z"},
                    {"text": "John works at Google on the AI team", "mentioned_at": "2024-01-14T09:00:00Z"},
                ],
            }
        }
    )

    entity_id: str = Field(description="Unique identifier for the entity")
    canonical_name: str = Field(description="Canonical name of the entity")
    observations: list[EntityObservation] = Field(
        default_factory=list, description="List of observations about this entity"
    )


class MentalModel(BaseModel):
    """
    A manually configured mental model for tracking specific topics/areas.

    Mental models are user-defined focus areas that the agent should track
    and maintain summaries for, unlike auto-extracted entities.
    """

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "id": "team-dynamics",
                "name": "Team Dynamics",
                "description": "Track how the team collaborates, communication patterns, conflicts, and resolutions",
                "summary": "The team has strong collaboration...",
                "summary_updated_at": "2024-01-15T10:30:00Z",
                "created_at": "2024-01-10T08:00:00Z",
            }
        }
    )

    id: str = Field(description="Unique identifier (alphanumeric lowercase)")
    name: str = Field(description="Display name for the mental model")
    description: str = Field(description="Prompt/directions for what to track and summarize")
    summary: str | None = Field(None, description="Generated summary based on relevant facts")
    summary_updated_at: str | None = Field(None, description="ISO format date when summary was last updated")
    created_at: str = Field(description="ISO format date when the mental model was created")