Added Firework.ai as provider. Better streamlit ui.

README.md

@@ -15,7 +15,7 @@ license: apache-2.0
 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.
 
-The Streamlit GUI hosted on 🤗 Sapces is [available to test](https://huggingface.co/spaces/CyranoB/search_agent)
+The Streamlit GUI hosted on 🤗 Spaces is [available to test](https://huggingface.co/spaces/CyranoB/search_agent)
 
 This Python script and Streamli GUI are a basic 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.
 

dotenv.sample

@@ -6,6 +6,7 @@ LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"
 
 OPENAI_API_KEY=sk-XXXXXXXXXXXXXXXXXXX
 ANTHROPIC_API_KEY=sk-ant-api03-XXXXXXXXXXXXXXXXXXX
+FIREWORKS_API_KEY=XXXXXXXXXXXXXXXXXXX
 GROQ_API_KEY=gsk_XXXXXXXXXXXXXXXXXXX
 CREDENTIALS_PROFILE_NAME=XXXXXXXXXXXXXXXXXXX
 COHERE_API_KEY=XXXXXXXXXXXXXXXXXXX

requirements.txt

@@ -8,6 +8,7 @@ pdfplumber
 python-dotenv
 langchain
 langchain-cohere
+langchain-fireworks
 langchain_core
 langchain_community
 langchain_experimental

search_agent.py

@@ -16,7 +16,7 @@ Options:
     --version                           Show version.
     -d domain --domain=domain           Limit search to a specific domain
     -t temp --temperature=temp          Set the temperature of the LLM [default: 0.0]
-    -p provider --provider=provider     Use a specific LLM (choices: bedrock,openai,groq,ollama,cohere) [default: openai]
+    -p provider --provider=provider     Use a specific LLM (choices: bedrock,openai,groq,ollama,cohere,fireworks) [default: openai]
     -m model --model=model              Use a specific model
     -n num --max_pages=num              Max number of pages to retrieve [default: 10]
     -o text --output=text               Output format (choices: text, markdown) [default: markdown]
@@ -78,8 +78,8 @@ if __name__ == '__main__':
     output=arguments["--output"]
     query = arguments["SEARCH_QUERY"]
 
-    chat = wr.get_chat_llm(provider, model, temperature)
-    console.log(f"Using {chat.model} on {provider} with temperature {temperature}")
+    chat, embedding_model = wr.get_models(provider, model, temperature)
+    #console.log(f"Using {chat.model_name} on {provider}")
 
     with console.status(f"[bold green]Optimizing query for search: {query}"):
         optimize_search_query = wr.optimize_search_query(chat, query, callbacks=callbacks)
@@ -98,7 +98,7 @@ if __name__ == '__main__':
     console.log(f"Managed to extract content from {len(contents)} sources")
 
     with console.status(f"[bold green]Embeddubg {len(contents)} sources for content", spinner="growVertical"):
-        vector_store = wc.vectorize(contents)
+        vector_store = wc.vectorize(contents, embedding_model)
 
     with console.status("[bold green]Querying LLM relevant context", spinner='dots8Bit'):
         respomse = wr.query_rag(chat, query, optimize_search_query, vector_store, top_k = 5, callbacks=callbacks)

search_agent_ui.py

@@ -1,4 +1,5 @@
 import datetime
+import os
 
 import dotenv
 import streamlit as st
@@ -13,11 +14,13 @@ import web_crawler as wc
 dotenv.load_dotenv()
 
 ls_tracer = LangChainTracer(
-    project_name="Search Agent UI",
+    project_name=os.getenv("LANGSMITH_PROJECT_NAME"),
     client=Client()
 )
 
+
 class StreamHandler(BaseCallbackHandler):
+    """Stream handler that appends tokens to container."""
     def __init__(self, container, initial_text=""):
         self.container = container
         self.text = initial_text
@@ -26,16 +29,34 @@ class StreamHandler(BaseCallbackHandler):
         self.text += token
         self.container.markdown(self.text)
 
-chat = wr.get_chat_llm(provider="cohere")
-
 st.title("🔍 Simple Search Agent 💬")
 
+if "providers" not in st.session_state:
+    providers = []
+    if os.getenv("COHERE_API_KEY"):
+        providers.append("cohere")
+    if os.getenv("OPENAI_API_KEY"):
+        providers.append("openai")
+    if os.getenv("GROQ_API_KEY"):
+        providers.append("groq")
+    if os.getenv("OLLAMA_API_KEY"):
+        providers.append("ollama")
+    if os.getenv("FIREWORKS_API_KEY"):
+        providers.append("fireworks")        
+    if os.getenv("CREDENTIALS_PROFILE_NAME"):
+        providers.append("bedrock")
+    st.session_state["providers"] = providers
+
+with st.sidebar:
+    st.write("Options")
+    model_provider = st.selectbox("🧠 Model provider 🧠", st.session_state["providers"])
+    temperature = st.slider("🌡️ Model temperature 🌡️", 0.0, 1.0, 0.1, help="The higher the more creative")
+    max_pages = st.slider("🔍 Max pages to retrieve 🔍", 1, 20, 15, help="How many web pages to retrive from the internet")
+    top_k_documents = st.slider("📄 How many document extracts to consider 📄", 1, 20, 5, help="How many of the top extracts to consider")
+
 if "messages" not in st.session_state:
     st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
     
-if "input_disabled" not in st.session_state:
-    st.session_state["input_disabled"] = False
-
 for message in st.session_state.messages:
     st.chat_message(message["role"]).write(message["content"])
     if message["role"] == "assistant" and 'message_id' in message:
@@ -46,26 +67,25 @@ for message in st.session_state.messages:
             mime="text/plain"
         )
 
-if prompt := st.chat_input("Enter you instructions...", disabled=st.session_state["input_disabled"] ):
-    
-    st.session_state["input_disabled"] = True
-    
+if prompt := st.chat_input("Enter you instructions..." ):   
     st.chat_message("user").write(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
+    chat, embedding_model = wr.get_models(model_provider, temperature=temperature)
+
     with st.status("Thinking", expanded=True):
         st.write("I first need to do some research")
         
         optimize_search_query = wr.optimize_search_query(chat, query=prompt, callbacks=[ls_tracer])
         st.write(f"I should search the web for: {optimize_search_query}")
         
-        sources = wc.get_sources(optimize_search_query, max_pages=20)
+        sources = wc.get_sources(optimize_search_query, max_pages=max_pages)
        
         st.write(f"I'll now retrieve the {len(sources)} webpages and documents I found")
         contents = wc.get_links_contents(sources)
 
         st.write( f"Reading through the {len(contents)} sources I managed to retrieve")
-        vector_store = wc.vectorize(contents)
+        vector_store = wc.vectorize(contents, embedding_model=embedding_model)
         st.write(f"I collected {vector_store.index.ntotal} chunk of data and I can now answer")
 
     rag_prompt = wr.build_rag_prompt(prompt, optimize_search_query, vector_store, top_k=5, callbacks=[ls_tracer])  
@@ -82,5 +102,3 @@ if prompt := st.chat_input("Enter you instructions...", disabled=st.session_stat
                 file_name=f"{message_id}.txt",
                 mime="text/plain"
             )
-    st.session_state["input_disabled"] = False
-    

web_crawler.py

@@ -124,7 +124,7 @@ def get_links_contents(sources, get_driver_func=None):
                 result['page_content'] = main_content
     return results
 
-def vectorize(contents):
+def vectorize(contents, embedding_model):
     documents = []
     for content in contents:
         try:
@@ -135,7 +135,7 @@ def vectorize(contents):
                 documents.append(doc)
         except Exception as e:
             print(f"[gray]Error processing content for {content['link']}: {e}")
-    semantic_chunker = SemanticChunker(OpenAIEmbeddings(model="text-embedding-3-large"), breakpoint_threshold_type="percentile")
+    semantic_chunker = SemanticChunker(embedding_model, breakpoint_threshold_type="percentile")
     docs = semantic_chunker.split_documents(documents)
     embeddings = OpenAIEmbeddings()
     store = FAISS.from_documents(docs, embeddings)

web_rag.py

@@ -29,40 +29,58 @@ from langchain.prompts.prompt import PromptTemplate
 from langchain.retrievers.multi_query import MultiQueryRetriever
 
 from langchain_cohere.chat_models import ChatCohere
+from langchain_cohere.embeddings import CohereEmbeddings
+from langchain_fireworks.chat_models import ChatFireworks
 from langchain_groq import ChatGroq
 from langchain_openai import ChatOpenAI
+from langchain_openai.embeddings import OpenAIEmbeddings
 from langchain_community.chat_models.bedrock import BedrockChat
+from langchain_community.embeddings.bedrock import BedrockEmbeddings
 from langchain_community.chat_models.ollama import ChatOllama
 
-def get_chat_llm(provider, model=None, temperature=0.0):
+def get_models(provider, model=None, temperature=0.0):
     match provider:
         case 'bedrock':
+            credentials_profile_name=os.getenv('CREDENTIALS_PROFILE_NAME')
             if model is None:
                 model = "anthropic.claude-3-sonnet-20240229-v1:0"
             chat_llm = BedrockChat(
-                credentials_profile_name=os.getenv('CREDENTIALS_PROFILE_NAME'),
+                credentials_profile_name=credentials_profile_name,
                 model_id=model,
                 model_kwargs={"temperature": temperature },
             )
+            embedding_model = BedrockEmbeddings(
+                model_id='cohere.embed-multilingual-v3',
+                credentials_profile_name=credentials_profile_name
+            )
         case 'openai':
             if model is None:
                 model = "gpt-3.5-turbo"
             chat_llm = ChatOpenAI(model_name=model, temperature=temperature)
+            embedding_model = OpenAIEmbeddings(model='text-embedding-3-small')
         case 'groq':
             if model is None:
                 model = 'mixtral-8x7b-32768'
             chat_llm = ChatGroq(model_name=model, temperature=temperature)
+            embedding_model = OpenAIEmbeddings(model='text-embedding-3-small')
         case 'ollama':
             if model is None:
                 model = 'llama2'
             chat_llm = ChatOllama(model=model, temperature=temperature)
+            embedding_model = OpenAIEmbeddings(model='text-embedding-3-small')
         case 'cohere':
             if model is None:
                 model = 'command-r-plus'
             chat_llm = ChatCohere(model=model, temperature=temperature)
+            embedding_model = CohereEmbeddings(model="embed-english-light-v3.0")
+        case 'fireworks':
+            if model is None:
+                model = 'accounts/fireworks/models/mixtral-8x22b-instruct-preview'
+            chat_llm = ChatFireworks(model_name=model, temperature=temperature)
+            embedding_model = OpenAIEmbeddings(model='text-embedding-3-small')
         case _:
             raise ValueError(f"Unknown LLM provider {provider}")
-    return chat_llm
+    return chat_llm, embedding_model
 
 
 def get_optimized_search_messages(query):