"""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