banko-ai-assistant 1.0.0__py3-none-any.whl

banko_ai/utils/database.py

@@ -0,0 +1,202 @@
+"""
+Database management utilities.
+
+This module provides database schema creation and management functionality.
+"""
+
+import os
+from typing import Optional
+
+
+class DatabaseManager:
+    """Database management utilities."""
+    
+    def __init__(self, database_url: Optional[str] = None):
+        """Initialize database manager."""
+        self.database_url = database_url or os.getenv('DATABASE_URL', "cockroachdb://root@localhost:26257/banko_ai?sslmode=disable")
+        self._engine = None
+    
+    @property
+    def engine(self):
+        """Get SQLAlchemy engine (lazy import)."""
+        if self._engine is None:
+            from sqlalchemy import create_engine
+            from sqlalchemy.dialects.postgresql.base import PGDialect
+            
+            # Monkey patch version parsing to handle CockroachDB
+            original_get_server_version_info = PGDialect._get_server_version_info
+            
+            def patched_get_server_version_info(self, connection):
+                try:
+                    return original_get_server_version_info(self, connection)
+                except Exception:
+                    return (25, 3, 0)  # Return compatible version tuple
+            
+            PGDialect._get_server_version_info = patched_get_server_version_info
+            
+            # Convert cockroachdb:// to postgresql:// for SQLAlchemy compatibility
+            database_url = self.database_url.replace("cockroachdb://", "postgresql://")
+            
+            self._engine = create_engine(
+                database_url,
+                connect_args={
+                    "options": "-c default_transaction_isolation=serializable"
+                },
+                pool_pre_ping=True,
+                pool_recycle=300
+            )
+        return self._engine
+    
+    def create_tables(self) -> bool:
+        """Create all required tables."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                # Create expenses table with vector support
+                conn.execute(text("""
+                    CREATE TABLE IF NOT EXISTS expenses (
+                        expense_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+                        user_id UUID NOT NULL,
+                        expense_date DATE NOT NULL,
+                        expense_amount DECIMAL(10,2) NOT NULL,
+                        shopping_type STRING NOT NULL,
+                        description STRING,
+                        merchant STRING,
+                        payment_method STRING NOT NULL,
+                        recurring BOOL DEFAULT false,
+                        tags STRING[],
+                        embedding VECTOR(384),
+                        created_at TIMESTAMP DEFAULT now()
+                    )
+                """))
+                
+                # Create vector index for general search
+                conn.execute(text("""
+                    CREATE VECTOR INDEX IF NOT EXISTS idx_expenses_embedding 
+                    ON expenses (embedding)
+                """))
+                
+                # Create user-specific vector index
+                conn.execute(text("""
+                    CREATE VECTOR INDEX IF NOT EXISTS idx_expenses_user_embedding 
+                    ON expenses (user_id, embedding)
+                """))
+                
+                # Create additional indexes for common queries
+                conn.execute(text("""
+                    CREATE INDEX IF NOT EXISTS idx_expenses_user_date 
+                    ON expenses (user_id, expense_date DESC)
+                """))
+                
+                conn.execute(text("""
+                    CREATE INDEX IF NOT EXISTS idx_expenses_merchant 
+                    ON expenses (merchant)
+                """))
+                
+                conn.execute(text("""
+                    CREATE INDEX IF NOT EXISTS idx_expenses_shopping_type 
+                    ON expenses (shopping_type)
+                """))
+                
+                conn.commit()
+                return True
+                
+        except Exception as e:
+            print(f"Error creating tables: {e}")
+            return False
+    
+    def drop_tables(self) -> bool:
+        """Drop all tables."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                conn.execute(text("DROP TABLE IF EXISTS expenses CASCADE"))
+                conn.commit()
+                return True
+        except Exception as e:
+            print(f"Error dropping tables: {e}")
+            return False
+    
+    def table_exists(self, table_name: str) -> bool:
+        """Check if a table exists."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                result = conn.execute(text("""
+                    SELECT EXISTS (
+                        SELECT FROM information_schema.tables 
+                        WHERE table_name = :table_name
+                    )
+                """), {"table_name": table_name})
+                return result.scalar()
+        except Exception as e:
+            print(f"Error checking table existence: {e}")
+            return False
+    
+    def get_table_info(self, table_name: str) -> dict:
+        """Get table information."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                # Get column information
+                columns_result = conn.execute(text("""
+                    SELECT column_name, data_type, is_nullable, column_default
+                    FROM information_schema.columns
+                    WHERE table_name = :table_name
+                    ORDER BY ordinal_position
+                """), {"table_name": table_name})
+                
+                columns = []
+                for row in columns_result:
+                    columns.append({
+                        "name": row[0],
+                        "type": row[1],
+                        "nullable": row[2] == "YES",
+                        "default": row[3]
+                    })
+                
+                # Get index information
+                indexes_result = conn.execute(text("""
+                    SELECT indexname, indexdef
+                    FROM pg_indexes
+                    WHERE tablename = :table_name
+                """), {"table_name": table_name})
+                
+                indexes = []
+                for row in indexes_result:
+                    indexes.append({
+                        "name": row[0],
+                        "definition": row[1]
+                    })
+                
+                return {
+                    "table_name": table_name,
+                    "columns": columns,
+                    "indexes": indexes
+                }
+                
+        except Exception as e:
+            print(f"Error getting table info: {e}")
+            return {}
+    
+    def get_record_count(self, table_name: str) -> int:
+        """Get record count for a table."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                result = conn.execute(text(f"SELECT COUNT(*) FROM {table_name}"))
+                return result.scalar()
+        except Exception as e:
+            print(f"Error getting record count: {e}")
+            return 0
+    
+    def test_connection(self) -> bool:
+        """Test database connection."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                conn.execute(text("SELECT 1"))
+                return True
+        except Exception as e:
+            print(f"Database connection failed: {e}")
+            return False

