worked on getting the summary from Brave Search, got the web search results and the summary output working but still need to work on converting the resulting webpages to documents.

brave_ai.py

@@ -0,0 +1,88 @@
+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,
+        }
+        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

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()