vector_memory.py

"""
Vector Memory System for Chrome Extension Idea Generator
Uses Weaviate as a vector database to store and retrieve extension ideas and insights
"""

import os
import json
import uuid
import time
from datetime import datetime
from typing import List, Dict, Any, Optional, Union, Tuple
from tqdm import tqdm

import weaviate
import weaviate.classes as wvc
from weaviate.classes.config import ConsistencyLevel, Property, DataType, Tokenization
import openai
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

class VectorMemory:
    """
    Vector Memory System for storing and retrieving Chrome extension ideas and insights
    Uses Weaviate as a vector database and OpenAI for embeddings
    """
    
    def __init__(self, 
                 weaviate_url: Optional[str] = None, 
                 api_key: Optional[str] = None,
                 openai_api_key: Optional[str] = None,
                 schema_name: str = "ChromeExtensionIdea"):
        """
        Initialize the Vector Memory system
        
        Args:
            weaviate_url: URL for the Weaviate instance (default: use env var WEAVIATE_URL)
            api_key: API key for Weaviate (default: use env var WEAVIATE_API_KEY)
            openai_api_key: API key for OpenAI (default: use env var OPENAI_API_KEY)
            schema_name: Name of the schema class to use in Weaviate
        """
        # Use environment variables if not provided
        self.weaviate_url = weaviate_url or os.getenv("WEAVIATE_URL")
        self.api_key = api_key or os.getenv("WEAVIATE_API_KEY")
        self.openai_api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
        self.schema_name = schema_name
        
        # Set up OpenAI client
        self.openai_client = openai.OpenAI(api_key=self.openai_api_key)
        
        # Connect to Weaviate
        try:
            if self.weaviate_url and self.api_key:
                self.client = self._setup_weaviate_client()
                # Check if schema exists, create if not
                self._ensure_schema_exists()
                self.is_available = True
                print(f"Connected to Weaviate at {self.weaviate_url}")
            else:
                print("Weaviate URL or API key not provided. Vector memory will be unavailable.")
                self.is_available = False
                self.client = None
        except Exception as e:
            print(f"Failed to connect to Weaviate: {str(e)}")
            self.is_available = False
            self.client = None
    
    def _setup_weaviate_client(self):
        """Set up and return a Weaviate client with the configured credentials"""
        auth_config = weaviate.AuthApiKey(api_key=self.api_key)
        client = weaviate.Client(
            url=self.weaviate_url,
            auth_client_secret=auth_config,
            additional_headers={
                "X-OpenAI-Api-Key": self.openai_api_key
            }
        )
        return client
    
    def _ensure_schema_exists(self):
        """Check if the schema exists, create it if not"""
        if not self.client:
            return False
        
        # Check if class exists
        try:
            schema = self.client.schema.get()
            class_names = [c['class'] for c in schema['classes']] if 'classes' in schema else []
            
            if self.schema_name not in class_names:
                # Create the class
                print(f"Creating schema class '{self.schema_name}'...")
                
                # Define properties
                properties = [
                    Property(name="name", data_type=DataType.TEXT, tokenization=Tokenization.FIELD),
                    Property(name="description", data_type=DataType.TEXT, tokenization=Tokenization.FIELD),
                    Property(name="problem", data_type=DataType.TEXT, tokenization=Tokenization.FIELD),
                    Property(name="features", data_type=DataType.TEXT_ARRAY),
                    Property(name="keywords", data_type=DataType.TEXT_ARRAY),
                    Property(name="audience", data_type=DataType.TEXT_ARRAY),
                    Property(name="competition", data_type=DataType.TEXT),
                    Property(name="monetization", data_type=DataType.TEXT_ARRAY),
                    Property(name="originalityScore", data_type=DataType.NUMBER),
                    Property(name="trendScore", data_type=DataType.NUMBER),
                    Property(name="isHybrid", data_type=DataType.BOOLEAN),
                    Property(name="notionPageId", data_type=DataType.TEXT),
                    Property(name="createdAt", data_type=DataType.DATE),
                    Property(name="source", data_type=DataType.TEXT),
                    Property(name="insights", data_type=DataType.TEXT),
                ]
                
                # Create class configuration
                config = wvc.Configure.new()
                
                # Set up vectorizer for OpenAI embeddings
                config = config.with_vectorizer("text2vec-openai", {
                    "model": "ada",
                    "modelVersion": "002",
                    "type": "text"
                })
                
                # Create class with properties
                self.client.collections.create(
                    name=self.schema_name, 
                    properties=properties,
                    vectorizer_config=config,
                    consistency_level=ConsistencyLevel.QUORUM
                )
                
                print(f"Schema class '{self.schema_name}' created successfully.")
                return True
            else:
                print(f"Schema class '{self.schema_name}' already exists.")
                return True
        except Exception as e:
            print(f"Error ensuring schema exists: {str(e)}")
            return False
    
    def store_idea(self, idea: Dict[str, Any]) -> Optional[str]:
        """
        Store a Chrome extension idea in the vector database
        
        Args:
            idea: Dictionary containing the idea data
            
        Returns:
            UUID of the stored idea, or None if storage failed
        """
        if not self.is_available or not self.client:
            print("Vector memory not available. Idea not stored.")
            return None
        
        try:
            # Generate a UUID for the idea
            idea_id = str(uuid.uuid4())
            
            # Prepare data for storage
            data_obj = {
                "name": idea.get("name", "Untitled Idea"),
                "description": idea.get("description", ""),
                "problem": idea.get("problem", ""),
                "features": idea.get("features", []),
                "keywords": idea.get("keywords", []),
                "audience": idea.get("suggestedAudience", []),
                "competition": idea.get("competition", ""),
                "monetization": idea.get("monetization", []),
                "originalityScore": idea.get("originalityScore", 0),
                "trendScore": idea.get("trendScore", 0),
                "isHybrid": idea.get("isHybrid", False),
                "notionPageId": idea.get("notionPageId", ""),
                "createdAt": datetime.now().isoformat(),
                "source": idea.get("source", "AI Generator"),
                "insights": idea.get("insights", "")
            }
            
            # Store in Weaviate
            with self.client.batch.dynamic() as batch:
                batch.add_data_object(
                    data_obj,
                    self.schema_name,
                    idea_id
                )
            
            print(f"Stored idea '{data_obj['name']}' with ID {idea_id}")
            return idea_id
        except Exception as e:
            print(f"Error storing idea: {str(e)}")
            return None
    
    def store_multiple_ideas(self, ideas: List[Dict[str, Any]]) -> List[str]:
        """
        Store multiple Chrome extension ideas in the vector database
        
        Args:
            ideas: List of idea dictionaries to store
            
        Returns:
            List of UUIDs for the stored ideas
        """
        if not self.is_available or not self.client:
            print("Vector memory not available. Ideas not stored.")
            return []
        
        stored_ids = []
        try:
            with self.client.batch.dynamic() as batch:
                for idea in tqdm(ideas, desc="Storing ideas"):
                    # Generate a UUID for the idea
                    idea_id = str(uuid.uuid4())
                    
                    # Prepare data for storage
                    data_obj = {
                        "name": idea.get("name", "Untitled Idea"),
                        "description": idea.get("description", ""),
                        "problem": idea.get("problem", ""),
                        "features": idea.get("features", []),
                        "keywords": idea.get("keywords", []),
                        "audience": idea.get("suggestedAudience", []),
                        "competition": idea.get("competition", ""),
                        "monetization": idea.get("monetization", []),
                        "originalityScore": idea.get("originalityScore", 0),
                        "trendScore": idea.get("trendScore", 0),
                        "isHybrid": idea.get("isHybrid", False),
                        "notionPageId": idea.get("notionPageId", ""),
                        "createdAt": datetime.now().isoformat(),
                        "source": idea.get("source", "AI Generator"),
                        "insights": idea.get("insights", "")
                    }
                    
                    # Store in Weaviate
                    batch.add_data_object(
                        data_obj,
                        self.schema_name,
                        idea_id
                    )
                    stored_ids.append(idea_id)
            
            print(f"Stored {len(stored_ids)} ideas in vector database")
            return stored_ids
        except Exception as e:
            print(f"Error storing multiple ideas: {str(e)}")
            return stored_ids
    
    def search_similar_ideas(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
        """
        Search for ideas similar to the query text
        
        Args:
            query: Text query to search for
            limit: Maximum number of results to return
            
        Returns:
            List of similar ideas
        """
        if not self.is_available or not self.client:
            print("Vector memory not available. Cannot search.")
            return []
        
        try:
            # Get collection
            collection = self.client.collections.get(self.schema_name)
            
            # Perform vector search
            result = collection.query.near_text(
                query=query,
                limit=limit
            ).do()
            
            # Extract and return the objects
            if "data" in result and "Get" in result["data"]:
                ideas = result["data"]["Get"].get(self.schema_name, [])
                return ideas
            return []
        except Exception as e:
            print(f"Error searching similar ideas: {str(e)}")
            return []
    
    def search_by_keywords(self, keywords: List[str], limit: int = 5) -> List[Dict[str, Any]]:
        """
        Search for ideas by keywords
        
        Args:
            keywords: List of keywords to search for
            limit: Maximum number of results to return
            
        Returns:
            List of matching ideas
        """
        if not self.is_available or not self.client:
            print("Vector memory not available. Cannot search.")
            return []
        
        try:
            # Get collection
            collection = self.client.collections.get(self.schema_name)
            
            # Prepare where filter for keywords
            where_filter = {
                "path": ["keywords"],
                "operator": "ContainsAny",
                "valueStringArray": keywords
            }
            
            # Perform search
            result = collection.query.get(
                limit=limit
            ).with_where(
                where_filter
            ).do()
            
            # Extract and return the objects
            if "data" in result and "Get" in result["data"]:
                ideas = result["data"]["Get"].get(self.schema_name, [])
                return ideas
            return []
        except Exception as e:
            print(f"Error searching by keywords: {str(e)}")
            return []
    
    def get_context_for_prompt(self, query: str, max_ideas: int = 3) -> str:
        """
        Get a formatted context string from similar ideas for use in AI prompts
        
        Args:
            query: The query to find similar ideas for
            max_ideas: Maximum number of ideas to include
            
        Returns:
            Formatted context string
        """
        similar_ideas = self.search_similar_ideas(query, limit=max_ideas)
        
        if not similar_ideas:
            return "No similar ideas found."
        
        # Format the ideas as context
        context = "Related ideas from memory:\n\n"
        
        for i, idea in enumerate(similar_ideas):
            context += f"IDEA {i+1}: {idea.get('name', 'Untitled')}\n"
            context += f"Problem: {idea.get('problem', 'N/A')}\n"
            
            if idea.get('features'):
                context += "Features:\n"
                for feat in idea.get('features', [])[:3]:  # Limit to 3 features
                    context += f"- {feat}\n"
            
            context += f"Keywords: {', '.join(idea.get('keywords', []))}\n"
            
            if idea.get('insights'):
                context += f"Insights: {idea.get('insights')}\n"
            
            context += "\n"
        
        return context
    
    def get_stats(self) -> Dict[str, Any]:
        """
        Get statistics about the vector memory
        
        Returns:
            Dictionary of statistics
        """
        if not self.is_available or not self.client:
            return {"available": False}
        
        try:
            # Get collection
            collection = self.client.collections.get(self.schema_name)
            
            # Get total count
            try:
                count_result = collection.query.aggregate().with_meta_count().do()
                total_count = count_result["data"]["Aggregate"][self.schema_name][0]["meta"]["count"]
            except:
                total_count = 0
            
            # Get hybrid count
            try:
                hybrid_filter = {
                    "path": ["isHybrid"],
                    "operator": "Equal",
                    "valueBoolean": True
                }
                hybrid_result = collection.query.aggregate().with_where(hybrid_filter).with_meta_count().do()
                hybrid_count = hybrid_result["data"]["Aggregate"][self.schema_name][0]["meta"]["count"]
            except:
                hybrid_count = 0
            
            # Get high originality count
            try:
                originality_filter = {
                    "path": ["originalityScore"],
                    "operator": "GreaterThan",
                    "valueNumber": 7.0
                }
                originality_result = collection.query.aggregate().with_where(originality_filter).with_meta_count().do()
                high_originality_count = originality_result["data"]["Aggregate"][self.schema_name][0]["meta"]["count"]
            except:
                high_originality_count = 0
            
            return {
                "available": True,
                "totalIdeas": total_count,
                "hybridIdeas": hybrid_count,
                "highOriginalityIdeas": high_originality_count
            }
        except Exception as e:
            print(f"Error getting stats: {str(e)}")
            return {"available": False, "error": str(e)}
    
    def add_insight(self, idea_id: str, insight: str) -> bool:
        """
        Add an insight to an existing idea
        
        Args:
            idea_id: UUID of the idea
            insight: Insight text to add
            
        Returns:
            True if successful, False otherwise
        """
        if not self.is_available or not self.client:
            print("Vector memory not available. Cannot add insight.")
            return False
        
        try:
            # Get collection
            collection = self.client.collections.get(self.schema_name)
            
            # Update the insight
            collection.data.update(
                uuid=idea_id,
                properties={
                    "insights": insight
                }
            )
            
            return True
        except Exception as e:
            print(f"Error adding insight: {str(e)}")
            return False

# Singleton instance
_vector_memory_instance = None

def get_vector_memory() -> VectorMemory:
    """
    Get the singleton instance of VectorMemory
    
    Returns:
        The VectorMemory instance
    """
    global _vector_memory_instance
    if _vector_memory_instance is None:
        _vector_memory_instance = VectorMemory()
    return _vector_memory_instance