banko_ai/utils/migration.py

@@ -0,0 +1,123 @@
+"""
+Database migration utilities.
+
+This module provides migration scripts to update the database schema
+for user-specific vector indexing and other enhancements.
+"""
+
+import os
+from typing import Optional
+
+
+class DatabaseMigration:
+    """Database migration utilities."""
+    
+    def __init__(self, database_url: Optional[str] = None):
+        """Initialize migration manager."""
+        self.database_url = database_url or os.getenv('DATABASE_URL', "cockroachdb://root@localhost:26257/defaultdb?sslmode=disable")
+        self._engine = None
+    
+    @property
+    def engine(self):
+        """Get SQLAlchemy engine (lazy import)."""
+        if self._engine is None:
+            from sqlalchemy import create_engine
+            self._engine = create_engine(self.database_url)
+        return self._engine
+    
+    def migrate_to_user_specific_indexing(self) -> bool:
+        """Migrate database to support user-specific vector indexing."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                # Check if user_id column exists
+                result = conn.execute(text("""
+                    SELECT column_name 
+                    FROM information_schema.columns 
+                    WHERE table_name = 'expenses' AND column_name = 'user_id'
+                """))
+                
+                if not result.fetchone():
+                    # Add user_id column if it doesn't exist
+                    conn.execute(text("""
+                        ALTER TABLE expenses 
+                        ADD COLUMN user_id UUID DEFAULT gen_random_uuid()
+                    """))
+                    print("Added user_id column to expenses table")
+                
+                # Create user-specific vector index
+                conn.execute(text("""
+                    CREATE INDEX IF NOT EXISTS idx_expenses_user_embedding 
+                    ON expenses (user_id, embedding) 
+                    USING ivfflat (embedding vector_cosine_ops) 
+                    WITH (lists = 100)
+                """))
+                print("Created user-specific vector index")
+                
+                # Create regional index if supported
+                try:
+                    conn.execute(text("""
+                        CREATE INDEX IF NOT EXISTS idx_expenses_user_embedding_regional 
+                        ON expenses (user_id, embedding) 
+                        LOCALITY REGIONAL BY ROW AS region
+                    """))
+                    print("Created regional user-specific vector index")
+                except Exception as e:
+                    print(f"Regional indexing not supported: {e}")
+                
+                # Create additional indexes for user queries
+                conn.execute(text("""
+                    CREATE INDEX IF NOT EXISTS idx_expenses_user_date 
+                    ON expenses (user_id, expense_date DESC)
+                """))
+                print("Created user date index")
+                
+                conn.commit()
+                return True
+                
+        except Exception as e:
+            print(f"Migration failed: {e}")
+            return False
+    
+    def add_created_at_column(self) -> bool:
+        """Add created_at timestamp column."""
+        try:
+            from sqlalchemy import text
+            with self.engine.connect() as conn:
+                # Check if created_at column exists
+                result = conn.execute(text("""
+                    SELECT column_name 
+                    FROM information_schema.columns 
+                    WHERE table_name = 'expenses' AND column_name = 'created_at'
+                """))
+                
+                if not result.fetchone():
+                    conn.execute(text("""
+                        ALTER TABLE expenses 
+                        ADD COLUMN created_at TIMESTAMP DEFAULT now()
+                    """))
+                    print("Added created_at column to expenses table")
+                    conn.commit()
+                    return True
+                else:
+                    print("created_at column already exists")
+                    return True
+                    
+        except Exception as e:
+            print(f"Failed to add created_at column: {e}")
+            return False
+    
+    def run_all_migrations(self) -> bool:
+        """Run all pending migrations."""
+        print("Running database migrations...")
+        
+        success = True
+        success &= self.add_created_at_column()
+        success &= self.migrate_to_user_specific_indexing()
+        
+        if success:
+            print("All migrations completed successfully")
+        else:
+            print("Some migrations failed")
+        
+        return success

banko_ai/vector_search/__init__.py

@@ -0,0 +1,18 @@
+"""Vector search functionality for Banko AI Assistant."""
+
+def get_data_enricher():
+    """Get DataEnricher (lazy import)."""
+    from .enrichment import DataEnricher
+    return DataEnricher
+
+def get_vector_search_engine():
+    """Get VectorSearchEngine (lazy import)."""
+    from .search import VectorSearchEngine
+    return VectorSearchEngine
+
+def get_enhanced_expense_generator():
+    """Get EnhancedExpenseGenerator (lazy import)."""
+    from .generator import EnhancedExpenseGenerator
+    return EnhancedExpenseGenerator
+
+__all__ = ["get_data_enricher", "get_vector_search_engine", "get_enhanced_expense_generator"]

banko_ai/vector_search/enrichment.py

@@ -0,0 +1,278 @@
+"""
+Data enrichment module for improving vector search accuracy.
+
+This module enriches expense descriptions with merchant context and other relevant
+information to improve vector search accuracy and relevance.
+"""
+
+from typing import Dict, Any, Optional
+from datetime import datetime
+import re
+
+
+class DataEnricher:
+    """Enriches expense data with contextual information for better vector search."""
+    
+    def __init__(self):
+        """Initialize the data enricher."""
+        self.merchant_categories = {
+            "grocery": ["Whole Foods Market", "Trader Joe's", "Kroger", "Safeway", "Publix", "Walmart", "Target"],
+            "retail": ["Amazon", "Best Buy", "Apple Store", "Home Depot", "Costco", "Target", "Walmart"],
+            "dining": ["Starbucks", "McDonald's", "Chipotle", "Subway", "Pizza Hut", "Domino's"],
+            "transportation": ["Shell Gas Station", "Exxon", "Uber", "Lyft", "Metro", "Parking"],
+            "healthcare": ["CVS Pharmacy", "Walgreens", "Rite Aid", "Hospital", "Clinic"],
+            "entertainment": ["Netflix", "Spotify", "Movie Theater", "Concert", "Gaming"],
+            "utilities": ["Electric Company", "Internet Provider", "Phone Company", "Water Company"]
+        }
+        
+        self.amount_ranges = {
+            "low": (0, 25),
+            "medium": (25, 100),
+            "high": (100, 500),
+            "very_high": (500, float('inf'))
+        }
+    
+    def enrich_expense_description(
+        self, 
+        description: str, 
+        merchant: str, 
+        amount: float, 
+        category: str,
+        payment_method: str,
+        date: datetime,
+        **kwargs
+    ) -> str:
+        """
+        Enrich expense description with contextual information.
+        
+        Args:
+            description: Original expense description
+            merchant: Merchant name
+            amount: Expense amount
+            category: Expense category
+            payment_method: Payment method used
+            date: Expense date
+            **kwargs: Additional metadata
+            
+        Returns:
+            Enriched description string
+        """
+        # Start with the original description
+        enriched_parts = [description]
+        
+        # Add merchant name and amount prominently
+        enriched_parts.append(f"at {merchant} for ${amount:.2f}")
+        
+        # Add merchant context
+        merchant_context = self._get_merchant_context(merchant, amount)
+        if merchant_context:
+            enriched_parts.append(merchant_context)
+        
+        # Add amount context
+        amount_context = self._get_amount_context(amount)
+        if amount_context:
+            enriched_parts.append(amount_context)
+        
+        # Add category context
+        category_context = self._get_category_context(category, merchant)
+        if category_context:
+            enriched_parts.append(category_context)
+        
+        # Add payment method context
+        payment_context = self._get_payment_context(payment_method)
+        if payment_context:
+            enriched_parts.append(payment_context)
+        
+        # Add temporal context
+        temporal_context = self._get_temporal_context(date)
+        if temporal_context:
+            enriched_parts.append(temporal_context)
+        
+        # Add merchant category context
+        merchant_category = self._get_merchant_category(merchant)
+        if merchant_category:
+            enriched_parts.append(f"at {merchant_category} store")
+        
+        # Combine all parts
+        enriched_description = " ".join(enriched_parts)
+        
+        # Clean up and format
+        enriched_description = self._clean_description(enriched_description)
+        
+        return enriched_description
+    
+    def _get_merchant_context(self, merchant: str, amount: float) -> Optional[str]:
+        """Get merchant-specific context."""
+        merchant_lower = merchant.lower()
+        
+        # Gas stations
+        if any(gas in merchant_lower for gas in ["shell", "exxon", "chevron", "bp", "gas"]):
+            return f"fuel purchase at {merchant}"
+        
+        # Grocery stores
+        if any(grocery in merchant_lower for grocery in ["whole foods", "trader joe", "kroger", "safeway"]):
+            return f"grocery shopping at {merchant}"
+        
+        # Online retailers
+        if merchant_lower == "amazon":
+            return f"online purchase from {merchant}"
+        
+        # Coffee shops
+        if any(coffee in merchant_lower for coffee in ["starbucks", "dunkin", "peet", "coffee"]):
+            return f"coffee and food at {merchant}"
+        
+        # Fast food
+        if any(fast in merchant_lower for fast in ["mcdonald", "burger", "pizza", "chipotle", "subway"]):
+            return f"fast food at {merchant}"
+        
+        # Pharmacies
+        if any(pharmacy in merchant_lower for pharmacy in ["cvs", "walgreens", "rite aid", "pharmacy"]):
+            return f"pharmacy visit at {merchant}"
+        
+        # Home improvement
+        if any(home in merchant_lower for home in ["home depot", "lowes", "ace hardware"]):
+            return f"home improvement at {merchant}"
+        
+        return None
+    
+    def _get_amount_context(self, amount: float) -> Optional[str]:
+        """Get amount-based context."""
+        if amount < 10:
+            return "small purchase"
+        elif amount < 50:
+            return "moderate expense"
+        elif amount < 200:
+            return "significant purchase"
+        elif amount < 500:
+            return "major expense"
+        else:
+            return "large transaction"
+    
+    def _get_category_context(self, category: str, merchant: str) -> Optional[str]:
+        """Get category-specific context."""
+        category_lower = category.lower()
+        
+        if category_lower == "groceries":
+            return "food and household items"
+        elif category_lower == "transportation":
+            return "travel and commuting"
+        elif category_lower == "dining":
+            return "restaurant and food service"
+        elif category_lower == "entertainment":
+            return "leisure and recreation"
+        elif category_lower == "healthcare":
+            return "medical and wellness"
+        elif category_lower == "shopping":
+            return "retail and consumer goods"
+        elif category_lower == "utilities":
+            return "essential services"
+        
+        return None
+    
+    def _get_payment_context(self, payment_method: str) -> Optional[str]:
+        """Get payment method context."""
+        payment_lower = payment_method.lower()
+        
+        if "credit" in payment_lower:
+            return "paid with credit card"
+        elif "debit" in payment_lower:
+            return "paid with debit card"
+        elif "cash" in payment_lower:
+            return "paid with cash"
+        elif "mobile" in payment_lower:
+            return "paid with mobile payment"
+        elif "bank" in payment_lower:
+            return "paid via bank transfer"
+        
+        return None
+    
+    def _get_temporal_context(self, date: datetime) -> Optional[str]:
+        """Get temporal context based on date."""
+        from datetime import date as date_type
+        # Convert both to datetime for consistent comparison
+        if isinstance(date, date_type):
+            date = datetime.combine(date, datetime.min.time())
+        now = datetime.now()
+        days_ago = (now - date).days
+        
+        if days_ago == 0:
+            return "today"
+        elif days_ago == 1:
+            return "yesterday"
+        elif days_ago <= 7:
+            return "this week"
+        elif days_ago <= 30:
+            return "this month"
+        elif days_ago <= 90:
+            return "recently"
+        else:
+            return "in the past"
+    
+    def _get_merchant_category(self, merchant: str) -> Optional[str]:
+        """Get merchant category for additional context."""
+        merchant_lower = merchant.lower()
+        
+        for category, merchants in self.merchant_categories.items():
+            if any(m in merchant_lower for m in merchants):
+                return category
+        
+        return None
+    
+    def _clean_description(self, description: str) -> str:
+        """Clean and format the enriched description."""
+        # Remove extra spaces
+        description = re.sub(r'\s+', ' ', description)
+        
+        # Remove duplicate words
+        words = description.split()
+        seen = set()
+        unique_words = []
+        for word in words:
+            if word.lower() not in seen:
+                unique_words.append(word)
+                seen.add(word.lower())
+        
+        return ' '.join(unique_words).strip()
+    
+    def create_searchable_text(
+        self, 
+        description: str, 
+        merchant: str, 
+        amount: float, 
+        category: str,
+        **kwargs
+    ) -> str:
+        """
+        Create a comprehensive searchable text block for embedding.
+        
+        This creates a rich text representation that includes all relevant
+        information for vector search.
+        """
+        # Extract required parameters from kwargs
+        payment_method = kwargs.get('payment_method', '')
+        date = kwargs.get('date', datetime.now())
+        
+        # Remove these from kwargs to avoid conflicts
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ['payment_method', 'date']}
+        
+        enriched_description = self.enrich_expense_description(
+            description, merchant, amount, category, payment_method, date, **filtered_kwargs
+        )
+        
+        # Create a comprehensive searchable text
+        searchable_parts = [
+            f"Spent ${amount:.2f} on {enriched_description}",
+            f"Merchant: {merchant}",
+            f"Category: {category}",
+            f"Amount: ${amount:.2f}"
+        ]
+        
+        # Add any additional context
+        if kwargs.get('payment_method'):
+            searchable_parts.append(f"Payment: {kwargs['payment_method']}")
+        
+        if kwargs.get('tags'):
+            tags = ', '.join(kwargs['tags'])
+            searchable_parts.append(f"Tags: {tags}")
+        
+        return " | ".join(searchable_parts)