Nighthawk/VDM
1
0
Fork 0
mirror of https://github.com/Nighthawk42/VDM.git synced 2026-08-30 04:30:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added Chonkie

Browse Source
This commit is contained in:
Nighthawk
2025-09-05 03:14:24 -04:00
parent 716d7121a7
commit af31ae8757
5 changed files with 447 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
requirements.txt
+350 -25
View File
@@ -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
# =========================
# 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
server/config.py
+3 -3
View File
@@ -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
server/memory_manager.py
+79 -111
View File
@@ -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:
server/models.py
+11 -2
View File
@@ -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."""
settings.yml
+4 -1
View File
@@ -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.