signalwire-agents 0.1.13__py3-none-any.whl → 1.0.17.dev4__py3-none-any.whl

signalwire_agents/search/index_builder.py

@@ -35,29 +35,117 @@ logger = logging.getLogger(__name__)
 class IndexBuilder:
     """Build searchable indexes from document directories"""
     
-    def __init__(self, model_name: str = 'sentence-transformers/all-mpnet-base-v2', 
-                 chunking_strategy: str = 'sentence',
-                 max_sentences_per_chunk: int = 50,
-                 chunk_size: int = 50, 
-                 chunk_overlap: int = 10,
-                 split_newlines: Optional[int] = None,
-                 verbose: bool = False):
+    def __init__(
+        self,
+        model_name: str = 'sentence-transformers/all-mpnet-base-v2',
+        chunking_strategy: str = 'sentence',
+        max_sentences_per_chunk: int = 5,
+        chunk_size: int = 50,
+        chunk_overlap: int = 10,
+        split_newlines: Optional[int] = None,
+        index_nlp_backend: str = 'nltk',
+        verbose: bool = False,
+        semantic_threshold: float = 0.5,
+        topic_threshold: float = 0.3,
+        backend: str = 'sqlite',
+        connection_string: Optional[str] = None
+    ):
+        """
+        Initialize the index builder
+        
+        Args:
+            model_name: Name of the sentence transformer model to use
+            chunking_strategy: Strategy for chunking documents ('sentence', 'sliding', 'paragraph', 'page', 'semantic', 'topic', 'qa', 'json')
+            max_sentences_per_chunk: For sentence strategy (default: 5)
+            chunk_size: For sliding strategy - words per chunk (default: 50)
+            chunk_overlap: For sliding strategy - overlap in words (default: 10)
+            split_newlines: For sentence strategy - split on multiple newlines (optional)
+            index_nlp_backend: NLP backend for indexing (default: 'nltk')
+            verbose: Whether to enable verbose logging (default: False)
+            semantic_threshold: Similarity threshold for semantic chunking (default: 0.5)
+            topic_threshold: Similarity threshold for topic chunking (default: 0.3)
+            backend: Storage backend ('sqlite' or 'pgvector') (default: 'sqlite')
+            connection_string: PostgreSQL connection string for pgvector backend
+        """
         self.model_name = model_name
         self.chunking_strategy = chunking_strategy
         self.max_sentences_per_chunk = max_sentences_per_chunk
         self.chunk_size = chunk_size
         self.chunk_overlap = chunk_overlap
         self.split_newlines = split_newlines
+        self.index_nlp_backend = index_nlp_backend
         self.verbose = verbose
+        self.semantic_threshold = semantic_threshold
+        self.topic_threshold = topic_threshold
+        self.backend = backend
+        self.connection_string = connection_string
         self.model = None
+        
+        # Validate backend
+        if self.backend not in ['sqlite', 'pgvector']:
+            raise ValueError(f"Invalid backend '{self.backend}'. Must be 'sqlite' or 'pgvector'")
+        
+        # Validate NLP backend
+        if self.index_nlp_backend not in ['nltk', 'spacy']:
+            logger.warning(f"Invalid index_nlp_backend '{self.index_nlp_backend}', using 'nltk'")
+            self.index_nlp_backend = 'nltk'
+        
         self.doc_processor = DocumentProcessor(
             chunking_strategy=chunking_strategy,
             max_sentences_per_chunk=max_sentences_per_chunk,
             chunk_size=chunk_size,
-            overlap_size=chunk_overlap,
-            split_newlines=split_newlines
+            chunk_overlap=chunk_overlap,
+            split_newlines=split_newlines,
+            index_nlp_backend=self.index_nlp_backend,
+            verbose=self.verbose,
+            semantic_threshold=self.semantic_threshold,
+            topic_threshold=self.topic_threshold
         )
     
+    def _extract_metadata_from_json_content(self, content: str) -> tuple[Dict[str, Any], str]:
+        """
+        Extract metadata from JSON content if present
+        
+        Returns:
+            (metadata_dict, metadata_text)
+        """
+        metadata_dict = {}
+        
+        # Try to extract metadata from JSON structure in content
+        if '"metadata":' in content:
+            try:
+                # Look for metadata object in content
+                import re
+                # Find all metadata objects
+                pattern = r'"metadata"\s*:\s*(\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\})'
+                matches = re.finditer(pattern, content)
+                
+                for match in matches:
+                    try:
+                        json_metadata = json.loads(match.group(1))
+                        # Merge all found metadata
+                        if isinstance(json_metadata, dict):
+                            metadata_dict.update(json_metadata)
+                    except:
+                        pass
+            except Exception as e:
+                logger.debug(f"Error extracting JSON metadata: {e}")
+        
+        # Create searchable text from all metadata keys and values
+        metadata_text_parts = []
+        for key, value in metadata_dict.items():
+            # Add key
+            metadata_text_parts.append(str(key))
+            # Add value(s)
+            if isinstance(value, list):
+                metadata_text_parts.extend(str(v) for v in value)
+            else:
+                metadata_text_parts.append(str(value))
+        
+        metadata_text = ' '.join(metadata_text_parts).lower()
+        
+        return metadata_dict, metadata_text
+    
     def _load_model(self):
         """Load embedding model (lazy loading)"""
         if self.model is None:
