By default now use spacy for retrieval and augmentation (vs embeddings)

nlp_rag.py

@@ -0,0 +1,144 @@
+import spacy
+from itertools import groupby
+from operator import itemgetter
+from langsmith import traceable
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import numpy as np
+
+def get_nlp_model():
+    if not spacy.util.is_package("en_core_web_md"):
+        print("Downloading en_core_web_md model...")
+        spacy.cli.download("en_core_web_md")
+        print("Model downloaded successfully!")
+    nlp = spacy.load("en_core_web_md")
+    return nlp
+
+
+def recursive_split_documents(contents, max_chunk_size=1000, overlap=100):
+    from langchain_core.documents.base import Document
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+    documents = []
+    for content in contents:
+        try:
+            page_content = content['page_content']
+            if page_content:
+                metadata = {'title': content['title'], 'source': content['link']}
+                doc = Document(page_content=content['page_content'], metadata=metadata)
+                documents.append(doc)
+        except Exception as e:
+            print(f"Error processing content for {content['link']}: {e}")
+
+    # Initialize recursive text splitter
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=max_chunk_size, chunk_overlap=overlap)
+
+    # Split documents
+    split_documents = text_splitter.split_documents(documents)
+
+    # Convert split documents to the same format as recursive_split
+    chunks = []
+    for doc in split_documents:
+        chunk = {
+            'text': doc.page_content,
+            'metadata': {
+                'title': doc.metadata.get('title', ''),
+                'source': doc.metadata.get('source', '')
+            }
+        }
+        chunks.append(chunk)
+
+    return chunks
+
+
+def semantic_search(query, chunks, nlp, similarity_threshold=0.5, top_n=10):
+    # Precompute query vector and its norm
+    query_vector = nlp(query).vector
+    query_norm = np.linalg.norm(query_vector) + 1e-8  # Add epsilon to avoid division by zero
+
+    # Check if chunks have precomputed vectors; if not, compute them
+    if 'vector' not in chunks[0]:
+        texts = [chunk['text'] for chunk in chunks]
+
+        # Process texts in batches using nlp.pipe()
+        batch_size = 1000  # Adjust based on available memory
+        with nlp.disable_pipes(*[pipe for pipe in nlp.pipe_names if pipe != 'tok2vec']):
+            docs = nlp.pipe(texts, batch_size=batch_size)
+
+        # Add vectors to chunks
+        for chunk, doc in zip(chunks, docs):
+            chunk['vector'] = doc.vector
+
+    # Prepare chunk vectors and norms
+    chunk_vectors = np.array([chunk['vector'] for chunk in chunks])
+    chunk_norms = np.linalg.norm(chunk_vectors, axis=1) + 1e-8  # Add epsilon to avoid division by zero
+
+    # Compute similarities
+    similarities = np.dot(chunk_vectors, query_vector) / (chunk_norms * query_norm)
+
+    # Filter and sort results
+    relevant_chunks = [
+        (chunk, sim) for chunk, sim in zip(chunks, similarities) if sim > similarity_threshold
+    ]
+    relevant_chunks.sort(key=lambda x: x[1], reverse=True)
+
+    return relevant_chunks[:top_n]
+
+
+# Perform semantic search using spaCy
+def semantic_search(query, chunks, nlp, similarity_threshold=0.5, top_n=10):
+    import numpy as np
+    from concurrent.futures import ThreadPoolExecutor
+    
+    # Precompute query vector and its norm with epsilon to prevent division by zero
+    with nlp.disable_pipes(*[pipe for pipe in nlp.pipe_names if pipe != 'tok2vec']):
+        query_vector = nlp(query).vector
+    query_norm = np.linalg.norm(query_vector) + 1e-8  # Add epsilon
+    
+    # Prepare texts from chunks
+    texts = [chunk['text'] for chunk in chunks]
+    
+    # Function to process each text and compute its vector
+    def compute_vector(text):
+        with nlp.disable_pipes(*[pipe for pipe in nlp.pipe_names if pipe != 'tok2vec']):
+            doc = nlp(text)
+            vector = doc.vector
+        return vector
+    
+    # Process texts in parallel using ThreadPoolExecutor
+    with ThreadPoolExecutor() as executor:
+        chunk_vectors = list(executor.map(compute_vector, texts))
+    
+    chunk_vectors = np.array(chunk_vectors)
+    chunk_norms = np.linalg.norm(chunk_vectors, axis=1) + 1e-8  # Add epsilon
+    
+    # Compute similarities using vectorized operations
+    similarities = np.dot(chunk_vectors, query_vector) / (chunk_norms * query_norm)
+    
+    # Filter and sort results
+    relevant_chunks = [
+        (chunk, sim) for chunk, sim in zip(chunks, similarities) if sim > similarity_threshold
+    ]
+    relevant_chunks.sort(key=lambda x: x[1], reverse=True)
+    
+    return relevant_chunks[:top_n]
+
+
+@traceable(run_type="llm", name="nlp_rag")
+def query_rag(chat_llm, query, relevant_results):
+    import web_rag as wr
+
+    formatted_chunks = ""
+    for chunk, similarity in relevant_results:
+        formatted_chunk = f"""
+        <source>
+        <url>{chunk['metadata']['source']}</url>
+        <title>{chunk['metadata']['title']}</title>
+        <text>{chunk['text']}</text>
+        </source>
+        """
+        formatted_chunks += formatted_chunk
+
+    prompt = wr.get_rag_prompt_template().format(query=query, context=formatted_chunks)  
+
+    draft = chat_llm.invoke(prompt).content
+    return draft

