Spaces:
Sleeping
Sleeping
File size: 2,596 Bytes
46bd2ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | from __future__ import annotations
import os
from functools import lru_cache
import chromadb
import yaml
from chromadb.utils import embedding_functions
from config import CONFIG
def load_profile() -> dict:
"""Load the profile from disk, or {} when there is none."""
if not os.path.exists(CONFIG.profile_path):
return {}
with open(CONFIG.profile_path, "r", encoding="utf-8") as fh:
return yaml.safe_load(fh) or {}
def parse_profile(text: str) -> dict:
"""Parse a YAML profile typed in the UI, tolerating empty or broken input."""
if not text or not text.strip():
return {}
try:
data = yaml.safe_load(text)
except yaml.YAMLError:
return {}
return data if isinstance(data, dict) else {}
def profile_to_prompt(profile: dict) -> str:
"""Render the profile into a system-prompt block, in the active language."""
if not profile:
return ""
lines = [CONFIG.pack["profile_header"]]
for key, value in profile.items():
if isinstance(value, list):
value = ", ".join(str(v) for v in value)
lines.append(f"- {key}: {value}")
return "\n".join(lines)
@lru_cache(maxsize=1)
def _collection():
"""Open the persistent Chroma collection once and cache it."""
client = chromadb.PersistentClient(path=CONFIG.chroma_dir)
embed = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name=CONFIG.embed_model
)
return client.get_or_create_collection(
name="his_life", embedding_function=embed
)
def warmup() -> None:
"""Load the embedding model before the first turn so it isn't cold."""
_collection()
def remember(text: str, source: str = "conversation") -> None:
"""Store a memory: a story, a fact, or a moment from a chat."""
text = text.strip()
if not text:
return
col = _collection()
col.add(
documents=[text],
metadatas=[{"source": source}],
ids=[f"{source}-{col.count()}"],
)
def recall(query: str) -> list[str]:
"""Return the stored memories most similar to the query."""
col = _collection()
if col.count() == 0:
return []
res = col.query(query_texts=[query], n_results=CONFIG.rag_top_k)
docs = res.get("documents") or [[]]
return docs[0]
def recall_block(query: str) -> str:
"""Render recalled memories as a prompt block, or '' when there are none."""
hits = recall(query)
if not hits:
return ""
bullets = "\n".join(f"- {h}" for h in hits)
return f"{CONFIG.pack['recall_header']}\n{bullets}"
|