🤖 RAG Chatbot Development

# Production-Grade RAG Implementation
import asyncio
import logging
from typing import List, Dict, Optional, Tuple
from datetime import datetime
import numpy as np
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
import openai

class ProductionRAGChatbot:
    def __init__(
        self, 
        embedding_model: str = "all-MiniLM-L6-v2",
        llm_model: str = "gpt-3.5-turbo",
        vector_db_path: str = "./chroma_db",
        enable_reranking: bool = True
    ):
        # Initialize embeddings
        self.embeddings = HuggingFaceEmbeddings(
            model_name=embedding_model,
            model_kwargs={'device': 'cuda'},
            encode_kwargs={'normalize_embeddings': True}
        )
        
        # Initialize vector store
        self.vectorstore = Chroma(
            persist_directory=vector_db_path,
            embedding_function=self.embeddings
        )
        
        # Initialize LLM
        self.llm_model = llm_model
        
        # Advanced text splitting
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
            separators=["\n\n", "\n", " ", ""]
        )
        
        # Re-ranking model for better results
        self.enable_reranking = enable_reranking
        if enable_reranking:
            self.reranker = SentenceTransformer('cross-encoder/ms-marco-MiniLM-L-6-v2')
        
        # Setup logging
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)
        
        # Query cache for performance
        self.query_cache = {}
        self.cache_ttl = 3600  # 1 hour
    
    async def ingest_documents(
        self, 
        documents: List[Dict],
        batch_size: int = 100
    ) -> Dict[str, int]:
        """Advanced document ingestion with metadata"""
        processed_chunks = []
        metadata_list = []
        
        for doc in documents:
            content = doc.get('content', '')
            metadata = doc.get('metadata', {})
            
            # Split document into chunks
            chunks = self.text_splitter.split_text(content)
            
            for i, chunk in enumerate(chunks):
                processed_chunks.append(chunk)
                chunk_metadata = {
                    **metadata,
                    'chunk_id': f"{metadata.get('doc_id', 'unknown')}_{i}",
                    'chunk_index': i,
                    'total_chunks': len(chunks),
                    'created_at': datetime.now().isoformat()
                }
                metadata_list.append(chunk_metadata)
        
        # Add to vector store in batches
        for i in range(0, len(processed_chunks), batch_size):
            batch_chunks = processed_chunks[i:i + batch_size]
            batch_metadata = metadata_list[i:i + batch_size]
            
            self.vectorstore.add_texts(
                texts=batch_chunks,
                metadatas=batch_metadata
            )
            
            self.logger.info(f"Processed batch {i//batch_size + 1}")
        
        # Persist the vector store
        self.vectorstore.persist()
        
        return {
            'total_documents': len(documents),
            'total_chunks': len(processed_chunks),
            'avg_chunk_size': sum(len(chunk) for chunk in processed_chunks) / len(processed_chunks)
        }
    
    def _rerank_results(self, query: str, results: List[Tuple]) -> List[Tuple]:
        """Re-rank search results for better relevance"""
        if not self.enable_reranking or len(results) <= 1:
            return results
        
        # Prepare query-document pairs
        pairs = [(query, result[0].page_content) for result in results]
        
        # Get relevance scores
        scores = self.reranker.predict(pairs)
        
        # Sort by relevance score
        scored_results = list(zip(results, scores))
        scored_results.sort(key=lambda x: x[1], reverse=True)
        
        return [result[0] for result in scored_results]
    
    def _build_context_prompt(
        self, 
        query: str, 
        context_chunks: List[Tuple],
        max_context_length: int = 4000
    ) -> str:
        """Build optimized prompt with context"""
        context_parts = []
        current_length = 0
        
        for i, (doc, score) in enumerate(context_chunks):
            chunk_text = doc.page_content
            source = doc.metadata.get('source', 'Unknown')
            
            context_part = f"[Source {i+1}: {source}]\n{chunk_text}\n"
            
            if current_length + len(context_part) > max_context_length:
                break
            
            context_parts.append(context_part)
            current_length += len(context_part)
        
        context = "\n---\n".join(context_parts)
        
        prompt = f"""You are a helpful assistant that answers questions based on the provided context. 
Use only the information from the context to answer the question. If the answer cannot be found in the context, say so.

Context:
{context}

---

Question: {query}

Answer (cite sources using [Source X] format):"""
        
        return prompt
    
    async def query(
        self, 
        question: str, 
        k: int = 5,
        filter_metadata: Optional[Dict] = None,
        use_cache: bool = True
    ) -> Dict:
        """Enhanced query with caching and filtering"""
        start_time = datetime.now()
        
        # Check cache first
        cache_key = f"{question}_{k}_{str(filter_metadata)}"
        if use_cache and cache_key in self.query_cache:
            cached_result = self.query_cache[cache_key]
            if (datetime.now() - cached_result['timestamp']).seconds < self.cache_ttl:
                return cached_result['result']
        
        # Retrieve relevant documents
        search_kwargs = {'k': k * 2}  # Get more for re-ranking
        if filter_metadata:
            search_kwargs['filter'] = filter_metadata
        
        raw_results = self.vectorstore.similarity_search_with_score(
            question, **search_kwargs
        )
        
        # Re-rank results
        reranked_results = self._rerank_results(question, raw_results)[:k]
        
        # Build context prompt
        prompt = self._build_context_prompt(question, reranked_results)
        
        # Generate response with OpenAI
        try:
            response = await openai.ChatCompletion.acreate(
                model=self.llm_model,
                messages=[
                    {"role": "user", "content": prompt}
                ],
                temperature=0.1,
                max_tokens=1000
            )
            
            answer = response.choices[0].message.content
            
        except Exception as e:
            self.logger.error(f"LLM generation error: {e}")
            answer = "I encountered an error while generating a response. Please try again."
        
        # Prepare response
        end_time = datetime.now()
        result = {
            'answer': answer,
            'sources': [
                {
                    'content': doc.page_content[:200] + '...',
                    'metadata': doc.metadata,
                    'relevance_score': float(score)
                } for doc, score in reranked_results
            ],
            'query_time': (end_time - start_time).total_seconds(),
            'retrieved_docs': len(reranked_results)
        }
        
        # Cache the result
        if use_cache:
            self.query_cache[cache_key] = {
                'result': result,
                'timestamp': datetime.now()
            }
        
        return result
    
    def get_analytics(self) -> Dict:
        """Get system analytics"""
        collection = self.vectorstore._collection
        
        return {
            'total_documents': collection.count(),
            'cache_hits': len(self.query_cache),
            'embedding_model': self.embeddings.model_name,
            'vector_db_path': self.vectorstore.persist_directory
        }