requirements.txt

@@ -30,4 +30,5 @@ tiktoken
 transformers >= 4.44.2
 rich >= 13.8.1
 trafilatura >= 1.12.2
-watchdog >= 2.1.5, < 5.0.0
+watchdog >= 2.1.5, < 5.0.0
+spacy >= 3.6.1, < 4.0.0

search_agent.py

@@ -10,7 +10,7 @@ Usage:
         [--copywrite]
         [--max_pages=num]
         [--max_extracts=num]
-        [--use_selenium]
+        [--use_browser]
         [--output=text]
         [--verbose]
         SEARCH_QUERY
@@ -23,10 +23,10 @@ Options:
     -d domain --domain=domain           Limit search to a specific domain
     -t temp --temperature=temp          Set the temperature of the LLM [default: 0.0]
     -m model --model=model              Use a specific model [default: openai/gpt-4o-mini]
-    -e model --embedding_model=model    Use a specific embedding model [default: same provider as model]
+    -e model --embedding_model=model    Use an embedding model
     -n num --max_pages=num              Max number of pages to retrieve [default: 10]
     -x num --max_extracts=num           Max number of page extract to consider [default: 7]
-    -s --use_selenium                   Use selenium to fetch content from the web [default: False]
+    -b --use_browser                    Use browser to fetch content from the web [default: False]
     -o text --output=text               Output format (choices: text, markdown) [default: markdown]
     -v --verbose                        Print verbose output [default: False]
 
@@ -49,6 +49,7 @@ import web_rag as wr
 import web_crawler as wc
 import copywriter as cw
 import models as md
+import nlp_rag as nr
 
 console = Console()
 dotenv.load_dotenv()
@@ -91,32 +92,35 @@ def main(arguments):
     max_pages=int(arguments["--max_pages"])
     max_extract=int(arguments["--max_extracts"])
     output=arguments["--output"]
-    use_selenium=arguments["--use_selenium"]
+    use_selenium=arguments["--use_browser"]
     query = arguments["SEARCH_QUERY"]
 
     chat = md.get_model(model, temperature)
-    if embedding_model.lower() == "same provider as model":
-        provider = model.split(':')[0]
-        embedding_model = md.get_embedding_model(f"{provider}")
+    if embedding_model is None:
+        use_nlp = True
+        nlp = nr.get_nlp_model()
     else:
         embedding_model = md.get_embedding_model(embedding_model)
+        use_nlp = False
 
     if verbose:
         model_name = getattr(chat, 'model_name', None) or getattr(chat, 'model', None) or getattr(chat, 'model_id', None) or str(chat)
