Update main.py

main.py

@@ -310,7 +310,7 @@ class QueryModel(BaseModel):
             }
         }
 
-def assistant_api(query, data_type):
+def 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)
     
@@ -321,7 +321,7 @@ def assistant_api(query, data_type):
     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):
+        for content in chat_with_llama_stream(messages, model=model):
             full_response += content
             yield content
         logger.info(f"Completed {data_type} assistant response for query: {query}")
@@ -335,7 +335,7 @@ 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")
+    return assistant_api(query.query, "news",model=query.model_id)
 
 @app.post("/search-assistant")
 async def search_assistant(query: QueryModel, api_key: str = Depends(verify_api_key)):
@@ -343,7 +343,7 @@ 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")
+    return assistant_api(query.query, "search",model=query.model_id)
 
 from pydantic import BaseModel, Field
 import yaml