cleaning code -> making phase 1 pipeline ready

config.yaml

@@ -4,7 +4,7 @@ project_name: "arxiv_cyber_advisor"
 # Stage 1: Data Acquisition
 data_ingestion:
   category: "cs.AI"
-  limit: 20
+  limit: 5
   save_local: true
   raw_data_path: "data/raw_arxiv.csv"
 

data_loader.py

@@ -35,7 +35,7 @@ def fetch_arxiv_data(category="cs.AI", limit=5):
             "id": r.entry_id.split('/')[-1],
             "title": r.title,
             "abstract": r.summary.replace('\n', ' '),
-            "full_text": full_text, # <--- NEW FIELD
+            "full_text": full_text, # <--- Main part of the data
             "url": r.pdf_url
         })
     return pd.DataFrame(results)

main.py

@@ -1,28 +1,95 @@
-import yaml
-from data_processor import fetch_arxiv_data, process_to_chunks
+import os
+from dotenv import load_dotenv
 
-def load_config():
-    with open("config.yaml", "r") as f:
-        return yaml.safe_load(f)
+from vector_db import get_pinecone_index, refresh_pinecone_index
+from retriever.retriever import HybridRetriever
+from retriever.generator import RAGGenerator
+from retriever.processor import ChunkProcessor
+import data_loader as dl
+
+from models.llama_3_8b import Llama3_8B
+from models.mistral_7b import Mistral_7b
+from models.qwen_2_5 import Qwen2_5
+from models.deepseek_v3 import DeepSeek_V3
+from models.tiny_aya import TinyAya
+
+load_dotenv()
 
 def main():
-    config = load_config()
-    
-    # Run Stage 1
-    raw_data = fetch_arxiv_data(
-        category=config['data_ingestion']['category'], 
-        limit=config['data_ingestion']['limit']
-    )
 
