Update main.py

main.py

@@ -12,6 +12,7 @@ import sqlite3
 import time
 from datetime import datetime, timedelta
 import asyncio
+import requests
 
 app = FastAPI()
 
@@ -48,10 +49,21 @@ class QueryModel(BaseModel):
             }
         }
 
+class NewsQueryModel(BaseModel):
+    query: str = Field(..., description="News topic to search for")
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "query": "Latest developments in AI"
+            }
+        }
+
 @lru_cache()
 def get_api_keys():
     return {
-        "OPENROUTER_API_KEY": f"sk-or-v1-{os.environ['OPENROUTER_API_KEY']}"
+        "OPENROUTER_API_KEY": f"sk-or-v1-{os.environ['OPENROUTER_API_KEY']}",
+        "BRAVE_API_KEY": os.environ['BRAVE_API_KEY']
     }
 
 api_keys = get_api_keys()
@@ -152,7 +164,6 @@ async def startup_event():
 async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
     """
     Coding assistant endpoint that provides programming help based on user queries.
-
     Available models:
     - meta-llama/llama-3-70b-instruct (default)
     - anthropic/claude-3.5-sonnet
@@ -161,7 +172,6 @@ async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks,
     - openai/gpt-3.5-turbo-instruct
     - qwen/qwen-72b-chat
     - google/gemma-2-27b-it
-
     Requires API Key authentication via X-API-Key header.
     """
     if query.conversation_id not in conversations:
@@ -184,6 +194,69 @@ async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks,
 
     return StreamingResponse(process_response(), media_type="text/event-stream")
 
+# New functions for news assistant
+def fetch_news(query, num_results=20):
+    url = "https://api.search.brave.com/res/v1/news/search"
+    headers = {
+        "Accept": "application/json",
+        "Accept-Encoding": "gzip",
+        "X-Subscription-Token": api_keys["BRAVE_API_KEY"]
+    }
+    params = {"q": query}
+
+    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:
+        return []
+
+def analyze_news(query):
+    news_data = fetch_news(query)
+    
+    if not news_data:
+        return "Failed to fetch news data.", []
+
+    # Prepare the prompt for the AI
+    prompt = f"Based on the following recent news about '{query}', Provide a well summarized and formatted answer using markdown, give importance to the latest news:\n\n"
+    for item in news_data:
+        prompt += f"Title: {item['title']}\n"
+        prompt += f"Snippet: {item['snippet']}\n"
+        prompt += f"Last Updated: {item['last_updated']}\n\n"
+
+    messages = [
+        {"role": "system", "content": "You are a knowledgeable Assistant capable of providing insightful answers on various topics."},
+        {"role": "user", "content": prompt}
+    ]
+
+    return messages
+
+@app.post("/news-assistant")
+async def news_assistant(query: NewsQueryModel, api_key: str = Depends(verify_api_key)):
+    """
+    News assistant endpoint that provides summaries and analysis of recent news based on user queries.
+    Requires API Key authentication via X-API-Key header.
+    """
+    messages = analyze_news(query.query)
+    
+    if not messages:
+        raise HTTPException(status_code=500, detail="Failed to fetch news data")
+
+    def process_response():
+        for content in chat_with_llama_stream(messages, model="meta-llama/llama-3-70b-instruct"):
+            yield content
+
+    return StreamingResponse(process_response(), media_type="text/event-stream")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)