Error Handling

{"type": "LLM_RETRY", "message": "Retrying after upstream rate limit"}

{"type": "TOOL_ERROR", "tool_name": "search", "error": "Upstream search timed out"}

{"type": "ERROR", "error": "Request failed: invalid checkpoint id"}

import json
import logging
import requests

log = logging.getLogger(__name__)


class ResearchAgentError(Exception):
    """Raised for unrecoverable Research Agent failures."""


def stream_with_error_handling(api_key: str, message: str) -> str:
    endpoint = "https://agents.bigdata.com/v1/research-agent"
    headers = {"X-API-KEY": api_key, "Content-Type": "application/json"}
    payload = {"message": message, "research_effort": "standard"}

    answer = ""
    tool_errors: dict[str, int] = {}

    with requests.post(endpoint, headers=headers, json=payload, stream=True, timeout=300) as r:
        # HTTP-level errors: convert to typed exceptions before streaming.
        if r.status_code == 401:
            raise ResearchAgentError("API key is missing or invalid")
        if r.status_code == 403:
            raise ResearchAgentError(f"Not entitled: {r.text}")
        if r.status_code == 404:
            raise ResearchAgentError(f"Resource not found: {r.text}")
        if r.status_code == 429:
            # The API does not currently set a Retry-After header on 429.
            # Use exponential backoff with jitter (see stream_with_retries below).
            raise ResearchAgentError("Rate limited")
        if r.status_code >= 500:
            raise ResearchAgentError(f"Upstream failure: {r.status_code}")
        r.raise_for_status()

        for raw_line in r.iter_lines(decode_unicode=True):
            if not raw_line or not raw_line.startswith("data: "):
                continue
            try:
                event = json.loads(raw_line[6:])
            except json.JSONDecodeError:
                log.debug("Skipping malformed SSE line")
                continue

            msg = event.get("message", {})
            msg_type = msg.get("type")

            if msg_type == "ANSWER":
                answer += msg.get("content", "")

            elif msg_type == "LLM_RETRY":
                log.info("Agent retrying: %s", msg.get("message"))

            elif msg_type == "TOOL_ERROR":
                tool = msg.get("tool_name") or "unknown"
                tool_errors[tool] = tool_errors.get(tool, 0) + 1
                log.warning("Tool error from %s: %s", tool, msg.get("error"))

            elif msg_type == "ERROR":
                raise ResearchAgentError(msg.get("error", "Unknown error"))

            elif msg_type == "COMPLETE":
                if tool_errors:
                    log.warning(
                        "Completed with tool errors: %s",
                        ", ".join(f"{k}={v}" for k, v in tool_errors.items()),
                    )
                return answer

    # Reached end of stream without COMPLETE: treat as truncated.
    raise ResearchAgentError("Stream ended without COMPLETE event")

import random
import time


def stream_with_retries(api_key: str, message: str, max_attempts: int = 4) -> str:
    delay = 1.0
    last_exc: Exception | None = None
    for attempt in range(1, max_attempts + 1):
        try:
            return stream_with_error_handling(api_key, message)
        except ResearchAgentError as exc:
            text = str(exc)
            if "Rate limited" not in text and "Upstream failure" not in text:
                raise
            last_exc = exc
            if attempt == max_attempts:
                break
            sleep_for = delay + random.uniform(0, delay)
            log.info("Retry %d/%d after %.1fs", attempt, max_attempts, sleep_for)
            time.sleep(sleep_for)
            delay *= 2

    assert last_exc is not None
    raise last_exc