Merge pull request #3 from leowalker89/EnableBraveSearchAI

app.py

@@ -1,20 +1,32 @@
 import streamlit as st
 import random
-from helpers import query_you_com, query_tavily, query_perplexity
-from mongod_db import MongoDBHandler
-from swarms.utils.loguru_logger import logger
+from helpers import query_you_com, query_tavily, query_perplexity, query_brave
+from provider_info import search_providers
+# from mongod_db import MongoDBHandler
+# from swarms.utils.loguru_logger import logger
 import time
 
-mongo = MongoDBHandler()
+# mongo = MongoDBHandler()
 
 # Set Streamlit to wide mode
 st.set_page_config(layout="wide")
 
+# Add information to sidebar
+st.sidebar.title("About the App")
+st.sidebar.write("""
+This app allows you to compare responses from different search engines.
+Submit a question, and you'll receive answers from two randomly selected search engines.
+You can then vote on which response you prefer.
+""")
+st.sidebar.write("""
+**[GitHub](https://github.com/leowalker89/SearchArena)**
+
+""")
 
 # Define the function to process the question
 def ProcessQuestion(question):
-    # Randomly select two out of the three functions
-    functions = [query_you_com, query_tavily, query_perplexity]
+    # Randomly select two out of the four functions
+    functions = [query_you_com, query_tavily, query_perplexity, query_brave]
     selected_functions = random.sample(functions, 2)
 
     # Get answers from the selected functions
@@ -22,98 +34,157 @@ def ProcessQuestion(question):
     answer_b = selected_functions[1](question)
     
     # Log into mongodb
-    try: 
-        logger.info(f"Logging question: {question}")
-        mongo.add(
-            {
-                "question": question,
-                "answer_a": answer_a,
-                "answer_b": answer_b,
-                "selected_functions": [f.__name__ for f in selected_functions],
-                "query_time": time.time(),
-            }
-        )
-        logger.info("Successfully logged into mongodb")
-    except Exception as e:
-        logger.error(f"Error logging into mongodb: {e}")
-
-    return answer_a, answer_b
+    # try: 
+    #     logger.info(f"Logging question: {question}")
+    #     mongo.add(
+    #         {
+    #             "question": question,
+    #             "answer_a": answer_a,
+    #             "answer_b": answer_b,
+    #             "selected_functions": [f.__name__ for f in selected_functions],
+    #             "query_time": time.time(),
+    #         }
+    #     )
+    #     logger.info("Successfully logged into mongodb")
+    # except Exception as e:
+    #     logger.error(f"Error logging into mongodb: {e}")
+
+    return answer_a, answer_b, selected_functions
 
 
 # Initialize session state if not already done
-if "results_displayed" not in st.session_state:
-    st.session_state["results_displayed"] = False
-if "answer_a" not in st.session_state:
-    st.session_state["answer_a"] = ""
-if "answer_b" not in st.session_state:
-    st.session_state["answer_b"] = ""
-if "question" not in st.session_state:
-    st.session_state["question"] = ""
+default_values = {
+    "state": "arena_ready",
+    "question": "",
+    "answer_a": "",
+    "answer_b": "",
+    "source_a": "",
+    "source_b": "",
+    "winner": "",
+    "selected_button": ""
+}
+
+for key, value in default_values.items():
+    if key not in st.session_state:
+        st.session_state[key] = value
 
 # Streamlit app layout
 st.title("Search Engine Agent Comparison")
 
-# Create columns for the input and model selection
-input_col, control_col = st.columns([4, 1])
-
-with input_col:
-    # Text box for user input with character limit
-    question = st.text_area(
-        "Enter your question here (max 1000 characters):", max_chars=1000
-    )
-
-with control_col:
-    # Submit button
-    submit_button = st.button("Submit")
-
-if submit_button:
+def on_submit():
+    question = st.session_state["question_input"]
     if question:
