feat: Add Oghma RAG Proxy for SkyrimNet lore injection
RAG proxy that intercepts SkyrimNet LLM requests and enriches them with relevant Tamrielic lore from CHIM's Oghma Infinium database. Features: - FastAPI proxy compatible with OpenAI API - ChromaDB semantic search for lore retrieval - NPC profile extraction from SkyrimNet prompts - Google Sheets ingestion for CHIM's Oghma data - Kubernetes deployment manifests - Debug endpoint for RAG operation monitoring Collections ingested to iris-dev ChromaDB: - oghma_lore: 1951 entries (scholar knowledge) - oghma_basic: 1949 entries (commoner knowledge) - oghma_visual: 1151 entries (Omnisight perception) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
dafit
173
oghma-proxy/src/oghma_proxy/retriever.py
Normal file
173
oghma-proxy/src/oghma_proxy/retriever.py
Normal file
@@ -0,0 +1,173 @@
|
||||
"""Oghma Lore Retriever - Queries ChromaDB for relevant Tamrielic lore."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import chromadb
|
||||
import structlog
|
||||
from chromadb.config import Settings
|
||||
|
||||
from .models import EducationLevel, LoreEntry, NPCProfile
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from chromadb import Collection
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
|
||||
class OghmaRetriever:
|
||||
"""Retrieves relevant lore from Oghma ChromaDB collections."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "iris-dev.eachpath.local",
|
||||
port: int = 35000,
|
||||
collection_lore: str = "oghma_lore",
|
||||
collection_basic: str = "oghma_basic",
|
||||
max_results: int = 5,
|
||||
min_score: float = 0.55,
|
||||
):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.collection_lore_name = collection_lore
|
||||
self.collection_basic_name = collection_basic
|
||||
self.max_results = max_results
|
||||
self.min_score = min_score
|
||||
|
||||
self._client: chromadb.HttpClient | None = None
|
||||
self._collection_lore: Collection | None = None
|
||||
self._collection_basic: Collection | None = None
|
||||
|
||||
def _get_client(self) -> chromadb.HttpClient:
|
||||
"""Get or create ChromaDB client."""
|
||||
if self._client is None:
|
||||
self._client = chromadb.HttpClient(
|
||||
host=self.host,
|
||||
port=self.port,
|
||||
settings=Settings(anonymized_telemetry=False),
|
||||
)
|
||||
logger.info("Connected to ChromaDB", host=self.host, port=self.port)
|
||||
return self._client
|
||||
|
||||
def _get_collection(self, education_level: EducationLevel) -> Collection:
|
||||
"""Get the appropriate collection based on education level."""
|
||||
client = self._get_client()
|
||||
|
||||
if education_level == EducationLevel.SCHOLAR:
|
||||
if self._collection_lore is None:
|
||||
self._collection_lore = client.get_collection(self.collection_lore_name)
|
||||
return self._collection_lore
|
||||
else:
|
||||
if self._collection_basic is None:
|
||||
self._collection_basic = client.get_collection(self.collection_basic_name)
|
||||
return self._collection_basic
|
||||
|
||||
def retrieve(
|
||||
self,
|
||||
query: str,
|
||||
npc_profile: NPCProfile,
|
||||
) -> tuple[list[LoreEntry], float]:
|
||||
"""
|
||||
Retrieve relevant lore entries for an NPC.
|
||||
|
||||
Args:
|
||||
query: Conversation context to search for
|
||||
npc_profile: NPC profile for knowledge filtering
|
||||
|
||||
Returns:
|
||||
Tuple of (lore entries, query time in ms)
|
||||
"""
|
||||
if not query.strip():
|
||||
return [], 0.0
|
||||
|
||||
start_time = time.perf_counter()
|
||||
|
||||
try:
|
||||
collection = self._get_collection(npc_profile.education_level)
|
||||
|
||||
# Build metadata filter for knowledge classes
|
||||
# NOTE: Currently disabled because CHIM's Oghma data doesn't have
|
||||
# knowledge_class populated consistently. Enable when data is enriched.
|
||||
where_filter = None
|
||||
# TODO: Re-enable when knowledge_class data is available
|
||||
# if npc_profile.knowledge_classes:
|
||||
# if len(npc_profile.knowledge_classes) == 1:
|
||||
# where_filter = {"knowledge_classes": {"$contains": npc_profile.knowledge_classes[0]}}
|
||||
# else:
|
||||
# where_filter = {
|
||||
# "$or": [
|
||||
# {"knowledge_classes": {"$contains": kc}}
|
||||
# for kc in npc_profile.knowledge_classes
|
||||
# ]
|
||||
# }
|
||||
|
||||
# Query ChromaDB
|
||||
results = collection.query(
|
||||
query_texts=[query],
|
||||
n_results=self.max_results,
|
||||
where=where_filter,
|
||||
include=["documents", "metadatas", "distances"],
|
||||
)
|
||||
|
||||
# Parse results
|
||||
entries = []
|
||||
if results and results["documents"] and results["documents"][0]:
|
||||
for i, doc in enumerate(results["documents"][0]):
|
||||
metadata = results["metadatas"][0][i] if results["metadatas"] else {}
|
||||
distance = results["distances"][0][i] if results["distances"] else 1.0
|
||||
|
||||
# Convert distance to similarity score (ChromaDB uses L2 distance)
|
||||
# Lower distance = higher similarity
|
||||
score = 1.0 / (1.0 + distance)
|
||||
|
||||
if score >= self.min_score:
|
||||
entries.append(
|
||||
LoreEntry(
|
||||
topic=metadata.get("topic", "Unknown"),
|
||||
content=doc,
|
||||
category=metadata.get("category", "Unknown"),
|
||||
score=score,
|
||||
knowledge_classes=metadata.get("knowledge_classes", "").split(","),
|
||||
)
|
||||
)
|
||||
|
||||
query_time = (time.perf_counter() - start_time) * 1000
|
||||
|
||||
logger.info(
|
||||
"Retrieved lore entries",
|
||||
query_preview=query[:100],
|
||||
npc_name=npc_profile.name,
|
||||
education=npc_profile.education_level.value,
|
||||
entries_found=len(entries),
|
||||
query_time_ms=round(query_time, 2),
|
||||
)
|
||||
|
||||
return entries, query_time
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to retrieve lore", error=str(e))
|
||||
query_time = (time.perf_counter() - start_time) * 1000
|
||||
return [], query_time
|
||||
|
||||
def health_check(self) -> bool:
|
||||
"""Check if ChromaDB is reachable."""
|
||||
try:
|
||||
client = self._get_client()
|
||||
client.heartbeat()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("ChromaDB health check failed", error=str(e))
|
||||
return False
|
||||
|
||||
|
||||
# Cached retriever instance
|
||||
@lru_cache(maxsize=1)
|
||||
def get_retriever(
|
||||
host: str = "iris-dev.eachpath.local",
|
||||
port: int = 35000,
|
||||
) -> OghmaRetriever:
|
||||
"""Get cached retriever instance."""
|
||||
return OghmaRetriever(host=host, port=port)
|
||||
Reference in New Issue
Block a user