vector_store.py

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