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

banko_ai/__init__.py

@@ -0,0 +1,19 @@
+"""
+Banko AI Assistant - AI-powered expense analysis and RAG system.
+
+A modern Python package for AI-powered expense analysis using CockroachDB vector search
+and multi-provider AI support (OpenAI, AWS Bedrock, IBM Watsonx, Google Gemini).
+"""
+
+__version__ = "1.0.0"
+__author__ = "Virag Tripathi"
+__email__ = "virag.tripathi@gmail.com"
+
+from .config.settings import Config
+
+def create_app():
+    """Create Flask application (lazy import)."""
+    from .web.app import create_app as _create_app
+    return _create_app()
+
+__all__ = ["Config", "create_app", "__version__"]

banko_ai/__main__.py

@@ -0,0 +1,10 @@
+"""
+Main entry point for Banko AI Assistant.
+
+This module allows the package to be run as a module: python -m banko_ai
+"""
+
+from .cli import main
+
+if __name__ == '__main__':
+    main()

banko_ai/ai_providers/__init__.py

@@ -0,0 +1,18 @@
+"""AI provider implementations for Banko AI Assistant."""
+
+from .base import AIProvider, AIProviderError
+from .openai_provider import OpenAIProvider
+from .aws_provider import AWSProvider
+from .watsonx_provider import WatsonxProvider
+from .gemini_provider import GeminiProvider
+from .factory import AIProviderFactory
+
+__all__ = [
+    "AIProvider",
+    "AIProviderError", 
+    "OpenAIProvider",
+    "AWSProvider",
+    "WatsonxProvider",
+    "GeminiProvider",
+    "AIProviderFactory"
+]

banko_ai/ai_providers/aws_provider.py

