nimmersky/oghma-proxy/src/oghma_proxy/models.py

"""Data models for Oghma RAG Proxy."""

from __future__ import annotations

from enum import Enum
from typing import Any

from pydantic import BaseModel, Field


class EducationLevel(str, Enum):
    """NPC education level determines lore depth."""

    SCHOLAR = "scholar"  # Full lore access
    COMMONER = "commoner"  # Basic summaries only


class NPCProfile(BaseModel):
    """Extracted NPC profile from SkyrimNet prompts."""

    name: str = "Unknown"
    race: str = "Unknown"
    gender: str = "Unknown"
    profession: str | None = None
    factions: list[str] = Field(default_factory=list)
    location: str | None = None
    traits: list[str] = Field(default_factory=list)

    # Computed
    knowledge_classes: list[str] = Field(default_factory=list)
    education_level: EducationLevel = EducationLevel.COMMONER

    def compute_knowledge_classes(self) -> None:
        """Compute knowledge classes from profile attributes."""
        classes = set()

        # Race-based knowledge
        race_map = {
            "nord": ["nord"],
            "dunmer": ["darkelf", "dunmer"],
            "altmer": ["highelf", "altmer"],
            "bosmer": ["woodelf", "bosmer"],
            "argonian": ["argonian"],
            "khajiit": ["khajiit"],
            "breton": ["breton"],
            "redguard": ["redguard"],
            "orsimer": ["orc", "orsimer"],
            "orc": ["orc", "orsimer"],
            "imperial": ["imperial"],
        }
        race_lower = self.race.lower()
        if race_lower in race_map:
            classes.update(race_map[race_lower])

        # Profession-based knowledge
        profession_map = {
            "priest": ["priest"],
            "mage": ["mage", "scholar"],
            "wizard": ["mage", "scholar"],
            "scholar": ["scholar"],
            "blacksmith": ["blacksmith"],
            "guard": ["guard", "warrior"],
            "soldier": ["warrior", "guard"],
            "warrior": ["warrior"],
            "thief": ["thief"],
            "merchant": ["merchant"],
            "innkeeper": ["innkeeper"],
            "hunter": ["hunter"],
            "farmer": ["peasant"],
            "peasant": ["peasant"],
            "noble": ["noble"],
            "jarl": ["noble"],
            "bard": ["bard"],
            "alchemist": ["alchemist"],
        }
        if self.profession:
            prof_lower = self.profession.lower()
            if prof_lower in profession_map:
                classes.update(profession_map[prof_lower])

        # Location-based knowledge
        location_map = {
            "whiterun": ["whiterun"],
            "windhelm": ["eastmarch"],
            "solitude": ["haafingar"],
            "riften": ["rift"],
            "markarth": ["reach"],
            "morthal": ["hjaalmarch"],
            "dawnstar": ["pale"],
            "winterhold": ["winterhold"],
            "falkreath": ["falkreath"],
            "solstheim": ["solstheim"],
        }
        if self.location:
            loc_lower = self.location.lower()
            if loc_lower in location_map:
                classes.update(location_map[loc_lower])

        # Faction-based knowledge
        faction_map = {
            "companions": ["companions"],
            "college of winterhold": ["college", "mage"],
            "college": ["college", "mage"],
            "thieves guild": ["thieves"],
            "dark brotherhood": ["darkbrotherhood"],
            "stormcloaks": ["stormcloak"],
            "stormcloak": ["stormcloak"],
            "imperial legion": ["imperial"],
            "legion": ["imperial"],
            "thalmor": ["thalmor"],
            "dawnguard": ["dawnguard"],
            "volkihar": ["vampire", "volkihar"],
        }
        for faction in self.factions:
            faction_lower = faction.lower()
            if faction_lower in faction_map:
                classes.update(faction_map[faction_lower])

        self.knowledge_classes = list(classes)

        # Determine education level
        educated_professions = {"mage", "wizard", "scholar", "priest", "noble", "bard"}
        educated_factions = {"college of winterhold", "thalmor", "college"}

        if self.profession and self.profession.lower() in educated_professions:
            self.education_level = EducationLevel.SCHOLAR
        elif any(f.lower() in educated_factions for f in self.factions):
            self.education_level = EducationLevel.SCHOLAR
        else:
            self.education_level = EducationLevel.COMMONER


class LoreEntry(BaseModel):
    """A retrieved lore entry from Oghma."""

    topic: str
    content: str
    category: str
    score: float
    knowledge_classes: list[str] = Field(default_factory=list)


class ChatMessage(BaseModel):
    """OpenRouter-compatible chat message."""

    role: str
    content: str
    name: str | None = None


class ChatCompletionRequest(BaseModel):
    """OpenRouter-compatible chat completion request."""

    model: str
    messages: list[ChatMessage]
    temperature: float | None = None
    max_tokens: int | None = None
    stream: bool = False
    # Allow additional fields to pass through
    model_config = {"extra": "allow"}


class InjectionResult(BaseModel):
    """Result of lore injection."""

    npc_profile: NPCProfile
    lore_entries: list[LoreEntry]
    injection_text: str
    query_time_ms: float