Update main.py

main.py

@@ -310,7 +310,7 @@ class QueryModel(BaseModel):
             }
         }
 
-def assistant_api(query, data_type,model = "openai/gpt-4o-mini"):
+def search_assistant_api(query, data_type, model="openai/gpt-4o-mini"):
     logger.info(f"Received {data_type} assistant query: {query}")
     messages = analyze_data(query, data_type)
     
@@ -326,8 +326,11 @@ def assistant_api(query, data_type,model = "openai/gpt-4o-mini"):
             yield content
         logger.info(f"Completed {data_type} assistant response for query: {query}")
         logger.info(f"LLM Response: {full_response}")
+    
+    return process_response
 
-    return StreamingResponse(process_response(), media_type="text/event-stream")
+def create_streaming_response(generator):
+    return StreamingResponse(generator(), media_type="text/event-stream")
 
 @app.post("/news-assistant")
 async def news_assistant(query: QueryModel, api_key: str = Depends(verify_api_key)):
@@ -335,7 +338,8 @@ async def news_assistant(query: QueryModel, api_key: str = Depends(verify_api_ke
     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",model=query.model_id)
+    response_generator = search_assistant_api(query.query, "news", model=query.model_id)
+    return create_streaming_response(response_generator)
 
 @app.post("/search-assistant")
 async def search_assistant(query: QueryModel, api_key: str = Depends(verify_api_key)):
@@ -343,7 +347,8 @@ async def search_assistant(query: QueryModel, api_key: str = Depends(verify_api_
     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",model=query.model_id)
+    response_generator = search_assistant_api(query.query, "web", model=query.model_id)
+    return create_streaming_response(response_generator)
 
 from pydantic import BaseModel, Field
 import yaml
@@ -547,15 +552,13 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
         logger.info(f"Completed followup agent response for query: {query.query}, send result: {result}")
 
     return StreamingResponse(process_response(), media_type="text/event-stream")
-
 @app.post("/v2/followup-tools-agent")
-async def followup_agent(query: FollowupQueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
+def followup_agent(query: FollowupQueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
     """
     Followup agent endpoint that provides helpful responses or generates clarifying questions based on user queries.
     Requires API Key authentication via X-API-Key header.
     """
     logger.info(f"Received followup agent query: {query.query}")
-
     if query.conversation_id not in conversations:
         conversations[query.conversation_id] = [
             {"role": "system", "content": MULTI_AGENT_PROMPT}
@@ -566,19 +569,19 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
     
     # Limit tokens in the conversation history
     limited_conversation = conversations[query.conversation_id]
-
+    
     def process_response():
         full_response = ""
         for content in chat_with_llama_stream(limited_conversation, model=query.model_id):
             full_response += content
             yield content
-
+        
         logger.info(f"LLM RAW response for query: {query.query}: {full_response}")
         response_content, interact, tools = parse_followup_and_tools(full_response)
         
         result = {
             "clarification": interact,
-            "tools":tools
+            "tools": tools
         }
         
         yield "<json>" + json.dumps(result)
@@ -586,9 +589,19 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
         # Add the assistant's response to the conversation history
         conversations[query.conversation_id].append({"role": "assistant", "content": full_response})
         
+        # Process tool if present
+        if tools and len(tools) > 0:
+            tool = tools[0]  # Assume only one tool is present
+            if tool["name"] in ["news", "web"]:
+                search_query = tool["input"]
+                search_response = search_assistant_api(search_query, tool["name"], model=query.model_id)
+                
+                for content in search_response():
+                    yield content
+        
         background_tasks.add_task(update_db, query.user_id, query.conversation_id, query.query, full_response)
         logger.info(f"Completed followup agent response for query: {query.query}, send result: {result}")
-
+    
     return StreamingResponse(process_response(), media_type="text/event-stream")