Update main.py

main.py

@@ -12,7 +12,7 @@ import aiohttp
 from bs4 import BeautifulSoup
 
 # --- Configuration ---
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 load_dotenv()
@@ -21,22 +21,24 @@ LLM_API_KEY = os.getenv("LLM_API_KEY")
 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]}...).")
+    logger.info(f"LLM API Key loaded successfully.")
 
-# API URLs, Models, and context size limit
+# ***** CHANGE 1: Update constants to match your new API provider *****
 SNAPZION_API_URL = "https://search.snapzion.com/get-snippets"
 LLM_API_URL = "https://api.typegpt.net/v1/chat/completions"
-LLM_MODEL = "gpt-4.1-mini" # Corrected model name from previous attempts
+LLM_MODEL = "gpt-4.1-mini"
 MAX_CONTEXT_CHAR_LENGTH = 120000
 
 # Headers for external services
 SNAPZION_HEADERS = { 'Content-Type': 'application/json', 'User-Agent': 'AI-Deep-Research-Agent/1.0' }
 SCRAPING_HEADERS = { 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36' }
-# ***** CHANGE 1: Add a User-Agent to the LLM headers *****
+
+# ***** CHANGE 2: Create more standard and robust headers for the LLM call *****
 LLM_HEADERS = {
     "Authorization": f"Bearer {LLM_API_KEY}",
     "Content-Type": "application/json",
-    "User-Agent": "AI-Deep-Research-Client/2.2"
+    "Accept": "application/json", # Explicitly request a JSON response
+    "User-Agent": "AI-Deep-Research-Client/2.3"
 }
 
 # --- Pydantic Models ---
@@ -47,7 +49,7 @@ class DeepResearchRequest(BaseModel):
 app = FastAPI(
     title="AI Deep Research API",
     description="Provides streaming deep research completions.",
-    version="2.2.0" # Version bump for critical bug fix
+    version="2.3.0" # Version bump for advanced error handling
 )
 
 # --- Core Service Functions (Unchanged) ---
@@ -94,6 +96,7 @@ async def run_deep_research_stream(query: str) -> AsyncGenerator[str, None]:
     def format_sse(data: dict) -> str:
         return f"data: {json.dumps(data)}\n\n"
 
+    raw_response_text_for_debugging = "" # Variable to hold response text for logging
     try:
         async with aiohttp.ClientSession() as session:
             # Step 1: Generate Sub-Questions
@@ -104,63 +107,63 @@ async def run_deep_research_stream(query: str) -> AsyncGenerator[str, None]:
                 "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. Respond ONLY with the JSON array. Example: [\"Question 1?\", \"Question 2?\"]" }]
             }
 
-            # ***** CHANGE 2: Implement robust parsing for the API call *****
+            # ***** CHANGE 3: The most critical fix. Heavily reinforced error handling. *****
             try:
+                logger.info(f"Sending request to LLM for planning. Model: {LLM_MODEL}, URL: {LLM_API_URL}")
                 async with session.post(LLM_API_URL, headers=LLM_HEADERS, json=sub_question_prompt, timeout=20) as response:
+                    raw_response_text_for_debugging = await response.text()
+                    
                     if response.status != 200:
-                        error_text = await response.text()
-                        logger.error(f"LLM API for planning failed with status {response.status}: {error_text}")
-                        raise Exception(f"LLM API returned non-200 status: {response.status}")
+                        logger.error(f"LLM API for planning failed! Status: {response.status}, Headers: {response.headers}, Body: {raw_response_text_for_debugging}")
+                        raise Exception(f"LLM provider returned non-200 status: {response.status}")
                     
-                    raw_response_text = await response.text()
-                    if not raw_response_text:
-                        raise Exception("LLM API returned an empty response.")
+                    if not raw_response_text_for_debugging:
+                        raise Exception("LLM provider returned an empty response body.")
                         
-                    result = json.loads(raw_response_text)
-                    llm_content = result['choices'][0]['message']['content']
+                    result = json.loads(raw_response_text_for_debugging)
+                    llm_content = result.get('choices', [{}])[0].get('message', {}).get('content', '')
+                    
+                    if not llm_content or not llm_content.strip().startswith('['):
+                        logger.error(f"LLM did not return a valid JSON array string. Received: {llm_content}")
+                        raise Exception("LLM failed to generate a valid research plan.")
+
                     sub_questions = json.loads(llm_content)
+
             except Exception as e:
-                logger.error(f"Failed to generate or parse research plan: {e}")
+                # This will now catch the JSON error and log the problematic text
+                logger.error(f"Failed to generate/parse research plan. Error: {e}. Raw API Response: '{raw_response_text_for_debugging}'")
                 yield format_sse({"event": "error", "data": f"Could not generate research plan. Reason: {e}"})
-                return # Stop the process if planning fails
+                return
 
             yield format_sse({"event": "plan", "data": sub_questions})
 
-            # (The rest of the logic remains the same)
-            # Step 2: Concurrently research all sub-questions
+            # (The rest of the logic remains the same, as it was not the point of failure)
             research_tasks = [search_and_scrape(session, sq) for sq in sub_questions]
-            all_research_results = []
+            yield format_sse({"event": "status", "data": f"Starting research on {len(sub_questions)} topics..."})
             
-            for i, task in enumerate(asyncio.as_completed(research_tasks)):
-                yield format_sse({"event": "status", "data": f"Researching: \"{sub_questions[i]}\""})
-                result = await task
-                all_research_results.append(result)
-
-            # Step 3: Consolidate all context and sources
-            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())
+            consolidated_context = ""
+            all_sources = []
+            for task in asyncio.as_completed(research_tasks):
+                context, sources = await task
+                if context: consolidated_context += context + "\n\n---\n\n"
+                if sources: all_sources.extend(sources)
             
-            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]
+            yield format_sse({"event": "status", "data": "Consolidating research..."})
+            if len(consolidated_context) > MAX_CONTEXT_CHAR_LENGTH:
+                consolidated_context = consolidated_context[:MAX_CONTEXT_CHAR_LENGTH]
             
-            if not full_context.strip():
+            if not consolidated_context.strip():
                 yield format_sse({"event": "error", "data": "Failed to gather any research context."})
                 return
 
-            # 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'Synthesize the provided context into a comprehensive report on "{query}". Use markdown. Context:\n{consolidated_context}'
             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:
                 if response.status != 200:
                     error_text = await response.text()
                     raise Exception(f"LLM API Error for final report: {response.status}, {error_text}")
-
                 async for line in response.content:
                     if line.strip():
                         line_str = line.decode('utf-8').strip()
@@ -171,11 +174,12 @@ async def run_deep_research_stream(query: str) -> AsyncGenerator[str, None]:
                             content = chunk.get("choices", [{}])[0].get("delta", {}).get("content")
                             if content: yield format_sse({"event": "chunk", "data": content})
                         except json.JSONDecodeError: continue
-
+            
+            unique_sources = list({s['link']: s for s in all_sources}.values())
             yield format_sse({"event": "sources", "data": unique_sources})
             
     except Exception as e:
-        logger.error(f"An error occurred during deep research: {e}")
+        logger.error(f"A critical error occurred in the main research stream: {e}")
         yield format_sse({"event": "error", "data": str(e)})
     finally:
         yield format_sse({"event": "done", "data": "Deep research complete."})