-        if len(question) <= 1000:
-            # Process the question and get answers
-            answer_a, answer_b = ProcessQuestion(question)
-
-            # Save answers and state to session state
-            st.session_state["answer_a"] = answer_a
-            st.session_state["answer_b"] = answer_b
-            st.session_state["question"] = question
-            st.session_state["results_displayed"] = True
-        else:
-            st.error(
-                "Your question exceeds the 1,000 character limit. Please shorten your question."
-            )
-    else:
-        st.error("Please enter a question.")
-
-# Display results if available in session state
-if st.session_state["results_displayed"]:
+        answer_a, answer_b, selected_functions = ProcessQuestion(question)
+        st.session_state["question"] = question
+        st.session_state["answer_a"] = answer_a
+        st.session_state["answer_b"] = answer_b
+        st.session_state["source_a"] = selected_functions[0].__name__
+        st.session_state["source_b"] = selected_functions[1].__name__
+        st.session_state["state"] = "arena_review"
+
+def handle_vote(vote):
+    st.session_state["winner"] = vote
+    st.session_state["state"] = "arena_results"
+
+def get_provider_info(provider_function_name):
+    provider_name_map = {
+        'query_you_com': 'You.com',
+        'query_tavily': 'Tavily',
+        'query_perplexity': 'Perplexity AI',
+        'query_brave': 'Brave Search'
+    }
+    provider_name = provider_name_map.get(provider_function_name)
+    return next((provider for provider in search_providers if provider['company_name'] == provider_name), {})
+
+def render_ready_state():
+    st.text_area("Enter your question here (max 1000 characters):", 
+                 max_chars=1000, 
+                 key="question_input", 
+                 on_change=on_submit)
+    st.button("Submit", on_click=on_submit)
+
+def render_review_state():
+    st.write("## Results")
     col1, col2 = st.columns(2)
     with col1:
-        st.write("### Output A")
+        st.write("### Answer A")
         st.write(st.session_state["answer_a"])
-        a_feedback_grid = st.columns(1)
     with col2:
-        st.write("### Output B")
+        st.write("### Answer B")
         st.write(st.session_state["answer_b"])
-        b_feedback_grid = st.columns(2)
-
-    # Create a placeholder for the feedback div
-    feedback_placeholder = st.empty()
-
-    def display_feedback(message):
-        feedback_placeholder.markdown(
-            f'<div style="position: fixed; bottom: 10px; left: 10px; background-color: #f0f0f0; padding: 10px; border-radius: 5px;">{message}</div>',
-            unsafe_allow_html=True,
-        )
-
-    with a_feedback_grid[0]:
-        if st.button("A is better 🥇"):
-            display_feedback("You selected: A is better")
-    with b_feedback_grid[0]:
-        if st.button("B is better 💪"):
-            display_feedback("You selected: B is better")
-    with a_feedback_grid[0]:
-        if st.button("It's a Tie 🤝"):
-            display_feedback("You selected: It's a Tie")
-    with b_feedback_grid[0]:
-        if st.button("Both are bad 👎"):
-            display_feedback("You selected: Both are bad")
+    st.write("### Vote for the Best Answer")
+    col1, col2, col3, col4 = st.columns(4)
+    if col1.button("It's a Tie 🤝"):
+        handle_vote("Tie")
+    if col2.button("A is better 💪"):
+        handle_vote("A")
+    if col3.button("B is better 🥇"):
+        handle_vote("B")
+    if col4.button("Both are bad 👎"):
+        handle_vote("Both are bad")
+
+def render_results_state():
+    st.write("## Results")
+    st.write(f"### Question: {st.session_state['question']}")
+    
+    provider_info_a = get_provider_info(st.session_state["source_a"])
+    provider_info_b = get_provider_info(st.session_state["source_b"])
+    
+    col1, col2 = st.columns(2)
+    with col1:
+        if st.session_state["winner"] == "A":
+            st.write(f"### ⭐ {provider_info_a['company_name']} 🥇")
+        elif st.session_state["winner"] == "Tie":
+            st.write(f"### 🤝 {provider_info_a['company_name']} 🤝")
+        elif st.session_state["winner"] == "Both are bad":
+            st.write(f"### 👎 {provider_info_a['company_name']} 👎")
+        else:
+            st.write(f"### {provider_info_a['company_name']} 🥈")
+        st.write("**Response:**")
+        st.markdown(f"<div style='padding: 10px; border: 1px solid #ddd;'>{st.session_state['answer_a']}</div>", unsafe_allow_html=True)
+
+    with col2:
+        if st.session_state["winner"] == "B":
+            st.write(f"### ⭐ {provider_info_b['company_name']} 🥇")
+        elif st.session_state["winner"] == "Tie":
+            st.write(f"### 🤝 {provider_info_b['company_name']} 🤝")
+        elif st.session_state["winner"] == "Both are bad":
+            st.write(f"### 👎 {provider_info_b['company_name']} 👎")
+        else:
+            st.write(f"### {provider_info_b['company_name']} 🥈")
+        st.write("**Response:**")
+        st.markdown(f"<div style='padding: 10px; border: 1px solid #ddd;'>{st.session_state['answer_b']}</div>", unsafe_allow_html=True)
+    
+    st.write("### Feedback")
+    st.text_area("Please provide feedback on why you chose the winner:", key="feedback")
+    st.write("### About the search providers:")
+    col1, col2 = st.columns(2)
+    with col1:
+        st.write(f"**Website:** [{provider_info_a['website']}]({provider_info_a['website']})")
+        st.write(f"**Overview:** {provider_info_a['overview']}")
+    with col2:
+        st.write(f"**Website:** [{provider_info_b['website']}]({provider_info_b['website']})")
+        st.write(f"**Overview:** {provider_info_b['overview']}")
+
+if st.session_state["state"] == "arena_ready":
+    render_ready_state()
+elif st.session_state["state"] == "arena_review":
+    render_review_state()
+elif st.session_state["state"] == "arena_results":
+    render_results_state()
+
+# Apply custom CSS to highlight the selected button
+selected_button = st.session_state.get("selected_button", "")
+
+if selected_button:
+    st.markdown(
+        f"""
+        <style>
+        button[kind="primary"]{{
+            background-color: #4CAF50 !important;
+            color: white !important;
+        }}
+        </style>
+        """,
+        unsafe_allow_html=True
+    )

