Module 7: Advanced Flows

Evolving Project: Retry Logic and Branching (v2)

Project overview

In Module 6, you built v1 of the AI Research Assistant: an agent that breaks a topic down into sub-queries, searches 3 sources, synthesizes findings, and generates a structured report with Pydantic. It works end to end. But it has a fundamental problem: it's fragile.

v1 searches 3 sources for each sub-query using futures. That looks parallel — and technically it is, inside search_all_sources. But the per-sub-query searches are fired off as independent futures with no protection: if the API fails, the pipeline crashes. If one source takes 30 seconds, everything waits. If two sources return the same finding, it gets duplicated in the report. There are no retries, no graceful degradation, and no visibility into what happened when something breaks.

The v2 you build in this module solves all of that. It isn't a rewrite — it's an evolution. Your v1 code keeps working. This module adds layers of robustness on top: retry with exponential backoff so transient errors recover on their own, protected parallel execution so the 3 sources are searched simultaneously but each one handles its own errors, merge with deduplication to consolidate results intelligently, and structured logging so you know exactly what happened at each step.

The difference is tangible. If you run v1 against an API that fails intermittently, your agent crashes. If you run v2 against the same API, your agent retries, recovers, and if the API is still down after 3 attempts, it generates the report with the sources that did respond. v1 is a demo. v2 is a system you can put in production.

Project goal

Evolve the AI Research Assistant from v1 (functional but fragile) to v2 (robust and production-ready), adding retry logic, protected parallel search, merge with deduplication, graceful degradation, and structured logging.

By the time you finish this project:

🔧 You'll implement retry with exponential backoff and jitter for every search
🔧 You'll protect each search source with its own error handling
🔧 You'll run searches in parallel with futures and collect results with fault tolerance
🔧 You'll implement smart merging with deduplication of findings
🔧 You'll add graceful degradation: a minimum of 1 source for a valid report
🔧 You'll wire in structured logging for traceability of every operation

Before and after

v1 (Module 6): functional but fragile

Topic → Decompose → [search_all_sources(q1), search_all_sources(q2), ...] → Synthesize → Report

Problems:
- If search_web fails → Exception, the pipeline crashes
- If search_academic takes 30s → everything waits 30s
- If web and news return the same finding → duplicated in the report
- If something fails → "Error during the research" (no details)

v2 (this module): robust and production-ready

Topic → Decompose → [search_with_retry(web), search_with_retry(academic), search_with_retry(news)] (parallel)
                      ↓ each one with 3x retry, backoff 1s→2s→4s
                      ↓ if it fails after 3 retries → graceful skip
                    → Merge with deduplication
                    → Synthesize (with the available sources)
                    → Report (includes the status of each source)

Improvements:
- Transient error → auto-recovery (retry with backoff)
- Slow API → timeout + retry, doesn't block the others
- Duplicate findings → deduplication by similarity
- Permanently down source → carry on with the rest
- Any failure → structured log with request_id, node, duration

Technical specifications

Tech stack

Component	Version	Purpose
Python	3.11+	Runtime
LangChain	v1.2+	LLM framework
LangGraph	v1.0+	Functional API (@entrypoint, @task)
langchain-openai	latest	Model provider
pydantic	v2+	Structured data models
python-dotenv	latest	Environment variables

Project structure (evolved from v1)

research-assistant/
├── .env
├── requirements.txt
├── agents/
│   └── researcher.py           # MODIFIED — v2 @entrypoint with retry + parallel
├── tools/
│   ├── web_search.py           # MODIFIED — search with simulated failures
│   └── calculator.py           # UNCHANGED
├── state/
│   └── research_state.py       # EXTENDED — new fields for source status
├── config/
│   └── settings.py             # EXTENDED — retry and timeout config
├── utils/                      # NEW
│   ├── retry.py                # Retry with exponential backoff
│   └── logger.py               # Structured logging
└── main.py                     # MODIFIED — shows source availability

The new files are utils/retry.py and utils/logger.py. The rest are extensions of existing files. Your v1 code doesn't get deleted — it gets better.

Step 1: Retry with exponential backoff (utils/retry.py)

The first building block. When an API fails with a transient error (timeout, 429, 503), you don't want to retry immediately — that makes the problem worse. You want to wait a growing amount of time: 1s, 2s, 4s. And you add jitter (random variation) to keep 100 clients from retrying at the same instant (thundering herd).

"""
utils/retry.py
Retry with exponential backoff and jitter for the AI Research Assistant v2.
"""

import time
import random


class RetryConfig:
    """Retry configuration with exponential backoff."""

    def __init__(
        self,
        max_retries: int = 3,
        base_delay: float = 1.0,
        max_delay: float = 10.0,
        jitter: bool = True,
    ):
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.jitter = jitter

    def get_delay(self, attempt: int) -> float:
        """Computes the delay for a specific attempt."""
        delay = min(self.base_delay * (2 ** attempt), self.max_delay)
        if self.jitter:
            delay = delay * (0.5 + random.random() * 0.5)
        return delay


def retry_with_backoff(
    func,
    args: tuple = (),
    kwargs: dict = None,
    config: RetryConfig = None,
    on_retry=None,
    retriable_exceptions: tuple = (Exception,),
) -> dict:
    """
    Runs func with retry and exponential backoff.

    Returns a dict with the result or the final error:
    - {"status": "success", "result": ..., "attempts": N}
    - {"status": "failed", "error": "...", "attempts": N}
    """
    if kwargs is None:
        kwargs = {}
    if config is None:
        config = RetryConfig()

    last_error = None

    for attempt in range(config.max_retries + 1):
        try:
            result = func(*args, **kwargs)
            return {
                "status": "success",
                "result": result,
                "attempts": attempt + 1,
            }
        except retriable_exceptions as e:
            last_error = e

            if attempt < config.max_retries:
                delay = config.get_delay(attempt)
                if on_retry:
                    on_retry(attempt + 1, config.max_retries, delay, str(e))
                time.sleep(delay)
            else:
                return {
                    "status": "failed",
                    "error": str(e),
                    "attempts": attempt + 1,
                }

    return {
        "status": "failed",
        "error": str(last_error),
        "attempts": config.max_retries + 1,
    }

