Update main.py

main.py

@@ -18,13 +18,17 @@ logger = logging.getLogger(__name__)
 load_dotenv()
 LLM_API_KEY = os.getenv("LLM_API_KEY")
 
+# ***** CHANGE 1: Add API Key loading confirmation *****
 if not LLM_API_KEY:
     raise RuntimeError("LLM_API_KEY must be set in a .env file.")
+else:
+    logger.info(f"LLM API Key loaded successfully (starts with: {LLM_API_KEY[:4]}...).")
 
-# API URLs and Models
+# API URLs, Models, and a new constant for context size
 SNAPZION_API_URL = "https://search.snapzion.com/get-snippets"
 LLM_API_URL = "https://api.inference.net/v1/chat/completions"
 LLM_MODEL = "mistralai/mistral-nemo-12b-instruct/fp-8"
+MAX_CONTEXT_CHAR_LENGTH = 120000 # Safeguard: roughly 30k tokens
 
 # Headers for external services
 SNAPZION_HEADERS = { 'Content-Type': 'application/json', 'User-Agent': 'AI-Deep-Research-Agent/1.0' }
@@ -39,10 +43,10 @@ class DeepResearchRequest(BaseModel):
 app = FastAPI(
     title="AI Deep Research API",
     description="Provides single-shot AI search and streaming deep research completions.",
-    version="2.0.0"
+    version="2.1.0" # Version bump for new robustness feature
 )
 
-# --- Core Service Functions (Reused and New) ---
+# --- Core Service Functions (Unchanged) ---
 
 async def call_snapzion_search(session: aiohttp.ClientSession, query: str) -> list:
     try:
@@ -52,7 +56,7 @@ async def call_snapzion_search(session: aiohttp.ClientSession, query: str) -> li
             return data.get("organic_results", [])
     except Exception as e:
         logger.error(f"Snapzion search failed for query '{query}': {e}")
-        return [] # Return empty list on failure instead of crashing
+        return []
 
 async def scrape_url(session: aiohttp.ClientSession, url: str) -> str:
     if url.lower().endswith('.pdf'): return "Error: PDF content cannot be scraped."
@@ -69,16 +73,15 @@ async def scrape_url(session: aiohttp.ClientSession, url: str) -> str:
         return f"Error: {e}"
 
 async def search_and_scrape(session: aiohttp.ClientSession, query: str) -> tuple[str, list]:
-    """Performs the search and scrape pipeline for a given query."""
     search_results = await call_snapzion_search(session, query)
-    sources = search_results[:4] # Use top 4 sources per sub-query
+    sources = search_results[:4]
     if not sources: return "", []
 
     scrape_tasks = [scrape_url(session, source["link"]) for source in sources]
     scraped_contents = await asyncio.gather(*scrape_tasks)
 
     context = "\n\n".join(
-        f"Source [{i+1}] (from {sources[i]['link']}):\n{content}"
+        f"Source Details: Title '{sources[i]['title']}', URL '{sources[i]['link']}'\nContent:\n{content}"
         for i, content in enumerate(scraped_contents) if not content.startswith("Error:")
     )
     return context, sources
@@ -86,35 +89,26 @@ async def search_and_scrape(session: aiohttp.ClientSession, query: str) -> tuple
 # --- Streaming Deep Research Logic ---
 
 async def run_deep_research_stream(query: str) -> AsyncGenerator[str, None]:
-    """The main async generator for the deep research process."""
     
     def format_sse(data: dict) -> str:
-        """Formats a dictionary as a Server-Sent Event string."""
         return f"data: {json.dumps(data)}\n\n"
 
     try:
         async with aiohttp.ClientSession() as session:
-            # Step 1: Generate Sub-Questions
+            # Step 1: Generate Sub-Questions (Unchanged)
             yield format_sse({"event": "status", "data": "Generating research plan..."})
             sub_question_prompt = {
                 "model": LLM_MODEL,
-                "messages": [{
-                    "role": "user",
-                    "content": f"You are a research planner. Based on the user's query '{query}', generate a list of 3 to 4 crucial sub-questions that would form the basis of a comprehensive research report. Respond with ONLY a JSON array of strings. Example: [\"Question 1?\", \"Question 2?\"]"
-                }]
+                "messages": [{ "role": "user", "content": f"You are a research planner. For the topic '{query}', create a JSON array of 3-4 key sub-questions for a research report. Example: [\"Question 1?\", \"Question 2?\"]" }]
             }
             async with session.post(LLM_API_URL, headers=LLM_HEADERS, json=sub_question_prompt) as response:
                 response.raise_for_status()
                 result = await response.json()
-                try:
-                    sub_questions = json.loads(result['choices'][0]['message']['content'])
-                except (json.JSONDecodeError, IndexError):
-                    yield format_sse({"event": "error", "data": "Failed to parse sub-questions from LLM."})
-                    return
+                sub_questions = json.loads(result['choices'][0]['message']['content'])
             
             yield format_sse({"event": "plan", "data": sub_questions})
 
