ai-agent-book/chapter3/agentic-rag-for-user-memory/indexer.py

"""RAG Indexer for User Memory Conversations

This module handles indexing of conversation chunks using the retrieval pipeline service.
Interfaces with the existing retrieval pipeline on port 4242.
"""

import os
import re
import math
import json
import logging
import requests
from collections import Counter
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path

from config import IndexConfig, IndexMode
from chunker import ConversationChunk


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def _tokenize(text: str) -> List[str]:
    """Lightweight tokenizer shared by the local backend (alphanumeric + CJK)."""
    return re.findall(r"[a-zA-Z0-9]+|[一-鿿]", text.lower())


class LocalBM25Backend:
    """A dependency-free, in-process BM25 index.

    This is the offline fallback for the external retrieval pipeline: it lets the
    whole store/retrieve path run without any network service or API key, which is
    what makes the experiment reproducible on a laptop. Sparse (BM25) retrieval is
    the same lexical scoring the pipeline exposes as its "sparse"/"hybrid" modes.
    """

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.doc_ids: List[str] = []
        self.doc_tokens: List[List[str]] = []
        self.idf: Dict[str, float] = {}
        self.avgdl: float = 0.0
        self._built: bool = False

    def clear(self):
        self.doc_ids = []
        self.doc_tokens = []
        self.idf = {}
        self.avgdl = 0.0
        self._built = False

    def add(self, doc_id: str, text: str):
        self.doc_ids.append(doc_id)
        self.doc_tokens.append(_tokenize(text))
        self._built = False

    def _build(self):
        n_docs = len(self.doc_tokens)
        df: Counter = Counter()
        for tokens in self.doc_tokens:
            for term in set(tokens):
                df[term] += 1
        # Standard BM25 idf with +1 smoothing so it stays non-negative.
        self.idf = {
            term: math.log(1 + (n_docs - freq + 0.5) / (freq + 0.5))
            for term, freq in df.items()
        }
        self.avgdl = (sum(len(t) for t in self.doc_tokens) / n_docs) if n_docs else 0.0
        self._built = True

    def search(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
        if not self._built:
            self._build()
        query_terms = _tokenize(query)
        scored: List[Tuple[str, float]] = []
        for idx, tokens in enumerate(self.doc_tokens):
            tf = Counter(tokens)
            dl = len(tokens)
            score = 0.0
            for term in query_terms:
                freq = tf.get(term, 0)
                if not freq:
                    continue
                idf = self.idf.get(term, 0.0)
                denom = freq + self.k1 * (1 - self.b + self.b * dl / (self.avgdl or 1))
                score += idf * (freq * (self.k1 + 1)) / denom
            if score > 0:
                scored.append((self.doc_ids[idx], score))
        scored.sort(key=lambda item: item[1], reverse=True)
        return scored[:top_k]


@dataclass
class SearchResult:
    """Result from searching the index"""
    chunk_id: str
    score: float
    chunk: ConversationChunk
    match_type: str  # "dense", "sparse", or "hybrid"
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "chunk_id": self.chunk_id,
            "score": self.score,
            "match_type": self.match_type,
            "conversation_id": self.chunk.conversation_id,
            "test_id": self.chunk.test_id,
            "rounds": f"{self.chunk.start_round}-{self.chunk.end_round}",
            "text": self.chunk.to_text()
        }