brave_ai.py

@@ -0,0 +1,89 @@
+import json
+import requests
+from typing import List, Optional
+# from langchain_community.document_transformers import MarkdownifyTransformer
+from langchain_core.documents import Document
+from langchain_core.pydantic_v1 import BaseModel, Field
+
+class BraveAIWrapper(BaseModel):
+    api_key: str = Field(..., description="API key for Brave Search")
+    base_search_url: str = Field("https://api.search.brave.com/res/v1/web/search", const=True)
+    base_summarize_url: str = Field("https://api.search.brave.com/res/v1/summarizer/search", const=True)
+    headers: dict = Field(default_factory=dict, description="HTTP headers for API requests")
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.headers = {
+            "X-Subscription-Token": self.api_key,
+            "Accept": "application/json",
+            "Accept-Encoding": "gzip",
+        }
+
+    def get_brave_results(self, query: str, count: int = 3, safe_search: str = 'off') -> Optional[dict]:
+        """
+        Get search results from Brave Search.
+
+        Args:
+            query (str): The search query.
+            count (int): Number of results to return.
+            safe_search (str): Safe search filter (off, moderate, strict).
+
+        Returns:
+            Optional[dict]: JSON response from Brave Search API or None if an error occurs.
+        """
+        params = {
+            "q": query,
+            "count": count,
+            "summary": True,
+            "safe_search": safe_search,
+            "extra_snippets": True,
+        }
+        try:
+            response = requests.get(self.base_search_url, headers=self.headers, params=params)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            print(f"Error: {e}")
+            return None
+
+    def summarize_results(self, summarizer_key: str) -> Optional[dict]:
+        """
+        Summarize search results using Brave Summarizer.
+
+        Args:
+            summarizer_key (str): The key for the summarizer.
+
+        Returns:
+            Optional[dict]: JSON response from Brave Summarizer API or None if an error occurs.
+        """
+        params = {"key": summarizer_key}
+        try:
+            response = requests.get(self.base_summarize_url, headers=self.headers, params=params)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            print(f"Error: {e}")
+            return None
+
+    def get_and_summarize(self, query: str, count: int = 3, safe_search: str = 'off') -> Optional[str]:
+        """
+        Get and summarize search results from Brave Search.
+
+        Args:
+            query (str): The search query.
+            count (int): Number of results to return.
+            safe_search (str): Safe search filter (off, moderate, strict).
+
+        Returns:
+            Optional[str]: Summarized result or None if an error occurs.
+        """
+        results = self.get_brave_results(query, count, safe_search)
+        if results and 'summarizer' in results:
+            summarizer_key = results['summarizer']['key']
+            summary = self.summarize_results(summarizer_key)
+            if summary and 'summary' in summary and len(summary['summary']) > 0:
+                return summary['summary'][0]['data']
+        return None

