pvanand commited on
Commit
dd91c47
1 Parent(s): cfc59d1

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +47 -0
main.py CHANGED
@@ -281,6 +281,53 @@ async def news_assistant(query: NewsQueryModel, api_key: str = Depends(verify_ap
281
  #meta-llama/llama-3-70b-instruct google/gemini-pro-1.5
282
  return StreamingResponse(process_response(), media_type="text/event-stream")
283
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
284
  if __name__ == "__main__":
285
  import uvicorn
286
  uvicorn.run(app, host="0.0.0.0", port=7860)
 
281
  #meta-llama/llama-3-70b-instruct google/gemini-pro-1.5
282
  return StreamingResponse(process_response(), media_type="text/event-stream")
283
 
284
+ class SearchQueryModel(BaseModel):
285
+ query: str = Field(..., description="Search query")
286
+ model_id: ModelID = Field(
287
+ default="meta-llama/llama-3-70b-instruct",
288
+ description="ID of the model to use for response generation"
289
+ )
290
+ class Config:
291
+ schema_extra = {
292
+ "example": {
293
+ "query": "What are the latest advancements in quantum computing?",
294
+ "model_id": "meta-llama/llama-3-70b-instruct"
295
+ }
296
+ }
297
+
298
+ def analyze_search_results(query):
299
+ search_data = internet_search(query, type="web")
300
+
301
+ if not search_data:
302
+ return "Failed to fetch search data.", []
303
+
304
+ # Prepare the prompt for the AI
305
+ prompt = generate_search_prompt(query, search_data)
306
+
307
+ messages = [
308
+ {"role": "system", "content": SEARCH_ASSISTANT_PROMPT},
309
+ {"role": "user", "content": prompt}
310
+ ]
311
+
312
+ return messages
313
+
314
+ @app.post("/search-assistant")
315
+ async def search_assistant(query: SearchQueryModel, api_key: str = Depends(verify_api_key)):
316
+ """
317
+ Search assistant endpoint that provides summaries and analysis of web search results based on user queries.
318
+ Requires API Key authentication via X-API-Key header.
319
+ """
320
+ messages = analyze_search_results(query.query)
321
+
322
+ if not messages:
323
+ raise HTTPException(status_code=500, detail="Failed to fetch search data")
324
+
325
+ def process_response():
326
+ for content in chat_with_llama_stream(messages, model=query.model_id):
327
+ yield content
328
+
329
+ return StreamingResponse(process_response(), media_type="text/event-stream")
330
+
331
  if __name__ == "__main__":
332
  import uvicorn
333
  uvicorn.run(app, host="0.0.0.0", port=7860)