File size: 2,637 Bytes
0eb4b73 | 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 | import chromadb
from chromadb.utils import embedding_functions
import uuid
from typing import List, Dict
class InsightVectorStore:
def __init__(self, persist_dir="./chroma_db"):
self.client = chromadb.PersistentClient(path=persist_dir)
self.embed_fn = embedding_functions.DefaultEmbeddingFunction()
self.collection = self.client.get_or_create_collection(
name="insights",
embedding_function=self.embed_fn
)
def add_insight(self, insight_doc: Dict) -> str:
doc_id = str(uuid.uuid4())
text = f"{insight_doc['subreddit']} {insight_doc['post_title']} {insight_doc['insight']['burning_pain']} {insight_doc['insight']['startup_idea']}"
metadata = {
"subreddit": insight_doc['subreddit'],
"post_title": insight_doc['post_title'],
"post_url": insight_doc['post_url'],
"timestamp": insight_doc['timestamp'],
"confidence_score": insight_doc['insight']['confidence_score'],
"problem_severity": insight_doc['insight']['problem_severity_score'],
"market_timing": insight_doc['insight']['market_timing_score'],
"pain_intensity": insight_doc['insight'].get('pain_intensity', 0),
"founder_persona": insight_doc['insight'].get('founder_persona', 'unknown'),
"investor_signal": insight_doc['insight'].get('investor_signal', False),
"early_startup_signal": insight_doc['insight'].get('early_startup_signal', False)
}
self.collection.upsert(
ids=[doc_id],
documents=[text],
metadatas=[metadata]
)
return doc_id
def query(self, query_text: str, n_results: int = 5) -> List[Dict]:
results = self.collection.query(query_texts=[query_text], n_results=n_results)
docs = []
if results['ids'] and results['ids'][0]:
for i, doc_id in enumerate(results['ids'][0]):
docs.append({
"id": doc_id,
"text": results['documents'][0][i],
"metadata": results['metadatas'][0][i],
"distance": results['distances'][0][i] if results.get('distances') else None
})
return docs
def get_all_insights(self) -> List[Dict]:
# Chroma doesn't have a direct get all, but we can query with empty text
# However, that's not efficient. We'll keep a separate list or use the collection.get()
# This method is not implemented; we'll use a separate cache.
pass |