helpers.py

@@ -1,5 +1,6 @@
 import requests
 from dotenv import load_dotenv
+from typing import Optional
 import os
 
 # Load environment variables from .env file
@@ -9,8 +10,7 @@ load_dotenv()
 YOU_COM_API_KEY = os.getenv("YOU_API_KEY")
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 PERPLEXITY_API_KEY = os.getenv("PPLX_API_KEY")
-BRAVE_API_KEY = os.getenv("BRAVESEARCH_API_KEY")
-
+BRAVE_API_KEY = os.getenv("BRAVE_AI_API_KEY")
 
 def query_you_com(query):
     headers = {"X-API-Key": YOU_COM_API_KEY}
@@ -76,71 +76,6 @@ def query_perplexity(query):
         return f"Request failed with status code: {response.status_code}"
 
 
-# def query_brave(query):
-#     headers = {"X-API-Key": BRAVE_API_KEY}
-#     params = {
-#         "q": query,
-#         "count": 1,
-#         "summary": True
-#     }
-#     response = requests.get("https://api.search.brave.com/res/v1/web/search", params=params, headers=headers)
-#     if response.status_code == 200:
-#         return response.json().get("summary", "No summary available.")
-#     else:
-#         return f"Request failed with status code: {response}"
-
-
-# def brave_search_summarization(query):
-#     # Endpoint for web search with summary
-#     web_search_url = "https://api.search.brave.com/res/v1/web/search"
-#     summarizer_url = "https://api.search.brave.com/res/v1/summarizer/search"
-
-#     # Headers for the requests
-#     headers = {
-#         "Accept": "application/json",
-#         "Accept-Encoding": "gzip",
-#         "X-Subscription-Token": BRAVE_API_KEY
-#     }
-
-#     # Parameters for the initial web search request
-#     web_search_params = {
-#         "q": query,
-#         "summary": 1
-#     }
-
-#     # Make the initial request to the web search endpoint
-#     web_search_response = requests.get(web_search_url, headers=headers, params=web_search_params)
-
-#     # Check if the request was successful
-#     if web_search_response.status_code != 200:
-#         raise Exception(f"Web search request failed with status code {web_search_response.status_code}")
-
-#     web_search_data = web_search_response.json()
-
-#     # Extract the summarizer key from the response
-#     summarizer_key = web_search_data.get('summarizer', {}).get('key')
-#     if not summarizer_key:
-#         raise Exception("No summarizer key found in the web search response")
-
-#     # Parameters for the summarizer request
-#     summarizer_params = {
-#         "key": summarizer_key,
-#         "entity_info": 1
-#     }
-
-#     # Make the request to the summarizer endpoint
-#     summarizer_response = requests.get(summarizer_url, headers=headers, params=summarizer_params)
-
-#     # Check if the request was successful
-#     if summarizer_response.status_code != 200:
-#         raise Exception(f"Summarizer request failed with status code {summarizer_response.status_code}")
-
-#     summarizer_data = summarizer_response.json()
-
-#     # Return the summarized content
-#     return summarizer_data
-
-
 def ProcessQuestion(question, model):
     if model == "You.com":
         return query_you_com(question)
@@ -152,3 +87,21 @@ def ProcessQuestion(question, model):
         return query_brave(question)
     else:
         return "Model not supported"