@@ -75,7 +163,8 @@ class IndexBuilder:
     
     def build_index_from_sources(self, sources: List[Path], output_file: str, 
                                 file_types: List[str], exclude_patterns: Optional[List[str]] = None,
-                                languages: List[str] = None, tags: Optional[List[str]] = None):
+                                languages: List[str] = None, tags: Optional[List[str]] = None,
+                                overwrite: bool = False):
         """
         Build complete search index from multiple sources (files and directories)
         
@@ -99,6 +188,7 @@ class IndexBuilder:
         
         # Process documents
         chunks = []
+        print(f"Processing {len(files)} files...")
         for file_path in files:
             try:
                 # For individual files, use the file's parent as the base directory
@@ -106,8 +196,8 @@ class IndexBuilder:
                 base_dir = self._get_base_directory_for_file(file_path, sources)
                 file_chunks = self._process_file(file_path, base_dir, tags)
                 chunks.extend(file_chunks)
-                if self.verbose:
-                    print(f"Processed {file_path}: {len(file_chunks)} chunks")
+                if self.verbose or file_path.suffix == '.json':
+                    print(f"  {file_path}: {len(file_chunks)} chunks")
             except Exception as e:
                 logger.error(f"Error processing {file_path}: {e}")
                 if self.verbose:
@@ -123,26 +213,47 @@ class IndexBuilder:
         # Generate embeddings
         self._load_model()
         if self.verbose:
-            print("Generating embeddings...")
+            print(f"Generating embeddings for {len(chunks)} chunks...")
+        else:
+            print(f"Generating embeddings for {len(chunks)} chunks...")
         
         for i, chunk in enumerate(chunks):
             try:
                 # Preprocess content for better search
                 processed = preprocess_document_content(
                     chunk['content'], 
-                    language=chunk.get('language', 'en')
+                    language=chunk.get('language', 'en'),
+                    index_nlp_backend=self.index_nlp_backend
                 )
                 
                 chunk['processed_content'] = processed['enhanced_text']
-                chunk['keywords'] = processed.get('keywords', [])
+                
+                # Include tags in keywords for better search matching
+                keywords = processed.get('keywords', [])
+                chunk_tags = chunk.get('tags', [])
+                if chunk_tags:
+                    # Add tags to keywords list for FTS matching
+                    keywords.extend(chunk_tags)
+                    # Remove duplicates while preserving order
+                    keywords = list(dict.fromkeys(keywords))
+                
+                chunk['keywords'] = keywords
+                
+                # For embedding, include tags in the text for better semantic matching
+                embedding_text = processed['enhanced_text']
+                if chunk_tags:
+                    # Append tags to the text for embedding generation
+                    embedding_text += " " + " ".join(chunk_tags)
                 
                 # Generate embedding (suppress progress bar)
-                embedding = self.model.encode(processed['enhanced_text'], show_progress_bar=False)
+                embedding = self.model.encode(embedding_text, show_progress_bar=False)
                 chunk['embedding'] = embedding.tobytes()
                 
-                if self.verbose and (i + 1) % 50 == 0:
+                # Show progress more frequently
+                show_every = 50 if len(chunks) > 500 else max(10, len(chunks) // 10)
+                if (i + 1) % show_every == 0 or (i + 1) == len(chunks):
                     progress_pct = ((i + 1) / len(chunks)) * 100
-                    print(f"Generated embeddings: {i + 1}/{len(chunks)} chunks ({progress_pct:.1f}%)")
+                    print(f"  Progress: {i + 1}/{len(chunks)} chunks ({progress_pct:.1f}%)")
                     
             except Exception as e:
                 logger.error(f"Error processing chunk {i}: {e}")
@@ -156,19 +267,24 @@ class IndexBuilder:
                 else:
                     chunk['embedding'] = b''
         
-        # Create SQLite database
-        sources_info = [str(s) for s in sources]
-        self._create_database(output_file, chunks, languages or ['en'], sources_info, file_types)
-        
-        if self.verbose:
-            print(f"Index created: {output_file}")
-            print(f"Total chunks: {len(chunks)}")
+        # Store chunks based on backend
+        if self.backend == 'sqlite':
+            # Create SQLite database
+            sources_info = [str(s) for s in sources]
+            self._create_database(output_file, chunks, languages or ['en'], sources_info, file_types)
+            
+            if self.verbose:
+                print(f"Index created: {output_file}")
+                print(f"Total chunks: {len(chunks)}")
+        else:
+            # Use pgvector backend
+            self._store_chunks_pgvector(chunks, output_file, languages or ['en'], overwrite)
 
     def build_index(self, source_dir: str, output_file: str, 
                    file_types: List[str], exclude_patterns: Optional[List[str]] = None,
                    languages: List[str] = None, tags: Optional[List[str]] = None):
         """
