| 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]:
|
|
|
|
|
|
|
| pass |