+
+
+from brave_ai import BraveAIWrapper
+
+def query_brave(query: str) -> Optional[str]:
+    """
+    Get a summary for the given query using BraveAIWrapper.
+
+    Args:
+        query (str): The search query.
+        api_key (str): The API key for Brave Search.
+
+    Returns:
+        Optional[str]: Summarized result or None if an error occurs.
+    """
+    brave_ai = BraveAIWrapper(api_key=BRAVE_API_KEY)
+    summary = brave_ai.get_and_summarize(query)
+    return summary

provider_info.py

@@ -0,0 +1,18 @@
+search_providers = [
+    { "company_name": "You.com",
+        "website": "you.com",
+        "overview": "You.com is an AI-powered search engine and productivity tool that aims to provide users with more control, privacy, and customization over their online search experience. Founded in 2020 by Richard Socher, former Chief Scientist at Salesforce, and Bryan McCann, You.com has evolved from a privacy-focused search engine into an AI assistant that supports chat-based search.\n\nKey features of You.com include:\n\n- AI Modes: You.com offers different AI interaction modes tailored to user needs, such as Smart mode for quick tasks and summarization, Genius mode for complex queries requiring multi-step reasoning, and Research mode for in-depth analysis with extensive citations.\n\n- Multimodal search: Beyond text, You.com can generate AI art, data visualizations, and even run Python code to solve math and science questions. It was one of the first search engines to introduce multimodal AI chat capabilities.\n\n- Customization: Users can personalize their search experience by selecting preferred sources, adjusting AI model parameters, and leveraging You.com's developer platform to build custom apps and integrations.\n\n- Privacy: You.com does not track users or sell their personal data. It provides a more secure and private alternative to mainstream search engines.\n\n- Subscription model: While offering a free tier, You.com provides a premium subscription called YouPro that unlocks advanced AI features, unlimited access, and an ad-free experience.\n\nYou.com has gained significant traction, amassing millions of users and raising substantial funding from prominent investors like Salesforce CEO Marc Benioff. In 2023, the company raised $45 million at a $700-900 million valuation.\n\nUnder the leadership of Richard Socher, who was named to Time's TIME100 AI list in 2023, You.com aims to make information more accessible and enhance user productivity through cutting-edge AI technologies. Its unique blend of search and AI assistance capabilities positions You.com as a major player in the rapidly evolving search and productivity tools market.\n\nIn summary, You.com is a next-generation AI search assistant that combines the power of large language models with live web access to deliver highly relevant, contextual, and multimodal search results while prioritizing user privacy and customization. By making information discovery more intuitive and empowering, You.com is reshaping the future of search."
+    },
+    { "company_name": "Tavily",
+        "website": "tavily.com",
+        "overview": "Tavily is an AI-powered search API and research platform designed specifically for large language models (LLMs) and AI agents. It aims to provide efficient, accurate, and factual search results to enhance the capabilities of AI applications.\n\nKey features of Tavily include:\n\n- Optimized for LLMs: Tavily's search API is tailored for AI agents and delivers real-time, relevant results to reduce hallucinations and bias in AI-generated content.\n\n- In-depth research: It simplifies data gathering by providing aggregated and curated results from multiple sources in a single API call. Tavily searches over 20+ sites per query and uses AI to score, filter, and rank the most relevant content.\n\n- Intelligent query suggestions: Tavily equips AI agents with the ability to iteratively deepen their knowledge through automated, nuanced answers and follow-up query suggestions.\n\n- Customization: The API offers customizable search depths, domain management, and HTML content parsing controls. It can be integrated with any LLM or used with partners like Langchain and LlamaIndex.\n\n- Versatility: Beyond fetching results, Tavily can answer questions, enable cross-agent communication, and complete comprehensive research tasks in seconds by leveraging proprietary data sources.\n\nTavily caters to a range of users, from new creators and researchers to solopreneurs, small teams, and enterprise-level needs. By automating the tedious aspects of online research like searching, scraping, filtering, and extracting relevant information, Tavily enables users to save time and make more informed decisions.\n\nIn summary, Tavily is a powerful AI search and research tool that connects LLMs and AI applications to trusted, real-time knowledge to enhance their performance and capabilities. It streamlines the research workflow and delivers comprehensive, accurate results to your inbox within minutes."
+    },
+    {"company_name": "Perplexity AI",
+        "website": "perplexity.ai",
+        "overview": "Perplexity AI is an AI-powered conversational search engine that provides accurate, concise answers to user queries by searching the web in real-time and leveraging advanced natural language processing capabilities. It aims to make information more accessible and understandable.\n\nKey features of Perplexity AI include:\n\n- Conversational interface: Users can engage with Perplexity AI using natural language, making information discovery as easy as having a conversation. It understands context and nuance to provide relevant answers.\n\n- Transparent sourcing: Perplexity AI cites the sources it uses to generate answers, enabling users to verify information and explore topics further. It aims to build trust through transparency.\n\n- Personalized results: The more a user interacts with Perplexity AI, the more it learns about their interests and preferences, allowing it to tailor future search results and recommendations.\n\n- Multi-model architecture: Perplexity AI utilizes multiple large language models, including GPT-4, Claude, and its own proprietary models, to process queries and generate human-like responses.\n\n- Versatility: Beyond answering questions, Perplexity AI can assist with a wide range of tasks such as writing, coding, analyzing data, and creative projects. It adapts to the user's needs.\n\n- Continuous learning: By indexing the web daily, Perplexity AI stays up-to-date with the latest information to provide relevant, timely answers.\n\nPerplexity AI caters to a broad user base, including researchers, writers, programmers, and curious individuals seeking to learn and explore. It has gained significant traction, amassing over 10 million monthly users as of early 2024.\n\nFounded in 2022 by a team with experience from Google, OpenAI, and other leading tech companies, Perplexity AI has raised over $165 million from prominent investors. Its mission is to democratize access to knowledge and empower people to learn in their own way.\n\nIn summary, Perplexity AI is a powerful AI search engine that combines advanced language understanding, real-time web indexing, and an intuitive conversational interface to deliver accurate, relevant answers and support a wide range of information needs. By making knowledge more accessible and engaging, it is transforming how people discover, learn, and create."
+    },
+    {"company_name": "Brave Search",
+    "website": "search.brave.com",
+    "overview": "Brave Search is a privacy-focused search engine that offers an AI API for developers to integrate search capabilities and generate summaries from search results. The API leverages Brave's own search index and AI models to provide relevant, accurate, and concise information.\n\nKey features of Brave Search's AI API include:\n\n- Two-step process: To use the API, developers first perform a search using the '/search' endpoint, which returns a list of relevant web pages. They can then pass the search results to the '/summarize' endpoint to generate a summary of the information.\n\n- Search quality: Brave Search maintains its own independent search index, ensuring that results are not influenced by biased algorithms or user tracking. The search functionality of the API delivers high-quality, relevant web pages for a given query.\n\n- Summarization: The '/summarize' endpoint uses advanced AI language models to analyze the content of the web pages returned by the search and generate a concise, coherent summary. This saves developers time and effort in extracting key information from search results.\n\n- Customization: Developers can fine-tune the search and summarization process by adjusting parameters such as the number of search results to consider, the length of the summary, and the focus of the summary (e.g., key facts, main ideas, or specific details).\n\n- Privacy: In line with Brave's commitment to privacy, the AI API does not track users or collect personal data. All searches and summarization requests are processed anonymously.\n\nBrave Search's AI API is suitable for a wide range of applications, such as content aggregation platforms, research tools, news summarization services, and chatbots. By providing access to Brave's search capabilities and AI-powered summarization, the API enables developers to create innovative, privacy-respecting applications that deliver valuable insights to users.\n\nThe API offers a free tier with limited requests and a paid tier for higher usage and additional features. Brave Search is committed to expanding the capabilities of its AI API and supporting developers in building next-generation search and language AI applications.\n\nIn summary, Brave Search's AI API combines the power of privacy-focused search with advanced AI summarization capabilities, allowing developers to integrate high-quality search results and concise summaries into their applications. By providing a two-step process of searching and summarizing, the API streamlines information retrieval and synthesis, while prioritizing user privacy."
+    }
+]