-        Build complete search index from a single directory (legacy method)
+        Build complete search index from a single directory
         
         Args:
             source_dir: Directory to scan for documents
@@ -332,16 +448,57 @@ class IndexBuilder:
                      global_tags: Optional[List[str]] = None) -> List[Dict[str, Any]]:
         """Process single file into chunks"""
         try:
-            # Try to read as text first
-            try:
-                content = file_path.read_text(encoding='utf-8')
-            except UnicodeDecodeError:
+            relative_path = str(file_path.relative_to(source_dir))
+            file_extension = file_path.suffix.lower()
+            
+            # Handle different file types appropriately
+            if file_extension == '.pdf':
+                # Use document processor for PDF extraction
+                content_result = self.doc_processor._extract_text_from_file(str(file_path))
+                if isinstance(content_result, str) and content_result.startswith('{"error"'):
+                    if self.verbose:
+                        print(f"Skipping PDF file (extraction failed): {file_path}")
+                    return []
+                content = content_result
+            elif file_extension in ['.docx', '.xlsx', '.pptx']:
+                # Use document processor for Office documents
+                content_result = self.doc_processor._extract_text_from_file(str(file_path))
+                if isinstance(content_result, str) and content_result.startswith('{"error"'):
+                    if self.verbose:
+                        print(f"Skipping office document (extraction failed): {file_path}")
+                    return []
+                content = content_result
+            elif file_extension == '.html':
+                # Use document processor for HTML
+                content_result = self.doc_processor._extract_text_from_file(str(file_path))
+                if isinstance(content_result, str) and content_result.startswith('{"error"'):
+                    if self.verbose:
+                        print(f"Skipping HTML file (extraction failed): {file_path}")
+                    return []
+                content = content_result
+            elif file_extension == '.rtf':
+                # Use document processor for RTF
+                content_result = self.doc_processor._extract_text_from_file(str(file_path))
+                if isinstance(content_result, str) and content_result.startswith('{"error"'):
+                    if self.verbose:
+                        print(f"Skipping RTF file (extraction failed): {file_path}")
+                    return []
+                content = content_result
+            else:
+                # Try to read as text file (markdown, txt, code, etc.)
+                try:
+                    content = file_path.read_text(encoding='utf-8')
+                except UnicodeDecodeError:
+                    if self.verbose:
+                        print(f"Skipping binary file: {file_path}")
+                    return []
+            
+            # Validate content
+            if not content or (isinstance(content, str) and len(content.strip()) == 0):
                 if self.verbose:
-                    print(f"Skipping binary file: {file_path}")
+                    print(f"Skipping empty file: {file_path}")
                 return []
             
-            relative_path = str(file_path.relative_to(source_dir))
-            
             # Create chunks using document processor - pass content directly, not file path
             chunks = self.doc_processor.create_chunks(
                 content=content,  # Pass the actual content, not the file path
@@ -390,6 +547,7 @@ class IndexBuilder:
                     end_line INTEGER,
                     tags TEXT,
                     metadata TEXT,
+                    metadata_text TEXT,  -- Searchable text representation of all metadata
                     chunk_hash TEXT UNIQUE,
                     created_at TEXT DEFAULT CURRENT_TIMESTAMP
                 )
@@ -399,6 +557,7 @@ class IndexBuilder:
                 CREATE VIRTUAL TABLE chunks_fts USING fts5(
                     processed_content,
                     keywords,
+                    metadata_text,
                     content='chunks',
                     content_rowid='id'
                 )
@@ -460,13 +619,47 @@ class IndexBuilder:
                 # Prepare data
                 keywords_json = json.dumps(chunk.get('keywords', []))
                 tags_json = json.dumps(chunk.get('tags', []))
