diff --git a/requirements.txt b/requirements.txt index 07ff5cd..b2d331f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,30 +1,355 @@ -# Core server -fastapi -uvicorn -pydantic -starlette -orjson -httpx +# A good many of these are likely not needed. But exist due to ongoing testing and development. -# Vector & embeddings -chromadb -sentence-transformers -# NumPy 2 + SciPy that plays nice with it -numpy==2.0 -scipy +# You MUST use a .venv for this project. uv or conda is HIGHLY recommended. -# Audio stack (modern, NumPy 2 compatible) -librosa -soundfile -soxr +# Not using one will likely break your system-wide Python library. -# TTS / GPU bits you were already using (pin to your known-good set) -#torch -#torchaudio -#torchvision +# ========================= +# Server / API framework +# ========================= +fastapi==0.116.1 +starlette==0.47.3 +uvicorn==0.35.0 +httptools==0.6.4 +orjson==3.11.3 +watchfiles==1.1.0 -# Kokoro TTS (your current) -kokoro +# ========================= +# HTTP / Networking / Async +# ========================= +aiohappyeyeballs==2.6.1 +aiohttp==3.12.15 +aiolimiter==1.2.1 +aiosignal==1.4.0 +anyio==4.10.0 +h11==0.16.0 +h2==4.3.0 +httpcore==1.0.9 +httpx==0.28.1 +httpx-sse==0.4.0 +hyperframe==6.1.0 +multidict==6.6.4 +rfc3986==1.5.0 +sniffio==1.3.1 +websocket-client==1.8.0 +websockets==15.0.1 +yarl==1.20.1 -# Misc already in your tree -rich \ No newline at end of file +# ========================= +# Auth / Security +# ========================= +authlib==1.6.3 +bcrypt==3.2.0 +cryptography==45.0.6 +google-auth==2.40.3 +oauthlib==3.3.1 +pyasn1==0.6.1 +pyasn1-modules==0.4.2 +pyjwt==2.10.1 +requests-oauthlib==2.0.0 +rsa==4.9.1 + +# ========================= +# Cloud SDKs / Telemetry / Infra +# ========================= +azure-core==1.35.0 +azure-identity==1.24.0 +googleapis-common-protos==1.70.0 +grpcio==1.74.0 +kubernetes==33.1.0 +opentelemetry-api==1.36.0 +opentelemetry-exporter-otlp-proto-common==1.36.0 +opentelemetry-exporter-otlp-proto-grpc==1.36.0 +opentelemetry-proto==1.36.0 +opentelemetry-sdk==1.36.0 +opentelemetry-semantic-conventions==0.57b0 +posthog==5.4.0 + +# ========================= +# Data / Serialization / Validation +# ========================= +annotated-types==0.7.0 +jsonpatch==1.33 +jsonpointer==3.0.0 +jsonschema==4.25.1 +jsonschema-specifications==2025.4.1 +pydantic==2.11.7 +pydantic-core==2.33.2 +pydantic-settings==2.10.1 +protobuf==6.32.0 +python-dateutil==2.9.0.post0 +python-dotenv==1.1.1 +pyyaml==6.0.2 +referencing==0.36.2 +rpds-py==0.27.1 +uritemplate==4.2.0 + +# ========================= +# Core ML / Math / Utils +# ========================= +llvmlite==0.43.0 +numba==0.60.0 +numpy==2.0.0 +scipy==1.16.1 +scikit-learn==1.7.1 +threadpoolctl==3.6.0 +sympy==1.14.0 +tqdm==4.67.1 + +# ========================= +# PyTorch stack +# ========================= +torch==2.8.0+cu126 +torchvision==0.23.0+cu126 +torchaudio==2.8.0 +tensorboardx==2.6.4 +torchcrepe==0.0.20 + +# ========================= +# NLP / LLM tooling +# ========================= +accelerate==1.10.1 +huggingface-hub==0.34.4 +sentence-transformers==5.1.0 +tokenizers==0.22.0 +transformers==4.56.0 +tiktoken==0.11.0 +sacrebleu==2.5.1 +regex==2025.8.29 +openai==1.102.0 +cohere==5.17.0 +google-genai==1.32.0 +voyageai==0.3.4 +langchain-core==0.3.75 +langchain-text-splitters==0.3.11 +langsmith==0.4.21 + +# ========================= +# spaCy + curated transformers +# ========================= +spacy==3.8.7 +en-core-web-sm==3.8.0 +thinc==8.3.4 +catalogue==2.0.10 +confection==0.1.5 +cymem==2.0.11 +murmurhash==1.0.13 +preshed==3.0.10 +srsly==2.5.1 +spacy-legacy==3.0.12 +spacy-loggers==1.0.5 +langcodes==3.5.0 +language-data==1.3.0 +marisa-trie==1.3.1 +wasabi==1.1.3 +spacy-curated-transformers==0.3.1 +curated-tokenizers==0.0.9 +curated-transformers==0.1.1 +misaki==0.9.4 + +# ========================= +# Vectors / Datastores / Retrieval +# ========================= +chromadb==1.0.20 +faiss-cpu==1.10.0 +qdrant-client==1.15.1 +pinecone==7.3.0 +pinecone-plugin-assistant==1.8.0 +pinecone-plugin-interface==0.0.7 +weaviate-client==4.16.9 +turbopuffer==1.0.0 +propcache==0.3.2 +smart-open==7.3.0.post1 + +# ========================= +# Audio / Speech / Media +# ========================= +audioread==3.0.1 +av==15.1.0 +edge-tts==7.2.3 +einops==0.8.1 +espeakng-loader==0.2.4 +ffmpeg-python==0.2.0 +kokoro==0.9.4 +librosa==0.11.0 +praat-parselmouth==0.4.6 +pyworld==0.3.5 +resampy==0.4.3 +soundfile==0.13.1 +soxr==0.5.0.post1 +tts-with-rvc==0.1.9.1 + +# ========================= +# Imaging / Files / Parsing +# ========================= +pillow==11.0.0 +lxml==6.0.1 +rdflib==7.1.4 +csvw==3.5.1 +isodate==0.7.2 +language-tags==1.2.0 +markdown-it-py==4.0.0 +mdurl==0.1.2 +tabulate==0.9.0 +validators==0.35.0 + +# ========================= +# ONNX / Inference backends +# ========================= +onnxruntime==1.22.1 +onnxruntime-gpu==1.22.0 + +# ========================= +# DB / Storage / Locking +# ========================= +pymongo==4.14.1 +portalocker==3.2.0 + +# ========================= +# CLI / Logging / Dev ergonomics +# ========================= +backoff==2.2.1 +click==8.2.1 +colorama==0.4.6 +coloredlogs==15.0.1 +colorlog==6.9.0 +humanfriendly==10.0 +loguru==0.7.3 +pygments==2.19.2 +rich==14.1.0 +tenacity==9.1.2 +typer==0.17.3 +shellingham==1.5.4 + +# ========================= +# Tree-sitter / code tooling +# ========================= +tree-sitter==0.25.1 +tree-sitter-c-sharp==0.23.1 +tree-sitter-embedded-template==0.25.0 +tree-sitter-language-pack==0.9.0 +tree-sitter-yaml==0.7.1 + +# ========================= +# Azure / MSAL helpers +# ========================= +msal==1.33.0 +msal-extensions==1.3.1 + +# ========================= +# Fairseq & friends +# ========================= +fairseq-built==0.12.3 + +# ========================= +# Misc libraries (as installed) +# ========================= +accelerate==1.10.1 +addict==2.4.0 +antlr4-python3-runtime==4.9.3 +attrs==25.3.0 +autotiktokenizer==0.2.2 +babel==2.17.0 +bitarray==3.7.1 +blis==1.2.1 +build==1.3.0 +cachetools==5.5.2 +certifi==2025.8.3 +cffi==1.17.1 +charset-normalizer==3.4.3 +chonkie==1.2.1 +chromadb==1.0.20 +cloudflared==1.0.0.2 +cloudpathlib==0.22.0 +cohere==5.17.0 +confection==0.1.5 +csvw==3.5.1 +cython==3.1.3 +decorator==5.2.1 +deprecation==2.1.0 +distro==1.9.0 +dlinfo==2.0.0 +dnspython==2.7.0 +docopt==0.6.2 +durationpy==0.10 +fastavro==1.12.0 +filelock==3.19.1 +flatbuffers==25.2.10 +frozenlist==1.7.0 +fsspec==2025.7.0 +future==1.0.0 +hpack==4.1.0 +hydra-core==1.3.2 +hyper-connections==0.2.1 +idna==3.10 +importlib-metadata==8.7.0 +importlib-resources==6.5.2 +isodate==0.7.2 +jinja2==3.1.6 +jiter==0.10.0 +joblib==1.5.2 +kokoro==0.9.4 +langcodes==3.5.0 +language-data==1.3.0 +lazy-loader==0.4 +librosa==0.11.0 +local-attention==1.11.2 +lxml==6.0.1 +magika==0.6.2 +marisa-trie==1.3.1 +markupsafe==3.0.2 +mdurl==0.1.2 +misaki==0.9.4 +mmh3==5.2.0 +model2vec==0.6.0 +mpmath==1.3.0 +msgpack==1.1.1 +nest-asyncio==1.6.0 +networkx==3.5 +num2words==0.5.14 +omegaconf==2.3.0 +overrides==7.7.0 +packaging==25.0 +passlib==1.7.4 +phonemizer-fork==3.3.2 +pinecone-plugin-assistant==1.8.0 +pinecone-plugin-interface==0.0.7 +pip==25.2 +platformdirs==4.4.0 +pooch==1.8.2 +portalocker==3.2.0 +praat-parselmouth==0.4.6 +pybase64==1.4.2 +pycparser==2.22 +pyparsing==3.2.3 +pypika==0.48.9 +pyproject-hooks==1.2.0 +pyreadline3==3.5.4 +pywin32==311 +qdrant-client==1.15.1 +rdflib==7.1.4 +requests==2.32.5 +requests-toolbelt==1.0.0 +rich==14.1.0 +safetensors==0.6.2 +segments==2.3.0 +setuptools==80.9.0 +setuptools-scm==9.2.0 +six==1.17.0 +smart-open==7.3.0.post1 +soundfile==0.13.1 +thinc==8.3.4 +tree-sitter==0.25.1 +types-requests==2.32.4.20250809 +typing-extensions==4.15.0 +typing-inspection==0.4.1 +urllib3==2.5.0 +validators==0.35.0 +wasabi==1.1.3 +weasel==0.4.1 +weaviate-client==4.16.9 +wheel==0.45.1 +win32-setctime==1.2.0 +wrapt==1.17.3 +zipp==3.23.0 +zstandard==0.24.0 diff --git a/server/config.py b/server/config.py index b42af73..b5be4f3 100644 --- a/server/config.py +++ b/server/config.py @@ -35,11 +35,11 @@ class AudioSettings(BaseModel): class MemorySettings(BaseModel): """ Settings for memory and database persistence. - UPDATED: Now includes separate paths for user and session databases. """ + chunker: Literal["simple", "chonkie"] embedding_model: str - sessions_db_file: str # Path to the sessions/rooms database. - users_db_file: str # Path to the users/accounts database. + sessions_db_file: str + users_db_file: str class PathsSettings(BaseModel): prompts_file: str diff --git a/server/memory_manager.py b/server/memory_manager.py index 742d766..e3f7315 100644 --- a/server/memory_manager.py +++ b/server/memory_manager.py @@ -1,184 +1,152 @@ # server/memory_manager.py +""" +Manages the long-term memory of the AI Game Master using a vector database. +""" from __future__ import annotations -from pathlib import Path -from typing import List, Iterable, Dict, Any, Optional, TypedDict, cast import re import time import uuid -import numpy as np +from pathlib import Path +from typing import Any, Dict, List, Optional, TypedDict, cast, TYPE_CHECKING import chromadb +import numpy as np from chromadb.config import Settings as ChromaSettings from sentence_transformers import SentenceTransformer from .config import settings from .logger import logger -# ---------- Simple sentence splitting & chunking ---------- +# This block is for static analysis only (for Pylance). +# It is not executed at runtime, so it won't cause an ImportError. +if TYPE_CHECKING: + from chonkie import SemanticChunker + from chonkie.embeddings import SentenceTransformerEmbeddings + from chonkie.types import Chunk -def _sentence_split(text: str) -> List[str]: - """Lightweight sentence splitter; avoids heavy deps.""" - t = re.sub(r"\s+", " ", text).strip() - if not t: - return [] - parts = re.split(r"(?<=[.!?])\s+", t) - return [p.strip() for p in parts if p.strip()] - - -def _chunk_sentences( - sentences: List[str], - max_chars: int = 700, - overlap: int = 1, -) -> List[str]: - """Pack sentences into ~max_chars chunks with small overlap for recall.""" - chunks: List[str] = [] - buf: List[str] = [] - cur = 0 - for s in sentences: - if cur + len(s) + (1 if buf else 0) > max_chars and buf: - chunks.append(" ".join(buf)) - buf = buf[-overlap:] if overlap > 0 else [] - cur = sum(len(x) for x in buf) + (len(buf) - 1 if buf else 0) - if cur > 0: - cur += 1 - buf.append(s) - cur += len(s) - if buf: - chunks.append(" ".join(buf)) - return chunks - - -# ---------- Embedding wrapper ---------- class _STEmbedder: - """Sentence-Transformers wrapper that returns numpy arrays with the right dtype.""" + """A simple fallback embedder using sentence-transformers.""" def __init__(self, model_name: str) -> None: self.model = SentenceTransformer(model_name, device="cpu") - # NOTE: You may see a warning about a missing `config.json`. This is expected - # for some models like EmbeddingGemma and can be safely ignored. logger.info(f"SentenceTransformer loaded: {model_name}") - def embed(self, texts: Iterable[str]) -> np.ndarray: - """Encodes a list of texts into numpy embeddings.""" - # Note: EmbeddingGemma requires float32, this code ensures that. - arr = self.model.encode( - list(texts), - normalize_embeddings=True, - convert_to_numpy=True, - ) + def embed_batch(self, texts: List[str]) -> np.ndarray: + arr = self.model.encode(texts, normalize_embeddings=True, convert_to_numpy=True) if not isinstance(arr, np.ndarray): arr = np.asarray(arr) if arr.dtype != np.float32: arr = arr.astype(np.float32, copy=False) return arr - -# ---------- Metadata shape (all primitives) ---------- - +# --- (Metadata types are unchanged) --- Primitive = str | int | float | bool | None - class MemoryMeta(TypedDict, total=False): room_id: str ts: int len: int -# ---------- Memory manager ---------- - class MemoryManager: """Manages the vector database for the AI's long-term memory.""" def __init__(self): + """Initializes the ChromaDB client and the selected chunking/embedding strategy.""" self.memory_dir = Path(settings.paths.memory_dir) self.memory_dir.mkdir(parents=True, exist_ok=True) + self.chunker: Optional["SemanticChunker"] = None + self.embedder: Any = None try: chroma_settings = ChromaSettings(anonymized_telemetry=False) - self.chroma = chromadb.PersistentClient( - path=str(self.memory_dir / "chroma_db"), - settings=chroma_settings, - ) + self.chroma = chromadb.PersistentClient(path=str(self.memory_dir / "chroma_db"), settings=chroma_settings) logger.info("ChromaDB PersistentClient initialized (telemetry OFF).") - self.embedder = _STEmbedder(settings.memory.embedding_model) + + if settings.memory.chunker == "chonkie": + try: + from chonkie import SemanticChunker + from chonkie.embeddings import SentenceTransformerEmbeddings + except ImportError: + raise ImportError("Chonkie is configured but not installed correctly. Please run 'uv pip install \"chonkie[st]\"'.") + + logger.info("Initializing Chonkie with SemanticChunker...") + embedding_handler = SentenceTransformerEmbeddings(model_name=settings.memory.embedding_model, device="cpu") + self.chunker = SemanticChunker(embedding_model=embedding_handler) + self.embedder = embedding_handler + logger.info("Using 'chonkie' for text chunking.") + else: + self.embedder = _STEmbedder(settings.memory.embedding_model) + logger.info("Using 'simple' fallback for text chunking.") + except Exception: logger.critical("Failed to initialize MemoryManager.", exc_info=True) raise - def _collection_name(self, room_id: str) -> str: - """Generates a ChromaDB collection name for a given room.""" - return f"vdm_{room_id}" + @staticmethod + def _simple_chunker(text: str) -> List[str]: + # ... (implementation is unchanged) + t = re.sub(r"\s+", " ", text).strip() + if not t: return [] + sentences = [p.strip() for p in re.split(r"(?<=[.!?])\s+", t) if p.strip()] + if not sentences: return [] + chunks: List[str] = []; buf: List[str] = []; cur = 0 + for s in sentences: + if cur + len(s) + (1 if buf else 0) > 700 and buf: + chunks.append(" ".join(buf)); buf = buf[-1:]; cur = sum(len(x) for x in buf) + (len(buf) - 1 if buf else 0) + if cur > 0: cur += 1 + buf.append(s); cur += len(s) + if buf: chunks.append(" ".join(buf)) + return chunks - def _get_collection(self, room_id: str): - """Retrieves or creates a ChromaDB collection for a room.""" - return self.chroma.get_or_create_collection( - name=self._collection_name(room_id), - metadata={"hnsw:space": "cosine"}, - ) + def _get_collection(self, room_id: str) -> chromadb.Collection: + # ... (implementation is unchanged) + return self.chroma.get_or_create_collection(name=f"vdm_{room_id}", metadata={"hnsw:space": "cosine"}) def add_memory(self, room_id: str, text: str) -> None: - """ - Adds a piece of text to the long-term memory for a room. - It automatically formats the text for the selected embedding model. - """ - if not text or not text.strip(): - return + """Adds a piece of text to the long-term memory for a room.""" + if not text or not text.strip(): return try: - sentences = _sentence_split(text) - if not sentences: return - chunks: List[str] = _chunk_sentences(sentences, max_chars=700, overlap=1) + if settings.memory.chunker == "chonkie" and self.chunker: + # The runtime code knows what self.chunker(text) returns + chonkie_chunks = self.chunker(text) + chunks = [c.text for c in chonkie_chunks if c.text.strip()] # type: ignore (This might break, we're see) + else: + chunks = self._simple_chunker(text) + if not chunks: return - # NEW: Add model-specific prefixes for backward compatibility. - # This formats the text for EmbeddingGemma, but leaves other models unchanged. if "embeddinggemma" in settings.memory.embedding_model: - # This is a "document" for storage. - formatted_chunks = [f"text: {c}" for c in chunks] + formatted_chunks_for_embedding = [f"text: {c}" for c in chunks] else: - formatted_chunks = chunks + formatted_chunks_for_embedding = chunks - embeds_np: np.ndarray = self.embedder.embed(formatted_chunks) + embeds_np: np.ndarray = self.embedder.embed_batch(formatted_chunks_for_embedding) + col = self._get_collection(room_id) - ts = int(time.time()) ids: List[str] = [uuid.uuid4().hex for _ in chunks] - # We store the ORIGINAL, unprefixed text in the database. - metadatas: List[Dict[str, Primitive]] = [ - {"room_id": room_id, "ts": ts, "len": len(c)} for c in chunks - ] - col.add( - ids=ids, - documents=chunks, - embeddings=cast(Any, embeds_np), - metadatas=cast(Any, metadatas), - ) - logger.info(f"Added {len(chunks)} memory chunk(s) to room '{room_id}'.") + metadatas: List[Dict[str, Primitive]] = [{"room_id": room_id, "ts": ts, "len": len(c)} for c in chunks] + + col.add(ids=ids, documents=chunks, embeddings=cast(Any, embeds_np), metadatas=cast(Any, metadatas)) + logger.info(f"Added {len(chunks)} memory chunk(s) to room '{room_id}' using '{settings.memory.chunker}' chunker.") except Exception: logger.error(f"Failed to add memory to room '{room_id}'.", exc_info=True) def search_memory(self, room_id: str, query_text: str, k: int = 3) -> List[str]: - """ - Searches the long-term memory for relevant information. - It automatically formats the query for the selected embedding model. - """ - if not query_text or not query_text.strip(): - return [] + # ... (implementation is unchanged) + if not query_text or not query_text.strip(): return [] try: col = self._get_collection(room_id) - if col.count() == 0: - return [] + if col.count() == 0: return [] - # Add model-specific prefixes for backward compatibility. if "embeddinggemma" in settings.memory.embedding_model: - # This is a "query" for retrieval. formatted_query = f"task: search result | query: {query_text}" else: formatted_query = query_text - - q_np: np.ndarray = self.embedder.embed([formatted_query]) - result = col.query( - query_embeddings=cast(Any, q_np), - n_results=max(1, k), - ) + + q_np: np.ndarray = self.embedder.embed_batch([formatted_query]) + + result = col.query(query_embeddings=cast(Any, q_np), n_results=max(1, k)) docs = (result.get("documents") or [[]])[0] return [d for d in docs if d] except Exception: diff --git a/server/models.py b/server/models.py index 5fe70c6..26aeef7 100644 --- a/server/models.py +++ b/server/models.py @@ -5,7 +5,15 @@ from typing import Dict, List, Literal, Any, Optional # =================================================================== # Core Game & Application Models # =================================================================== -# These Pydantic models define the structure of our application's state. + +class PlayerSheet(BaseModel): + """Represents a player's character sheet with stats and inventory.""" + hp: int = 10 + max_hp: int = 10 + # Placeholders for future expansion + # attributes: Dict[str, int] = Field(default_factory=dict) + # inventory: List[str] = Field(default_factory=list) + class Player(BaseModel): """Represents a player within a game room.""" @@ -13,6 +21,8 @@ class Player(BaseModel): name: str avatar_style: str = "adventurer" is_active: bool = True + sheet: PlayerSheet = Field(default_factory=PlayerSheet) + class ChatMessage(BaseModel): """Represents a single message in the chat history.""" @@ -46,7 +56,6 @@ class LoginRequest(BaseModel): # =================================================================== # WebSocket Protocol Models # =================================================================== -# These models define the contract for messages sent over the WebSocket. class WSIncomingMessage(BaseModel): """A message received from a client.""" diff --git a/settings.yml b/settings.yml index 47c65ad..e4052a4 100644 --- a/settings.yml +++ b/settings.yml @@ -60,8 +60,11 @@ audio: # Configures the AI's long-term memory and how game sessions are saved. # ------------------------------------------------------------------------------ memory: + # The chunking strategy to use for breaking down text for the RAG system. + # Options: "simple" (original method), "chonkie" (advanced library). + chunker: "simple" + # The SentenceTransformer model used to create embeddings for the AI's - # long-term memory (RAG). This runs locally. embedding_model: "google/embeddinggemma-300m" # UPDATED: Split database paths for better organization.