-    # Run Stage 2 using YAML defaults
-    final_chunks = process_to_chunks(
+    # ------------------------------------------------------------------
+    # 0. Configuration
+    # ------------------------------------------------------------------
+    hf_token        = os.getenv("HF_TOKEN")
+    pinecone_api_key = os.getenv("PINECONE_API_KEY")
+    if not pinecone_api_key:
+        raise ValueError("PINECONE_API_KEY not found in environment variables")
+
+    query      = "How do transformers handle long sequences?"
+
+    # ------------------------------------------------------------------
+    # 1. Data Ingestion
+    # ------------------------------------------------------------------
+    raw_data = dl.fetch_arxiv_data(category="cs.AI", limit=5)
+
+    # ------------------------------------------------------------------
+    # 2. Chunking & Embedding
+    # ------------------------------------------------------------------
+    proc = ChunkProcessor(model_name='all-MiniLM-L6-v2', verbose=True)
+    final_chunks = proc.process(
         raw_data,
-        model_name=config['embedding']['model_name'],
-        chunk_size=config['chunking']['chunk_size'],
-        chunk_overlap=config['chunking']['chunk_overlap']
+        technique="sentence",   # options: fixed, recursive, character, sentence, semantic
+        chunk_size=500,
+        chunk_overlap=50
+    )
+
+    # ------------------------------------------------------------------
+    # 3. Vector DB
+    # ------------------------------------------------------------------
+    index_name = "arxiv-index"
+    index = get_pinecone_index(pinecone_api_key, index_name, dimension=384, metric="cosine")
+    refresh_pinecone_index(index, final_chunks, batch_size=100)
+
+    # ------------------------------------------------------------------
+    # 4. Retrieval
+    # ------------------------------------------------------------------
+    retriever = HybridRetriever(final_chunks, proc.encoder, verbose=True)
+    context_chunks = retriever.search(
+        query,
+        index,
+        mode="hybrid",   # options: bm25, semantic, hybrid
+        rerank_strategy="cross-encoder", # options: cross-encoder, rrf
+        use_mmr=True,
+        top_k=10,
+        final_k=5
     )
 
-    print(f"✅ Pipeline finished with {len(final_chunks)} chunks.")
+    if not context_chunks:
+        print("No context chunks retrieved. Check your index and query.")
+        return
+
+    # ------------------------------------------------------------------
+    # 5. Generation
+    # ------------------------------------------------------------------
+    rag_engine = RAGGenerator()
+
+    
+
+    models = {
+        "Llama-3-8B":  Llama3_8B(token=hf_token),
+        "Mistral-7B":  Mistral_7b(token=hf_token),
+        "Qwen-2.5":    Qwen2_5(token=hf_token),
+        "DeepSeek-V3": DeepSeek_V3(token=hf_token),
+        "TinyAya":     TinyAya(token=hf_token)
+    }
+
+    for name, model in models.items():
+        print(f"\n--- {name} ---")
+        try:
+            print(rag_engine.get_answer(model, query, context_chunks, temperature=0.1))
+        except Exception as e:
+            print(f"Error: {e}")
+
 
 if __name__ == "__main__":
     main()

models/deepseek_v3.py

@@ -5,7 +5,7 @@ class DeepSeek_V3:
         self.client = InferenceClient(token=token)
         self.model_id = "deepseek-ai/DeepSeek-V3"
 
-    def generate(self, prompt, max_tokens=500, temperature=0.15):
+    def generate(self, prompt, max_tokens=500, temperature=0.1):
         response = ""
         try:
             for message in self.client.chat_completion(
@@ -19,5 +19,5 @@ class DeepSeek_V3:
                     content = message.choices[0].delta.content
                     if content: response += content
         except Exception as e:
-            return f"⚠️ DeepSeek API Busy: {e}"
+            return f" DeepSeek API Busy: {e}"
         return response

models/llama_3_8b.py

@@ -5,13 +5,13 @@ class Llama3_8B:
         self.client = InferenceClient(token=token)
         self.model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 
-    def generate(self, prompt, max_tokens=500, temp=0.1):
+    def generate(self, prompt, max_tokens=500, temperature=0.1):
         response = ""
         for message in self.client.chat_completion(
             model=self.model_id,
             messages=[{"role": "user", "content": prompt}],
             max_tokens=max_tokens,
-            temperature=temp,
+            temperature=temperature,
             stream=True,
         ):
             if message.choices:

models/mistral_7b.py

@@ -1,20 +1,13 @@
-import os
 from huggingface_hub import InferenceClient
 
 class Mistral_7b:
     def __init__(self, token):
-        # Initializing with api_key as per latest documentation
         self.client = InferenceClient(api_key=token)
-        # Using the specific provider suffix
         self.model_id = "mistralai/Mistral-7B-Instruct-v0.2:featherless-ai"
 
-    def generate(self, prompt, max_tokens=500, **kwargs):
-        # Extract temperature, defaulting to 0.2 if not provided
-        temperature = kwargs.get('temperature', kwargs.get('temp', 0.2))
-        
+    def generate(self, prompt, max_tokens=500, temperature=0.1):
         response = ""
         try:
-            # Using the new .chat.completions.create syntax for Featherless
             stream = self.client.chat.completions.create(
                 model=self.model_id,
                 messages=[{"role": "user", "content": prompt}],
@@ -22,14 +15,12 @@ class Mistral_7b:
                 temperature=temperature,
                 stream=True,
             )
-
             for chunk in stream:
-                # Accessing content through the standard completion object structure
                 if chunk.choices and chunk.choices[0].delta.content:
                     content = chunk.choices[0].delta.content
                     response += content
                     
         except Exception as e:
-            return f"❌ Mistral Featherless Error: {e}"
+            return f" Mistral Featherless Error: {e}"
             
         return response

models/qwen_2_5.py

@@ -5,7 +5,7 @@ class Qwen2_5:
         self.client = InferenceClient(token=token)
         self.model_id = "Qwen/Qwen2.5-72B-Instruct"
 
-    def generate(self, prompt, max_tokens=500, temperature=0.3):
+    def generate(self, prompt, max_tokens=500, temperature=0.1):
         response = ""
         for message in self.client.chat_completion(
             model=self.model_id,

models/tiny_aya.py

@@ -3,16 +3,10 @@ from huggingface_hub import InferenceClient
 class TinyAya:
     def __init__(self, token):
         self.client = InferenceClient(token=token)
-        # 3.3B parameter model, great for multilingual/efficient RAG
         self.model_id = "CohereLabs/tiny-aya-global"
 
-    def generate(self, prompt, max_tokens=400, **kwargs):
-        """
-        Using **kwargs makes this compatible with calls using 'temp' or 'temperature'.
-        """
-        # This line looks for 'temperature', then 'temp', and defaults to 0.3 if neither exist
-        temperature = kwargs.get('temperature', kwargs.get('temp', 0.3))
-        
+    def generate(self, prompt, max_tokens=500, temperature=0.1):
+
         response = ""
         try:
             for message in self.client.chat_completion(
@@ -26,6 +20,6 @@ class TinyAya:
                     content = message.choices[0].delta.content
                     if content: response += content
         except Exception as e:
-            return f"❌ TinyAya Error: {e}"
+            return f" TinyAya Error: {e}"
         
         return response

retriever/generator.py

@@ -15,5 +15,5 @@ Answer:"""
 
     def get_answer(self, model_instance, query, retrieved_contexts, **kwargs):
         """Uses a specific model instance to generate the final answer."""
-        prompt = self.generate_prompt(query, retrieved_contexts)
+        prompt = self.generate_prompt(query, retrieved_contexts)    
         return model_instance.generate(prompt, **kwargs)

retriever/processor.py

@@ -1,169 +1,133 @@
 from langchain_text_splitters import (
     RecursiveCharacterTextSplitter, 
     CharacterTextSplitter,
-    SentenceTransformersTokenTextSplitter
+    SentenceTransformersTokenTextSplitter,
+    NLTKTextSplitter
 )
 from langchain_experimental.text_splitter import SemanticChunker
 from langchain_huggingface import HuggingFaceEmbeddings
 from sentence_transformers import SentenceTransformer
 from typing import List, Dict, Any, Optional
+import nltk
+nltk.download('punkt_tab', quiet=True)
 import pandas as pd
 
+
 class ChunkProcessor:
     def __init__(self, model_name='all-MiniLM-L6-v2', verbose: bool = True):
         self.model_name = model_name
         self.encoder = SentenceTransformer(model_name)
         self.verbose = verbose
-        # Required for Semantic Chunking
         self.hf_embeddings = HuggingFaceEmbeddings(model_name=model_name)
 
-    def _print(self, *args, **kwargs):
-        """Helper method to conditionally print"""
-        if self.verbose:
-            print(*args, **kwargs)
+    # ------------------------------------------------------------------
+    # Splitters
+    # ------------------------------------------------------------------
 
     def get_splitter(self, technique: str, chunk_size: int = 500, chunk_overlap: int = 50, **kwargs):
         """
         Factory method to return different chunking strategies.
-        
+
         Strategies:
-        - "fixed": Simple character-based splitting with empty separator (can split mid-sentence)
-        - "recursive": Recursive character splitting with hierarchical separators (preserves semantics)
-        - "character": Character-based splitting with paragraph separator
-        - "sentence": Recursive splitting optimized for sentence boundaries
-        - "semantic": Embedding-based semantic chunking
-        - "token": Token-based splitting for transformer models
+        - "fixed":     Character-based, may split mid-sentence
+        - "recursive": Recursive character splitting with hierarchical separators
+        - "character": Character-based splitting on paragraph boundaries
+        - "sentence":  Sliding window over NLTK sentences
+        - "semantic":  Embedding-based semantic chunking
         """
         if technique == "fixed":
-            # FIXED: Simple character-based splitter - WILL split mid-sentence
             return CharacterTextSplitter(
-                separator=kwargs.get('separator', ""), 
-                chunk_size=chunk_size, 
+                separator=kwargs.get('separator', ""),
+                chunk_size=chunk_size,
                 chunk_overlap=chunk_overlap,
                 length_function=len,
                 is_separator_regex=False
             )
-        
+
         elif technique == "recursive":
-            # FIXED: Proper recursive splitter with default separators that preserve semantics
-            separators = kwargs.get('separators', ["\n\n", "\n", ". ", "! ", "? ", "; ", ", ", " ", ""])
             return RecursiveCharacterTextSplitter(
-                chunk_size=chunk_size, 
+                chunk_size=chunk_size,
                 chunk_overlap=chunk_overlap,
-                separators=separators,
+                separators=kwargs.get('separators', ["\n\n", "\n", ". ", "! ", "? ", "; ", ", ", " ", ""]),
                 length_function=len,
                 keep_separator=kwargs.get('keep_separator', True)
             )
-        
+
         elif technique == "character":
-            # FIXED: Character splitter with paragraph separator
             return CharacterTextSplitter(
-                separator=kwargs.get('separator', "\n\n"), 
-                chunk_size=chunk_size, 
+                separator=kwargs.get('separator', "\n\n"),
+                chunk_size=chunk_size,
                 chunk_overlap=chunk_overlap,
                 length_function=len,
                 is_separator_regex=False
             )
-        
+
         elif technique == "sentence":
-            # FIXED: Using Recursive Splitter with comprehensive sentence boundaries
-            # This preserves full sentences whenever possible
-            return RecursiveCharacterTextSplitter(
+            # sentence-level chunking using NLTK
+            return NLTKTextSplitter(
                 chunk_size=chunk_size,
                 chunk_overlap=chunk_overlap,
-                separators=kwargs.get('separators', ["\n\n", "\n", ". ", "? ", "! ", ".\n", "?\n", "!\n", "; ", ": ", ", ", " ", ""]),
-                length_function=len,
-                keep_separator=kwargs.get('keep_separator', True)
-            )
-        
+                separator="\n"
+            )  
+
         elif technique == "semantic":
-            # FIXED: Semantic chunker with proper configuration
             return SemanticChunker(
-                self.hf_embeddings, 
+                self.hf_embeddings,
                 breakpoint_threshold_type=kwargs.get('breakpoint_threshold_type', "percentile"),
-                breakpoint_threshold_amount=kwargs.get('breakpoint_threshold_amount', 95),
-                min_chunk_size=kwargs.get('min_chunk_size', chunk_size // 10),
-                max_chunk_size=kwargs.get('max_chunk_size', chunk_size)
-            )
-        
-        elif technique == "token":
-            # FIXED: Token-based splitter with proper token counting
-            return SentenceTransformersTokenTextSplitter(
-                model_name=self.model_name, 
-                tokens_per_chunk=chunk_size,
-                chunk_overlap=chunk_overlap,
-                length_function=kwargs.get('length_function', lambda x: len(self.encoder.encode(x)))
+                breakpoint_threshold_amount=kwargs.get('breakpoint_threshold_amount', 95)
             )
-        
+
         else:
-            raise ValueError(f"Technique '{technique}' is not supported. Choose from: fixed, recursive, character, sentence, semantic, token")
+            raise ValueError(f"Technique '{technique}' is not supported. Choose from: fixed, recursive, character, sentence, semantic")
+
+    # ------------------------------------------------------------------
+    # Processing
+    # ------------------------------------------------------------------
 
-    def process(self, df: pd.DataFrame, technique: str = "recursive", chunk_size: int = 500, 
-                chunk_overlap: int = 50, max_docs: Optional[int] = 5, verbose: Optional[bool] = None, 
-                **kwargs) -> List[Dict[str, Any]]:
+    def process(self, df: pd.DataFrame, technique: str = "recursive", chunk_size: int = 500,
+                chunk_overlap: int = 50, max_docs: Optional[int] = 5,
+                verbose: Optional[bool] = None, **kwargs) -> List[Dict[str, Any]]:
         """
-        Processes a DataFrame into vector-ready chunks with full output for documents.
-        
+        Processes a DataFrame into vector-ready chunks.
+
         Args:
-            df: DataFrame containing documents with columns: id, title, url, full_text
-            technique: Chunking strategy to use
-            chunk_size: Maximum size of each chunk (characters for most, tokens for token splitter)
+            df:            DataFrame with columns: id, title, url, full_text
+            technique:     Chunking strategy to use
+            chunk_size:    Maximum size of each chunk in characters
             chunk_overlap: Overlap between consecutive chunks
-            max_docs: Maximum number of documents to process (None for all)
-            verbose: Override the instance's verbose setting (if None, uses instance setting)
-            **kwargs: Additional arguments to pass to splitter
-            
+            max_docs:      Number of documents to process (None for all)
+            verbose:       Override instance verbose setting
+            **kwargs:      Additional arguments passed to the splitter
+
         Returns:
-            List of processed chunks with embeddings and metadata
+            List of chunk dicts with embeddings and metadata
         """
-        # Determine if we should print
         should_print = verbose if verbose is not None else self.verbose
-        
-        splitter = self.get_splitter(technique, chunk_size, chunk_overlap, **kwargs)
-        processed_chunks = []
-        
-        # Select documents to process
-        if max_docs:
-            subset_df = df.head(max_docs)
-        else:
-            subset_df = df
-        
-        # Validate required columns exist
+
         required_cols = ['id', 'title', 'url', 'full_text']
-        missing_cols = [col for col in required_cols if col not in subset_df.columns]
+        missing_cols = [col for col in required_cols if col not in df.columns]
         if missing_cols:
             raise ValueError(f"DataFrame missing required columns: {missing_cols}")
-        
+
+        splitter = self.get_splitter(technique, chunk_size, chunk_overlap, **kwargs)
+        subset_df = df.head(max_docs) if max_docs else df
+        processed_chunks = []
+
         for _, row in subset_df.iterrows():
             if should_print:
-                self._print("\n" + "="*80)
-                self._print(f"📄 DOCUMENT: {row['title']}")
-                self._print(f"🔗 URL: {row['url']}")
-                self._print(f"📏 Technique: {technique.upper()} | Chunk Size: {chunk_size} | Overlap: {chunk_overlap}")
-                self._print("-" * 80)
-            
-            # Split the text
+                self._print_document_header(row['title'], row['url'], technique, chunk_size, chunk_overlap)
+
             raw_chunks = splitter.split_text(row['full_text'])
-            
-            if should_print:
-                self._print(f"🎯 Total Chunks Generated: {len(raw_chunks)}")
-            
+
             for i, text in enumerate(raw_chunks):
-                # Standardize output (handle both string and Document objects)
                 content = text.page_content if hasattr(text, 'page_content') else text
-                
+
                 if should_print:
-                    # Print chunk preview
-                    self._print(f"\n[Chunk {i}] ({len(content)} chars):")
-                    preview = content[:200] + "..." if len(content) > 200 else content
-                    self._print(f"   {preview}")
-                
-                # Generate embedding
-                embedding = self.encoder.encode(content).tolist()
-                
+                    self._print_chunk(i, content)
+
                 processed_chunks.append({
                     "id": f"{row['id']}-chunk-{i}",
-                    "values": embedding,
+                    "values": self.encoder.encode(content).tolist(),
                     "metadata": {
                         "title": row['title'],
                         "text": content,
@@ -174,67 +138,37 @@ class ChunkProcessor:
                         "total_chunks": len(raw_chunks)
                     }
                 })
-            
+
             if should_print:
-                self._print("="*80)
-        
+                self._print_document_summary(len(raw_chunks))
+
         if should_print:
-            self._print(f"\n✅ Finished processing {len(subset_df)} documents into {len(processed_chunks)} chunks.")
-            if len(processed_chunks) > 0:
-                self._print(f"📊 Average chunk size: {sum(c['metadata']['chunk_size'] for c in processed_chunks) / len(processed_chunks):.0f} chars")
-        
+            self._print_processing_summary(len(subset_df), processed_chunks)
+
         return processed_chunks
 
-    def compare_strategies(self, df: pd.DataFrame, text_column: str = 'full_text', 
-                          chunk_size: int = 500, max_docs: int = 1, 
-                          verbose: Optional[bool] = None) -> Dict[str, Any]:
-        """
-        Compare different chunking strategies on the same document.
-        
-        Returns:
-            Dictionary with comparison metrics for each strategy
-        """
-        # Determine if we should print
-        should_print = verbose if verbose is not None else self.verbose
-        
-        strategies = ['fixed', 'recursive', 'character', 'sentence', 'semantic', 'token']
-        results = {}
-        
-        # Get first document
-        sample_text = df.iloc[0][text_column]
-        
-        for technique in strategies:
-            try:
-                if should_print:
-                    self._print(f"\n🔍 Testing {technique.upper()} strategy...")
-                
-                splitter = self.get_splitter(technique, chunk_size=chunk_size)
-                chunks = splitter.split_text(sample_text)
-                
-                # Analyze chunks
-                chunk_lengths = [len(c.page_content if hasattr(c, 'page_content') else c) for c in chunks]
-                avg_chunk_size = sum(chunk_lengths) / len(chunk_lengths) if chunk_lengths else 0
-                
-                # Count how many chunks end with sentence boundaries
-                sentence_enders = ['.', '!', '?', '"', "'"]
-                complete_sentences = sum(1 for c in chunks 
-                                        if (c.page_content if hasattr(c, 'page_content') else c).strip()[-1] in sentence_enders)
-                
-                results[technique] = {
-                    'num_chunks': len(chunks),
-                    'avg_chunk_size': avg_chunk_size,
-                    'min_chunk_size': min(chunk_lengths) if chunk_lengths else 0,
-                    'max_chunk_size': max(chunk_lengths) if chunk_lengths else 0,
-                    'complete_sentences_ratio': complete_sentences / len(chunks) if chunks else 0,
-                    'chunk_lengths': chunk_lengths
-                }
-                
-                if should_print:
-                    self._print(f"   ✓ Generated {len(chunks)} chunks, avg size: {avg_chunk_size:.0f} chars")
-                    
-            except Exception as e:
-                results[technique] = {'error': str(e)}
-                if should_print:
-                    self._print(f"   ✗ Error: {str(e)}")
-        
-        return results
+
+    # ------------------------------------------------------------------
+    # Printing
+    # ------------------------------------------------------------------
+
+    def _print_document_header(self, title, url, technique, chunk_size, chunk_overlap):
+        print("\n" + "="*80)
+        print(f"DOCUMENT: {title}")
+        print(f"URL: {url}")
+        print(f"Technique: {technique.upper()} | Chunk Size: {chunk_size} | Overlap: {chunk_overlap}")
+        print("-" * 80)
+
+    def _print_chunk(self, index, content):
+        print(f"\n[Chunk {index}] ({len(content)} chars):")
+        print(f"   {content}")
+
+    def _print_document_summary(self, num_chunks):
+        print(f"Total Chunks Generated: {num_chunks}")
+        print("="*80)
+
+    def _print_processing_summary(self, num_docs, processed_chunks):
+        print(f"\nFinished processing {num_docs} documents into {len(processed_chunks)} chunks.")
+        if processed_chunks:
+            avg = sum(c['metadata']['chunk_size'] for c in processed_chunks) / len(processed_chunks)
+            print(f"Average chunk size: {avg:.0f} chars")

retriever/retriever.py

@@ -2,194 +2,166 @@ import numpy as np
 from rank_bm25 import BM25Okapi
 from sentence_transformers import CrossEncoder
 from sklearn.metrics.pairwise import cosine_similarity
-from typing import Optional 
+from typing import Optional, List
 
 class HybridRetriever:
     def __init__(self, final_chunks, embed_model, rerank_model_name='cross-encoder/ms-marco-MiniLM-L-6-v2', verbose: bool = True):
-        """
-        :param final_chunks: The list of chunk dictionaries with metadata.
-        :param embed_model: The SentenceTransformer model used for query and chunk embedding.
-        :param verbose: Whether to print retrieval details and final results.
-        """
         self.final_chunks = final_chunks
         self.embed_model = embed_model
         self.rerank_model = CrossEncoder(rerank_model_name)
         self.verbose = verbose
         
-        # Initialize BM25 corpus
         self.tokenized_corpus = [chunk['metadata']['text'].lower().split() for chunk in final_chunks]
         self.bm25 = BM25Okapi(self.tokenized_corpus)
 
-    def _print(self, *args, **kwargs):
-        """Helper method to conditionally print"""
-        if self.verbose:
-            print(*args, **kwargs)
+    # ------------------------------------------------------------------
+    # Retrieval
+    # ------------------------------------------------------------------
+
+    def _semantic_search(self, query, index, top_k) -> tuple[np.ndarray, List[str]]:
+        query_vector = self.embed_model.encode(query)
+        res = index.query(vector=query_vector.tolist(), top_k=top_k, include_metadata=True)
+        chunks = [match['metadata']['text'] for match in res['matches']]
+        return query_vector, chunks
+
+    def _bm25_search(self, query, top_k) -> List[str]:
+        tokenized_query = query.lower().split()
+        scores = self.bm25.get_scores(tokenized_query)
+        top_indices = np.argsort(scores)[::-1][:top_k]
+        return [self.final_chunks[i]['metadata']['text'] for i in top_indices]
+
+    # ------------------------------------------------------------------
+    # Fusion
+    # ------------------------------------------------------------------
 
-    def _rrf_score(self, semantic_results, bm25_results, k=60):
-        """Reciprocal Rank Fusion (RRF) Implementation."""
+    def _rrf_score(self, semantic_results, bm25_results, k=60) -> List[str]:
         scores = {}
         for rank, chunk in enumerate(semantic_results):
             scores[chunk] = scores.get(chunk, 0) + 1 / (k + rank + 1)
         for rank, chunk in enumerate(bm25_results):
             scores[chunk] = scores.get(chunk, 0) + 1 / (k + rank + 1)
-        
-        sorted_chunks = sorted(scores.items(), key=lambda x: x[1], reverse=True)
-        return [item[0] for item in sorted_chunks]
+        return [chunk for chunk, _ in sorted(scores.items(), key=lambda x: x[1], reverse=True)]
 
-    def _maximal_marginal_relevance(self, query_embedding, chunk_texts, lambda_param=0.5, top_k=3):
-        """
-        MMR Re-ranking to balance relevance and diversity.
-        """
-        if not chunk_texts: return []
-        
-        chunk_embeddings = self.embed_model.encode(chunk_texts)
-        query_embedding = query_embedding.reshape(1, -1)
+    # ------------------------------------------------------------------
+    # Reranking
+    # ------------------------------------------------------------------
+
+    def _cross_encoder_rerank(self, query, chunks, final_k) -> List[str]:
+        pairs = [[query, chunk] for chunk in chunks]
+        scores = self.rerank_model.predict(pairs)
+        ranked = sorted(zip(chunks, scores), key=lambda x: x[1], reverse=True)
+        return [chunk for chunk, _ in ranked[:final_k]]
+
+    # ------------------------------------------------------------------
+    # MMR (applied after reranking as a diversity filter)
+    # ------------------------------------------------------------------
+
+    def _maximal_marginal_relevance(self, query_vector, chunks, lambda_param=0.5, top_k=3) -> List[str]:
+        if not chunks:
+            return []
         
-        # Initial relevance scores
+        chunk_embeddings = self.embed_model.encode(chunks)
+        query_embedding = query_vector.reshape(1, -1)
         relevance_scores = cosine_similarity(query_embedding, chunk_embeddings)[0]
         
-        selected_indices = []
-        unselected_indices = list(range(len(chunk_texts)))
+        selected, unselected = [], list(range(len(chunks)))
         
-        # First pick: most relevant
-        idx = np.argmax(relevance_scores)
-        selected_indices.append(idx)
-        unselected_indices.remove(idx)
+        first = int(np.argmax(relevance_scores))
+        selected.append(first)
+        unselected.remove(first)
         
-        while len(selected_indices) < min(top_k, len(chunk_texts)):
-            mmr_scores = []
-            for un_idx in unselected_indices:
-                # Similarity to query
-                rel = relevance_scores[un_idx]
-                # Max similarity to already selected chunks (redundancy)
-                sim_to_selected = max([cosine_similarity(chunk_embeddings[un_idx].reshape(1, -1), 
-                                                        chunk_embeddings[sel_idx].reshape(1, -1))[0][0] 
-                                       for sel_idx in selected_indices])
-                
-                mmr_score = lambda_param * rel - (1 - lambda_param) * sim_to_selected
-                mmr_scores.append((un_idx, mmr_score))
-            
-            next_idx = max(mmr_scores, key=lambda x: x[1])[0]
-            selected_indices.append(next_idx)
-            unselected_indices.remove(next_idx)
-            
-        return [chunk_texts[i] for i in selected_indices]
-
-    def search(self, query, index, top_k=10, final_k=3, mode="hybrid", rerank_strategy="cross-encoder", 
-               verbose: Optional[bool] = None):
+        while len(selected) < min(top_k, len(chunks)):
+            mmr_scores = [
+                (i, lambda_param * relevance_scores[i] - (1 - lambda_param) * max(
+                    cosine_similarity(chunk_embeddings[i].reshape(1, -1),
+                                      chunk_embeddings[s].reshape(1, -1))[0][0]
+                    for s in selected
+                ))
+                for i in unselected
+            ]
+            best = max(mmr_scores, key=lambda x: x[1])[0]
+            selected.append(best)
+            unselected.remove(best)
+        
+        return [chunks[i] for i in selected]
+
+    # ------------------------------------------------------------------
+    # Main search
+    # ------------------------------------------------------------------
+
+    def search(self, query, index, top_k=10, final_k=3, mode="hybrid",
+               rerank_strategy="cross-encoder", use_mmr=True, lambda_param=0.5,
+               verbose: Optional[bool] = None) -> List[str]:
         """
-        :param mode: "semantic", "bm25", or "hybrid"
-        :param rerank_strategy: "cross-encoder", "rrf", "mmr", or "none"
-        :param verbose: Override the instance's verbose setting (if None, uses instance setting)
+        :param mode:             "semantic", "bm25", or "hybrid"
+        :param rerank_strategy:  "cross-encoder", "rrf", or "none"
+        :param use_mmr:          Whether to apply MMR diversity filter after reranking
+        :param lambda_param:     MMR trade-off between relevance (1.0) and diversity (0.0)
         """
-        # Determine if we should print
         should_print = verbose if verbose is not None else self.verbose
         
         if should_print:
-            self._print("\n" + "="*80)
-            self._print(f"🔍 SEARCH QUERY: {query}")
-            self._print(f"📊 Mode: {mode.upper()} | Rerank: {rerank_strategy.upper()}")
-            self._print(f"🎯 Top-K: {top_k} | Final-K: {final_k}")
-            self._print("-" * 80)
-        
-        semantic_chunks = []
-        bm25_chunks = []
+            self._print_search_header(query, mode, rerank_strategy, top_k, final_k)
+
+        # 1. Retrieve candidates
         query_vector = None
+        semantic_chunks, bm25_chunks = [], []
 
-        # 1. Fetch Candidates
         if mode in ["semantic", "hybrid"]:
+            query_vector, semantic_chunks = self._semantic_search(query, index, top_k)
             if should_print:
-                self._print(f"📚 Semantic Search: Retrieving top {top_k} candidates...")
-            
-            query_vector = self.embed_model.encode(query)
-            res = index.query(vector=query_vector.tolist(), top_k=top_k, include_metadata=True)
-            semantic_chunks = [match['metadata']['text'] for match in res['matches']]
-            
-            if should_print:
-                self._print(f"   ✓ Retrieved {len(semantic_chunks)} semantic candidates")
-                for i, chunk in enumerate(semantic_chunks[:3]):  # Show first 3
-                    preview = chunk[:100] + "..." if len(chunk) > 100 else chunk
-                    self._print(f"      [{i}] {preview}")
+                self._print_candidates("Semantic Search", semantic_chunks)
 
         if mode in ["bm25", "hybrid"]:
+            bm25_chunks = self._bm25_search(query, top_k)
             if should_print:
-                self._print(f"📚 BM25 Search: Retrieving top {top_k} candidates...")
-            
-            tokenized_query = query.lower().split()
-            bm25_scores = self.bm25.get_scores(tokenized_query)
-            top_indices = np.argsort(bm25_scores)[::-1][:top_k]
-            bm25_chunks = [self.final_chunks[i]['metadata']['text'] for i in top_indices]
-            
-            if should_print:
-                self._print(f"   ✓ Retrieved {len(bm25_chunks)} BM25 candidates")
-                for i, chunk in enumerate(bm25_chunks[:3]):  # Show first 3
-                    preview = chunk[:100] + "..." if len(chunk) > 100 else chunk
-                    self._print(f"      [{i}] {preview}")
+                self._print_candidates("BM25 Search", bm25_chunks)
 
-        # 2. Re-Ranking / Fusion
-        if mode == "hybrid" and rerank_strategy == "rrf":
-            if should_print:
-                self._print(f"🔄 Applying Reciprocal Rank Fusion (RRF)...")
-            
-            results = self._rrf_score(semantic_chunks, bm25_chunks)[:final_k]
-            
-            if should_print:
-                self._print(f"✅ Final {final_k} Results:")
-                for i, chunk in enumerate(results):
-                    preview = chunk[:150] + "..." if len(chunk) > 150 else chunk
-                    self._print(f"   [{i+1}] {preview}")
-                self._print("="*80)
-            
-            return results
-
-        # Standard combination for other methods
-        combined = list(dict.fromkeys(semantic_chunks + bm25_chunks)) # Deduplicate keep order
-        
-        if should_print:
-            self._print(f"🔄 Combined unique candidates: {len(combined)}")
-            self._print(f"🔄 Applying {rerank_strategy.upper()} reranking...")
-
-        if rerank_strategy == "cross-encoder" and combined:
-            
-            pairs = [[query, chunk] for chunk in combined]
-            scores = self.rerank_model.predict(pairs)
-            results = sorted(zip(combined, scores), key=lambda x: x[1], reverse=True)
-            results = [res[0] for res in results[:final_k]]
-            
-            if should_print:
-                self._print(f"\n✅ Final {final_k} Results (Cross-Encoder Reranked):")
-                for i, chunk in enumerate(results):
-                    preview = chunk[:150] + "..." if len(chunk) > 150 else chunk
-                    self._print(f"   [{i+1}] {preview}")
-                self._print("="*80)
-            
-            return results
-
-        elif rerank_strategy == "mmr" and combined:
-            if should_print:
-                self._print(f"   Using MMR with λ=0.5 to balance relevance and diversity")
-            
-            if query_vector is None: 
+        # 2. Fuse / rerank
+        if rerank_strategy == "rrf":
+            candidates = self._rrf_score(semantic_chunks, bm25_chunks)[:final_k]
+            label = "RRF"
+        elif rerank_strategy == "cross-encoder":
+            combined = list(dict.fromkeys(semantic_chunks + bm25_chunks))
+            candidates = self._cross_encoder_rerank(query, combined, final_k)
+            label = "Cross-Encoder"
+        else:  # "none"
+            candidates = list(dict.fromkeys(semantic_chunks + bm25_chunks))[:final_k]
+            label = "No Reranking"
+
+        # 3. MMR diversity filter (applied after reranking)
+        if use_mmr and candidates:
+            if query_vector is None:
                 query_vector = self.embed_model.encode(query)
-            results = self._maximal_marginal_relevance(query_vector, combined, top_k=final_k)
-            
-            if should_print:
-                self._print(f"\n✅ Final {final_k} Results (MMR Reranked):")
-                for i, chunk in enumerate(results):
-                    preview = chunk[:150] + "..." if len(chunk) > 150 else chunk
-                    self._print(f"   [{i+1}] {preview}")
-                self._print("="*80)
-            
-            return results
-        
-        else:  # "none" or fallback
-            results = combined[:final_k]
-            
-            if should_print:
-                self._print(f"\n✅ Final {final_k} Results (No Reranking):")
-                for i, chunk in enumerate(results):
-                    preview = chunk[:150] + "..." if len(chunk) > 150 else chunk
-                    self._print(f"   [{i+1}] {preview}")
-                self._print("="*80)
-            
-            return results
+            candidates = self._maximal_marginal_relevance(query_vector, candidates,
+                                                          lambda_param=lambda_param, top_k=3)
+            label += " + MMR"
+
+        if should_print:
+            self._print_final_results(candidates, label)
+
+        return candidates
+
+    # ------------------------------------------------------------------
+    # Printing
+    # ------------------------------------------------------------------
+
+    def _print_search_header(self, query, mode, rerank_strategy, top_k, final_k):
+        print("\n" + "="*80)
+        print(f" SEARCH QUERY: {query}")
+        print(f"Mode: {mode.upper()} | Rerank: {rerank_strategy.upper()}")
+        print(f"Top-K: {top_k} | Final-K: {final_k}")
+        print("-" * 80)
+
+    def _print_candidates(self, label, chunks, preview_n=3):
+        print(f"{label}: Retrieved {len(chunks)} candidates")
+        for i, chunk in enumerate(chunks[:preview_n]):
+            preview = chunk[:100] + "..." if len(chunk) > 100 else chunk
+            print(f"   [{i}] {preview}")
+
+    def _print_final_results(self, results, strategy_label):
+        print(f"\n Final {len(results)} Results ({strategy_label}):")
+        for i, chunk in enumerate(results):
+            preview = chunk[:150] + "..." if len(chunk) > 150 else chunk
+            print(f"   [{i+1}] {preview}")
+        print("="*80)

vector_db.py

@@ -22,12 +22,81 @@ def get_pinecone_index(api_key, index_name, dimension=384, metric="cosine"):
     
     return pc.Index(index_name)
 
+def prepare_vectors_for_upsert(final_chunks):
+    """Convert final_chunks to the format expected by Pinecone upsert"""
+    vectors = []
+    for chunk in final_chunks:
+        vectors.append({
+            'id': chunk['id'],
+            'values': chunk['values'],  # The embedding vector
+            'metadata': {
+                'text': chunk['metadata']['text'],
+                'title': chunk['metadata']['title'],
+                'url': chunk['metadata']['url'],
+                'chunk_index': chunk['metadata']['chunk_index'],
+                'technique': chunk['metadata']['technique'],
+                'chunk_size': chunk['metadata']['chunk_size'],
+                'total_chunks': chunk['metadata']['total_chunks']
+            }
+        })
+    return vectors
+
 def upsert_to_pinecone(index, chunks, batch_size=100):
-    """Upserts chunks to Pinecone in manageable batches."""
+    """Upserts chunks to Pinecone in manageable batches.
+    
+    Args:
+        index: Pinecone index object
+        chunks: List of chunk dictionaries (as returned by prepare_vectors_for_upsert)
+        batch_size: Number of vectors to upsert in each batch
+    """
     print(f"Uploading {len(chunks)} chunks to Pinecone...")
     
     for i in range(0, len(chunks), batch_size):
         batch = chunks[i : i + batch_size]
         index.upsert(vectors=batch)
+        print(f"   Uploaded batch {i//batch_size + 1}/{(len(chunks)-1)//batch_size + 1} ({len(batch)} vectors)")
+        
+    print(" Upsert complete.")
+
+def refresh_pinecone_index(index, final_chunks, batch_size=100):
+    """Helper function to refresh index with new chunks.
+    
+    This function checks if the index has the expected number of vectors,
+    and upserts if necessary.
+    
+    Args:
+        index: Pinecone index object
+        final_chunks: List of chunk dictionaries with embeddings
+        batch_size: Batch size for upsert
+    
+    Returns:
+        Boolean indicating if upsert was performed
+    """
+    try:
+        stats = index.describe_index_stats()
+        current_vector_count = stats.get('total_vector_count', 0)
+        expected_vector_count = len(final_chunks)
+        
+        print(f"\n Current vectors in index: {current_vector_count}")
+        print(f" Expected vectors: {expected_vector_count}")
         
-    print("✅ Upsert complete.")
+        if current_vector_count == 0:
+            print(" Index is empty. Preparing vectors for upsert...")
+            vectors_to_upsert = prepare_vectors_for_upsert(final_chunks)
+            upsert_to_pinecone(index, vectors_to_upsert, batch_size)
+            
+            # Verify upsert
+            stats = index.describe_index_stats()
+            print(f" After upsert - Total vectors: {stats.get('total_vector_count', 0)}")
+            return True
+        elif current_vector_count != expected_vector_count:
+            print(f"  Vector count mismatch. Expected {expected_vector_count}, got {current_vector_count}")
+            print("   Consider recreating the index if you want to refresh.")
+            return False
+        else:
+            print(f"ℹ Index already has {current_vector_count} vectors. Ready for search.")
+            return False
+            
+    except Exception as e:
+        print(f"Error checking index stats: {e}")
+        return False