The backoff pattern: attempt 0 = no wait, attempt 1 = ~1s, attempt 2 = ~2s, attempt 3 = ~4s. The jitter makes the actual delay vary between 50% and 100% of the computed value. That way, 100 clients that fail at the same moment don't retry at the same moment.

Let's check that it works:

from utils.retry import retry_with_backoff, RetryConfig

call_count = 0

def flaky_function():
    """Fails the first 2 times, works on the third."""
    global call_count
    call_count += 1
    if call_count <= 2:
        raise ConnectionError(f"Attempt {call_count}: API unavailable")
    return "Success on attempt 3!"


config = RetryConfig(max_retries=3, base_delay=0.5, jitter=False)
result = retry_with_backoff(
    flaky_function,
    config=config,
    on_retry=lambda attempt, max_r, delay, err: print(
        f"  Retry {attempt}/{max_r} in {delay:.1f}s — {err}"
    ),
)
print(f"Status: {result['status']}")
print(f"Result: {result.get('result', result.get('error'))}")
print(f"Attempts: {result['attempts']}")
# Expected output:
#   Retry 1/3 in 0.5s — Attempt 1: API unavailable
#   Retry 2/3 in 1.0s — Attempt 2: API unavailable
# Status: success
# Result: Success on attempt 3!
# Attempts: 3

Step 2: Structured logging (utils/logger.py)

Every operation the agent runs needs to be traceable. The logger emits JSON with a request_id, the node name, duration, status, and any relevant metadata.

"""
utils/logger.py
Structured logging for the AI Research Assistant v2.
"""

import json
import time
import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
)


class ResearchLogger:
    """Structured logger with a correlation ID for the Research Assistant."""

    def __init__(self, name: str = "research_agent"):
        self._logger = logging.getLogger(name)
        self._request_id = None

    def set_request_id(self, request_id: str):
        self._request_id = request_id

    def _emit(self, level: str, node: str, event: str, **kwargs):
        entry = {
            "ts": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "level": level,
            "request_id": self._request_id,
            "node": node,
            "event": event,
        }
        entry.update(kwargs)
        self._logger.info(json.dumps(entry, ensure_ascii=False))

    def start(self, node: str, **kwargs):
        self._emit("INFO", node, "start", **kwargs)

    def end(self, node: str, duration_ms: float, **kwargs):
        self._emit("INFO", node, "end", duration_ms=round(duration_ms, 1), **kwargs)

    def retry(self, node: str, attempt: int, max_retries: int, delay: float, error: str, **kwargs):
        self._emit("WARN", node, "retry", attempt=attempt, max_retries=max_retries, delay_s=round(delay, 2), error=error, **kwargs)

    def error(self, node: str, error: str, duration_ms: float, **kwargs):
        self._emit("ERROR", node, "error", error=error, duration_ms=round(duration_ms, 1), **kwargs)

    def skip(self, node: str, reason: str, **kwargs):
        self._emit("WARN", node, "skip", reason=reason, **kwargs)

Step 3: Extend the data models (state/research_state.py)

The v2 report needs information about which sources responded and which failed. Extend the Pydantic models — without breaking v1:

"""
state/research_state.py
Data models for the AI Research Assistant v2.
Extends v1 with source availability fields.
"""

from pydantic import BaseModel, Field
from datetime import datetime


class Source(BaseModel):
    """An information source found during the research."""
    name: str = Field(description="Name of the source")
    source_type: str = Field(description="Type: web, academic, news")
    content: str = Field(description="Content extracted from the source")


class KeyFinding(BaseModel):
    """A key finding identified in the research."""
    title: str = Field(description="Title of the finding")
    description: str = Field(description="Description of the finding")
    confidence: float = Field(
        description="Confidence in the finding (0.0 to 1.0)",
        ge=0.0,
        le=1.0,
    )


class SourceStatus(BaseModel):
    """State of a search source (v2)."""
    source_type: str = Field(description="Type of source")
    status: str = Field(description="ok, failed, timeout, skipped")
    attempts: int = Field(default=1, description="Number of attempts made")
    error: str = Field(default="", description="Error message if it failed")
    duration_ms: float = Field(default=0, description="Total duration including retries")


class ResearchReport(BaseModel):
    """Structured research report (v2)."""
    topic: str = Field(description="Topic researched")
    summary: str = Field(description="Executive summary (2-3 sentences)")
    key_findings: list[KeyFinding] = Field(
        description="Main findings",
        min_length=1,
    )
    sources: list[Source] = Field(
        description="Sources consulted successfully",
        min_length=1,
    )
    sub_queries: list[str] = Field(
        description="Sub-questions generated for the research",
    )
    confidence: float = Field(
        description="Overall confidence of the report (0.0 to 1.0)",
        ge=0.0,
        le=1.0,
    )
    generated_at: str = Field(
        default_factory=lambda: datetime.now().isoformat(),
        description="Generation timestamp",
    )
    source_availability: list[SourceStatus] = Field(
        default_factory=list,
        description="(v2) State of each search source",
    )
    version: str = Field(default="v2", description="Agent version")


class SubQuery(BaseModel):
    """A sub-question generated from the main topic."""
    query: str = Field(description="The sub-question")
    rationale: str = Field(description="Why this question is relevant")

The changes:

✅ SourceStatus — a new model to track what happened with each source
✅ source_availability in ResearchReport — a list of SourceStatus with an empty default (compatible with v1)
✅ version — distinguishes v1 reports from v2

Every change is additive. A v1 report is still a valid ResearchReport — the new fields have defaults.

Step 4: Extend the configuration (config/settings.py)

Add retry and timeout configuration:

"""
config/settings.py
Configuration for the AI Research Assistant v2.
"""

from dotenv import load_dotenv
load_dotenv()

MODEL_NAME = "openai:gpt-4.1-mini"
MODEL_TEMPERATURE = 0.2

MAX_SUB_QUERIES = 4
MAX_SOURCES_PER_QUERY = 3

SEARCH_SOURCES = ["web", "academic", "news"]

# v2: Retry configuration
RETRY_MAX_ATTEMPTS = 3
RETRY_BASE_DELAY = 1.0
RETRY_MAX_DELAY = 10.0
RETRY_JITTER = True

# v2: Minimum sources for a valid report
MIN_SOURCES_FOR_REPORT = 1

# v2: Search timeout
SEARCH_TIMEOUT_SECONDS = 5.0