@@ -0,0 +1,337 @@
+"""
+AWS Bedrock AI provider implementation.
+
+This module provides AWS Bedrock integration for vector search and RAG responses.
+"""
+
+import os
+import json
+from typing import List, Dict, Any, Optional
+import boto3
+from sentence_transformers import SentenceTransformer
+from sqlalchemy import create_engine, text
+
+from .base import AIProvider, SearchResult, RAGResponse, AIConnectionError, AIAuthenticationError
+
+
+class AWSProvider(AIProvider):
+    """AWS Bedrock AI provider implementation."""
+    
+    def __init__(self, config: Dict[str, Any], cache_manager=None):
+        """Initialize AWS provider."""
+        self.access_key_id = config.get("access_key_id")
+        self.secret_access_key = config.get("secret_access_key")
+        self.region = config.get("region", "us-east-1")
+        self.bedrock_client = None
+        self.embedding_model = None
+        self.db_engine = None
+        self.cache_manager = cache_manager
+        super().__init__(config)
+    
+    def _validate_config(self) -> None:
+        """Validate AWS configuration."""
+        if not self.access_key_id or not self.secret_access_key:
+            raise AIAuthenticationError("AWS access key ID and secret access key are required")
+        
+        # Initialize Bedrock client
+        try:
+            self.bedrock_client = boto3.client(
+                'bedrock-runtime',
+                aws_access_key_id=self.access_key_id,
+                aws_secret_access_key=self.secret_access_key,
+                region_name=self.region
+            )
+        except Exception as e:
+            raise AIConnectionError(f"Failed to initialize AWS Bedrock client: {str(e)}")
+    
+    def get_default_model(self) -> str:
+        """Get the default AWS model."""
+        return "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
+    
+    def get_available_models(self) -> List[str]:
+        """Get available AWS models."""
+        return [
+            "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
+            "us.anthropic.claude-3-5-haiku-20241022-v1:0",
+            "us.anthropic.claude-3-opus-20240229-v1:0",
+            "us.anthropic.claude-3-sonnet-20240229-v1:0",
+            "us.anthropic.claude-3-haiku-20240307-v1:0"
+        ]
+    
+    def _get_embedding_model(self) -> SentenceTransformer:
+        """Get or create the embedding model."""
+        if self.embedding_model is None:
+            try:
+                self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+            except Exception as e:
+                raise AIConnectionError(f"Failed to load embedding model: {str(e)}")
+        return self.embedding_model
+    
+    def _get_db_engine(self):
+        """Get database engine."""
+        if self.db_engine is None:
+            database_url = os.getenv("DATABASE_URL", "cockroachdb://root@localhost:26257/defaultdb?sslmode=disable")
+            try:
+                self.db_engine = create_engine(database_url)
+            except Exception as e:
+                raise AIConnectionError(f"Failed to connect to database: {str(e)}")
+        return self.db_engine
+    
+    def search_expenses(
+        self, 
+        query: str, 
+        user_id: Optional[str] = None,
+        limit: int = 10,
+        threshold: float = 0.7
+    ) -> List[SearchResult]:
+        """Search for expenses using vector similarity."""
+        try:
+            # Generate query embedding
+            embedding_model = self._get_embedding_model()
+            query_embedding = embedding_model.encode([query])[0]
+            
+            # Convert to PostgreSQL vector format
+            search_embedding = json.dumps(query_embedding.tolist())
+            
+            # Build SQL query
+            sql = """
+            SELECT 
+                expense_id,
+                user_id,
+                description,
+                merchant,
+                expense_amount,
+                expense_date,
+                1 - (embedding <-> %s) as similarity_score
+            FROM expenses
+            WHERE 1 - (embedding <-> %s) > %s
+            """
+            
+            params = [search_embedding, search_embedding, threshold]
+            
+            if user_id:
+                sql += " AND user_id = %s"
+                params.append(user_id)
+            
+            sql += " ORDER BY similarity_score DESC LIMIT %s"
+            params.append(limit)
+            
+            # Execute query
+            engine = self._get_db_engine()
+            with engine.connect() as conn:
+                result = conn.execute(text(sql), params)
+                rows = result.fetchall()
+            
+            # Convert to SearchResult objects
+            results = []
+            for row in rows:
+                results.append(SearchResult(
+                    expense_id=str(row[0]),
+                    user_id=str(row[1]),
+                    description=row[2] or "",
+                    merchant=row[3] or "",
+                    amount=float(row[4]),
+                    date=str(row[5]),
+                    similarity_score=float(row[6]),
+                    metadata={}
+                ))
+            
+            return results
+            
+        except Exception as e:
+            raise AIConnectionError(f"Search failed: {str(e)}")
+    
+    def generate_rag_response(
+        self, 
+        query: str, 
+        context: List[SearchResult],
+        user_id: Optional[str] = None,
+        language: str = "en"
+    ) -> RAGResponse:
+        """Generate RAG response using AWS Bedrock."""
+        try:
+            print(f"\n🤖 AWS BEDROCK RAG (with caching):")
+            print(f"1. Query: '{query[:60]}...'")
+            
+            # Check for cached response first
+            if self.cache_manager:
+                # Convert SearchResult objects to dict format for cache lookup
+                search_results_dict = []
+                for result in context:
+                    search_results_dict.append({
+                        'expense_id': result.expense_id,
+                        'user_id': result.user_id,
+                        'description': result.description,
+                        'merchant': result.merchant,
+                        'expense_amount': result.amount,
+                        'expense_date': result.date,
+                        'similarity_score': result.similarity_score,
+                        'shopping_type': result.metadata.get('shopping_type'),
+                        'payment_method': result.metadata.get('payment_method'),
+                        'recurring': result.metadata.get('recurring'),
+                        'tags': result.metadata.get('tags')
+                    })
+                
+                cached_response = self.cache_manager.get_cached_response(
+                    query, search_results_dict, "aws"
+                )
+                if cached_response:
+                    print(f"2. ✅ Response cache HIT! Returning cached response")
+                    return RAGResponse(
+                        response=cached_response,
+                        sources=context,
+                        metadata={
+                            'provider': 'aws',
+                            'model': self.get_default_model(),
+                            'user_id': user_id,
+                            'language': language,
+                            'cached': True
+                        }
+                    )
+                print(f"2. ❌ Response cache MISS, generating fresh response")
+            else:
+                print(f"2. No cache manager available, generating fresh response")
+            
+            # Prepare context
+            context_text = self._prepare_context(context)
+            
+            # Prepare the prompt
+            prompt = f"""You are Banko, a financial assistant. Answer based on this expense data:
+
+Q: {query}
+
+Data:
+{context_text}
+
+Provide helpful insights with numbers, markdown formatting, and actionable advice."""
+            
+            # Define input parameters for Claude
+            payload = {
+                "anthropic_version": "bedrock-2023-05-31",
+                "max_tokens": 1000,
+                "top_k": 250,
+                "stop_sequences": [],
+                "temperature": 1,
+                "top_p": 0.999,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "text",
+                                "text": prompt
+                            }
+                        ]
+                    }
+                ]
+            }
+            
+            # Convert to JSON format
+            body = json.dumps(payload)
+            
+            # Use current model
+            model_id = self.current_model
+            
+            # Invoke model
+            response = self.bedrock_client.invoke_model(
+                modelId=model_id,
+                contentType="application/json",
+                accept="application/json",
+                body=body
+            )
+            
+            # Parse response
+            response_body = json.loads(response['body'].read())
+            ai_response = response_body['content'][0]['text']
+            
+            # Cache the response for future similar queries
+            if self.cache_manager and ai_response:
+                # Convert SearchResult objects to dict format for caching
+                search_results_dict = []
+                for result in context:
+                    search_results_dict.append({
+                        'expense_id': result.expense_id,
+                        'user_id': result.user_id,
+                        'description': result.description,
+                        'merchant': result.merchant,
+                        'expense_amount': result.amount,
+                        'expense_date': result.date,
+                        'similarity_score': result.similarity_score,
+                        'shopping_type': result.metadata.get('shopping_type'),
+                        'payment_method': result.metadata.get('payment_method'),
+                        'recurring': result.metadata.get('recurring'),
+                        'tags': result.metadata.get('tags')
+                    })
+                
+                # Estimate token usage (rough approximation for AWS)
+                prompt_tokens = len(query.split()) * 1.3  # ~1.3 tokens per word
+                response_tokens = len(ai_response.split()) * 1.3
+                
+                self.cache_manager.cache_response(
+                    query, ai_response, search_results_dict, "aws",
+                    int(prompt_tokens), int(response_tokens)
+                )
+                print(f"3. ✅ Cached response (est. {int(prompt_tokens + response_tokens)} tokens)")
+            
+            return RAGResponse(
+                response=ai_response,
+                sources=context,
+                metadata={
+                    "model": "claude-3-5-sonnet",
+                    "region": self.region,
+                    "language": language
+                }
+            )
+            
+        except Exception as e:
+            raise AIConnectionError(f"RAG response generation failed: {str(e)}")
+    
+    def generate_embedding(self, text: str) -> List[float]:
+        """Generate embedding for text."""
+        try:
+            embedding_model = self._get_embedding_model()
+            embedding = embedding_model.encode([text])[0]
+            return embedding.tolist()
+        except Exception as e:
+            raise AIConnectionError(f"Embedding generation failed: {str(e)}")
+    
+    def test_connection(self) -> bool:
+        """Test AWS Bedrock connection."""
+        try:
+            # Test with a simple completion
+            payload = {
+                "anthropic_version": "bedrock-2023-05-31",
+                "max_tokens": 5,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": [{"type": "text", "text": "Hello"}]
+                    }
+                ]
+            }
+            
+            response = self.bedrock_client.invoke_model(
+                modelId=self.current_model,
+                contentType="application/json",
+                accept="application/json",
+                body=json.dumps(payload)
+            )
+            
+            response_body = json.loads(response['body'].read())
+            return response_body['content'][0]['text'] is not None
+        except Exception:
+            return False
+    
+    def _prepare_context(self, context: List[SearchResult]) -> str:
+        """Prepare context text from search results."""
+        if not context:
+            return "No relevant expense data found."
+        
+        context_parts = []
+        for i, result in enumerate(context, 1):
+            context_parts.append(
+                f"• **{result.description}** at {result.merchant}: ${result.amount:.2f} "
+                f"({result.date}) - similarity: {result.similarity_score:.3f}"
+            )
+        
+        return "\n".join(context_parts)

