Spaces:
Configuration error
Configuration error
| """ | |
| Vector memory layer for the Decision Debate Agent. | |
| Design choice: we use a local TF-IDF embedding function instead of a | |
| downloaded neural embedding model (e.g. all-MiniLM-L6-v2). This is a | |
| deliberate engineering decision, not a shortcut: | |
| 1. Reproducibility: a grader running this on a clean machine should not | |
| depend on being able to reach an external model host at eval time. | |
| 2. No heavy ML runtime dependency (torch) needed just to embed short | |
| decision-query text. | |
| 3. TF-IDF is genuinely adequate here: our corpus is short, keyword-rich | |
| decision statements (e.g. "should I take the higher paying job or | |
| stay for growth"), where lexical overlap is a strong relevance signal. | |
| Trade-off (documented honestly for the changelog): TF-IDF will not | |
| capture pure semantic paraphrase the way a neural embedding would. | |
| This is called out as a known limitation, not hidden. | |
| """ | |
| import os | |
| import pickle | |
| from typing import List, Dict, Any | |
| import chromadb | |
| from chromadb import EmbeddingFunction, Documents, Embeddings | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| import numpy as np | |
| class TfidfEmbeddingFunction(EmbeddingFunction): | |
| """A local, dependency-light embedding function for Chroma. | |
| Fits a TF-IDF vectorizer over whatever text has been seen so far. | |
| Vocabulary grows as more decisions are stored. Vectors are padded / | |
| truncated to a fixed dimensionality so Chroma's index stays stable. | |
| """ | |
| def __init__(self, vocab_path: str, max_features: int = 512): | |
| self.vocab_path = vocab_path | |
| self.max_features = max_features | |
| self.vectorizer = TfidfVectorizer(max_features=max_features) | |
| self._fitted = False | |
| self._load_if_exists() | |
| def _load_if_exists(self): | |
| if os.path.exists(self.vocab_path): | |
| with open(self.vocab_path, "rb") as f: | |
| self.vectorizer = pickle.load(f) | |
| self._fitted = True | |
| def _save(self): | |
| with open(self.vocab_path, "wb") as f: | |
| pickle.dump(self.vectorizer, f) | |
| def _refit(self, all_texts: List[str]): | |
| self.vectorizer = TfidfVectorizer(max_features=self.max_features) | |
| self.vectorizer.fit(all_texts) | |
| self._fitted = True | |
| self._save() | |
| def __call__(self, input: Documents) -> Embeddings: | |
| if not self._fitted: | |
| self._refit(list(input)) | |
| try: | |
| vecs = self.vectorizer.transform(input).toarray() | |
| except Exception: | |
| self._refit(list(input)) | |
| vecs = self.vectorizer.transform(input).toarray() | |
| # pad to fixed width so Chroma's collection dimensionality never changes | |
| width = self.max_features | |
| if vecs.shape[1] < width: | |
| pad = np.zeros((vecs.shape[0], width - vecs.shape[1])) | |
| vecs = np.hstack([vecs, pad]) | |
| return vecs.tolist() | |
| class DecisionMemory: | |
| """Stores past decision debates and retrieves similar ones.""" | |
| def __init__(self, persist_dir: str = "./chroma_db"): | |
| os.makedirs(persist_dir, exist_ok=True) | |
| self.client = chromadb.PersistentClient(path=persist_dir) | |
| self.embed_fn = TfidfEmbeddingFunction( | |
| vocab_path=os.path.join(persist_dir, "tfidf_vocab.pkl") | |
| ) | |
| self.collection = self.client.get_or_create_collection( | |
| name="decisions", | |
| embedding_function=self.embed_fn, | |
| metadata={"hnsw:space": "cosine"}, | |
| ) | |
| def add_decision(self, decision_id: str, query: str, metadata: Dict[str, Any]): | |
| self.collection.add( | |
| documents=[query], | |
| ids=[decision_id], | |
| metadatas=[metadata], | |
| ) | |
| def retrieve_similar( | |
| self, | |
| query: str, | |
| n_results: int = 3, | |
| min_similarity: float = 0.35, | |
| ) -> List[Dict[str, Any]]: | |
| """Retrieve past decisions, but only ones that are ACTUALLY similar. | |
| Fix for a discovered failure mode (see README Hot Take): earlier | |
| versions returned the top-N nearest neighbors unconditionally, | |
| even when nothing in the store was genuinely related to the | |
| current query. With a small, unscoped memory store, "nearest" | |
| can still mean "barely related" -- e.g. an unrelated car-purchase | |
| query being retrieved for a career-switch question just because | |
| both contain common decision-framing words. That produced | |
| outputs referencing information the user never gave. | |
| min_similarity is a cosine-similarity floor (0-1). Chroma returns | |
| cosine *distance* (1 - similarity) when the collection is | |
| configured with hnsw:space="cosine", so we filter on | |
| (1 - distance) >= min_similarity. Nothing below the floor is | |
| returned, even if it's the "closest" thing in the store -- | |
| empty results are the correct, honest answer when there's | |
| nothing truly similar yet. | |
| """ | |
| count = self.collection.count() | |
| if count == 0: | |
| return [] | |
| n_results = min(n_results, count) | |
| try: | |
| results = self.collection.query( | |
| query_texts=[query], | |
| n_results=n_results, | |
| include=["documents", "metadatas", "distances"], | |
| ) | |
| except Exception: | |
| return [] | |
| out = [] | |
| docs = results.get("documents", [[]])[0] | |
| metas = results.get("metadatas", [[]])[0] | |
| distances = results.get("distances", [[]])[0] | |
| for doc, meta, dist in zip(docs, metas, distances): | |
| similarity = 1.0 - dist | |
| if similarity >= min_similarity: | |
| out.append({"query": doc, "metadata": meta, "similarity": round(similarity, 3)}) | |
| return out |