Step 5: Evolve the search tools (tools/web_search.py)

The v1 tools never failed — they were deterministic mocks. v2 simulates real errors so the retry logic has something to work against. But the mock stays controllable for testing.

"""
tools/web_search.py
Mock web search tool for the AI Research Assistant v2.
Simulates transient errors and variable latency for testing retry logic.
"""

import time
import random
import hashlib
from langgraph.func import task

from utils.retry import retry_with_backoff, RetryConfig
from utils.logger import ResearchLogger


MOCK_RESULTS = {
    "web": {
        "default": (
            "Multiple web sources agree that {query} is a topic of growing "
            "interest. Experts highlight significant advances over the last 2 years. "
            "Practical applications include automation, data analysis "
            "and content generation."
        ),
    },
    "academic": {
        "default": (
            "Recent research (2025-2026) shows that {query} has solid theoretical "
            "foundations backed by multiple peer-reviewed studies. "
            "Experimental results demonstrate 40-60% improvements on key "
            "metrics compared to traditional methods."
        ),
    },
    "news": {
        "default": (
            "Recent news reports that {query} is having an impact on the "
            "industry. Leading companies like Google, Microsoft and emerging startups "
            "are investing significantly in this area. Important developments "
            "are expected by the end of 2026."
        ),
    },
}

FAILURE_CONFIG = {
    "web": {"fail_rate": 0.0, "latency_range": (0.1, 0.3)},
    "academic": {"fail_rate": 0.0, "latency_range": (0.2, 0.5)},
    "news": {"fail_rate": 0.0, "latency_range": (0.1, 0.4)},
}


def configure_failures(source: str, fail_rate: float = 0.0, latency_range: tuple = None):
    """Configures the failure rate of a source for testing."""
    if latency_range is None:
        latency_range = (0.1, 0.3)
    FAILURE_CONFIG[source] = {
        "fail_rate": fail_rate,
        "latency_range": latency_range,
    }


def _generate_deterministic_score(query: str, source_type: str) -> float:
    hash_input = f"{query}:{source_type}"
    hash_value = int(hashlib.md5(hash_input.encode()).hexdigest()[:8], 16)
    return round(0.5 + (hash_value % 50) / 100, 2)


def _raw_search(query: str, source_type: str) -> dict:
    """Raw search that can fail according to FAILURE_CONFIG."""
    config = FAILURE_CONFIG.get(source_type, {"fail_rate": 0.0, "latency_range": (0.1, 0.3)})

    latency = random.uniform(*config["latency_range"])
    time.sleep(latency)

    if random.random() < config["fail_rate"]:
        error_types = [
            ConnectionError(f"{source_type} API: connection refused"),
            TimeoutError(f"{source_type} API: request timed out"),
            RuntimeError(f"{source_type} API: 503 Service Unavailable"),
        ]
        raise random.choice(error_types)

    content = MOCK_RESULTS[source_type]["default"].format(query=query)
    return {
        "source_name": f"{source_type.title()} Search",
        "source_type": source_type,
        "content": content,
        "relevance": _generate_deterministic_score(query, source_type),
    }


@task
def search_with_retry(query: str, source_type: str, logger: ResearchLogger = None) -> dict:
    """
    Search with retry and exponential backoff (v2).
    Returns a successful result or an error with attempt metadata.
    """
    start = time.time()
    retry_config = RetryConfig(
        max_retries=3,
        base_delay=1.0,
        max_delay=10.0,
        jitter=True,
    )

    def on_retry(attempt, max_retries, delay, error):
        if logger:
            logger.retry(
                f"search_{source_type}",
                attempt=attempt,
                max_retries=max_retries,
                delay=delay,
                error=error,
                query=query[:50],
            )

    outcome = retry_with_backoff(
        func=_raw_search,
        args=(query, source_type),
        config=retry_config,
        on_retry=on_retry,
        retriable_exceptions=(ConnectionError, TimeoutError, RuntimeError),
    )

    duration_ms = (time.time() - start) * 1000

    if outcome["status"] == "success":
        if logger:
            logger.end(
                f"search_{source_type}",
                duration_ms=duration_ms,
                attempts=outcome["attempts"],
                status="success",
                query=query[:50],
            )
        return {
            **outcome["result"],
            "search_status": "ok",
            "attempts": outcome["attempts"],
            "duration_ms": round(duration_ms, 1),
        }
    else:
        if logger:
            logger.error(
                f"search_{source_type}",
                error=outcome["error"],
                duration_ms=duration_ms,
                attempts=outcome["attempts"],
                query=query[:50],
            )
        return {
            "source_name": f"{source_type.title()} Search",
            "source_type": source_type,
            "content": "",
            "relevance": 0.0,
            "search_status": "failed",
            "error": outcome["error"],
            "attempts": outcome["attempts"],
            "duration_ms": round(duration_ms, 1),
        }


SEARCH_FUNCTIONS = {
    "web": search_with_retry,
    "academic": search_with_retry,
    "news": search_with_retry,
}

Changes from v1:

✅ _raw_search simulates configurable transient errors with FAILURE_CONFIG
✅ search_with_retry wraps _raw_search with retry + backoff + logging
✅ It never raises exceptions — it returns {"search_status": "ok"} or {"search_status": "failed"}
✅ configure_failures() lets you turn failures on and off for testing

Step 6: Evolve the main agent (agents/researcher.py)

This is the central change. The @entrypoint now runs searches with retry, merges with deduplication, and degrades gracefully.

"""
agents/researcher.py
Main agent of the AI Research Assistant (v2).
Evolution of v1: retry, protected parallelism, merge, graceful degradation.
"""

import json
import time
import uuid
from langchain.chat_models import init_chat_model
from langgraph.func import entrypoint, task
from langgraph.checkpoint.memory import MemorySaver

import sys
sys.path.insert(0, ".")

from config.settings import (
    MODEL_NAME,
    MODEL_TEMPERATURE,
    MAX_SUB_QUERIES,
    SEARCH_SOURCES,
    MIN_SOURCES_FOR_REPORT,
)
from state.research_state import (
    ResearchReport,
    Source,
    KeyFinding,
    SubQuery,
    SourceStatus,
)
from tools.web_search import search_with_retry
from tools.calculator import calculate_confidence
from utils.logger import ResearchLogger