class MemoryIndexer:
    """Indexes and searches conversation chunks using the retrieval pipeline service"""
    
    def __init__(self, config: Optional[IndexConfig] = None):
        """
        Initialize the indexer
        
        Args:
            config: Index configuration
        """
        self.config = config or IndexConfig()
        self.chunks: Dict[str, ConversationChunk] = {}
        self.chunk_texts: Dict[str, str] = {}  # Map chunk_id to prepared text
        self.doc_id_mapping: Dict[str, str] = {}  # Map generated doc_id to our chunk_id

        # Retrieval pipeline URL
        self.retrieval_url = getattr(self.config, "retrieval_url", "http://localhost:4242")

        # Local, in-process fallback backend (no external service required)
        self.local_backend = LocalBM25Backend()

        # Create directories
        Path(self.config.index_path).parent.mkdir(parents=True, exist_ok=True)
        Path(self.config.chunk_store_path).parent.mkdir(parents=True, exist_ok=True)

        # Decide which backend to use: "local", "pipeline", or "auto"
        backend = getattr(self.config, "retrieval_backend", "auto")
        if backend == "local":
            self.use_local = True
        elif backend == "pipeline":
            self.use_local = False
            self._check_retrieval_pipeline()
        else:  # auto
            self.use_local = not self._check_retrieval_pipeline()

        if self.use_local:
            logger.info("Using built-in local BM25 backend (offline mode, no port 4242 needed)")
        else:
            logger.info("Using external retrieval pipeline backend")

        logger.info(f"Initialized indexer with mode: {self.config.mode}")

    def _check_retrieval_pipeline(self) -> bool:
        """Check if the retrieval pipeline service is available. Returns True if reachable."""
        try:
            response = requests.get(f"{self.retrieval_url}/health", timeout=2)
            if response.status_code == 200:
                logger.info("✓ Retrieval pipeline service is available")
                return True
            logger.warning(f"Retrieval pipeline returned status {response.status_code}")
            return False
        except requests.exceptions.RequestException as e:
            logger.warning(f"Retrieval pipeline service not available at {self.retrieval_url}: {e}")
            logger.info("Note: falling back to the built-in local BM25 backend (offline).")
            logger.info("To use the external pipeline instead, start it with:")
            logger.info("  cd ../retrieval-pipeline && python api_server.py")
            return False
    
    def add_chunks(self, chunks: List[ConversationChunk], rebuild: bool = True):
        """
        Add conversation chunks to the index
        
        Args:
            chunks: List of conversation chunks to index
            rebuild: Whether to rebuild indexes after adding (for retrieval pipeline)
        """
        documents = []
        
        for chunk in chunks:
            chunk_id = chunk.chunk_id
            
            # Store chunk locally
            self.chunks[chunk_id] = chunk
            
            # Prepare text for indexing
            chunk_text = self._prepare_chunk_text(chunk)
            self.chunk_texts[chunk_id] = chunk_text
            
            # Prepare document for retrieval pipeline
            doc = {
                "text": chunk_text,
                "metadata": {
                    "doc_id": chunk_id,
                    "test_id": chunk.test_id,
                    "conversation_id": chunk.conversation_id,
                    "chunk_index": chunk.chunk_index,
                    "start_round": chunk.start_round,
                    "end_round": chunk.end_round,
                    **chunk.metadata
                }
            }
            documents.append(doc)
            logger.debug(f"Added chunk {chunk_id} to index")
        
        if rebuild and documents:
            self._index_documents(documents)
        
        logger.info(f"Added {len(chunks)} chunks to index. Total chunks: {len(self.chunks)}")
    
    def _prepare_chunk_text(self, chunk: ConversationChunk) -> str:
        """
        Prepare chunk text for indexing with contextual enrichment
        
        Args:
            chunk: Conversation chunk
            
        Returns:
            Enriched text for indexing
        """
        if not self.config.enable_contextual:
            return chunk.to_text()
        
        # Build enriched text with contextual information
        lines = []
        
        # Add test case context
        lines.append(f"Test Case: {chunk.test_id}")
        lines.append(f"Conversation: {chunk.conversation_id}")
        
        # Add metadata as searchable text
        if chunk.metadata:
            for key, value in chunk.metadata.items():
                lines.append(f"{key}: {value}")
        
        # Add the main chunk content
        lines.append(chunk.to_text())
        
        # Add semantic tags for better retrieval
        lines.append(self._generate_semantic_tags(chunk))
        
        return "\n".join(lines)
    
    def _generate_semantic_tags(self, chunk: ConversationChunk) -> str:
        """
        Generate semantic tags for better retrieval
        
        Args:
            chunk: Conversation chunk
            
        Returns:
            Semantic tags as string
        """
        tags = []
        
        # Analyze content for common topics
        content = chunk.to_text().lower()
        
        # Financial topics
        if any(word in content for word in ["account", "bank", "credit", "loan", "payment"]):
            tags.append("financial")
        
        # Insurance topics
        if any(word in content for word in ["insurance", "claim", "policy", "coverage"]):
            tags.append("insurance")
        
        # Medical topics
        if any(word in content for word in ["medical", "doctor", "appointment", "prescription"]):
            tags.append("medical")
        
        # Travel topics
        if any(word in content for word in ["flight", "hotel", "travel", "booking", "reservation"]):
            tags.append("travel")
        
        # Add position tags
        if chunk.chunk_index == 0:
            tags.append("conversation_start")
        
        # Add round count tags
        round_count = chunk.end_round - chunk.start_round + 1
        if round_count < 10:
            tags.append("short_segment")
        elif round_count > 30:
            tags.append("long_segment")
        
        return f"Tags: {', '.join(tags)}" if tags else ""
    
    def _index_documents(self, documents: List[Dict[str, Any]]):
        """Index documents into the active backend (local BM25 or the external pipeline)."""
        if self.use_local:
            self.local_backend.clear()
            for doc in documents:
                chunk_id = doc.get("metadata", {}).get("doc_id")
                if chunk_id:
                    self.local_backend.add(chunk_id, doc["text"])
                    self.doc_id_mapping[chunk_id] = chunk_id
            logger.info(f"Indexed {len(documents)} documents into local BM25 backend")
            return
        try:
            # First, clear existing index
            clear_response = requests.post(f"{self.retrieval_url}/clear", timeout=30)
            if clear_response.status_code == 200:
                logger.info("Cleared existing index")
            
            # Index documents one by one (retrieval pipeline expects individual documents)
            indexed_count = 0
            failed_count = 0
            
            for doc in documents:
                try:
                    response = requests.post(
                        f"{self.retrieval_url}/index",
                        json=doc  # Send individual document, timeout=30
                    )
                    
                    if response.status_code == 200:
                        result = response.json()
                        generated_doc_id = result.get("doc_id")
                        our_chunk_id = doc.get("metadata", {}).get("doc_id")
                        
                        # Store the mapping between generated doc_id and our chunk_id
                        if generated_doc_id and our_chunk_id:
                            self.doc_id_mapping[generated_doc_id] = our_chunk_id
                            
                        indexed_count += 1
                    else:
                        failed_count += 1
                        logger.warning(f"Failed to index document: {doc.get('metadata', {}).get('doc_id', 'unknown')}")
                        
                except requests.exceptions.RequestException as e:
                    failed_count += 1
                    logger.warning(f"Error indexing document: {e}")
            
            logger.info(f"Indexed {indexed_count} documents successfully ({failed_count} failed)")
            
        except requests.exceptions.RequestException as e:
            logger.error(f"Error connecting to retrieval pipeline: {e}")
            logger.info("Make sure the retrieval pipeline is running on port 4242")
    
    def build_indexes(self):
        """Build or rebuild indexes by sending all chunks to retrieval pipeline"""
        if not self.chunks:
            logger.warning("No chunks to index")
            return
        
        # Prepare all documents
        documents = []
        for chunk_id, chunk in self.chunks.items():
            chunk_text = self.chunk_texts.get(chunk_id) or self._prepare_chunk_text(chunk)
            doc = {
                "text": chunk_text,
                "metadata": {
                    "doc_id": chunk_id,
                    "test_id": chunk.test_id,
                    "conversation_id": chunk.conversation_id,
                    "chunk_index": chunk.chunk_index,
                    "start_round": chunk.start_round,
                    "end_round": chunk.end_round,
                    **chunk.metadata
                }
            }
            documents.append(doc)
        
        # Send to retrieval pipeline
        self._index_documents(documents)
        logger.info("Index building complete")
    
    def search(self,
              query: str,
              top_k: int = 3,
              mode: Optional[IndexMode] = None) -> List[SearchResult]:
        """
        Search the index for relevant chunks using retrieval pipeline
        
        Args:
            query: Search query
            top_k: Number of results to return
            mode: Search mode (uses config default if not specified)
            
        Returns:
            List of search results
        """
        mode = mode or self.config.mode
        
        # Map IndexMode to retrieval pipeline mode strings
        mode_map = {
            IndexMode.DENSE: "dense",
            IndexMode.SPARSE: "sparse", 
            IndexMode.HYBRID: "hybrid"
        }
        
        search_mode = mode_map.get(mode, "hybrid")

        if not top_k or top_k < 1:
            top_k = 3

        # Offline path: score against the in-process BM25 index.
        if self.use_local:
            results = []
            for chunk_id, score in self.local_backend.search(query, top_k=top_k):
                chunk = self.chunks.get(chunk_id)
                if chunk:
                    results.append(SearchResult(
                        chunk_id=chunk_id,
                        score=float(score),
                        chunk=chunk,
                        match_type="local_bm25"
                    ))
            logger.info(f"Search returned {len(results)} results from local BM25 backend")
            return results

        try:
            # Query the retrieval pipeline
            # Note: The pipeline has two top_k parameters:
            # - top_k: for initial retrieval (we set to max(20, top_k))
            # - rerank_top_k: for final results (we set to the requested top_k)
            response = requests.post(
                f"{self.retrieval_url}/search",
                json={
                    "query": query,
                    "mode": search_mode,
                    "top_k": max(20, top_k),  # Retrieve more candidates for better reranking
                    "rerank_top_k": top_k,    # Return the requested number of results
                    "skip_reranking": False    # Always use reranking for better quality
                }, timeout=30
            )
            response.raise_for_status()
            
            data = response.json()
            
            # Get results based on mode
            if search_mode == "hybrid" and "reranked_results" in data:
                search_results = data["reranked_results"]
            elif search_mode == "dense" and "dense_results" in data:
                search_results = data["dense_results"]
            elif search_mode == "sparse" and "sparse_results" in data:
                search_results = data["sparse_results"]
            else:
                # Fallback to any available results
                search_results = (data.get("reranked_results", []) or 
                                data.get("dense_results", []) or 
                                data.get("sparse_results", []))
            
            # Convert to SearchResult objects
            results = []
            for item in search_results:
                # Try to get our chunk_id from different sources
                chunk_id = None
                
                # First, check if metadata contains our doc_id
                metadata = item.get("metadata", {})
                if metadata.get("doc_id"):
                    chunk_id = metadata.get("doc_id")
                else:
                    # Fall back to doc_id mapping
                    generated_doc_id = item.get("doc_id", "")
                    chunk_id = self.doc_id_mapping.get(generated_doc_id)
                
                # Get chunk from local storage
                if chunk_id and chunk_id in self.chunks:
                    chunk = self.chunks[chunk_id]
                    
                    # Get score based on result type
                    score = item.get("rerank_score", item.get("score", 0.0))
                    
                    results.append(SearchResult(
                        chunk_id=chunk_id,
                        score=float(score),
                        chunk=chunk,
                        match_type=search_mode
                    ))
                else:
                    # Log warning but don't fail
                    doc_id = item.get("doc_id", "unknown")
                    if chunk_id:
                        logger.debug(f"Chunk {chunk_id} not found in local storage")
                    else:
                        logger.debug(f"No mapping found for doc_id {doc_id}")
            
            logger.info(f"Search returned {len(results)} results from retrieval pipeline")
            return results
            
        except requests.exceptions.RequestException as e:
            logger.error(f"Error searching via retrieval pipeline: {e}")
            logger.info("Falling back to empty results. Ensure retrieval pipeline is running.")
            return []
    
    def save_index(self, path: Optional[str] = None):
        """
        Save the chunks and metadata to disk
        
        Args:
            path: Path to save index (uses config default if not specified)
        """
        path = path or self.config.index_path
        
        # Save chunks
        chunks_data = {
            chunk_id: chunk.to_dict() 
            for chunk_id, chunk in self.chunks.items()
        }
        
        with open(f"{path}_chunks.json", 'w', encoding='utf-8') as f:
            json.dump(chunks_data, f, ensure_ascii=False, indent=2)
        
        # Save chunk texts
        with open(f"{path}_texts.json", 'w', encoding='utf-8') as f:
            json.dump(self.chunk_texts, f, ensure_ascii=False, indent=2)
        
        logger.info(f"Chunks saved to {path}. Total chunks: {len(self.chunks)}")
    
    def load_index(self, path: Optional[str] = None):
        """
        Load chunks from disk and re-index in retrieval pipeline
        
        Args:
            path: Path to load index from (uses config default if not specified)
        """
        path = path or self.config.index_path
        
        try:
            # Load chunks
            with open(f"{path}_chunks.json", 'r', encoding='utf-8') as f:
                chunks_data = json.load(f)
            
            self.chunks = {}
            for chunk_id, chunk_dict in chunks_data.items():
                # Convert messages
                from chunker import ConversationMessage
                messages = []
                for msg_data in chunk_dict.get('messages', []):
                    messages.append(ConversationMessage(**msg_data))
                
                # Create chunk
                chunk = ConversationChunk(
                    chunk_id=chunk_dict['chunk_id'],
                    conversation_id=chunk_dict['conversation_id'],
                    test_id=chunk_dict['test_id'],
                    chunk_index=chunk_dict['chunk_index'],
                    start_round=chunk_dict['start_round'],
                    end_round=chunk_dict['end_round'],
                    messages=messages,
                    metadata=chunk_dict.get('metadata', {}),
                    context_before=chunk_dict.get('context_before'),
                    context_after=chunk_dict.get('context_after'),
                    created_at=chunk_dict.get('created_at', '')
                )
                self.chunks[chunk_id] = chunk
            
            # Load chunk texts if available
            texts_path = f"{path}_texts.json"
            if Path(texts_path).exists():
                with open(texts_path, 'r', encoding='utf-8') as f:
                    self.chunk_texts = json.load(f)
            else:
                # Regenerate texts if not saved
                self.chunk_texts = {}
                for chunk_id, chunk in self.chunks.items():
                    self.chunk_texts[chunk_id] = self._prepare_chunk_text(chunk)
            
            logger.info(f"Loaded {len(self.chunks)} chunks from {path}")
            
            # Re-index in retrieval pipeline
            self.build_indexes()
            
        except Exception as e:
            logger.error(f"Error loading index: {e}")
            raise