banko_ai/ai_providers/base.py

@@ -0,0 +1,175 @@
+"""
+Base AI provider interface and common functionality.
+
+This module defines the abstract base class for AI providers and common error handling.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Optional
+from dataclasses import dataclass
+
+
+class AIProviderError(Exception):
+    """Base exception for AI provider errors."""
+    pass
+
+
+class AIConnectionError(AIProviderError):
+    """Exception raised when AI service connection fails."""
+    pass
+
+
+class AIAuthenticationError(AIProviderError):
+    """Exception raised when AI service authentication fails."""
+    pass
+
+
+class AIQuotaExceededError(AIProviderError):
+    """Exception raised when AI service quota is exceeded."""
+    pass
+
+
+@dataclass
+class SearchResult:
+    """Result from vector similarity search."""
+    expense_id: str
+    user_id: str
+    description: str
+    merchant: str
+    amount: float
+    date: str
+    similarity_score: float
+    metadata: Dict[str, Any]
+
+
+@dataclass
+class RAGResponse:
+    """Response from RAG (Retrieval-Augmented Generation) query."""
+    response: str
+    sources: List[SearchResult]
+    metadata: Dict[str, Any]
+
+
+class AIProvider(ABC):
+    """Abstract base class for AI providers."""
+    
+    def __init__(self, config: Dict[str, Any]):
+        """Initialize the AI provider with configuration."""
+        self.config = config
+        self.current_model = config.get("model", self.get_default_model())
+        self._validate_config()
+    
+    @abstractmethod
+    def _validate_config(self) -> None:
+        """Validate provider-specific configuration."""
+        pass
+    
+    @abstractmethod
+    def search_expenses(
+        self, 
+        query: str, 
+        user_id: Optional[str] = None,
+        limit: int = 10,
+        threshold: float = 0.7
+    ) -> List[SearchResult]:
+        """
+        Search for expenses using vector similarity.
+        
+        Args:
+            query: Search query text
+            user_id: Optional user ID to filter results
+            limit: Maximum number of results to return
+            threshold: Minimum similarity score threshold
+            
+        Returns:
+            List of SearchResult objects
+        """
+        pass
+    
+    @abstractmethod
+    def generate_rag_response(
+        self, 
+        query: str, 
+        context: List[SearchResult],
+        user_id: Optional[str] = None,
+        language: str = "en"
+    ) -> RAGResponse:
+        """
+        Generate a RAG response using the provided context.
+        
+        Args:
+            query: User query
+            context: List of relevant search results
+            user_id: Optional user ID
+            language: Response language code
+            
+        Returns:
+            RAGResponse object
+        """
+        pass
+    
+    @abstractmethod
+    def generate_embedding(self, text: str) -> List[float]:
+        """
+        Generate embedding vector for the given text.
+        
+        Args:
+            text: Text to embed
+            
+        Returns:
+            List of float values representing the embedding
+        """
+        pass
+    
+    @abstractmethod
+    def test_connection(self) -> bool:
+        """
+        Test the connection to the AI service.
+        
+        Returns:
+            True if connection is successful, False otherwise
+        """
+        pass
+    
+    def get_provider_name(self) -> str:
+        """Get the name of this AI provider."""
+        return self.__class__.__name__.replace("Provider", "").lower()
+    
+    def get_provider_info(self) -> Dict[str, Any]:
+        """Get provider information and status."""
+        return {
+            "name": self.get_provider_name(),
+            "connected": self.test_connection(),
+            "current_model": self.current_model,
+            "config_keys": list(self.config.keys())
+        }
+    
+    @abstractmethod
+    def get_default_model(self) -> str:
+        """Get the default model for this provider."""
+        pass
+    
+    def get_available_models(self) -> List[str]:
+        """Get available models for this provider."""
+        # Default implementation - can be overridden by providers
+        return [self.current_model]
+    
+    def set_model(self, model: str) -> bool:
+        """
+        Switch to a different model.
+        
+        Args:
+            model: Model name to switch to
+            
+        Returns:
+            True if model was switched successfully, False otherwise
+        """
+        available_models = self.get_available_models()
+        if model in available_models:
+            self.current_model = model
+            return True
+        return False
+    
+    def get_current_model(self) -> str:
+        """Get the current model."""
+        return self.current_model

banko_ai/ai_providers/factory.py

@@ -0,0 +1,84 @@
+"""
+AI Provider Factory for creating provider instances.
+
+This module provides a factory pattern for creating AI provider instances
+based on configuration.
+"""
+
+from typing import Dict, Any, Type
+from .base import AIProvider, AIProviderError
+from .openai_provider import OpenAIProvider
+from .aws_provider import AWSProvider
+from .watsonx_provider import WatsonxProvider
+from .gemini_provider import GeminiProvider
+
+
+class AIProviderFactory:
+    """Factory for creating AI provider instances."""
+    
+    _providers: Dict[str, Type[AIProvider]] = {
+        "openai": OpenAIProvider,
+        "aws": AWSProvider,
+        "watsonx": WatsonxProvider,
+        "gemini": GeminiProvider,
+    }
+    
+    @classmethod
+    def create_provider(cls, service_name: str, config: Dict[str, Any], cache_manager=None) -> AIProvider:
+        """
+        Create an AI provider instance.
+        
+        Args:
+            service_name: Name of the AI service (openai, aws, watsonx, gemini)
+            config: Configuration dictionary for the provider
+            cache_manager: Optional cache manager instance
+            
+        Returns:
+            AIProvider instance
+            
+        Raises:
+            AIProviderError: If the service is not supported
+        """
+        service_name = service_name.lower()
+        
+        if service_name not in cls._providers:
+            available = ", ".join(cls._providers.keys())
+            raise AIProviderError(
+                f"Unsupported AI service: {service_name}. "
+                f"Available services: {available}"
+            )
+        
+        provider_class = cls._providers[service_name]
+        
+        try:
+            # Pass cache_manager to all providers that support it
+            if hasattr(provider_class, '__init__'):
+                # Check if the provider's __init__ method accepts cache_manager parameter
+                import inspect
+                sig = inspect.signature(provider_class.__init__)
+                if 'cache_manager' in sig.parameters:
+                    return provider_class(config, cache_manager)
+                else:
+                    return provider_class(config)
+            else:
+                return provider_class(config)
+        except Exception as e:
+            raise AIProviderError(
+                f"Failed to create {service_name} provider: {str(e)}"
+            )
+    
+    @classmethod
+    def get_available_providers(cls) -> list[str]:
+        """Get list of available AI providers."""
+        return list(cls._providers.keys())
+    
+    @classmethod
+    def register_provider(cls, name: str, provider_class: Type[AIProvider]) -> None:
+        """
+        Register a new AI provider.
+        
+        Args:
+            name: Name of the provider
+            provider_class: Provider class that implements AIProvider
+        """
+        cls._providers[name.lower()] = provider_class