Update main.py

main.py

@@ -13,7 +13,7 @@ import time
 from datetime import datetime, timedelta
 import asyncio
 import requests
-from prompts import CODING_ASSISTANT_PROMPT, NEWS_ASSISTANT_PROMPT, generate_news_prompt
+from prompts import CODING_ASSISTANT_PROMPT, NEWS_ASSISTANT_PROMPT, generate_news_prompt, SEARCH_ASSISTANT_PROMPT, generate_search_prompt
 from fastapi_cache import FastAPICache
 from fastapi_cache.backends.inmemory import InMemoryBackend
 from fastapi_cache.decorator import cache
@@ -205,8 +205,12 @@ async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks,
 
 # New functions for news assistant
 
-def fetch_news(query, num_results=20):
-    url = "https://api.search.brave.com/res/v1/news/search"
+def internet_search(query, type = "web", num_results=20):
+    if type == "web":
+        url = "https://api.search.brave.com/res/v1/web/search"
+    else:
+        url = "https://api.search.brave.com/res/v1/news/search"
+
     headers = {
         "Accept": "application/json",
         "Accept-Encoding": "gzip",
@@ -216,27 +220,35 @@ def fetch_news(query, num_results=20):
 
     response = requests.get(url, headers=headers, params=params)
 
-    if response.status_code == 200:
-        news_data = response.json()
-        return [
-            {
-                "title": item["title"],
-                "snippet": item["extra_snippets"][0] if "extra_snippets" in item and item["extra_snippets"] else "",
-                "last_updated": item.get("age", ""),
-            }
-            for item in news_data['results']
-            if "extra_snippets" in item and item["extra_snippets"]
-        ][:num_results]
-    else:
+    if response.status_code != 200:
         return []
+    
+    if type == "web":
+        search_data = response.json()["web"]["results"]
+    else:
+        search_data = response.json()["results"]
+    processed_results = []
+
+    for item in search_data:
+        if not item.get("extra_snippets"):
+            continue
+        
+        result = {
+            "title": item["title"],
+            "snippet": item["extra_snippets"][0],
+            "last_updated": item.get("age", "")
+        }
+        processed_results.append(result)
+
+    return processed_results[:num_results]
 
 @lru_cache(maxsize=100)
-def cached_fetch_news(query: str):
-    return fetch_news(query)
+def cached_internet_search(query: str):
+    return internet_search(query, type = "news")
 
 
 def analyze_news(query):
-    news_data = cached_fetch_news(query)
+    news_data = cached_internet_search(query)
     
     if not news_data:
         return "Failed to fetch news data.", []