Update main.py

main.py

@@ -55,11 +55,15 @@ class QueryModel(BaseModel):
 
 class NewsQueryModel(BaseModel):
     query: str = Field(..., description="News topic to search for")
-
+    model_id: ModelID = Field(
+        default="meta-llama/llama-3-70b-instruct",
+        description="ID of the model to use for response generation"
+    )
     class Config:
         schema_extra = {
             "example": {
-                "query": "Latest developments in AI"
+                "query": "Latest developments in AI",
+                "model_id": "meta-llama/llama-3-70b-instruct"
             }
         }
 
@@ -260,7 +264,7 @@ async def news_assistant(query: NewsQueryModel, api_key: str = Depends(verify_ap
         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"):
+        for content in chat_with_llama_stream(messages, model=query.model_id):
             yield content
 #meta-llama/llama-3-70b-instruct  google/gemini-pro-1.5
     return StreamingResponse(process_response(), media_type="text/event-stream")