requirements.txt

@@ -1,6 +1,92 @@
-swarms
-exa_py
-tavily-python
-openai
-exa_py
-pymongo
+aiohttp==3.9.5
+altair==5.3.0
+anyio==4.4.0
+attrs==23.2.0
+beautifulsoup4==4.12.3
+blinker==1.8.2
+cachetools==5.3.2
+certifi==2024.6.2
+charset-normalizer==3.3.2
+click==8.1.7
+dataclasses-json==0.6.7
+dnspython==2.6.1
+email_validator==2.1.2
+fastapi==0.111.0
+filelock==3.15.1
+frozenlist==1.4.1
+fsspec==2024.6.0
+gitdb==4.0.11
+GitPython==3.1.43
+google-api-core==2.15.0
+google-auth==2.26.2
+google-cloud-aiplatform==1.39.0
+google-cloud-bigquery==3.16.0
+google-cloud-core==2.4.1
+google-cloud-resource-manager==1.11.0
+google-cloud-storage==2.14.0
+google-crc32c==1.5.0
+google-resumable-media==2.7.0
+googleapis-common-protos==1.62.0
+grpc-google-iam-v1==0.13.0
+grpcio-status==1.60.0
+h11==0.14.0
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.4
+idna==3.7
+Jinja2==3.1.4
+jsonpatch==1.33
+jsonpointer==3.0.0
+jsonschema==4.22.0
+langchain==0.1.13
+loguru==0.7.2
+markdown-it-py==3.0.0
+markdownify==0.12.1
+MarkupSafe==2.1.5
+marshmallow==3.21.3
+multidict==6.0.5
+networkx==3.3
+numpy==1.26.4
+openai==1.34.0
+orjson==3.10.5
+packaging==23.2
+pandas==2.2.2
+pillow==10.3.0
+proto-plus==1.23.0
+protobuf==4.25.3
+pymongo==4.7.3
+pypdf==4.1.0
+python-dotenv==1.0.1
+pytz==2024.1
+PyYAML==6.0.1
+ratelimit==2.2.1
+referencing==0.35.1
+requests==2.32.3
+rich==13.7.1
+rsa==4.9
+safetensors==0.4.3
+sentry-sdk==2.5.1
+shapely==2.0.2
+six==1.16.0
+sniffio==1.3.1
+soupsieve==2.5
+SQLAlchemy==2.0.30
+starlette==0.37.2
+streamlit==1.35.0
+sympy==1.12.1
+tenacity==8.3.0
+termcolor==2.4.0
+tqdm==4.66.4
+transformers==4.41.2
+typer==0.12.3
+typing-inspect==0.9.0
+typing_extensions==4.8.0
+tzdata==2024.1
+ujson==5.10.0
+urllib3==2.2.1
+uvicorn==0.30.1
+uvloop==0.19.0
+watchfiles==0.22.0
+websockets==12.0
+yarl==1.9.4

