refactor

main.py

@@ -234,13 +234,10 @@ async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks,
 
 # New functions for news assistant
 
-def internet_search(query, type = "web", num_results=20):
-    logger.info(f"Performing internet search for query: {query}, type: {type}")
-    if type == "web":
-        url = "https://api.search.brave.com/res/v1/web/search"
-    else:
-        url = "https://api.search.brave.com/res/v1/news/search"
-
+def internet_search(query, search_type="web", num_results=20):
+    logger.info(f"Performing internet search for query: {query}, type: {search_type}")
+    url = f"https://api.search.brave.com/res/v1/{'web' if search_type == 'web' else 'news'}/search"
+    
     headers = {
         "Accept": "application/json",
         "Accept-Encoding": "gzip",
@@ -254,73 +251,52 @@ def internet_search(query, type = "web", num_results=20):
         logger.error(f"Failed to fetch search results. Status code: {response.status_code}")
         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 = {
+    search_data = response.json()["web"]["results"] if search_type == "web" else response.json()["results"]
+    
+    processed_results = [
+        {
             "title": item["title"],
             "snippet": item["extra_snippets"][0],
             "last_updated": item.get("age", "")
         }
-        processed_results.append(result)
+        for item in search_data
+        if item.get("extra_snippets")
+    ][:num_results]
 
     logger.info(f"Retrieved {len(processed_results)} search results")
-    return processed_results[:num_results]
+    return processed_results
 
 @lru_cache(maxsize=100)
 def cached_internet_search(query: str):
     logger.info(f"Performing cached internet search for query: {query}")
-    return internet_search(query, type = "news")
-
+    return internet_search(query, search_type="news")
 
-def analyze_news(query):
-    logger.info(f"Analyzing news for query: {query}")
-    news_data = cached_internet_search(query)
+def analyze_data(query, data_type="news"):
+    logger.info(f"Analyzing {data_type} for query: {query}")
     
-    if not news_data:
-        logger.error("Failed to fetch news data")
-        return "Failed to fetch news data.", []
-
-    # Prepare the prompt for the AI
-    # Use the imported function to generate the prompt (now includes today's date)
-    prompt = generate_news_prompt(query, news_data)
+    if data_type == "news":
+        data = cached_internet_search(query)
+        prompt_generator = generate_news_prompt
+        system_prompt = NEWS_ASSISTANT_PROMPT
+    else:
+        data = internet_search(query, search_type="web")
+        prompt_generator = generate_search_prompt
+        system_prompt = SEARCH_ASSISTANT_PROMPT
+    
+    if not data:
+        logger.error(f"Failed to fetch {data_type} data")
+        return None
 
+    prompt = prompt_generator(query, data)
     messages = [
-        {"role": "system", "content": NEWS_ASSISTANT_PROMPT},
+        {"role": "system", "content": system_prompt},
         {"role": "user", "content": prompt}
     ]
 
-    logger.info("News analysis completed")
+    logger.info(f"{data_type.capitalize()} analysis completed")
     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.
-    """
-    logger.info(f"Received news assistant query: {query.query}")
-    messages = analyze_news(query.query)
-    
-    if not messages:
-        logger.error("Failed to fetch news data")
-        raise HTTPException(status_code=500, detail="Failed to fetch news data")
-        
-    def process_response():
-        for content in chat_with_llama_stream(messages, model=query.model_id):
-            yield content
-        logger.info(f"Completed news assistant response for query: {query.query}")
-
-    return StreamingResponse(process_response(), media_type="text/event-stream")
-
-class SearchQueryModel(BaseModel):
+class QueryModel(BaseModel):
     query: str = Field(..., description="Search query")
     model_id: ModelID = Field(
         default="openai/gpt-4o-mini",
@@ -334,49 +310,40 @@ class SearchQueryModel(BaseModel):
             }
         }
 
-def analyze_search_results(query):
-    logger.info(f"Analyzing search results for query: {query}")
-    search_data = internet_search(query, type="web")
-    
-    if not search_data:
-        logger.error("Failed to fetch search data")
-        return "Failed to fetch search data.", []
-
-    # Prepare the prompt for the AI
-    prompt = generate_search_prompt(query, search_data)
-
-    messages = [
-        {"role": "system", "content": SEARCH_ASSISTANT_PROMPT},
-        {"role": "user", "content": prompt}
-    ]
-
-    logger.info("Search results analysis completed")
-    return messages
-
-@app.post("/search-assistant")
-async def search_assistant(query: SearchQueryModel, api_key: str = Depends(verify_api_key)):
-    """
-    Search assistant endpoint that provides summaries and analysis of web search results based on user queries.
-    Requires API Key authentication via X-API-Key header.
-    """
-    logger.info(f"Received search assistant query: {query.query}")
-    messages = analyze_search_results(query.query)
+def assistant_api(query, data_type):
+    logger.info(f"Received {data_type} assistant query: {query}")
+    messages = analyze_data(query, data_type)
     
     if not messages:
-        logger.error("Failed to fetch search data")
-        raise HTTPException(status_code=500, detail="Failed to fetch search data")
-        
+        logger.error(f"Failed to fetch {data_type} data")
+        raise HTTPException(status_code=500, detail=f"Failed to fetch {data_type} data")
+    
     def process_response():
         logger.info(f"Generating response using LLM: {messages}")
         full_response = ""
         for content in chat_with_llama_stream(messages, model=query.model_id):
-            full_response+=content
+            full_response += content
             yield content
-        logger.info(f"Completed search assistant response for query: {query.query}")
+        logger.info(f"Completed {data_type} assistant response for query: {query}")
         logger.info(f"LLM Response: {full_response}")
 
     return StreamingResponse(process_response(), media_type="text/event-stream")
 
+@app.post("/news-assistant")
+async def news_assistant(query: QueryModel, 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.
+    """
+    return assistant_api(query.query, "news")
+
+@app.post("/search-assistant")
+async def search_assistant(query: QueryModel, api_key: str = Depends(verify_api_key)):
+    """
+    Search assistant endpoint that provides summaries and analysis of web search results based on user queries.
+    Requires API Key authentication via X-API-Key header.
+    """
+    return assistant_api(query.query, "search")
 
 from pydantic import BaseModel, Field
 import yaml
@@ -624,6 +591,7 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
 
     return StreamingResponse(process_response(), media_type="text/event-stream")
 
+
 @app.post("/v3/followup-agent")
 async def followup_agent(query: FollowupQueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
     """