-        embedding_model_name = getattr(embedding_model, 'model_name', None) or getattr(embedding_model, 'model', None) or getattr(embedding_model, 'model_id', None) or str(embedding_model)
-        console.log(f"Using model: {model_name}")
         console.log(f"Using embedding model: {embedding_model_name}")
+        if not use_nlp:
+            embedding_model_name = getattr(embedding_model, 'model_name', None) or getattr(embedding_model, 'model', None) or getattr(embedding_model, 'model_id', None) or str(embedding_model)
+            console.log(f"Using model: {embedding_model_name}")
+        
 
     with console.status(f"[bold green]Optimizing query for search: {query}"):
-        optimize_search_query = wr.optimize_search_query(chat, query)
-        if len(optimize_search_query) < 3:
-            optimize_search_query = query
-    console.log(f"Optimized search query: [bold blue]{optimize_search_query}")
+        optimized_search_query = wr.optimize_search_query(chat, query)
+        if len(optimized_search_query) < 3:
+            optimized_search_query = query
+    console.log(f"Optimized search query: [bold blue]{optimized_search_query}")
 
     with console.status(
-            f"[bold green]Searching sources using the optimized query: {optimize_search_query}"
+            f"[bold green]Searching sources using the optimized query: {optimized_search_query}"
         ):
-        sources = wc.get_sources(optimize_search_query, max_pages=max_pages, domain=domain)
+        sources = wc.get_sources(optimized_search_query, max_pages=max_pages, domain=domain)
     console.log(f"Found {len(sources)} sources {'on ' + domain if domain else ''}")
 
     with console.status(
@@ -125,11 +129,28 @@ def main(arguments):
         contents = wc.get_links_contents(sources, get_selenium_driver, use_selenium=use_selenium)
     console.log(f"Managed to extract content from {len(contents)} sources")
 
-    with console.status(f"[bold green]Embedding {len(contents)} sources for content", spinner="growVertical"):
-        vector_store = wc.vectorize(contents, embedding_model)
+    if use_nlp:
+        with console.status(f"[bold green]Splitting {len(contents)} sources for content", spinner="growVertical"):            
+            chunks = nr.recursive_split_documents(contents)
+            #chunks = nr.chunk_contents(nlp, contents)
+            console.log(f"Split {len(contents)} sources into {len(chunks)} chunks")
+        with console.status(f"[bold green]Searching relevant chunks", spinner="growVertical"):
+            import time
+
+            start_time = time.time()
+            relevant_results = nr.semantic_search(optimized_search_query, chunks, nlp, top_n=max_extract)
+            end_time = time.time()
+            execution_time = end_time - start_time
+            console.log(f"Semantic search took {execution_time:.2f} seconds")
+            console.log(f"Found {len(relevant_results)} relevant chunks")
+        with console.status(f"[bold green]Writing content", spinner="growVertical"):
+            draft = nr.query_rag(chat, query, relevant_results)
+    else:
+        with console.status(f"[bold green]Embedding {len(contents)} sources for content", spinner="growVertical"):
+            vector_store = wc.vectorize(contents, embedding_model)
+        with console.status("[bold green]Writing content", spinner='dots8Bit'):
+            draft = wr.query_rag(chat, query, optimized_search_query, vector_store, top_k = max_extract)
 
-    with console.status("[bold green]Writing content", spinner='dots8Bit'):
-        draft = wr.query_rag(chat, query, optimize_search_query, vector_store, top_k = max_extract)
 
     console.rule(f"[bold green]Response")
     if output == "text":

web_rag.py

@@ -74,13 +74,13 @@ def get_optimized_search_messages(query):
                 chocolate chip cookies recipe from scratch**
             Example:
                 Question: I would like you to show me a timeline of Marie Curie's life. Show results as a markdown table
-                "Marie Curie" timeline**
+                Marie Curie timeline**
             Example:
                 Question: I would like you to write a long article on NATO vs Russia. Use known geopolitical frameworks.
                 geopolitics nato russia**
             Example:
                 Question: Write an engaging LinkedIn post about Andrew Ng
-                "Andrew Ng"**
+                Andrew Ng**
             Example:
                 Question: Write a short article about the solar system in the style of Carl Sagan
                 solar system**