""" Knowledge Base for Experience Retrieval This module provides indexing and retrieval capabilities for experiences. """ import json import logging import os from typing import Dict, Any, List, Optional, Tuple import numpy as np try: from sentence_transformers import SentenceTransformer except ImportError: SentenceTransformer = None try: import faiss except ImportError: faiss = None import pickle logger = logging.getLogger(__name__) class KnowledgeBase: """ Knowledge base for storing and retrieving experiences using semantic search. """ def __init__( self, index_path: str = "./kb_index", model_name: str = "all-MiniLM-L6-v2", embedding_dim: int = 384 ): """ Initialize the knowledge base. Args: index_path: Path to store the index files model_name: Name of the sentence transformer model embedding_dim: Dimension of the embeddings """ self.index_path = index_path self.model_name = model_name self.embedding_dim = embedding_dim if SentenceTransformer is None or faiss is None: self.encoder = None else: try: self.encoder = SentenceTransformer(model_name) # Derive the real embedding dimension from the loaded model so the # FAISS index matches it. A fixed 384 silently breaks any non-384 # model chosen via --embedding-model / config.yaml (e.g. # all-mpnet-base-v2 = 768): index.add() then raises, is swallowed, # and every search falls back to keyword-only for the whole KB. model_dim = self.encoder.get_sentence_embedding_dimension() if model_dim: self.embedding_dim = model_dim except Exception as e: logger.warning(f"Failed to load SentenceTransformer, falling back to simple search: {e}") self.encoder = None # Initialize FAISS index self.index = None self.documents = [] self.metadata = [] # Create index directory if it doesn't exist os.makedirs(index_path, exist_ok=True) # Load existing index if available self._load_index() def _load_index(self): """Load existing index from disk.""" index_file = os.path.join(self.index_path, "faiss.index") docs_file = os.path.join(self.index_path, "documents.pkl") meta_file = os.path.join(self.index_path, "metadata.pkl") # Documents remain useful for keyword search even when the optional # semantic-search dependencies (or the FAISS file) are unavailable. # Load them independently so a keyword-only run survives a restart. if os.path.exists(docs_file): try: with open(docs_file, 'rb') as f: self.documents = pickle.load(f) if not isinstance(self.documents, list): raise ValueError("Persisted documents must be a list") if os.path.exists(meta_file): with open(meta_file, 'rb') as f: self.metadata = pickle.load(f) if not isinstance(self.metadata, list): raise ValueError("Persisted metadata must be a list") else: self.metadata = [{}] * len(self.documents) # Keep one metadata entry per document and tolerate older or # partially-written metadata files. self.metadata = [ item if isinstance(item, dict) else {} for item in self.metadata[:len(self.documents)] ] self.metadata.extend( {} for _ in range(len(self.documents) - len(self.metadata)) ) except Exception as e: logger.error(f"Failed to load persisted documents: {e}") self.documents = [] self.metadata = [] if not self.encoder: logger.info(f"Loaded knowledge base with {len(self.documents)} documents") return rebuild_reason = None if os.path.exists(index_file): try: self.index = faiss.read_index(index_file) except Exception as e: logger.warning(f"Failed to load FAISS index: {e}") rebuild_reason = "FAISS index could not be loaded" else: if self.index.d != self.embedding_dim: rebuild_reason = ( f"FAISS dimension {self.index.d} != model dimension " f"{self.embedding_dim}" ) elif self.index.ntotal != len(self.documents): rebuild_reason = ( f"FAISS row count {self.index.ntotal} != document count " f"{len(self.documents)}" ) elif self.documents: rebuild_reason = "FAISS index is missing" if rebuild_reason: logger.warning(f"{rebuild_reason}; rebuilding from stored queries") self._rebuild_index_from_metadata() elif self.index is None: self._create_new_index() logger.info(f"Loaded knowledge base with {len(self.documents)} documents") def _create_new_index(self): """Create a new empty index.""" if self.encoder: self.index = faiss.IndexFlatL2(self.embedding_dim) else: self.index = None def _rebuild_index_from_metadata(self): """Rebuild the FAISS index at the current embedding dimension by re-encoding the query texts persisted in metadata (used when a loaded index was built with a different embedding model). One embedding per document, in order, so index rows stay aligned with self.documents.""" if not self.encoder or faiss is None: self.index = None return self.index = faiss.IndexFlatL2(self.embedding_dim) if not self.documents: return queries = [ (self.metadata[i].get('query', '') if i < len(self.metadata) else '') for i in range(len(self.documents)) ] try: embeddings = self.encoder.encode(queries) self.index.add(embeddings) self._save_index() logger.info(f"Rebuilt FAISS index with {self.index.ntotal} embeddings") except Exception as e: logger.error(f"Failed to rebuild FAISS index: {e}") def _save_index(self): """Save index to disk.""" try: if self.encoder and self.index is not None: index_file = os.path.join(self.index_path, "faiss.index") faiss.write_index(self.index, index_file) docs_file = os.path.join(self.index_path, "documents.pkl") with open(docs_file, 'wb') as f: pickle.dump(self.documents, f) meta_file = os.path.join(self.index_path, "metadata.pkl") with open(meta_file, 'wb') as f: pickle.dump(self.metadata, f) except Exception as e: logger.error(f"Failed to save index: {e}") def index_gaia_validation(self, validation_file: str): """ Index the GAIA validation file for experience retrieval. Args: validation_file: Path to gaia-validation.jsonl """ if not os.path.exists(validation_file): logger.error(f"Validation file not found: {validation_file}") return logger.info(f"Indexing GAIA validation data from {validation_file}") # The index is persisted and reloaded by __init__ (_load_index), so # without this every run appends another full copy of the dataset and # search() starts returning the same document top_k times. existing_ids = { doc.get('task_id') for doc in self.documents if doc.get('source') == 'gaia_validation' } skipped = 0 try: with open(validation_file, 'r', encoding='utf-8') as f: for line_num, line in enumerate(f, 1): try: data = json.loads(line) # Extract relevant information question = data.get('Question', '') answer = data.get('Final answer', '') level = data.get('Level', 0) metadata = data.get('Annotator Metadata', {}) task_id = data.get('task_id', f'gaia_{line_num}') if task_id in existing_ids: skipped += 1 continue # Create experience document experience = { 'task_id': task_id, 'question': question, 'answer': answer, 'level': level, 'approach': self._extract_approach(metadata), 'tools_used': self._extract_tools(metadata), 'steps': metadata.get('Steps', ''), 'num_steps': metadata.get('Number of steps', '0'), 'source': 'gaia_validation' } # Add to index self.add_experience(question, experience) existing_ids.add(task_id) except json.JSONDecodeError as e: logger.error(f"Failed to parse line {line_num}: {e}") except Exception as e: logger.error(f"Error processing line {line_num}: {e}") # Save the index after bulk indexing self._save_index() if skipped: logger.info(f"Skipped {skipped} GAIA records already present in the index") logger.info(f"Successfully indexed {len(self.documents)} experiences from GAIA validation") except Exception as e: logger.error(f"Failed to index validation file: {e}") def _extract_approach(self, metadata: Dict[str, Any]) -> str: """ Extract approach from metadata steps. Args: metadata: Annotator metadata Returns: Summarized approach """ steps = metadata.get('Steps', '') if not steps: return "" # Extract key actions from steps lines = steps.split('\n') key_actions = [] for line in lines[:3]: # Take first 3 steps as approach if line.strip(): # Remove step numbers clean_line = line.strip() if clean_line[0].isdigit(): clean_line = clean_line.split('.', 1)[-1].strip() key_actions.append(clean_line) return " → ".join(key_actions) if key_actions else steps[:200] def _extract_tools(self, metadata: Dict[str, Any]) -> List[str]: """ Extract tools used from metadata. Args: metadata: Annotator metadata Returns: List of tools used """ tools = metadata.get('Tools', '') if not tools: return [] # Parse tools string tool_list = [] # Handle numbered list format lines = tools.split('\n') for line in lines: line = line.strip() if line: # Remove numbering if '. ' in line: tool = line.split('. ', 1)[-1].strip() else: tool = line if tool and tool not in ['', 'None']: tool_list.append(tool) return tool_list def add_experience(self, query: str, experience: Dict[str, Any]): """ Add an experience to the knowledge base. Args: query: The query/question for indexing experience: The experience data """ # Store document self.documents.append(experience) self.metadata.append({ 'query': query, 'task_id': experience.get('task_id', ''), 'level': experience.get('level', 0) }) # Create embedding and add to index if encoder is available if self.encoder and self.index: try: embedding = self.encoder.encode([query]) self.index.add(embedding) except Exception as e: logger.error(f"Failed to create embedding: {e}") def search(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]: """ Search for relevant experiences. Args: query: The search query top_k: Number of top results to return Returns: List of relevant experiences """ if top_k <= 0 or not self.documents: return [] # If we have embeddings, use semantic search if self.encoder and self.index and self.index.ntotal > 0: try: query_embedding = self.encoder.encode([query]) distances, indices = self.index.search(query_embedding, min(top_k, len(self.documents))) results = [] for idx in indices[0]: if 0 <= idx < len(self.documents): results.append(self.documents[idx]) return results except Exception as e: logger.error(f"Semantic search failed, falling back to keyword search: {e}") # Fallback to simple keyword search return self._keyword_search(query, top_k) def _keyword_search(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]: """ Simple keyword-based search fallback. Args: query: The search query top_k: Number of results to return Returns: List of relevant experiences """ query_words = set(query.lower().split()) scored_docs = [] for doc in self.documents: tools = doc.get('tools_used') if tools is None: tools = [] elif isinstance(tools, str): tools = [tools] elif not isinstance(tools, (list, tuple, set)): tools = [tools] tools_str = ' '.join(str(t) for t in tools if t is not None) doc_text = f"{doc.get('question', '')} {doc.get('approach', '')} {tools_str}" doc_words = set(doc_text.lower().split()) # Calculate simple overlap score overlap = len(query_words & doc_words) if overlap > 0: scored_docs.append((overlap, doc)) # Sort by score and return top k scored_docs.sort(key=lambda x: x[0], reverse=True) return [doc for _, doc in scored_docs[:top_k]] def get_statistics(self) -> Dict[str, Any]: """ Get statistics about the knowledge base. Returns: Dictionary with statistics """ stats = { 'total_documents': len(self.documents), 'has_embeddings': self.encoder is not None, 'index_size': self.index.ntotal if self.encoder and self.index else 0, 'sources': {} } # Count by source for doc in self.documents: source = doc.get('source', 'unknown') stats['sources'][source] = stats['sources'].get(source, 0) + 1 return stats