class SmartChunker:
    def __init__(self):
        self.strategies = {
            'semantic': self.semantic_chunking,
            'sliding_window': self.sliding_window,
            'document_structure': self.structure_aware
        }
    
    def semantic_chunking(self, text: str, max_size: int = 500):
        """Split based on semantic boundaries"""
        # Use sentence transformers for semantic similarity
        sentences = text.split('.')
        chunks = []
        current_chunk = []
        
        for sentence in sentences:
            current_chunk.append(sentence)
            if len(' '.join(current_chunk)) > max_size:
                chunks.append(' '.join(current_chunk))
                current_chunk = []
        
        return chunks
    
    def sliding_window(self, text: str, window_size: int = 500, stride: int = 250):
        """Overlapping window approach for better context preservation"""
        chunks = []
        for i in range(0, len(text), stride):
            chunk = text[i:i + window_size]
            if chunk:
                chunks.append(chunk)
        return chunks

from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel, Field
import asyncio
from datetime import datetime

app = FastAPI(title="RAG Chatbot API")

class Query(BaseModel):
    question: str = Field(..., description="User question")
    context_size: int = Field(5, description="Number of context chunks")
    session_id: Optional[str] = None

class Response(BaseModel):
    answer: str
    sources: List[Dict]
    confidence: float
    response_time: float

@app.post("/query", response_model=Response)
async def query_rag(query: Query, background_tasks: BackgroundTasks):
    """Process user query through RAG pipeline"""
    start_time = datetime.now()
    
    try:
        # Get response from RAG chatbot
        answer, sources = await rag_chatbot.query(
            query.question, 
            k=query.context_size
        )
        
        # Calculate confidence score
        confidence = calculate_confidence(answer, sources)
        
        # Log for monitoring
        background_tasks.add_task(
            log_query, 
            query.session_id, 
            query.question, 
            answer
        )
        
        response_time = (datetime.now() - start_time).total_seconds()
        
        return Response(
            answer=answer,
            sources=sources,
            confidence=confidence,
            response_time=response_time
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "timestamp": datetime.now()}