Update app.py

app.py

@@ -1,142 +1,33 @@
-import os
 import asyncio
-import re
-import json
-from typing import Optional
-from datetime import datetime, timezone
 from fastapi import FastAPI
-from pydantic import BaseModel
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-import httpx
-import trafilatura
-import google.generativeai as genai
+from pydantic import BaseModel
 
 # --- FastAPI App with CORS ---
-app = FastAPI(title="AI Research Agent API")
+app = FastAPI(title="Streaming Test App")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
 )
 
-# --- Prompts (No changes) ---
-PROMPT_ULTRADEEP_PLANNER = """You are an expert research planner... Your output MUST be a valid JSON object: {"queries": ["query 1", ...]}..."""
-PROMPT_ULTRADEEP_SYNTHESIZER = """You are a master research analyst... If info is incomplete, output JSON: `{"status": "INCOMPLETE", "missing_queries": ["..."]}`. If complete, output the final markdown report..."""
-
-# --- Core Logic (No changes) ---
-async def search_web_logic(query: str, serper_api_key: str) -> str:
-    # ... (code from previous version, no changes)
-    if not serper_api_key: return "Error: Serper API Key is missing."
-    try:
-        headers = {"X-API-KEY": serper_api_key, "Content-Type": "application/json"}
-        async with httpx.AsyncClient(timeout=15) as client:
-            resp = await client.post("https://google.serper.dev/search", headers=headers, json={"q": query, "num": 5})
-        if resp.status_code == 401: return "Error: Invalid Serper API Key."
-        if resp.status_code != 200: return f"Error: Serper API returned status {resp.status_code}."
-        results = resp.json().get("organic", []);
-        if not results: return f"Error: No web results found for query '{query}'."
-        urls = [r["link"] for r in results]
-        async with httpx.AsyncClient(timeout=20, follow_redirects=True) as client:
-            tasks = [client.get(u) for u in urls]; responses = await asyncio.gather(*tasks, return_exceptions=True)
-        texts = [f"Source URL: {meta['link']}\nContent: {body.strip()}\n" for meta, response in zip(results, responses) if not isinstance(response, Exception) and (body := trafilatura.extract(response.text))]
-        if not texts: return "Error: Found web results, but could not extract content."
-        return "\n---\n".join(texts)
-    except Exception as e: return f"Error during web search: {str(e)}"
-
-async def call_gemini(prompt: str, gemini_key: str, model_name: str, json_mode: bool = False) -> str:
-    # ... (code from previous version, no changes)
-    if not gemini_key: return json.dumps({"error": "Gemini API Key is missing."})
-    try:
-        genai.configure(api_key=gemini_key)
-        model = genai.GenerativeModel(model_name)
-        generation_config = {"response_mime_type": "application/json"} if json_mode else None
-        response = await model.generate_content_async(prompt, generation_config=generation_config)
-        return response.text
-    except Exception as e: return json.dumps({"error": f"Error calling Gemini: {str(e)}"})
-
-# --- AI Agent with Heartbeat Logic ---
-async def ultradeep_research_agent_streamer(query: str, serper_api_key: str, gemini_key: str, model_name: str):
-    yield "STATUS: Initiating UltraDeep Research Agent... 🤖\n"
-    
-    # Step 1: Plan with Heartbeat
-    yield "STATUS: Creating a research plan with AI... 🧠"
-    planner_prompt = PROMPT_ULTRADEEP_PLANNER.format(query=query)
+# --- Dummy Streaming Function ---
+async def dummy_streamer():
+    yield "STATUS: Connection established! Starting test...\n"
+    await asyncio.sleep(2)
     
-    # <<< HEARTBEAT LOGIC START >>>
-    gemini_task = asyncio.create_task(call_gemini(planner_prompt, gemini_key, model_name, json_mode=True))
-    while not gemini_task.done():
-        yield "." # Send a heartbeat every 2 seconds
-        await asyncio.sleep(2)
-    plan_str = await gemini_task
-    yield "\n" # Newline after heartbeats
-    # <<< HEARTBEAT LOGIC END >>>
-
-    try:
-        match = re.search(r'\{.*\}', plan_str, re.DOTALL)
-        if not match: raise ValueError("No JSON object found in Gemini's planner response.")
-        plan_data = json.loads(match.group(0))
-        if "error" in plan_data:
-            yield f"FINAL: Error during planning phase: {plan_data['error']}"
-            return
-        search_queries = plan_data["queries"]
-        plan_display = "\n".join([f"  - {q}" for q in search_queries])
-        yield f"PLAN:\n{plan_display}\n\n"
-    except Exception as e:
-        yield f"FINAL:Error: Could not create a valid research plan. Details: {str(e)}\nRaw Response: {plan_str}"
-        return
-
-    # Step 2: Execute
-    all_scraped_text = ""
-    for i, sub_query in enumerate(search_queries):
-        yield f"STATUS: Searching for '{sub_query}' ({i+1}/{len(search_queries)})... 🕵️‍♂️\n"
-        # Search can also be slow, so we can add a heartbeat here too if needed, but let's skip for now
-        scraped_text = await search_web_logic(sub_query, serper_api_key)
-        if scraped_text.startswith("Error:"):
-            yield f"WARNING: Skipping search for '{sub_query}': {scraped_text}\n"
-        else:
-            all_scraped_text += scraped_text
+    for i in range(1, 6):
+        yield f"MESSAGE: Ping #{i} from server.\n"
         await asyncio.sleep(1)
+        
+    yield "FINAL: Test complete. Connection is working!"
 
-    if not all_scraped_text:
-        yield "FINAL: Error: Could not retrieve any web content. Please check Serper key and try again."
-        return
-
-    # Step 3: Synthesize with Heartbeat
-    yield "STATUS: All searches complete. Synthesizing the final report... ✍️"
-    current_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
-    synthesizer_prompt = PROMPT_ULTRADEEP_SYNTHESIZER.format(query=query, current_date=current_date, context_text=all_scraped_text)
-    
-    # <<< HEARTBEAT LOGIC START >>>
-    synthesis_task = asyncio.create_task(call_gemini(synthesizer_prompt, gemini_key, model_name))
-    while not synthesis_task.done():
-        yield "."
-        await asyncio.sleep(2)
-    final_report = await synthesis_task
-    yield "\n"
-    # <<< HEARTBEAT LOGIC END >>>
-    
-    yield f"FINAL:{final_report}"
-
-# --- FastAPI Endpoints and Server Startup (No Changes) ---
-class ResearchRequest(BaseModel):
-    query: str
-    serper_api_key: str
-    gemini_api_key: str
-    research_mode: str = "ultradeep"
-    gemini_model: str = "gemini-1.5-flash-latest"
-
-@app.post("/api/research-stream")
-async def api_research_stream(request: ResearchRequest):
-    # This endpoint is now only for ultradeep
-    return StreamingResponse(
-        ultradeep_research_agent_streamer(
-            request.query, request.serper_api_key, request.gemini_api_key, request.gemini_model
-        ), 
-        media_type="text/event-stream"
-    )
-
-# ... (We can add back the simple endpoint and Gradio UI later if needed)
+# --- Test Endpoint ---
+@app.post("/api/test-stream")
+async def api_test_stream():
+    return StreamingResponse(dummy_streamer(), media_type="text/event-stream")
 
+# --- Server Startup ---
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)