-            # Step 2: Concurrently research all sub-questions
+            # Step 2: Concurrently research all sub-questions (Unchanged)
             research_tasks = [search_and_scrape(session, sq) for sq in sub_questions]
             all_research_results = []
             
@@ -127,7 +121,13 @@ async def run_deep_research_stream(query: str) -> AsyncGenerator[str, None]:
             yield format_sse({"event": "status", "data": "Consolidating research..."})
             full_context = "\n\n---\n\n".join(res[0] for res in all_research_results if res[0])
             all_sources = [source for res in all_research_results for source in res[1]]
-            unique_sources = list({s['link']: s for s in all_sources}.values()) # Deduplicate sources
+            unique_sources = list({s['link']: s for s in all_sources}.values())
+            
+            # ***** CHANGE 2: Implement the context truncation safeguard *****
+            logger.info(f"Consolidated context size: {len(full_context)} characters.")
+            if len(full_context) > MAX_CONTEXT_CHAR_LENGTH:
+                logger.warning(f"Context is too long. Truncating from {len(full_context)} to {MAX_CONTEXT_CHAR_LENGTH} characters.")
+                full_context = full_context[:MAX_CONTEXT_CHAR_LENGTH]
             
             if not full_context.strip():
                 yield format_sse({"event": "error", "data": "Failed to gather any research context."})
@@ -135,38 +135,28 @@ async def run_deep_research_stream(query: str) -> AsyncGenerator[str, None]:
 
             # Step 4: Generate the final report with streaming
             yield format_sse({"event": "status", "data": "Generating final report..."})
+            final_report_prompt = f'Synthesize the provided context into a comprehensive report on "{query}". Use the context exclusively. Structure the report with markdown.\n\n## Research Context ##\n{full_context}'
             
-            final_report_prompt = f"""
-You are a research analyst. Your task is to synthesize the provided context into a comprehensive, well-structured report on the topic: "{query}".
-Use the context below exclusively. Do not use outside knowledge. Structure the report with markdown headings.
-
-## Research Context ##
-{full_context}
-"""
-            
-            final_report_payload = {
-                "model": LLM_MODEL,
-                "messages": [{"role": "user", "content": final_report_prompt}],
-                "stream": True # Enable streaming from the LLM
-            }
+            final_report_payload = {"model": LLM_MODEL, "messages": [{"role": "user", "content": final_report_prompt}], "stream": True}
 
             async with session.post(LLM_API_URL, headers=LLM_HEADERS, json=final_report_payload) as response:
-                response.raise_for_status()
+                # ***** CHANGE 3: More robust error handling for the streaming call *****
+                if response.status != 200:
+                    error_text = await response.text()
+                    logger.error(f"LLM API returned a non-200 status: {response.status} - {error_text}")
+                    raise Exception(f"LLM API Error: {response.status}, {error_text}")
+
                 async for line in response.content:
+                    # (Rest of the streaming logic is the same)
                     if line.strip():
-                        # The inference API might wrap its stream chunks in a 'data: ' prefix
                         line_str = line.decode('utf-8').strip()
-                        if line_str.startswith('data:'):
-                            line_str = line_str[5:].strip()
-                        if line_str == "[DONE]":
-                            break
+                        if line_str.startswith('data:'): line_str = line_str[5:].strip()
+                        if line_str == "[DONE]": break
                         try:
                             chunk = json.loads(line_str)
                             content = chunk.get("choices", [{}])[0].get("delta", {}).get("content")
-                            if content:
-                                yield format_sse({"event": "chunk", "data": content})
-                        except json.JSONDecodeError:
-                            continue # Ignore empty or malformed lines
+                            if content: yield format_sse({"event": "chunk", "data": content})
+                        except json.JSONDecodeError: continue
 
             yield format_sse({"event": "sources", "data": unique_sources})
             
@@ -178,25 +168,6 @@ Use the context below exclusively. Do not use outside knowledge. Structure the r
 
 
 # --- API Endpoints ---
-
-@app.get("/", include_in_schema=False)
-def root():
-    return {"message": "AI Deep Research API is active. See /docs for details."}
-
 @app.post("/v1/deepresearch/completions")
 async def deep_research_endpoint(request: DeepResearchRequest):
-    """
-    Performs a multi-step, streaming deep research task.
-
-    **Events Streamed:**
-    - `status`: Provides updates on the current stage of the process.
-    - `plan`: The list of sub-questions that will be researched.
-    - `chunk`: A piece of the final generated report.
-    - `sources`: The list of web sources used for the report.
-    - `error`: Indicates a fatal error occurred.
-    - `done`: Signals the end of the stream.
-    """
-    return StreamingResponse(
-        run_deep_research_stream(request.query),
-        media_type="text/event-stream"
-    )
+    return StreamingResponse(run_deep_research_stream(request.query), media_type="text/event-stream")