Update app.py

app.py

@@ -1,87 +1,95 @@
 import os
 import json
-import time
-import requests
+import asyncio
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from tavily import TavilyClient
+from duckduckgo_search import DDGS
+import requests
 
 app = FastAPI()
 
-class SearchRequest(BaseModel):
+class QueryRequest(BaseModel):
     query: str
 
-@app.post("/search")
-async def search_endpoint(request: SearchRequest):
-    tavily_key = os.environ.get("TAVILY_API_KEY")
-    openrouter_key = os.environ.get("OPENROUTER_API_KEY")
+def search_duckduckgo(query: str, max_results: int = 5):
+    try:
+        with DDGS() as ddgs:
+            results = [r for r in ddgs.text(query, max_results=max_results)]
+            return results
+    except Exception:
+        return []
+
+async def stream_search_and_llm(query: str):
+    OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
+    if not OPENROUTER_API_KEY:
+        yield f"data: {json.dumps({'error': 'API key missing'})}\n\n"
+        return
+
+    search_results = search_duckduckgo(query)
     
-    if not tavily_key or not openrouter_key:
-        raise HTTPException(status_code=500, detail="Missing API keys configuration.")
+    context = ""
+    sources = []
+    for idx, result in enumerate(search_results, 1):
+        title = result.get('title', '')
+        href = result.get('href', '')
+        body = result.get('body', '')
+        context += f"[{idx}] Source: {title}\nURL: {href}\nContent: {body}\n\n"
+        sources.append({"id": idx, "title": title, "url": href})
+
+    system_prompt = (
+        "You are a helpful AI assistant with access to real-time web search results. "
+        "Answer the user's question accurately based on the provided search results. "
+        "Cite sources using [1], [2], etc., matching the given source numbers. "
+        "If the search results do not contain the answer, use your knowledge but state so."
+    )
     
+    user_content = f"Search Results:\n{context}\n\nUser Query: {query}"
+
+    url = "https://openrouter.ai/api/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+        "Content-Type": "application/json",
+    }
+    payload = {
+        "model": "nvidia/nemotron-3-ultra-550b-a55b:free",
+        "stream": True,
+        "messages": [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_content}
+        ]
+    }
+
     try:
-        tavily_client = TavilyClient(api_key=tavily_key)
-        research_res = tavily_client.research(request.query)
-        request_id = research_res.get("request_id")
-        
-        context = ""
-        sources = []
-        
-        while True:
-            result = tavily_client.get_research_result(request_id)
-            if result.get("status") == "completed":
-                context = result.get("content", "")
-                sources = result.get("sources", [])
-                break
-            elif result.get("status") == "failed":
-                break
-            time.sleep(1)
-    except Exception:
-        context = ""
-        sources = []
+        response = requests.post(url, headers=headers, json=payload, stream=True)
+        if response.status_code != 200:
+            yield f"data: {json.dumps({'error': 'LLM request failed'})}\n\n"
+            return
+
+        for line in response.iter_lines():
+            if line:
+                line_str = line.decode('utf-8').strip()
+                if line_str.startswith("data:"):
+                    data_body = line_str[5:].strip()
+                    if data_body == "[DONE]":
+                        break
+                    try:
+                        data_json = json.loads(data_body)
+                        choices = data_json.get("choices", [])
+                        if choices:
+                            delta = choices[0].get("delta", {})
+                            content = delta.get("content", "")
+                            if content:
+                                yield f"data: {json.dumps({'type': 'answer', 'content': content})}\n\n"
+                    except Exception:
+                        continue
 
-    def event_generator():
-        url = "https://openrouter.ai/api/v1/chat/completions"
-        headers = {
-            "Authorization": f"Bearer {openrouter_key}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "model": "nvidia/nemotron-3-ultra-550b-a55b:free",
-            "stream": True,
-            "messages": [
-                {"role": "system", "content": f"Answer the user query based on this context:\n{context}"},
-                {"role": "user", "content": request.query}
-            ]
-        }
-        
-        try:
-            response = requests.post(url, headers=headers, json=payload, stream=True)
-            for line in response.iter_lines():
-                if line:
-                    decoded_line = line.decode('utf-8')
-                    if decoded_line.startswith("data: "):
-                        data_str = decoded_line[6:]
-                        if data_str.strip() == "[DONE]":
-                            break
-                        try:
-                            data_json = json.loads(data_str)
-                            choices = data_json.get("choices", [])
-                            if choices:
-                                delta = choices[0].get("delta", {})
-                                content = delta.get("content", "")
-                                if content:
-                                    yield f"data: {json.dumps({'type': 'content', 'value': content})}\n\n"
-                        except Exception:
-                            pass
-        except Exception as e:
-            yield f"data: {json.dumps({'type': 'error', 'value': str(e)})}\n\n"
-        
-        yield f"data: {json.dumps({'type': 'sources', 'value': sources})}\n\n"
+    except Exception:
+        yield f"data: {json.dumps({'error': 'Streaming failed'})}\n\n"
+        return
 
-    return StreamingResponse(event_generator(), media_type="text/event-stream")
+    yield f"data: {json.dumps({'type': 'sources', 'content': sources})}\n\n"
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+@app.post("/api/search")
+async def search_endpoint(request: QueryRequest):
+    return StreamingResponse(stream_search_and_llm(request.query), media_type="text/event-stream")