model = init_chat_model(MODEL_NAME, temperature=MODEL_TEMPERATURE)
agent_logger = ResearchLogger("research_agent_v2")


# =============================================================================
# TASK: Decompose the topic into sub-queries (unchanged from v1)
# =============================================================================

@task
def decompose_query(topic: str) -> list[dict]:
    """Breaks a research topic down into specific sub-questions."""
    response = model.invoke(
        f"You are an expert researcher. Break this topic down into "
        f"{MAX_SUB_QUERIES} specific, researchable sub-questions.\n\n"
        f"Topic: {topic}\n\n"
        f"Answer in JSON (no markdown, no ```json):\n"
        f'[{{"query": "sub-question", "rationale": "why it is relevant"}}]\n\n'
        f"Only the JSON, nothing else."
    )

    try:
        queries = json.loads(response.content)
        return queries[:MAX_SUB_QUERIES]
    except json.JSONDecodeError:
        return [
            {"query": topic, "rationale": "Original query as a fallback"},
            {"query": f"recent advances in {topic}", "rationale": "Current trends"},
            {"query": f"practical applications of {topic}", "rationale": "Real-world use"},
        ]


# =============================================================================
# TASK: Search every source with retry (v2)
# =============================================================================

@task
def search_all_sources_v2(query: str, logger: ResearchLogger) -> list[dict]:
    """
    Searches every source with per-source retry (v2).
    Fires all the searches in parallel with futures.
    Each search has its own retry — a failure in 'academic' doesn't affect 'web'.
    """
    logger.start("search_all", query=query[:50], sources=SEARCH_SOURCES)
    start = time.time()

    futures = []
    for source_type in SEARCH_SOURCES:
        logger.start(f"search_{source_type}", query=query[:50])
        futures.append(search_with_retry(query, source_type, logger))

    results = [f.result() for f in futures]

    duration_ms = (time.time() - start) * 1000
    ok_count = sum(1 for r in results if r["search_status"] == "ok")
    logger.end("search_all", duration_ms=duration_ms, sources_ok=ok_count, sources_total=len(results))

    return results


# =============================================================================
# TASK: Merge with deduplication (v2 — NEW)
# =============================================================================

@task
def merge_and_deduplicate(all_results: list[dict]) -> list[dict]:
    """
    Combines results from multiple searches and removes duplicates.
    Deduplication by: same source_type + very similar content.
    """
    seen_keys = set()
    unique_results = []

    for result in all_results:
        if result["search_status"] != "ok":
            continue

        dedup_key = f"{result['source_type']}:{result['content'][:100]}"

        if dedup_key not in seen_keys:
            seen_keys.add(dedup_key)
            unique_results.append(result)

    unique_results.sort(key=lambda r: r.get("relevance", 0), reverse=True)

    return unique_results


# =============================================================================
# TASK: Synthesize findings (unchanged from v1)
# =============================================================================

@task
def synthesize_findings(topic: str, all_results: list[dict]) -> list[dict]:
    """Synthesizes search results into key findings."""
    results_text = ""
    for i, result in enumerate(all_results, 1):
        results_text += (
            f"\nSource {i} ({result['source_type']}): {result['content']}\n"
        )

    response = model.invoke(
        f"You are a research analyst. Based on these sources, "
        f"identify 3-5 key findings about '{topic}'.\n\n"
        f"Sources:\n{results_text}\n\n"
        f"Answer in JSON (no markdown, no ```json):\n"
        f'[{{"title": "short title", "description": "1-2 sentence description", '
        f'"confidence": 0.8}}]\n\n'
        f"confidence goes from 0.0 to 1.0. Only the JSON, nothing else."
    )

    try:
        findings = json.loads(response.content)
        return findings[:5]
    except json.JSONDecodeError:
        return [{
            "title": "General finding",
            "description": f"Research on {topic} shows relevant results across multiple sources.",
            "confidence": 0.6,
        }]


# =============================================================================
# TASK: Generate the executive summary (unchanged from v1)
# =============================================================================

@task
def generate_summary(topic: str, findings: list[dict], source_info: str) -> str:
    """Generates an executive summary that includes info about source availability."""
    findings_text = "\n".join(
        f"- {f['title']}: {f['description']}" for f in findings
    )

    response = model.invoke(
        f"Generate a 2-3 sentence executive summary of the research "
        f"on the topic '{topic}'.\n\n"
        f"Main findings:\n{findings_text}\n\n"
        f"Note on sources: {source_info}\n\n"
        f"Only the summary, no title and no extra formatting."
    )
    return response.content.strip()


# =============================================================================
# ENTRYPOINT: Research agent v2
# =============================================================================

memory = MemorySaver()


