Update app.py

app.py

@@ -1,100 +1,72 @@
 import os
 import json
-import httpx
-from fastapi import FastAPI, Request
+import requests
+from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
-from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 from tavily import TavilyClient
 
 app = FastAPI()
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+class ChatRequest(BaseModel):
+    query: str
 
-@app.post("/chat")
-async def chat_endpoint(request: Request):
-    data = await request.json()
-    messages = data.get("messages", [])
+def generate_response(query: str):
+    tavily_api_key = os.environ.get("TAVILY_API_KEY", "")
+    tavily_client = TavilyClient(api_key=tavily_api_key)
     
-    user_query = ""
-    for msg in reversed(messages):
-        if msg.get("role") == "user":
-            user_query = msg.get("content", "")
-            break
+    tavily_res = tavily_client.search(query=query)
+    
+    context = "\n".join([r['content'] for r in tavily_res.get('results', [])])
+    sources = [{"title": r['title'], "url": r['url']} for r in tavily_res.get('results', [])]
 
-    tavily_key = os.getenv("TAVILY_API_KEY", "")
-    openrouter_key = os.getenv("OPENROUTER_API_KEY", "")
+    messages = [
+        {"role": "system", "content": f"Answer based on context:\n{context}"},
+        {"role": "user", "content": query}
+    ]
 
-    sources = []
-    context_text = ""
+    openrouter_api_key = os.environ.get("OPENROUTER_API_KEY", "")
+    
+    payload = {
+        "model": "nvidia/nemotron-3-ultra-550b-a55b:free",
+        "messages": messages,
+        "stream": True,
+        "reasoning": {"enabled": True}
+    }
 
-    if tavily_key and user_query:
-        try:
-            tavily_client = TavilyClient(api_key=tavily_key)
-            search_res = tavily_client.search(query=user_query, max_results=5)
-            
-            for item in search_res.get("results", []):
-                sources.append({
-                    "title": item.get("title", ""),
-                    "url": item.get("url", "")
-                })
-                context_text += f"\nTitle: {item.get('title')}\nURL: {item.get('url')}\nContent: {item.get('content')}\n"
-        except Exception:
-            pass
+    headers = {
+        "Authorization": f"Bearer {openrouter_api_key}",
+        "Content-Type": "application/json"
+    }
 
-    system_prompt = (
-        "You are a helpful AI assistant with real-time web search capabilities. "
-        "Use the following search context to answer the user's question accurately.\n"
-        f"Search Context:\n{context_text}"
+    response = requests.post(
+        "https://openrouter.ai/api/v1/chat/completions",
+        headers=headers,
+        json=payload,
+        stream=True
     )
 
-    llm_messages = [{"role": "system", "content": system_prompt}] + messages
-
-    async def event_generator():
-        async with httpx.AsyncClient(timeout=60.0) as client:
-            async with client.stream(
-                "POST",
-                "https://openrouter.ai/api/v1/chat/completions",
-                headers={
-                    "Authorization": f"Bearer {openrouter_key}",
-                    "Content-Type": "application/json",
-                },
-                json={
-                    "model": "nvidia/nemotron-3-ultra-550b-a55b:free",
-                    "messages": llm_messages,
-                    "stream": True,
-                    "reasoning": {"enabled": True}
-                }
-            ) as response:
-                async for line in response.aiter_lines():
-                    if line.startswith("data: "):
-                        data_str = line[6:].strip()
-                        if data_str == "[DONE]":
-                            break
-                        try:
-                            chunk = json.loads(data_str)
-                            delta = chunk.get("choices", [{}])[0].get("delta", {})
-
-                            reasoning = delta.get("reasoning_content") or delta.get("reasoning_details") or delta.get("reasoning")
-                            answer = delta.get("content")
-
-                            if reasoning:
-                                payload = json.dumps({"reasoning_content": reasoning}, ensure_ascii=False)
-                                yield f"data: {payload}\n\n"
-
-                            if answer:
-                                payload = json.dumps({"answer_content": answer}, ensure_ascii=False)
-                                yield f"data: {payload}\n\n"
-
-                        except Exception:
-                            pass
-
-        sources_payload = json.dumps({"sources": sources}, ensure_ascii=False)
-        yield f"data: {sources_payload}\n\n"
+    for line in response.iter_lines():
+        if line:
+            line_str = line.decode('utf-8')
+            if line_str.startswith("data: ") and line_str != "data: [DONE]":
+                try:
+                    data = json.loads(line_str[6:])
+                    delta = data['choices'][0]['delta']
+                    
+                    reasoning = delta.get('reasoning')
+                    if reasoning:
+                        yield f"data: {json.dumps({'type': 'reasoning', 'content': reasoning})}\n\n"
+                        
+                    content = delta.get('content')
+                    if content:
+                        yield f"data: {json.dumps({'type': 'answer', 'content': content})}\n\n"
+                except:
+                    continue
+                    
+    yield f"data: {json.dumps({'type': 'sources', 'content': sources})}\n\n"
+    yield "data: [DONE]\n\n"
 
-    return StreamingResponse(event_generator(), media_type="text/event-stream")
+@app.post("/api/search")
+async def search_api(request: ChatRequest):
+    return StreamingResponse(generate_response(request.query), media_type="text/event-stream")