Added support for PDF files.

Generated README using Claude 3

README.md

@@ -1,15 +1,33 @@
 # Simple Search Agent
 
-This is a simple search agent that (kind of) does what [Perplexity AI](https://www.perplexity.ai/) does.
+This Python project provides a search agent that can perform web searches, optimize search queries, fetch and process web content, and generate responses using a language model and the retrieved information.
+Does a bit what [Perplexity AI](https://www.perplexity.ai/) does.
 
-## How It Works
 
-1. The user asks the agent a question.
-2. The agent performs a web search using the question as the query.
-3. The agent extracts the most relevant snippets and information from the top search results. 
-4. The extracted web results are passed as context to a large language model.
-5. The LLM uses the web search context to generate a final answer to the original question.
-6. The agent returns the generated answer to the user.
+This Python script is a search agent that utilizes the LangChain library to perform optimized web searches, retrieve relevant content, and generate informative answers to user queries. The script supports multiple language models and providers, including OpenAI, Anthropic, and Groq.
+
+The main functionality of the script can be summarized as follows:
+
+1. **Query Optimization**: The user's input query is optimized for web search by identifying the key information requested and transforming it into a concise search string using the language model's capabilities.
+2. **Web Search**: The optimized search query is used to fetch search results from the Brave Search API. The script allows limiting the search to a specific domain and setting the maximum number of pages to retrieve.
+3. **Content Extraction**: The script fetches the content of the retrieved search results, handling both HTML and PDF documents. It extracts the main text content from web pages and text from PDF files.
+4. **Vectorization**: The extracted content is split into smaller text chunks and vectorized using OpenAI's text embeddings. The vectorized data is stored in a FAISS vector store for efficient retrieval.
+5. **Query Answering**: The user's original query is answered by retrieving the most relevant text chunks from the vector store using a Multi-Query Retriever. The language model generates an informative answer by synthesizing the retrieved information, citing the sources used, and formatting the response in Markdown.
+
+The script supports various options for customization, such as specifying the language model provider (OpenAI, Anthropic, Groq, or OllaMa), temperature for language model generation, and output format (text or Markdown).
+
+Additionally, the script integrates with the LangChain Tracing V2 feature, allowing users to monitor and analyze the execution of their LangChain applications using the LangChain Studio.
+
+To run the script, users need to provide their API keys for the desired language model provider and the Brave Search API in a `.env` file. The script can be executed from the command line, passing the desired options and the search query as arguments.
+
+## Features
+
+- Supports multiple language model providers (Bedrock, OpenAI, Groq, and Ollama)
+- Optimizes search queries using a language model
+- Fetches web pages and extracts main content (HTML and PDF)
+- Vectorizes the content for efficient retrieval
+- Queries the vectorized content using a Retrieval-Augmented Generation (RAG) approach
+- Generates markdown-formatted responses with references to the used sources
 
 ## Setup and Installation
 
@@ -24,6 +42,13 @@ This is a simple search agent that (kind of) does what [Perplexity AI](https://w
 
 ## Usage
 
+```
+python search_agent.py --query "your search query" --provider "provider_name" --model "model_name" --temperature 0.0
+```
+
+Replace `"your search query"` with your desired search query, `"provider_name"` with the language model provider (e.g., `bedrock`, `openai`, `groq`, `ollama`), `"model_name"` with the specific model name (optional), and `temperature` with the desired temperature value for the language model (optional).
+
+Example:
 ```
 ➜ python ./search_agent.py  --provider groq -o text "Write a linkedin post on how Sequoia Capital AI Ascent 2024 is interesting"
 [21:44:05] Using mixtral-8x7b-32768 on groq with temperature 0.0             search_agent.py:78

messages.py

@@ -38,7 +38,6 @@ def get_optimized_search_messages(query):
             - Remove lenght instruction (example: essay, article, letter, etc)
             
             Add "**" to the end of the search string to indicate the end of the query
-            Provide your output in this format: optimized search string**
             
             Example:
                 Question: How do I bake chocolate chip cookies from scratch?

requirements.txt

@@ -2,6 +2,7 @@ boto3
 bs4
 docopt
 faiss-cpu
+pdfplumber
 python-dotenv
 langchain
 langchain_core

search_agent.py

@@ -25,12 +25,14 @@ Options:
 
 import json
 import os
+import io
 from concurrent.futures import ThreadPoolExecutor
 from urllib.parse import quote
 
 from bs4 import BeautifulSoup
 from docopt import docopt
 import dotenv
+import pdfplumber
 
 from langchain_core.documents.base import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -77,7 +79,7 @@ def get_chat_llm(provider, model=None, temperature=0.0):
             chat_llm = ChatOllama(model=model, temperature=temperature)
         case _:
             raise ValueError(f"Unknown LLM provider {provider}")
-    
+
     console.log(f"Using {model} on {provider} with temperature {temperature}")
     return chat_llm
 
@@ -140,17 +142,35 @@ def extract_main_content(html):
         soup = BeautifulSoup(html, 'html.parser')
         for element in soup(["script", "style", "head", "nav", "footer", "iframe", "img"]):
             element.extract()
-        main_content = ' '.join(soup.body.get_text().split())
+        main_content = soup.get_text(separator='\n', strip=True)
         return main_content
     except Exception:
         return None
 
 def process_source(source):
     response = fetch_with_timeout(source['link'], 8)
+    console.log(f"Processing {source['link']}")
     if response:
-        html = response.text
-        main_content = extract_main_content(html)
-        return {**source, 'html': main_content}
+        content_type = response.headers.get('Content-Type')
+        if content_type == 'application/pdf':
+            # The response is a PDF file
+            pdf_content = response.content
+            # Create a file-like object from the bytes
+            pdf_file = io.BytesIO(pdf_content)
+            # Extract text from PDF using pdfplumber
+            with pdfplumber.open(pdf_file) as pdf:
+                text = ""
+                for page in pdf.pages:
+                    text += page.extract_text()
+            return {**source, 'pdf_content': text}
+        elif content_type.startswith('text/html'):
+            # The response is an HTML file
+            html = response.text
+            main_content = extract_main_content(html)
+            return {**source, 'html': main_content}
+        else:
+            console.log(f"Skipping {source['link']}! Unsupported content type: {content_type}")
+            return None
     return None
 
 def get_links_contents(sources):
@@ -163,14 +183,17 @@ def get_links_contents(sources):
 def vectorize(contents, text_chunk_size=1000,text_chunk_overlap=200,):
     documents = []
     for content in contents:
-        if content['html']:
-            try:
-                page_content = content['html']
-                metadata = {'title': content['title'], 'source': content['link']}
-                doc = Document(page_content=page_content, metadata=metadata)
-                documents.append(doc)
-            except Exception as e:
-                console.log(f"[gray]Error processing content for {content['link']}: {e}")
+        page_content = content['snippet']
+        if 'htlm' in content:
+            page_content = content['html']
+        if 'pdf_content' in content:
+            page_content = content['pdf_content']        
+        try:
+            metadata = {'title': content['title'], 'source': content['link']}
+            doc = Document(page_content=page_content, metadata=metadata)
+            documents.append(doc)
+        except Exception as e:
+            console.log(f"[gray]Error processing content for {content['link']}: {e}")
 
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=text_chunk_size,
@@ -195,9 +218,11 @@ def format_docs(docs):
 
 def query_rag(chat_llm, question, search_query, vectorstore):
     retriever_from_llm = MultiQueryRetriever.from_llm(
-        retriever=vectorstore.as_retriever(), llm=chat_llm,
+        retriever=vectorstore.as_retriever(), llm=chat_llm, include_original=True,
+    )
+    unique_docs = retriever_from_llm.get_relevant_documents(
+        query=search_query, callbacks=callbacks, verbose=True
     )
-    unique_docs = retriever_from_llm.get_relevant_documents(query=search_query, config={"callbacks": callbacks})
     context = format_docs(unique_docs)
     prompt = get_rag_prompt_template().format(query=question, context=context)
     response = chat_llm.invoke(prompt, config={"callbacks": callbacks})
@@ -249,7 +274,7 @@ if __name__ == '__main__':
         contents = get_links_contents(sources)
     console.log(f"Managed to extract content from {len(contents)} sources")
 
-    with console.status(f"[bold green]Embeddubg {len(sources)} sources", spinner="growVertical"):
+    with console.status(f"[bold green]Embeddubg {len(contents)} sources for content", spinner="growVertical"):
         vector_store = vectorize(contents)
 
     with console.status("[bold green]Querying LLM relevant context", spinner='dots8Bit'):