test_func.py

@@ -24,7 +24,43 @@ def test_queries():
     # print("\nTesting Brave.com API:")
     # brave_result = brave_search_summarization(test_query)
     # print(brave_result)
+    
+def test_brave_ai_wrapper():
+    # Initialize the BraveAIWrapper with your API key
+    api_key = "your_api_key_here"
+    brave_ai = BraveAIWrapper(api_key=api_key)
+
+    # Define the test query
+    query = "What is some of the best mountain biking near Crested Butte, CO?"
+
+    # Test get_brave_results
+    print("Testing get_brave_results...")
+    results = brave_ai.get_brave_results(query)
+    if results:
+        print("get_brave_results output:", json.dumps(results, indent=2))
+    else:
+        print("get_brave_results failed.")
+
+    # Test get_and_summarize
+    print("\nTesting get_and_summarize...")
+    summary = brave_ai.get_and_summarize(query)
+    if summary:
+        print("get_and_summarize output:", summary)
+    else:
+        print("get_and_summarize failed.")
+
+    # Test download_documents
+    print("\nTesting download_documents...")
+    documents = brave_ai.download_documents(query)
+    if documents:
+        for doc in documents:
+            print("Document metadata:", doc.metadata)
+            print("Document content:", doc.page_content[:200])  # Print first 200 characters
+            print("-" * 40)
+    else:
+        print("download_documents failed.")
 
 
 if __name__ == "__main__":
     test_queries()
+    test_brave_ai_wrapper()