ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/contextual-retrieval-for-user-memory/contextual_indexer.py

@@ -0,0 +1,672 @@
+"""Contextual RAG Indexer with Advanced Memory Cards
+
+This module combines contextual chunking for conversation histories with
+advanced JSON cards for structured user memory.
+"""
+
+import os
+import json
+import logging
+import requests
+import time
+from datetime import datetime
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass
+from pathlib import Path
+
+from config import IndexConfig, IndexMode, ChunkingConfig
+from chunker import ConversationChunk
+from contextual_chunking import ContextualConversationChunk, ContextualConversationChunker
+from advanced_memory_manager import AdvancedMemoryManager, AdvancedMemoryCard
+from indexer import SearchResult
+
+
+def _reasoning_safe_temperature(model, requested=1.0):
+    """Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
+    Return 1 for those; otherwise the requested value so non-reasoning
+    providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
+    m = str(model or "").lower().replace("/", "-")
+    return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass 
+class ContextualSearchResult(SearchResult):
+    """Enhanced search result with contextual information"""
+    context: str = ""
+    contextual_chunk: Optional[ContextualConversationChunk] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        result = super().to_dict()
+        result["context"] = self.context
+        if self.contextual_chunk:
+            result["contextual_info"] = {
+                "context": self.contextual_chunk.context,
+                "context_tokens": self.contextual_chunk.context_tokens,
+                "generation_time": self.contextual_chunk.generation_time
+            }
+        return result
+
+
+class ContextualMemoryIndexer:
+    """
+    Dual-layer memory indexer combining:
+    1. Contextual RAG for conversation chunks
+    2. Advanced JSON cards for structured facts
+    """
+    
+    def __init__(self,
+                 user_id: str,
+                 index_config: Optional[IndexConfig] = None,
+                 chunking_config: Optional[ChunkingConfig] = None,
+                 use_contextual: bool = True):
+        """
+        Initialize the contextual memory indexer.
+        
+        Args:
+            user_id: User identifier
+            index_config: Index configuration
+            chunking_config: Chunking configuration
+            use_contextual: Whether to use contextual chunking
+        """
+        self.user_id = user_id
+        self.index_config = index_config or IndexConfig()
+        self.chunking_config = chunking_config or ChunkingConfig()
+        self.use_contextual = use_contextual
+        
+        # Initialize contextual chunker
+        self.contextual_chunker = ContextualConversationChunker(
+            chunking_config=self.chunking_config,
+            use_contextual=use_contextual
+        )
+        
+        # Initialize advanced memory manager
+        self.memory_manager = AdvancedMemoryManager(
+            user_id=user_id,
+            storage_dir=self.index_config.chunk_store_path
+        )
+        
+        # Storage for chunks
+        self.contextual_chunks: Dict[str, ContextualConversationChunk] = {}
+        self.doc_id_mapping: Dict[str, str] = {}
+        
+        # Retrieval pipeline URL
+        self.retrieval_url = "http://localhost:4242"
+        
+        # Create directories
+        Path(self.index_config.index_path).parent.mkdir(parents=True, exist_ok=True)
+        Path(self.index_config.chunk_store_path).mkdir(parents=True, exist_ok=True)
+        
+        # Statistics
+        self.stats = {
+            "chunks_indexed": 0,
+            "memory_cards": 0,
+            "context_generation_time": 0.0,
+            "indexing_time": 0.0
+        }
+        
+        self._check_retrieval_pipeline()
+        logger.info(f"Initialized ContextualMemoryIndexer for user {user_id}")
+    
+    def _check_retrieval_pipeline(self):
+        """Check if the retrieval pipeline service is available"""
+        try:
+            response = requests.get(f"{self.retrieval_url}/health", timeout=2)
+            if response.status_code == 200:
+                logger.info("✓ Retrieval pipeline service is available")
+            else:
+                logger.warning(f"Retrieval pipeline returned status {response.status_code}")
+        except requests.exceptions.RequestException as e:
+            logger.warning(f"Retrieval pipeline service not available at {self.retrieval_url}: {e}")
+            logger.info("Note: The retrieval pipeline needs to be running. Start it with:")
+            logger.info("  cd projects/week3/retrieval-pipeline && python api_server.py")
+    
+    def process_conversation_history(self,
+                                    chunks: List[ConversationChunk],
+                                    conversation_id: str,
+                                    generate_summary_cards: bool = True) -> Dict[str, Any]:
+        """
+        Process conversation history with both contextual chunking and memory cards.
+        
+        Args:
+            chunks: Basic conversation chunks
+            conversation_id: Conversation identifier
+            generate_summary_cards: Whether to generate summary cards
+            
+        Returns:
+            Processing results
+        """
+        start_time = time.time()
+        results = {
+            "conversation_id": conversation_id,
+            "contextual_chunks": 0,
+            "memory_cards_before": len(self.memory_manager.categories),
+            "memory_cards_after": 0,
+            "processing_time": 0
+        }
+        
+        # Step 1: Generate contextual chunks
+        logger.info(f"Processing {len(chunks)} chunks for conversation {conversation_id}")
+        
+        if self.use_contextual:
+            contextual_chunks = self.contextual_chunker.contextualize_chunks(chunks)
+            logger.info(f"Generated {len(contextual_chunks)} contextual chunks")
+        else:
+            # Convert to contextual chunks without context
+            contextual_chunks = [
+                ContextualConversationChunk.from_basic_chunk(chunk)
+                for chunk in chunks
+            ]
+        
+        # Store contextual chunks
+        for chunk in contextual_chunks:
+            self.contextual_chunks[chunk.chunk_id] = chunk
+        
+        results["contextual_chunks"] = len(contextual_chunks)
+        
+        # Step 2: Index contextual chunks
+        self._index_contextual_chunks(contextual_chunks)
+        
+        # Step 3: Generate summary cards if requested
+        if generate_summary_cards:
+            summary_cards = self._generate_summary_cards(chunks, conversation_id)
+            for card in summary_cards:
+                self.memory_manager.add_card(card)
+            
+            logger.info(f"Generated {len(summary_cards)} summary cards")
+            results["new_summary_cards"] = len(summary_cards)
+        
+        # Step 4: Create conversation summary linking to memory cards
+        conversation_summary = self.memory_manager.summarize_for_conversation(conversation_id)
+        results["memory_cards_after"] = sum(len(cards) for cards in self.memory_manager.categories.values())
+        results["conversation_summary"] = conversation_summary
+        
+        # Update statistics
+        results["processing_time"] = time.time() - start_time
+        self.stats["chunks_indexed"] += len(contextual_chunks)
+        self.stats["memory_cards"] = results["memory_cards_after"]
+        self.stats["indexing_time"] += results["processing_time"]
+        
+        logger.info(f"Processing complete for conversation {conversation_id} in {results['processing_time']:.2f}s")
+        
+        return results
+    
+    def _index_contextual_chunks(self, chunks: List[ContextualConversationChunk]):
+        """Index contextual chunks in the retrieval pipeline"""
+        documents = []
+        
+        for chunk in chunks:
+            # Use contextualized text for indexing
+            doc = {
+                "text": chunk.contextualized_text,
+                "metadata": {
+                    "doc_id": chunk.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,
+                    "has_context": bool(chunk.context),
+                    "context_preview": chunk.context[:100] if chunk.context else "",
+                    **chunk.metadata
+                }
+            }
+            documents.append(doc)
+        
+        # Send to retrieval pipeline
+        try:
+            indexed_count = 0
+            for doc in documents:
+                try:
+                    response = requests.post(
+                        f"{self.retrieval_url}/index",
+                        json=doc, 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")
+                        
+                        if generated_doc_id and our_chunk_id:
+                            self.doc_id_mapping[generated_doc_id] = our_chunk_id
+                        
+                        indexed_count += 1
+                except Exception as e:
+                    logger.warning(f"Error indexing document: {e}")
+            
+            logger.info(f"Indexed {indexed_count}/{len(documents)} contextual chunks")
+            
+        except Exception as e:
+            logger.error(f"Error connecting to retrieval pipeline: {e}")
+    
+    def _generate_summary_cards(self, 
+                               chunks: List[ConversationChunk],
+                               conversation_id: str) -> List[AdvancedMemoryCard]:
+        """
+        Generate summary cards from conversation chunks using LLM extraction.
+        
+        Uses an LLM to analyze conversation chunks and extract structured memory cards
+        following the Advanced JSON Cards format from week2/user-memory.
+        """
+        cards = []
+        
+        if not chunks:
+            return cards
+        
+        # Combine all conversation text
+        full_text = "\n".join([chunk.to_text() for chunk in chunks])
+        
+        # Use LLM to extract structured memory cards
+        try:
+            # Initialize LLM client if needed
+            if not hasattr(self, '_llm_client'):
+                from openai import OpenAI
+                from config import Config
+                config = Config.from_env()
+                client_config, model = config.llm.get_client_config()
+                base_url = client_config.pop("base_url", None)
+                
+                if base_url:
+                    self._llm_client = OpenAI(base_url=base_url, **client_config)
+                else:
+                    self._llm_client = OpenAI(**client_config)
+                self._llm_model = model
+            
+            # Create extraction prompt
+            extraction_prompt = f"""Analyze this conversation and extract structured memory cards.
+
+Conversation:
+{full_text}
+
+Extract any important information into memory cards. Each card MUST include:
+- category: The type of information (e.g., 'personal', 'financial', 'medical', 'travel', 'family', 'work')
+- card_key: A unique identifier for this card
+- backstory: Context about when/why this information was learned (1-2 sentences)
+- date_created: Current timestamp (YYYY-MM-DD HH:MM:SS)
+- person: Who this relates to (e.g., "John Smith (primary)", "Sarah Smith (daughter)")
+- relationship: Role/relationship (e.g., "primary account holder", "family member")
+- Additional relevant fields based on the information type
+
+Respond with a JSON array of memory cards. If no significant information found, return empty array [].
+
+Example format:
+[{{
+    "category": "financial",
+    "card_key": "bank_account_primary",
+    "backstory": "User shared their banking details while setting up automatic bill payments",
+    "date_created": "2024-01-15 10:30:00",
+    "person": "John Smith (primary)",
+    "relationship": "primary account holder",
+    "bank_name": "Chase Bank",
+    "account_type": "checking",
+    "account_ending": "4567"
+}}]
+
+Extract memory cards:"""
+            
+            # Call LLM for extraction
+            response = self._llm_client.chat.completions.create(
+                model=self._llm_model,
+                messages=[
+                    {"role": "system", "content": "You are a memory extraction assistant. Extract structured information from conversations into memory cards."},
+                    {"role": "user", "content": extraction_prompt}
+                ],
+                temperature=_reasoning_safe_temperature(self._llm_model, 0.3),
+                response_format={"type": "json_object"}
+            )
+            
+            # Parse response
+            content = response.choices[0].message.content
+            extracted_data = json.loads(content) if content else {}
+            
+            # Handle both single card and array of cards
+            if isinstance(extracted_data, dict):
+                # Check if it's a wrapper with 'cards' key
+                if 'cards' in extracted_data:
+                    extracted_cards = extracted_data['cards']
+                elif 'memory_cards' in extracted_data:
+                    extracted_cards = extracted_data['memory_cards']
+                else:
+                    # Single card
+                    extracted_cards = [extracted_data] if extracted_data else []
+            elif isinstance(extracted_data, list):
+                extracted_cards = extracted_data
+            else:
+                extracted_cards = []
+            
+            # Convert extracted data to AdvancedMemoryCard objects
+            for card_data in extracted_cards:
+                try:
+                    # Ensure required fields
+                    if 'backstory' not in card_data:
+                        card_data['backstory'] = f"Information extracted from conversation {conversation_id}"
+                    if 'date_created' not in card_data:
+                        card_data['date_created'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+                    if 'person' not in card_data:
+                        card_data['person'] = 'User (primary)'
+                    if 'relationship' not in card_data:
+                        card_data['relationship'] = 'primary account holder'
+                    
+                    # Extract main fields
+                    category = card_data.pop('category', 'general')
+                    card_key = card_data.pop('card_key', f"{category}_{conversation_id[:8]}")
+                    backstory = card_data.pop('backstory')
+                    date_created = card_data.pop('date_created')
+                    person = card_data.pop('person')
+                    relationship = card_data.pop('relationship')
+                    
+                    # Create card with remaining fields as data
+                    card = AdvancedMemoryCard(
+                        category=category,
+                        card_key=card_key,
+                        backstory=backstory,
+                        date_created=date_created,
+                        person=person,
+                        relationship=relationship,
+                        data=card_data  # All remaining fields
+                    )
+                    cards.append(card)
+                    
+                except Exception as e:
+                    logger.warning(f"Error creating memory card: {e}")
+                    continue
+                    
+        except Exception as e:
+            logger.error(f"Error generating memory cards with LLM: {e}")
+            # Fallback to basic extraction if LLM fails
+            cards = self._fallback_extraction(chunks, conversation_id, full_text)
+        
+        return cards
+    
+    def _fallback_extraction(self, chunks: List[ConversationChunk], 
+                             conversation_id: str, full_text: str) -> List[AdvancedMemoryCard]:
+        """Fallback extraction method if LLM is unavailable"""
+        cards = []
+        
+        # Extract basic information based on keywords
+        categories_keywords = {
+            'financial': ['bank', 'account', 'money', 'payment', 'credit', 'loan', 'savings', 'checking'],
+            'medical': ['doctor', 'medical', 'health', 'appointment', 'prescription', 'hospital', 'clinic'],
+            'travel': ['flight', 'travel', 'trip', 'vacation', 'hotel', 'booking', 'destination'],
+            'family': ['family', 'child', 'parent', 'spouse', 'daughter', 'son', 'wife', 'husband'],
+            'work': ['job', 'work', 'employer', 'salary', 'office', 'meeting', 'project', 'colleague']
+        }
+        
+        text_lower = full_text.lower()
+        
+        for category, keywords in categories_keywords.items():
+            if any(keyword in text_lower for keyword in keywords):
+                card = AdvancedMemoryCard(
+                    category=category,
+                    card_key=f"{category}_{conversation_id[:8]}",
+                    backstory=f"Information about {category} topics discussed in conversation",
+                    date_created=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
+                    person="User (primary)",
+                    relationship="primary account holder",
+                    data={
+                        "conversation_id": conversation_id,
+                        "topics": [kw for kw in keywords if kw in text_lower][:3],
+                        "extracted_from": f"{len(chunks)} conversation chunks",
+                        "extraction_method": "fallback_keywords"
+                    }
+                )
+                cards.append(card)
+                break  # Only add one card per conversation in fallback mode
+        
+        return cards
+    
+    def search_with_context(self,
+                           query: str,
+                           top_k: int = 3,
+                           include_memory_cards: bool = True) -> Dict[str, Any]:
+        """
+        Search both contextual chunks and memory cards.
+        
+        Args:
+            query: Search query
+            top_k: Number of results
+            include_memory_cards: Whether to search memory cards too
+            
+        Returns:
+            Combined search results
+        """
+        results = {
+            "query": query,
+            "chunk_results": [],
+            "memory_card_results": [],
+            "combined_context": ""
+        }
+        
+        # Search contextual chunks via retrieval pipeline or fallback to local search
+        try:
+            response = requests.post(
+                f"{self.retrieval_url}/search",
+                json={
+                    "query": query,
+                    "mode": "hybrid",
+                    "top_k": max(20, top_k),
+                    "rerank_top_k": top_k,
+                    "skip_reranking": False
+                },
+                timeout=5
+            )
+            response.raise_for_status()
+            
+            data = response.json()
+            search_results = data.get("reranked_results", [])
+            
+            for item in search_results:
+                metadata = item.get("metadata", {})
+                chunk_id = metadata.get("doc_id")
+                
+                if chunk_id and chunk_id in self.contextual_chunks:
+                    contextual_chunk = self.contextual_chunks[chunk_id]
+                    
+                    # Create a basic ConversationChunk for compatibility
+                    from chunker import ConversationChunk
+                    basic_chunk = ConversationChunk(
+                        chunk_id=contextual_chunk.chunk_id,
+                        conversation_id=contextual_chunk.conversation_id,
+                        test_id=contextual_chunk.test_id,
+                        chunk_index=contextual_chunk.chunk_index,
+                        start_round=contextual_chunk.start_round,
+                        end_round=contextual_chunk.end_round,
+                        messages=contextual_chunk.messages,
+                        metadata=contextual_chunk.metadata,
+                        context_before=contextual_chunk.context_before,
+                        context_after=contextual_chunk.context_after,
+                        created_at=contextual_chunk.created_at
+                    )
+                    
+                    result = ContextualSearchResult(
+                        chunk_id=chunk_id,
+                        score=float(item.get("rerank_score", 0)),
+                        chunk=basic_chunk,
+                        match_type="hybrid",
+                        context=contextual_chunk.context,
+                        contextual_chunk=contextual_chunk
+                    )
+                    results["chunk_results"].append(result.to_dict())
+            
+        except Exception as e:
+            logger.error(f"Error searching via retrieval pipeline: {e}")
+            logger.info("Falling back to local search...")
+            
+            # Fallback: Local search through contextual chunks
+            query_lower = query.lower()
+            local_results = []
+            
+            for chunk_id, chunk in self.contextual_chunks.items():
+                # Search in contextualized text (skip chunks with empty text)
+                if query_lower and chunk.contextualized_text and query_lower in chunk.contextualized_text.lower():
+                    # Calculate a simple relevance score based on frequency
+                    score = chunk.contextualized_text.lower().count(query_lower) / len(chunk.contextualized_text)
+                    local_results.append((score, chunk_id, chunk))
+            
+            # Sort by score and take top_k
+            local_results.sort(key=lambda x: x[0], reverse=True)
+            
+            for score, chunk_id, contextual_chunk in local_results[:top_k]:
+                # Create a basic ConversationChunk for compatibility
+                from chunker import ConversationChunk
+                basic_chunk = ConversationChunk(
+                    chunk_id=contextual_chunk.chunk_id,
+                    conversation_id=contextual_chunk.conversation_id,
+                    test_id=contextual_chunk.test_id,
+                    chunk_index=contextual_chunk.chunk_index,
+                    start_round=contextual_chunk.start_round,
+                    end_round=contextual_chunk.end_round,
+                    messages=contextual_chunk.messages,
+                    metadata=contextual_chunk.metadata,
+                    context_before=contextual_chunk.context_before,
+                    context_after=contextual_chunk.context_after,
+                    created_at=contextual_chunk.created_at
+                )
+                
+                results["chunk_results"].append({
+                    "chunk_id": chunk_id,
+                    "score": score,
+                    "context": contextual_chunk.context,
+                    "conversation_id": contextual_chunk.conversation_id,
+                    "rounds": f"{contextual_chunk.start_round}-{contextual_chunk.end_round}",
+                    "text": contextual_chunk.original_text
+                })
+            
+            logger.info(f"Local search returned {len(results['chunk_results'])} results")
+        
+        # Search memory cards
+        if include_memory_cards:
+            card_results = self.memory_manager.search_cards(query)
+            for memory_id, card in card_results[:top_k]:
+                results["memory_card_results"].append({
+                    "memory_id": memory_id,
+                    "category": card.category,
+                    "backstory": card.backstory,
+                    "person": card.person,
+                    "data": card.data
+                })
+        
+        # Build combined context
+        context_parts = []
+        
+        if results["memory_card_results"]:
+            context_parts.append("=== RELEVANT MEMORY CARDS ===")
+            for card_result in results["memory_card_results"]:
+                context_parts.append(f"- {card_result['category']}: {card_result['backstory']}")
+        
+        if results["chunk_results"]:
+            context_parts.append("\n=== RELEVANT CONVERSATION CHUNKS ===")
+            for chunk_result in results["chunk_results"]:
+                context_parts.append(f"- {chunk_result.get('context', 'No context')}")
+        
+        results["combined_context"] = "\n".join(context_parts)
+        
+        return results
+    
+    def get_agent_context(self, max_memory_cards: int = 10) -> str:
+        """
+        Get the complete context for the agent including memory cards.
+        
+        Args:
+            max_memory_cards: Maximum number of memory cards to include
+            
+        Returns:
+            Formatted context string
+        """
+        return self.memory_manager.get_context_string(max_cards=max_memory_cards)
+    
+    def save_index(self, path: Optional[str] = None):
+        """Save the index and memory to disk"""
+        path = path or self.index_config.index_path
+        
+        # Save contextual chunks
+        chunks_data = {
+            chunk_id: chunk.to_dict()
+            for chunk_id, chunk in self.contextual_chunks.items()
+        }
+        
+        with open(f"{path}_contextual_chunks.json", 'w', encoding='utf-8') as f:
+            json.dump(chunks_data, f, ensure_ascii=False, indent=2)
+        
+        # Memory cards are automatically saved by the memory manager
+        
+        # Save statistics
+        with open(f"{path}_stats.json", 'w', encoding='utf-8') as f:
+            json.dump(self.stats, f, indent=2)
+        
+        logger.info(f"Saved index to {path}")
+    
+    def load_index(self, path: Optional[str] = None):
+        """Load the index from disk"""
+        path = path or self.index_config.index_path
+        
+        try:
+            # Load contextual chunks
+            with open(f"{path}_contextual_chunks.json", 'r', encoding='utf-8') as f:
+                chunks_data = json.load(f)
+            
+            self.contextual_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 contextual chunk
+                chunk = ContextualConversationChunk(
+                    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,
+                    original_text=chunk_dict.get('original_text', ''),
+                    context=chunk_dict.get('context', ''),
+                    contextualized_text=chunk_dict.get('contextualized_text', ''),
+                    context_tokens=chunk_dict.get('context_tokens', 0),
+                    generation_time=chunk_dict.get('generation_time', 0),
+                    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.contextual_chunks[chunk_id] = chunk
+            
+            # Memory cards are automatically loaded by the memory manager
+            
+            # Load statistics
+            stats_path = f"{path}_stats.json"
+            if Path(stats_path).exists():
+                with open(stats_path, 'r', encoding='utf-8') as f:
+                    self.stats = json.load(f)
+            
+            logger.info(f"Loaded {len(self.contextual_chunks)} contextual chunks from {path}")
+            
+            # Re-index in retrieval pipeline
+            self._index_contextual_chunks(list(self.contextual_chunks.values()))
+            
+        except Exception as e:
+            logger.error(f"Error loading index: {e}")
+            raise
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        """Get comprehensive statistics"""
+        stats = self.stats.copy()
+        
+        # Add chunker statistics
+        stats["chunker_stats"] = self.contextual_chunker.get_statistics()
+        
+        # Add memory manager statistics
+        stats["memory_stats"] = self.memory_manager.get_statistics()
+        
+        # Calculate averages
+        if stats["chunks_indexed"] > 0:
+            stats["avg_indexing_time"] = stats["indexing_time"] / stats["chunks_indexed"]
+        
+        return stats