@entrypoint(checkpointer=memory)
def research_agent(topic: str) -> dict:
    """
    AI Research Assistant v2.
    Flow: topic → decompose → search (parallel + retry) → merge → synthesize → report.
    """
    request_id = uuid.uuid4().hex[:8]
    agent_logger.set_request_id(request_id)
    agent_logger.start("pipeline", topic=topic, version="v2")

    pipeline_start = time.time()

    print(f"\n{'=' * 60}")
    print(f"  🔬 AI Research Assistant v2")
    print(f"  Topic: {topic}")
    print(f"  Request ID: {request_id}")
    print(f"{'=' * 60}")

    # --- Step 1: Decompose the topic ---
    print(f"\n📋 Step 1: Decomposing the topic into sub-queries...")
    agent_logger.start("decompose", topic=topic[:50])
    decompose_start = time.time()

    sub_queries_raw = decompose_query(topic).result()
    sub_queries = [SubQuery(**sq) for sq in sub_queries_raw]

    agent_logger.end("decompose", duration_ms=(time.time() - decompose_start) * 1000, count=len(sub_queries))
    print(f"   ✓ {len(sub_queries)} sub-queries generated:")
    for i, sq in enumerate(sub_queries, 1):
        print(f"     {i}. {sq.query}")

    # --- Step 2: Search in parallel with retry ---
    print(f"\n🔍 Step 2: Searching {len(SEARCH_SOURCES)} sources per sub-query (with retry)...")

    all_raw_results = []
    source_statuses = []

    search_futures = [
        search_all_sources_v2(sq.query, agent_logger) for sq in sub_queries
    ]

    for i, future in enumerate(search_futures):
        results = future.result()
        all_raw_results.extend(results)

        ok = sum(1 for r in results if r["search_status"] == "ok")
        failed = sum(1 for r in results if r["search_status"] == "failed")
        print(f"   Sub-query {i + 1}: {ok} OK, {failed} failed")

        for r in results:
            source_statuses.append(SourceStatus(
                source_type=r["source_type"],
                status=r["search_status"],
                attempts=r.get("attempts", 1),
                error=r.get("error", ""),
                duration_ms=r.get("duration_ms", 0),
            ))

    total_ok = sum(1 for r in all_raw_results if r["search_status"] == "ok")
    total_failed = sum(1 for r in all_raw_results if r["search_status"] == "failed")
    print(f"   Total: {total_ok} successful, {total_failed} failed out of {len(all_raw_results)}")

    # --- Check the minimum number of sources ---
    if total_ok < MIN_SOURCES_FOR_REPORT:
        pipeline_ms = (time.time() - pipeline_start) * 1000
        agent_logger.error("pipeline", error="Insufficient sources", duration_ms=pipeline_ms, sources_ok=total_ok)
        print(f"\n   ❌ Insufficient sources ({total_ok} < {MIN_SOURCES_FOR_REPORT}). Aborting.")

        return {
            "error": f"Only {total_ok} sources responded. Minimum required: {MIN_SOURCES_FOR_REPORT}",
            "request_id": request_id,
            "source_availability": [s.model_dump() for s in source_statuses],
            "version": "v2",
        }

    # --- Step 3: Merge with deduplication ---
    print(f"\n🔀 Step 3: Merging and deduplicating results...")
    unique_results = merge_and_deduplicate(all_raw_results).result()
    print(f"   ✓ {len(all_raw_results)} raw results → {len(unique_results)} unique after deduplication")

    # --- Step 4: Synthesize findings ---
    print(f"\n🧠 Step 4: Synthesizing findings...")
    findings_raw = synthesize_findings(topic, unique_results).result()
    findings = [KeyFinding(**f) for f in findings_raw]
    print(f"   ✓ {len(findings)} findings identified:")
    for i, f in enumerate(findings, 1):
        print(f"     {i}. [{f.confidence:.0%}] {f.title}")

    # --- Step 5: Generate the summary ---
    print(f"\n📝 Step 5: Generating the executive summary...")
    source_info = f"{total_ok} of {total_ok + total_failed} sources responded successfully"
    if total_failed > 0:
        failed_sources = set(r["source_type"] for r in all_raw_results if r["search_status"] == "failed")
        source_info += f". Unavailable sources: {', '.join(failed_sources)}"

    summary = generate_summary(topic, findings_raw, source_info).result()
    print(f"   ✓ Summary generated ({len(summary)} chars)")

    # --- Step 6: Compute confidence (adjusted for availability) ---
    print(f"\n📊 Step 6: Computing the report's confidence...")
    avg_relevance = (
        sum(r["relevance"] for r in unique_results) / len(unique_results)
        if unique_results else 0.5
    )

    base_confidence = calculate_confidence(
        num_sources=len(unique_results),
        avg_relevance=avg_relevance,
        num_findings=len(findings),
    ).result()

    availability_factor = total_ok / (total_ok + total_failed) if (total_ok + total_failed) > 0 else 0.5
    adjusted_confidence = round(base_confidence * (0.7 + 0.3 * availability_factor), 2)
    print(f"   ✓ Base confidence: {base_confidence:.0%}, adjusted for availability: {adjusted_confidence:.0%}")

    # --- Step 7: Build the structured report ---
    print(f"\n📄 Step 7: Building the v2 report...")
    sources = [
        Source(
            name=r["source_name"],
            source_type=r["source_type"],
            content=r["content"],
        )
        for r in unique_results
    ]

    report = ResearchReport(
        topic=topic,
        summary=summary,
        key_findings=findings,
        sources=sources,
        sub_queries=[sq.query for sq in sub_queries],
        confidence=adjusted_confidence,
        source_availability=source_statuses,
        version="v2",
    )

    pipeline_ms = (time.time() - pipeline_start) * 1000
    agent_logger.end("pipeline", duration_ms=pipeline_ms, sources_ok=total_ok, sources_failed=total_failed, confidence=adjusted_confidence)

    print(f"   ✓ v2 report generated successfully")
    print(f"   📊 Pipeline total: {pipeline_ms:.0f}ms")
    print(f"\n{'=' * 60}")

    return report.model_dump()

The key differences from v1:

✅ Request ID — every run has a unique ID for traceability
✅ search_all_sources_v2 — each source has its own retry, errors don't propagate
✅ merge_and_deduplicate — a new step that removes duplicates and ranks by relevance
✅ Graceful degradation — checks MIN_SOURCES_FOR_REPORT before synthesizing
✅ Adjusted confidence — confidence drops when sources are unavailable
✅ source_availability — the report includes the status of each source
✅ Structured logging — every step is logged with its duration and metadata

Step 7: Update the CLI (main.py)

The CLI now shows source availability information and the request ID:

"""
main.py
CLI entrypoint for the AI Research Assistant v2.
"""

import sys
import json
import uuid

sys.path.insert(0, ".")

from agents.researcher import research_agent


