main.py

import os
import asyncio
import json
import logging
from typing import AsyncGenerator

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from dotenv import load_dotenv
import aiohttp
from bs4 import BeautifulSoup

# --- Configuration ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

load_dotenv()
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]}...).")

# API URLs, Models, and context size limit
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
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 *****
LLM_HEADERS = {
    "Authorization": f"Bearer {LLM_API_KEY}",
    "Content-Type": "application/json",
    "User-Agent": "AI-Deep-Research-Client/2.2"
}

# --- Pydantic Models ---
class DeepResearchRequest(BaseModel):
    query: str

# --- FastAPI App ---
app = FastAPI(
    title="AI Deep Research API",
    description="Provides streaming deep research completions.",
    version="2.2.0" # Version bump for critical bug fix
)

# --- Core Service Functions (Unchanged) ---
async def call_snapzion_search(session: aiohttp.ClientSession, query: str) -> list:
    try:
        async with session.post(SNAPZION_API_URL, headers=SNAPZION_HEADERS, json={"query": query}, timeout=15) as response:
            response.raise_for_status()
            data = await response.json()
            return data.get("organic_results", [])
    except Exception as e:
        logger.error(f"Snapzion search failed for query '{query}': {e}")
        return []

async def scrape_url(session: aiohttp.ClientSession, url: str) -> str:
    if url.lower().endswith('.pdf'): return "Error: PDF content cannot be scraped."
    try:
        async with session.get(url, headers=SCRAPING_HEADERS, timeout=10, ssl=False) as response:
            if response.status != 200: return f"Error: HTTP status {response.status}"
            html = await response.text()
            soup = BeautifulSoup(html, "html.parser")
            for tag in soup(['script', 'style', 'nav', 'footer', 'header', 'aside']):
                tag.decompose()
            return " ".join(soup.stripped_strings)
    except Exception as e:
        logger.warning(f"Scraping failed for {url}: {e}")
        return f"Error: {e}"

async def search_and_scrape(session: aiohttp.ClientSession, query: str) -> tuple[str, list]:
    search_results = await call_snapzion_search(session, 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 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

# --- Streaming Deep Research Logic ---

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"

    try:
        async with aiohttp.ClientSession() as session:
            # Step 1: Generate Sub-Questions
            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. 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 *****
            try:
                async with session.post(LLM_API_URL, headers=LLM_HEADERS, json=sub_question_prompt, timeout=20) as response:
                    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}")
                    
                    raw_response_text = await response.text()
                    if not raw_response_text:
                        raise Exception("LLM API returned an empty response.")
                        
                    result = json.loads(raw_response_text)
                    llm_content = result['choices'][0]['message']['content']
                    sub_questions = json.loads(llm_content)
            except Exception as e:
                logger.error(f"Failed to generate or parse research plan: {e}")
                yield format_sse({"event": "error", "data": f"Could not generate research plan. Reason: {e}"})
                return # Stop the process if planning fails

            yield format_sse({"event": "plan", "data": sub_questions})

            # (The rest of the logic remains the same)
            # Step 2: Concurrently research all sub-questions
            research_tasks = [search_and_scrape(session, sq) for sq in sub_questions]
            all_research_results = []
            
            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())
            
            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."})
                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_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()
                        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

            yield format_sse({"event": "sources", "data": unique_sources})
            
    except Exception as e:
        logger.error(f"An error occurred during deep research: {e}")
        yield format_sse({"event": "error", "data": str(e)})
    finally:
        yield format_sse({"event": "done", "data": "Deep research complete."})

# --- API Endpoints ---
@app.post("/v1/deepresearch/completions")
async def deep_research_endpoint(request: DeepResearchRequest):
    return StreamingResponse(run_deep_research_stream(request.query), media_type="text/event-stream")