rag_service.py

"""
RAG (Retrieval-Augmented Generation) Pipeline for Financial Research
FAISS hybrid search with semantic and keyword matching
"""

import asyncio
import json
import logging
from typing import Dict, List, Optional, Tuple, Any
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
import pickle
import os

# Vector search and embeddings
import faiss
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# Web scraping and data sources
import requests
from bs4 import BeautifulSoup
import yfinance as yf
from dataclasses import dataclass

@dataclass
class Document:
    """Document structure for RAG pipeline"""
    content: str
    source: str
    timestamp: datetime
    metadata: Dict[str, Any]
    embedding: Optional[np.ndarray] = None
    relevance_score: float = 0.0

class FinancialDataCollector:
    """Collect and process financial data from various sources"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
    async def collect_news_data(self, tickers: List[str]) -> List[Document]:
        """Collect financial news for specified tickers"""
        documents = []
        
        try:
            # Using yfinance for news data
            for ticker in tickers:
                stock = yf.Ticker(ticker)
                news = stock.news
                
                for article in news[:5]:  # Limit to recent articles
                    doc = Document(
                        content=f"{article.get('title', '')} - {article.get('summary', '')}",
                        source=f"yahoo_finance_{ticker}",
                        timestamp=datetime.fromtimestamp(article.get('providerPublishTime', 0)),
                        metadata={
                            'ticker': ticker,
                            'url': article.get('link', ''),
                            'publisher': article.get('publisher', '')
                        }
                    )
                    documents.append(doc)
                    
        except Exception as e:
            self.logger.error(f"News collection failed: {e}")
            
        return documents
    
    async def collect_earnings_data(self, tickers: List[str]) -> List[Document]:
        """Collect earnings reports and financial data"""
        documents = []
        
        try:
            for ticker in tickers:
                stock = yf.Ticker(ticker)
                
                # Get earnings data
                earnings = stock.earnings
                if not earnings.empty:
                    earnings_summary = earnings.to_string()
                    
                    doc = Document(
                        content=f"Earnings data for {ticker}: {earnings_summary}",
                        source=f"earnings_{ticker}",
                        timestamp=datetime.now(),
                        metadata={
                            'ticker': ticker,
                            'data_type': 'earnings',
                            'period_count': len(earnings)
                        }
                    )
                    documents.append(doc)
                
                # Get financial statements
                financials = stock.financials
                if not financials.empty:
                    financials_summary = financials.head().to_string()
                    
                    doc = Document(
                        content=f"Financial statements for {ticker}: {financials_summary}",
                        source=f"financials_{ticker}",
                        timestamp=datetime.now(),
                        metadata={
                            'ticker': ticker,
                            'data_type': 'financials',
                            'metrics': list(financials.index)
                        }
                    )
                    documents.append(doc)
                    
        except Exception as e:
            self.logger.error(f"Earnings data collection failed: {e}")
            
        return documents
    
    async def collect_market_analysis(self) -> List[Document]:
        """Collect general market analysis and economic indicators"""
        documents = []
        
        try:
            # Economic indicators
            economic_data = {
                "VIX": "Volatility Index",
                "^TNX": "10-Year Treasury",
                "^GSPC": "S&P 500",
                "^DJI": "Dow Jones",
                "^IXIC": "NASDAQ"
            }
            
            for symbol, name in economic_data.items():
                try:
                    ticker = yf.Ticker(symbol)
                    hist = ticker.history(period="5d")
                    
                    if not hist.empty:
                        latest = hist.iloc[-1]
                        prev = hist.iloc[-2] if len(hist) > 1 else hist.iloc[-1]
                        
                        change = ((latest['Close'] - prev['Close']) / prev['Close']) * 100
                        
                        content = f"""
                        {name} ({symbol}) Market Update:
                        Current Level: {latest['Close']:.2f}
                        Daily Change: {change:.2f}%
                        Volume: {latest['Volume']:,.0f}
                        High: {latest['High']:.2f}
                        Low: {latest['Low']:.2f}
                        """
                        
                        doc = Document(
                            content=content,
                            source=f"market_data_{symbol}",
                            timestamp=datetime.now(),
                            metadata={
                                'symbol': symbol,
                                'name': name,
                                'data_type': 'market_indicator',
                                'current_price': float(latest['Close']),
                                'daily_change': float(change)
                            }
                        )
                        documents.append(doc)
                        
                except Exception as e:
                    self.logger.error(f"Failed to collect data for {symbol}: {e}")
                    continue
                    
        except Exception as e:
            self.logger.error(f"Market analysis collection failed: {e}")
            
        return documents

class HybridSearchEngine:
    """FAISS-based hybrid search combining semantic and keyword search"""
    
    def __init__(self, embedding_model: str = "all-MiniLM-L6-v2"):
        self.embedding_model = SentenceTransformer(embedding_model)
        self.dimension = 384  # Default for MiniLM
        
        # FAISS indices
        self.semantic_index = None
        self.documents = []
        
        # TF-IDF for keyword search
        self.tfidf_vectorizer = TfidfVectorizer(
            max_features=10000,
            stop_words='english',
            ngram_range=(1, 3)
        )
        self.tfidf_matrix = None
        
        self.logger = logging.getLogger(__name__)
        
    def build_index(self, documents: List[Document]):
        """Build FAISS index from documents"""
        try:
            self.documents = documents
            
            # Extract text content
            texts = [doc.content for doc in documents]
            
            # Generate embeddings
            embeddings = self.embedding_model.encode(texts, convert_to_numpy=True)
            
            # Store embeddings in documents
            for i, doc in enumerate(documents):
                doc.embedding = embeddings[i]
            
            # Build FAISS index
            self.dimension = embeddings.shape[1]
            self.semantic_index = faiss.IndexFlatIP(self.dimension)  # Inner product for cosine similarity
            
            # Normalize embeddings for cosine similarity
            faiss.normalize_L2(embeddings)
            self.semantic_index.add(embeddings.astype(np.float32))
            
            # Build TF-IDF matrix
            self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(texts)
            
            self.logger.info(f"Built search index with {len(documents)} documents")
            
        except Exception as e:
            self.logger.error(f"Index building failed: {e}")
            raise
    
    def search(
        self, 
        query: str, 
        k: int = 10, 
        semantic_weight: float = 0.7,
        keyword_weight: float = 0.3,
        filter_metadata: Optional[Dict] = None
    ) -> List[Document]:
        """Hybrid search combining semantic and keyword search"""
        
        try:
            if not self.semantic_index or not self.documents:
                return []
            
            # Semantic search
            query_embedding = self.embedding_model.encode([query], convert_to_numpy=True)
            faiss.normalize_L2(query_embedding)
            
            semantic_scores, semantic_indices = self.semantic_index.search(
                query_embedding.astype(np.float32), 
                min(k * 2, len(self.documents))
            )
            
            # Keyword search
            query_tfidf = self.tfidf_vectorizer.transform([query])
            keyword_similarities = cosine_similarity(query_tfidf, self.tfidf_matrix).flatten()
            
            # Combine scores
            results = []
            processed_indices = set()
            
            # Process semantic results
            for i, idx in enumerate(semantic_indices[0]):
                if idx in processed_indices:
                    continue
                    
                doc = self.documents[idx].copy() if hasattr(self.documents[idx], 'copy') else self.documents[idx]
                
                semantic_score = semantic_scores[0][i]
                keyword_score = keyword_similarities[idx] if idx < len(keyword_similarities) else 0
                
                # Combined score
                combined_score = (semantic_weight * semantic_score + 
                                keyword_weight * keyword_score)
                
                doc.relevance_score = combined_score
                results.append(doc)
                processed_indices.add(idx)
            
            # Add top keyword results not in semantic results
            keyword_indices = np.argsort(keyword_similarities)[::-1]
            for idx in keyword_indices[:k]:
                if idx not in processed_indices and len(results) < k * 2:
                    doc = self.documents[idx]
                    doc.relevance_score = keyword_weight * keyword_similarities[idx]
                    results.append(doc)
                    processed_indices.add(idx)
            
            # Apply metadata filters
            if filter_metadata:
                filtered_results = []
                for doc in results:
                    match = True
                    for key, value in filter_metadata.items():
                        if key not in doc.metadata or doc.metadata[key] != value:
                            match = False
                            break
                    if match:
                        filtered_results.append(doc)
                results = filtered_results
            
            # Sort by relevance score and limit results
            results.sort(key=lambda x: x.relevance_score, reverse=True)
            
            return results[:k]
            
        except Exception as e:
            self.logger.error(f"Search failed: {e}")
            return []
    
    def save_index(self, filepath: str):
        """Save FAISS index and metadata"""
        try:
            # Save FAISS index
            faiss.write_index(self.semantic_index, f"{filepath}.faiss")
            
            # Save documents and vectorizer
            with open(f"{filepath}.pkl", 'wb') as f:
                pickle.dump({
                    'documents': self.documents,
                    'tfidf_vectorizer': self.tfidf_vectorizer,
                    'tfidf_matrix': self.tfidf_matrix,
                    'dimension': self.dimension
                }, f)
                
            self.logger.info(f"Index saved to {filepath}")
            
        except Exception as e:
            self.logger.error(f"Index saving failed: {e}")
    
    def load_index(self, filepath: str):
        """Load FAISS index and metadata"""
        try:
            # Load FAISS index
            self.semantic_index = faiss.read_index(f"{filepath}.faiss")
            
            # Load documents and vectorizer
            with open(f"{filepath}.pkl", 'rb') as f:
                data = pickle.load(f)
                self.documents = data['documents']
                self.tfidf_vectorizer = data['tfidf_vectorizer']
                self.tfidf_matrix = data['tfidf_matrix']
                self.dimension = data['dimension']
                
            self.logger.info(f"Index loaded from {filepath}")
            
        except Exception as e:
            self.logger.error(f"Index loading failed: {e}")

class RAGPipeline:
    """Complete RAG pipeline for financial research"""
    
    def __init__(self, index_path: str = "./data/financial_rag_index"):
        self.data_collector = FinancialDataCollector()
        self.search_engine = HybridSearchEngine()
        self.index_path = index_path
        self.logger = logging.getLogger(__name__)
        
        # Create data directory if it doesn't exist
        os.makedirs(os.path.dirname(index_path), exist_ok=True)
        
    async def initialize(self, tickers: List[str] = None):
        """Initialize RAG pipeline with fresh data"""
        try:
            if tickers is None:
                tickers = ['AAPL', 'GOOGL', 'MSFT', 'TSLA', 'AMZN', 'NVDA']
            
            self.logger.info("Collecting financial data...")
            
            # Collect data from various sources
            all_documents = []
            
            # News data
            news_docs = await self.data_collector.collect_news_data(tickers)
            all_documents.extend(news_docs)
            
            # Earnings and financial data
            earnings_docs = await self.data_collector.collect_earnings_data(tickers)
            all_documents.extend(earnings_docs)
            
            # Market analysis
            market_docs = await self.data_collector.collect_market_analysis()
            all_documents.extend(market_docs)
            
            self.logger.info(f"Collected {len(all_documents)} documents")
            
            # Build search index
            if all_documents:
                self.search_engine.build_index(all_documents)
                self.search_engine.save_index(self.index_path)
                
            return len(all_documents)
            
        except Exception as e:
            self.logger.error(f"RAG initialization failed: {e}")
            raise
    
    def load_existing_index(self):
        """Load existing index if available"""
        try:
            if os.path.exists(f"{self.index_path}.faiss"):
                self.search_engine.load_index(self.index_path)
                return True
            return False
        except Exception as e:
            self.logger.error(f"Index loading failed: {e}")
            return False
    
    async def search_and_summarize(
        self, 
        query: str, 
        context: Optional[Dict] = None,
        max_results: int = 5
    ) -> Dict:
        """Search for relevant documents and create summary"""
        
        try:
            # Search for relevant documents
            results = self.search_engine.search(
                query, 
                k=max_results,
                filter_metadata=context
            )
            
            if not results:
                return {
                    "query": query,
                    "results": [],
                    "summary": "No relevant information found.",
                    "sources": [],
                    "timestamp": datetime.now().isoformat()
                }
            
            # Create context for LLM
            context_text = self._build_context(results, query)
            
            # Extract sources
            sources = [
                {
                    "source": doc.source,
                    "relevance_score": doc.relevance_score,
                    "timestamp": doc.timestamp.isoformat() if doc.timestamp else None,
                    "metadata": doc.metadata
                }
                for doc in results
            ]
            
            return {
                "query": query,
                "context": context_text,
                "results": [doc.content for doc in results],
                "sources": sources,
                "summary": self._create_summary(results, query),
                "total_documents": len(self.search_engine.documents),
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error(f"Search and summarize failed: {e}")
            return {
                "query": query,
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
    
    def _build_context(self, documents: List[Document], query: str) -> str:
        """Build context text from relevant documents"""
        context_parts = [f"Research Query: {query}\n\nRelevant Information:\n"]
        
        for i, doc in enumerate(documents[:5], 1):
            context_parts.append(f"\n{i}. Source: {doc.source}")
            context_parts.append(f"   Content: {doc.content[:500]}...")
            context_parts.append(f"   Relevance: {doc.relevance_score:.3f}")
            if doc.timestamp:
                context_parts.append(f"   Date: {doc.timestamp.strftime('%Y-%m-%d')}")
        
        return "\n".join(context_parts)
    
    def _create_summary(self, documents: List[Document], query: str) -> str:
        """Create a basic summary of search results"""
        if not documents:
            return "No relevant information found."
        
        # Group documents by source type
        source_counts = {}
        for doc in documents:
            source_type = doc.source.split('_')[0]
            source_counts[source_type] = source_counts.get(source_type, 0) + 1
        
        summary_parts = [
            f"Found {len(documents)} relevant documents for query: '{query}'",
            f"Source breakdown: {dict(source_counts)}",
            f"Most relevant source: {documents[0].source} (score: {documents[0].relevance_score:.3f})"
        ]
        
        return " | ".join(summary_parts)
    
    async def update_index(self, tickers: List[str] = None):
        """Update index with fresh data"""
        return await self.initialize(tickers)
    
    def get_index_stats(self) -> Dict:
        """Get statistics about the current index"""
        if not self.search_engine.documents:
            return {"total_documents": 0, "index_built": False}
        
        source_stats = {}
        for doc in self.search_engine.documents:
            source_type = doc.source.split('_')[0]
            source_stats[source_type] = source_stats.get(source_type, 0) + 1
        
        return {
            "total_documents": len(self.search_engine.documents),
            "index_built": self.search_engine.semantic_index is not None,
            "source_distribution": source_stats,
            "last_updated": datetime.now().isoformat()
        }