def format_report(report: dict) -> str:
    """Formats the v2 report for terminal display."""
    if "error" in report:
        lines = [
            "",
            "╔" + "═" * 58 + "╗",
            "║" + "  ❌ RESEARCH ERROR".center(58) + "║",
            "╚" + "═" * 58 + "╝",
            f"\n  {report['error']}",
            f"  Request ID: {report.get('request_id', 'N/A')}",
        ]
        if "source_availability" in report:
            lines.append(f"\n  Source availability:")
            for sa in report["source_availability"]:
                status_icon = "✅" if sa["status"] == "ok" else "❌"
                lines.append(f"    {status_icon} {sa['source_type']}: {sa['status']} ({sa['attempts']} attempts)")
        return "\n".join(lines)

    lines = []
    lines.append("")
    lines.append("╔" + "═" * 58 + "╗")
    lines.append("║" + f"  📄 RESEARCH REPORT ({report.get('version', 'v1')})".center(58) + "║")
    lines.append("╚" + "═" * 58 + "╝")

    lines.append(f"\n📌 Topic: {report['topic']}")
    lines.append(f"📅 Generated: {report['generated_at']}")
    lines.append(f"🎯 Confidence: {report['confidence']:.0%}")

    lines.append(f"\n{'─' * 60}")
    lines.append("📋 EXECUTIVE SUMMARY")
    lines.append(f"{'─' * 60}")
    lines.append(report["summary"])

    lines.append(f"\n{'─' * 60}")
    lines.append("🔍 SUB-QUERIES RESEARCHED")
    lines.append(f"{'─' * 60}")
    for i, sq in enumerate(report["sub_queries"], 1):
        lines.append(f"  {i}. {sq}")

    lines.append(f"\n{'─' * 60}")
    lines.append("💡 MAIN FINDINGS")
    lines.append(f"{'─' * 60}")
    for i, finding in enumerate(report["key_findings"], 1):
        conf = finding["confidence"]
        lines.append(f"\n  {i}. {finding['title']} [{conf:.0%} confidence]")
        lines.append(f"     {finding['description']}")

    lines.append(f"\n{'─' * 60}")
    lines.append(f"📚 SOURCES CONSULTED ({len(report['sources'])})")
    lines.append(f"{'─' * 60}")
    seen = set()
    for source in report["sources"]:
        key = f"{source['name']}:{source['source_type']}"
        if key not in seen:
            seen.add(key)
            lines.append(f"  • [{source['source_type'].upper()}] {source['name']}")

    if report.get("source_availability"):
        lines.append(f"\n{'─' * 60}")
        lines.append("🔌 SOURCE AVAILABILITY")
        lines.append(f"{'─' * 60}")
        for sa in report["source_availability"]:
            icon = "✅" if sa["status"] == "ok" else "❌"
            retry_info = f" ({sa['attempts']} attempts, {sa['duration_ms']:.0f}ms)" if sa["attempts"] > 1 else f" ({sa['duration_ms']:.0f}ms)"
            error_info = f" — {sa['error']}" if sa.get("error") else ""
            lines.append(f"  {icon} {sa['source_type']}: {sa['status']}{retry_info}{error_info}")

    lines.append(f"\n{'═' * 60}")

    return "\n".join(lines)


def run_interactive():
    """Interactive mode: the user types topics."""
    print("=" * 60)
    print("  🔬 AI Research Assistant v2")
    print("  Type a topic to research.")
    print("  Commands: 'exit' to quit")
    print("=" * 60)

    while True:
        try:
            topic = input("\n🔎 Topic: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\n\nSee you soon!")
            break

        if not topic:
            continue

        if topic.lower() in ("exit", "quit", "q"):
            print("\nSee you soon!")
            break

        thread_id = f"research-v2-{uuid.uuid4().hex[:8]}"

        try:
            report = research_agent.invoke(
                topic,
                config={"configurable": {"thread_id": thread_id}},
            )
            print(format_report(report))

        except Exception as e:
            print(f"\n❌ Unexpected error: {e}")
            print("   Try another topic.")