-                metadata_json = json.dumps(chunk.get('metadata', {}))
+                
+                # Extract metadata from JSON content and merge with chunk metadata
+                json_metadata, json_metadata_text = self._extract_metadata_from_json_content(chunk['content'])
+                chunk_metadata = chunk.get('metadata', {})
+                
+                # Merge metadata: chunk metadata takes precedence
+                merged_metadata = {**json_metadata, **chunk_metadata}
+                metadata_json = json.dumps(merged_metadata)
+                
+                # Create comprehensive metadata_text including tags
+                metadata_text_parts = []
+                
+                # Add metadata text from JSON content
+                if json_metadata_text:
+                    metadata_text_parts.append(json_metadata_text)
+                
+                # Add tags
+                tags = chunk.get('tags', [])
+                if tags:
+                    metadata_text_parts.extend(str(tag).lower() for tag in tags)
+                
+                # Add section if present
+                if chunk.get('section'):
+                    metadata_text_parts.append(chunk['section'].lower())
+                
+                # Add any additional metadata values
+                for key, value in chunk_metadata.items():
+                    if key not in json_metadata:  # Avoid duplicates
+                        metadata_text_parts.append(str(key).lower())
+                        if isinstance(value, list):
+                            metadata_text_parts.extend(str(v).lower() for v in value)
+                        else:
+                            metadata_text_parts.append(str(value).lower())
+                
+                metadata_text = ' '.join(metadata_text_parts)
                 
                 cursor.execute('''
                     INSERT OR IGNORE INTO chunks (
                         content, processed_content, keywords, language, embedding,
-                        filename, section, start_line, end_line, tags, metadata, chunk_hash
-                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                        filename, section, start_line, end_line, tags, metadata, metadata_text, chunk_hash
+                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                 ''', (
                     chunk['content'],
                     chunk.get('processed_content', chunk['content']),
@@ -479,6 +672,7 @@ class IndexBuilder:
                     chunk.get('end_line'),
                     tags_json,
                     metadata_json,
+                    metadata_text,
                     chunk_hash
                 ))
             
@@ -531,4 +725,80 @@ class IndexBuilder:
             }
             
         except Exception as e:
-            return {"valid": False, "error": str(e)} 
+            return {"valid": False, "error": str(e)}
+    
+    def _store_chunks_pgvector(self, chunks: List[Dict[str, Any]], collection_name: str,
+                              languages: List[str], overwrite: bool = False):
+        """
+        Store chunks in pgvector backend
+        
+        Args:
+            chunks: List of processed chunks
+            collection_name: Name for the collection (from output_file parameter)
+            languages: List of supported languages
+        """
+        from .pgvector_backend import PgVectorBackend
+        
+        # Extract collection name from the provided name
+        if collection_name.endswith('.swsearch'):
+            collection_name = collection_name[:-9]  # Remove .swsearch extension
+        
+        # Clean collection name for PostgreSQL
+        import re
+        collection_name = re.sub(r'[^a-zA-Z0-9_]', '_', collection_name)
+        
+        if self.verbose:
+            print(f"Storing chunks in pgvector collection: {collection_name}")
+        
+        # Create backend instance
+        backend = PgVectorBackend(self.connection_string)
+        
+        try:
+            # Get embedding dimensions from model
+            if self.model:
+                embedding_dim = self.model.get_sentence_embedding_dimension()
+            else:
+                embedding_dim = 768  # Default for all-mpnet-base-v2
+            
+            # Delete existing collection if overwrite is requested
+            if overwrite:
+                if self.verbose:
+                    print(f"Dropping existing collection: {collection_name}")
+                backend.delete_collection(collection_name)
+            
+            # Create schema
+            backend.create_schema(collection_name, embedding_dim)
+            
+            # Convert embeddings from bytes to numpy arrays
+            for chunk in chunks:
+                if chunk.get('embedding') and isinstance(chunk['embedding'], bytes):
+                    if np:
+                        chunk['embedding'] = np.frombuffer(chunk['embedding'], dtype=np.float32)
+                    else:
+                        # If numpy not available, leave as bytes
+                        pass
+            
+            # Prepare config
+            config = {
+                'model_name': self.model_name,
+                'embedding_dimensions': embedding_dim,
+                'chunking_strategy': self.chunking_strategy,
+                'languages': languages,
+                'metadata': {
+                    'max_sentences_per_chunk': self.max_sentences_per_chunk,
+                    'chunk_size': self.chunk_size,
+                    'chunk_overlap': self.chunk_overlap,
+                    'index_nlp_backend': self.index_nlp_backend
+                }
+            }
+            
+            # Store chunks
+            backend.store_chunks(chunks, collection_name, config)
+            
+            if self.verbose:
+                stats = backend.get_stats(collection_name)
+                print(f"Stored {stats['total_chunks']} chunks in pgvector")
+                print(f"Collection: {collection_name}")
+                
+        finally:
+            backend.close()