def run_single(topic: str):
    """Runs a single research pass."""
    thread_id = f"research-v2-{uuid.uuid4().hex[:8]}"

    report = research_agent.invoke(
        topic,
        config={"configurable": {"thread_id": thread_id}},
    )
    print(format_report(report))

    print("\n📦 Report JSON:")
    print(json.dumps(report, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    if len(sys.argv) > 1:
        run_single(" ".join(sys.argv[1:]))
    else:
        run_interactive()

Complete updated code

For quick reference, here are all the v2 project files consolidated.

utils/retry.py

"""
utils/retry.py
Retry with exponential backoff and jitter for the AI Research Assistant v2.
"""

import time
import random


class RetryConfig:
    def __init__(
        self,
        max_retries: int = 3,
        base_delay: float = 1.0,
        max_delay: float = 10.0,
        jitter: bool = True,
    ):
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.jitter = jitter

    def get_delay(self, attempt: int) -> float:
        delay = min(self.base_delay * (2 ** attempt), self.max_delay)
        if self.jitter:
            delay = delay * (0.5 + random.random() * 0.5)
        return delay


def retry_with_backoff(
    func,
    args: tuple = (),
    kwargs: dict = None,
    config: RetryConfig = None,
    on_retry=None,
    retriable_exceptions: tuple = (Exception,),
) -> dict:
    if kwargs is None:
        kwargs = {}
    if config is None:
        config = RetryConfig()

    last_error = None

    for attempt in range(config.max_retries + 1):
        try:
            result = func(*args, **kwargs)
            return {"status": "success", "result": result, "attempts": attempt + 1}
        except retriable_exceptions as e:
            last_error = e
            if attempt < config.max_retries:
                delay = config.get_delay(attempt)
                if on_retry:
                    on_retry(attempt + 1, config.max_retries, delay, str(e))
                time.sleep(delay)

    return {"status": "failed", "error": str(last_error), "attempts": config.max_retries + 1}

utils/logger.py

"""
utils/logger.py
Structured logging for the AI Research Assistant v2.
"""

import json
import time
import logging

logging.basicConfig(level=logging.INFO, format="%(message)s")


class ResearchLogger:
    def __init__(self, name: str = "research_agent"):
        self._logger = logging.getLogger(name)
        self._request_id = None

    def set_request_id(self, request_id: str):
        self._request_id = request_id

    def _emit(self, level: str, node: str, event: str, **kwargs):
        entry = {
            "ts": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "level": level,
            "request_id": self._request_id,
            "node": node,
            "event": event,
        }
        entry.update(kwargs)
        self._logger.info(json.dumps(entry, ensure_ascii=False))

    def start(self, node, **kw): self._emit("INFO", node, "start", **kw)
    def end(self, node, duration_ms, **kw): self._emit("INFO", node, "end", duration_ms=round(duration_ms, 1), **kw)
    def retry(self, node, attempt, max_retries, delay, error, **kw): self._emit("WARN", node, "retry", attempt=attempt, max_retries=max_retries, delay_s=round(delay, 2), error=error, **kw)
    def error(self, node, error, duration_ms, **kw): self._emit("ERROR", node, "error", error=error, duration_ms=round(duration_ms, 1), **kw)
    def skip(self, node, reason, **kw): self._emit("WARN", node, "skip", reason=reason, **kw)

state/research_state.py

"""
state/research_state.py
Data models for the AI Research Assistant v2.
"""

from pydantic import BaseModel, Field
from datetime import datetime


class Source(BaseModel):
    name: str = Field(description="Name of the source")
    source_type: str = Field(description="Type: web, academic, news")
    content: str = Field(description="Content extracted from the source")


class KeyFinding(BaseModel):
    title: str = Field(description="Title of the finding")
    description: str = Field(description="Description of the finding")
    confidence: float = Field(ge=0.0, le=1.0)


class SourceStatus(BaseModel):
    source_type: str = Field(description="Type of source")
    status: str = Field(description="ok, failed, timeout, skipped")
    attempts: int = Field(default=1)
    error: str = Field(default="")
    duration_ms: float = Field(default=0)


class ResearchReport(BaseModel):
    topic: str
    summary: str
    key_findings: list[KeyFinding] = Field(min_length=1)
    sources: list[Source] = Field(min_length=1)
    sub_queries: list[str]
    confidence: float = Field(ge=0.0, le=1.0)
    generated_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    source_availability: list[SourceStatus] = Field(default_factory=list)
    version: str = Field(default="v2")


class SubQuery(BaseModel):
    query: str
    rationale: str

config/settings.py

"""
config/settings.py
Configuration for the AI Research Assistant v2.
"""

from dotenv import load_dotenv
load_dotenv()

MODEL_NAME = "openai:gpt-4.1-mini"
MODEL_TEMPERATURE = 0.2

MAX_SUB_QUERIES = 4
MAX_SOURCES_PER_QUERY = 3
SEARCH_SOURCES = ["web", "academic", "news"]

RETRY_MAX_ATTEMPTS = 3
RETRY_BASE_DELAY = 1.0
RETRY_MAX_DELAY = 10.0
RETRY_JITTER = True

MIN_SOURCES_FOR_REPORT = 1
SEARCH_TIMEOUT_SECONDS = 5.0

Running it

Normal mode (no simulated failures)

cd research-assistant
python main.py "the impact of artificial intelligence on education"

Expected output: every source responds, 0 retries, high confidence.

Mode with simulated failures

To test the retry logic, add this at the top of main.py (before the invocation):

from tools.web_search import configure_failures

configure_failures("academic", fail_rate=0.7)
configure_failures("news", fail_rate=0.3)

Now academic fails 70% of the time and news 30%. The agent retries automatically.

Success criteria

Your project is complete when you meet these criteria:

✅ The searches run in parallel — the 3 sources are fired simultaneously (measurable: total time is ~equal to the slowest source, not the sum)
✅ Transient errors recover — with configure_failures("academic", fail_rate=0.5), the agent retries and eventually gets results
✅ Permanent failures don't crash it — with configure_failures("academic", fail_rate=1.0), the agent carries on with web and news
✅ The report includes source availability — you can see how many attempts each source took and which ones failed
✅ The results are deduplicated — there are no repeated findings from different sub-queries that hit the same source
✅ Confidence reflects availability — with every source, high confidence. With sources missing, confidence adjusted downward
✅ The logs are structured — every operation has a JSON log with request_id, node, duration, and status
✅ The v2 report JSON includes version and source_availability — the new fields are present and valid

Test scenarios

Test 1: Every source OK (happy path)

# Without configure_failures (default: 0% fail rate)
# python main.py "machine learning in medicine"

============================================================
  🔬 AI Research Assistant v2
  Topic: machine learning in medicine
  Request ID: a1b2c3d4
============================================================

📋 Step 1: Decomposing the topic into sub-queries...
   ✓ 4 sub-queries generated

🔍 Step 2: Searching 3 sources per sub-query (with retry)...
   Sub-query 1: 3 OK, 0 failed
   Sub-query 2: 3 OK, 0 failed
   Sub-query 3: 3 OK, 0 failed
   Sub-query 4: 3 OK, 0 failed
   Total: 12 successful, 0 failed out of 12

🔀 Step 3: Merging and deduplicating results...
   ✓ 12 raw results → 12 unique after deduplication

🧠 Step 4: Synthesizing findings...
   ✓ 4 findings identified

📝 Step 5: Generating the executive summary...
   ✓ Summary generated

📊 Step 6: Computing the report's confidence...
   ✓ Base confidence: 82%, adjusted for availability: 82%

╔══════════════════════════════════════════════════════════╗
║       📄 RESEARCH REPORT (v2)                            ║
╚══════════════════════════════════════════════════════════╝

🔌 SOURCE AVAILABILITY
────────────────────────────────────────────────────────────
  ✅ web: ok (305ms)
  ✅ academic: ok (420ms)
  ✅ news: ok (280ms)
  ...

Test 2: One source fails intermittently (retry success)

from tools.web_search import configure_failures
configure_failures("academic", fail_rate=0.7)

🔍 Step 2: Searching 3 sources per sub-query (with retry)...
   [Log: {"node": "search_academic", "event": "retry", "attempt": 1, "delay_s": 1.05, "error": "academic API: connection refused"}]
   [Log: {"node": "search_academic", "event": "retry", "attempt": 2, "delay_s": 2.13, "error": "academic API: 503 Service Unavailable"}]
   [Log: {"node": "search_academic", "event": "end", "duration_ms": 3850.2, "attempts": 3, "status": "success"}]
   Sub-query 1: 3 OK, 0 failed  ← academic recovered on attempt 3

🔌 SOURCE AVAILABILITY
  ✅ web: ok (305ms)
  ✅ academic: ok (3 attempts, 3850ms)  ← took longer but it worked
  ✅ news: ok (280ms)

Test 3: One source fails permanently (graceful degradation)

configure_failures("academic", fail_rate=1.0)

🔍 Step 2: Searching 3 sources per sub-query (with retry)...
   [Log: {"node": "search_academic", "event": "retry", "attempt": 1, ...}]
   [Log: {"node": "search_academic", "event": "retry", "attempt": 2, ...}]
   [Log: {"node": "search_academic", "event": "retry", "attempt": 3, ...}]
   [Log: {"node": "search_academic", "event": "error", "attempts": 4, "error": "academic API: 503"}]
   Sub-query 1: 2 OK, 1 failed  ← academic failed after 4 attempts, web and news OK

📊 Step 6: Computing the report's confidence...
   ✓ Base confidence: 78%, adjusted for availability: 72%  ← penalized

🔌 SOURCE AVAILABILITY
  ✅ web: ok (305ms)
  ❌ academic: failed (4 attempts, 7200ms) — academic API: 503 Service Unavailable
  ✅ news: ok (280ms)

Test 4: Every source fails

configure_failures("web", fail_rate=1.0)
configure_failures("academic", fail_rate=1.0)
configure_failures("news", fail_rate=1.0)

🔍 Step 2: Searching 3 sources per sub-query (with retry)...
   Sub-query 1: 0 OK, 3 failed
   Total: 0 successful, 12 failed out of 12

   ❌ Insufficient sources (0 < 1). Aborting.

╔══════════════════════════════════════════════════════════╗
║              ❌ RESEARCH ERROR                            ║
╚══════════════════════════════════════════════════════════╝

  Only 0 sources responded. Minimum required: 1
  Request ID: f4e5d6c7

  Source availability:
    ❌ web: failed (4 attempts)
    ❌ academic: failed (4 attempts)
    ❌ news: failed (4 attempts)

Test 5: Intermittent failures (stress test)

configure_failures("web", fail_rate=0.3)
configure_failures("academic", fail_rate=0.5)
configure_failures("news", fail_rate=0.2)

for topic in ["AI in education", "quantum computing", "renewable energy"]:
    result = research_agent.invoke(topic, config={"configurable": {"thread_id": f"test-{topic}"}})
    ok = sum(1 for sa in result.get("source_availability", []) if sa["status"] == "ok")
    total = len(result.get("source_availability", []))
    print(f"  {topic}: {ok}/{total} sources OK, confidence {result.get('confidence', 'N/A')}")

# Expected output (varies with the random):
#   AI in education: 10/12 sources OK, confidence 0.78
#   quantum computing: 11/12 sources OK, confidence 0.80
#   renewable energy: 9/12 sources OK, confidence 0.75

Common errors

1. `ModuleNotFoundError: No module named 'utils'`

Cause: The utils/ directory doesn't exist or has no __init__.py (though it isn't always required with sys.path.insert).

Fix: Run from the project root and check the structure:

cd research-assistant
ls utils/
# Should show: retry.py  logger.py
python main.py

2. Retry takes too long during testing

Cause: With base_delay=1.0 and 3 retries, each failed search takes ~7 seconds (1s + 2s + 4s of backoff). With 4 sub-queries and 3 sources, that's potentially 12 × 7s = 84 seconds.

Fix: For testing, shrink the delays:

from utils.retry import RetryConfig
test_config = RetryConfig(max_retries=2, base_delay=0.1, jitter=False)

Or change RETRY_BASE_DELAY in config/settings.py for development.

3. Deduplication doesn't remove similar findings from the LLM

Cause: merge_and_deduplicate compares by source_type + content[:100]. If two different sub-queries hit the same source, the mock content is different (because it embeds the query in the text), so they aren't detected as duplicates.

Fix: The current deduplication is conservative — it only removes exact duplicates. For semantic deduplication (findings that say the same thing in different words), you'd need embeddings or an LLM judge. That's the subject of Module 11.

4. `search_with_retry` doesn't use the config from `settings.py`

Cause: The search_with_retry function creates its own RetryConfig with hardcoded values.

Fix: Import the configuration from settings:

from config.settings import RETRY_MAX_ATTEMPTS, RETRY_BASE_DELAY, RETRY_MAX_DELAY, RETRY_JITTER

retry_config = RetryConfig(
    max_retries=RETRY_MAX_ATTEMPTS,
    base_delay=RETRY_BASE_DELAY,
    max_delay=RETRY_MAX_DELAY,
    jitter=RETRY_JITTER,
)

5. The retry logs are excessive in production

Cause: With many sub-queries and unstable sources, every retry generates a log. 4 sub-queries × 3 sources × 3 retries = 36 retry logs.

Fix: ResearchLogger already uses levels — retries are WARN. In production, set the minimum level:

logging.basicConfig(level=logging.WARNING)  # Only WARN and ERROR

6. `SourceStatus` doesn't show up in the report JSON

Cause: You're using the v1 ResearchReport without the source_availability field.

Fix: Check that state/research_state.py has the updated model with SourceStatus and source_availability. The field has default_factory=list, so it's compatible with v1 (it shows up as an empty list).

7. The adjusted confidence is always equal to the base

Cause: The availability_factor is total_ok / (total_ok + total_failed). If there are no failures, availability_factor = 1.0 and the formula becomes base * (0.7 + 0.3 * 1.0) = base * 1.0.

Fix: That's correct — confidence is only penalized when sources are missing. With every source available, the adjusted confidence equals the base.

8. `TypeError` when passing `logger` as a `@task` argument

Cause: The logger isn't JSON-serializable, and @task tries to serialize it for the checkpoint.

Fix: Use a global logger instead of passing it as an argument, or exclude it from serialization. In this project's code, search_with_retry receives the logger as an argument — make sure your checkpointer can handle that, or use the global logger pattern shown in capsule 07.

What's coming: Module 8 — Memory and Persistence

Your Research Assistant v2 is robust: it retries errors, searches in parallel, degrades gracefully, and logs everything. But it has one fundamental limitation: every research run starts from zero.

If you researched "artificial intelligence in medicine" yesterday, and today you research "AI applied to diagnosis", the agent doesn't know it already has relevant information. If the process is interrupted halfway through a long research run, you lose all the progress.

Module 8 adds two capabilities:

Persistent checkpointing with PostgresSaver — the agent saves its state at each step. If it goes down at step 4, it resumes from step 3 without repeating work
Long-term memory — the agent remembers earlier research, user preferences, and can reuse relevant previous findings

v3 isn't just more robust — it's smarter. An agent that remembers is fundamentally different from one that doesn't.

Project resources

LangGraph Functional API — Reference for @entrypoint and @task, including Futures
Exponential Backoff and Jitter — AWS blog on why jitter is essential in retry logic
Python logging Cookbook — Structured logging with JSON formatters
Pydantic v2 — Field Defaults — How to make models backwards-compatible with default_factory
LangGraph Persistence — Checkpointing and MemorySaver
Graceful Degradation in Distributed Systems — Chapter 8 of "Designing Data-Intensive Applications" (Kleppmann)

Module 7 — LangChain & LangGraph: From Chains to Agents

← PreviousProduction Patterns Next →Introduction: Why Agents Need Memory