Building a Multi-Agent Marketing Analytics Team: Part 1 — Architecture and Core Implementation

# marketing_schema.py
# Run once to set up the data warehouse

import psycopg2
from datetime import date, timedelta
import random
import decimal

DDL = """
CREATE TABLE IF NOT EXISTS campaigns (
    id VARCHAR PRIMARY KEY,
    name VARCHAR NOT NULL,
    channel VARCHAR NOT NULL,  -- 'google_search', 'meta', 'email', 'youtube'
    objective VARCHAR,         -- 'awareness', 'consideration', 'conversion'
    daily_budget DECIMAL(10,2),
    start_date DATE,
    end_date DATE,
    status VARCHAR DEFAULT 'active'
);

CREATE TABLE IF NOT EXISTS campaign_daily_metrics (
    campaign_id VARCHAR REFERENCES campaigns(id),
    date DATE,
    spend DECIMAL(10,2),
    impressions INTEGER,
    clicks INTEGER,
    conversions INTEGER,
    revenue DECIMAL(10,2),
    cpc DECIMAL(10,4),
    cpm DECIMAL(10,4),
    ctr DECIMAL(6,4),
    conversion_rate DECIMAL(6,4),
    roas DECIMAL(8,4),
    frequency DECIMAL(6,2),  -- Meta only, NULL for other channels
    PRIMARY KEY (campaign_id, date)
);

CREATE TABLE IF NOT EXISTS conversion_touchpoints (
    order_id VARCHAR,
    customer_id VARCHAR,
    touchpoint_sequence INTEGER,
    channel VARCHAR,
    campaign_id VARCHAR,
    touchpoint_timestamp TIMESTAMP,
    is_conversion_event BOOLEAN,
    PRIMARY KEY (order_id, touchpoint_sequence)
);

CREATE TABLE IF NOT EXISTS customers (
    id VARCHAR PRIMARY KEY,
    acquisition_channel VARCHAR,
    acquisition_date DATE,
    total_orders INTEGER DEFAULT 0,
    total_revenue DECIMAL(10,2) DEFAULT 0,
    ltv_tier VARCHAR,     -- 'high', 'medium', 'low'
    last_order_date DATE,
    rfm_segment VARCHAR   -- 'champions', 'at_risk', 'lost', 'new'
);
"""

def seed_sample_data(conn_str: str, num_campaigns: int = 8, days: int = 60):
    """Generate realistic-looking sample data for testing."""
    conn = psycopg2.connect(conn_str)
    cur = conn.cursor()
    cur.execute(DDL)

    channels = ['google_search', 'meta', 'email', 'youtube']
    campaign_ids = []

    for i, channel in enumerate(channels * 2):
        cid = f"camp_{channel}_{i:02d}"
        campaign_ids.append((cid, channel))
        cur.execute("""
            INSERT INTO campaigns (id, name, channel, objective, daily_budget, start_date, status)
            VALUES (%s, %s, %s, %s, %s, %s, 'active')
            ON CONFLICT (id) DO NOTHING
        """, (cid, f"{channel.replace('_',' ').title()} Campaign {i}",
              channel, 'conversion', random.uniform(500, 5000),
              date.today() - timedelta(days=days)))

    today = date.today()
    for days_back in range(days, 0, -1):
        d = today - timedelta(days=days_back)
        for cid, channel in campaign_ids:
            spend = random.uniform(300, 4000)
            impressions = int(spend * random.uniform(200, 600))
            clicks = int(impressions * random.uniform(0.01, 0.05))
            conversions = int(clicks * random.uniform(0.02, 0.08))
            revenue = conversions * random.uniform(80, 350)
            frequency = round(random.uniform(1.2, 5.5), 2) if channel == 'meta' else None

            cur.execute("""
                INSERT INTO campaign_daily_metrics
                    (campaign_id, date, spend, impressions, clicks, conversions,
                     revenue, cpc, cpm, ctr, conversion_rate, roas, frequency)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                ON CONFLICT (campaign_id, date) DO NOTHING
            """, (cid, d, spend, impressions, clicks, conversions, revenue,
                  spend/clicks if clicks else 0,
                  (spend/impressions)*1000 if impressions else 0,
                  clicks/impressions if impressions else 0,
                  conversions/clicks if clicks else 0,
                  revenue/spend if spend else 0,
                  frequency))

    conn.commit()
    cur.close()
    conn.close()
    print(f"Seeded {len(campaign_ids)} campaigns × {days} days of metrics")

# data_analyst_agent.py
from langchain_openai import AzureChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel, Field
import psycopg2
import psycopg2.extras
import re
from typing import Optional

SCHEMA_CONTEXT = """
You have access to a marketing data warehouse with these tables:

TABLE: campaigns
  id VARCHAR PRIMARY KEY
  name VARCHAR
  channel VARCHAR  -- values: 'google_search', 'meta', 'email', 'youtube'
  objective VARCHAR  -- values: 'awareness', 'consideration', 'conversion'
  daily_budget DECIMAL(10,2)
  start_date DATE
  end_date DATE
  status VARCHAR  -- values: 'active', 'paused', 'ended'

TABLE: campaign_daily_metrics
  campaign_id VARCHAR (FK to campaigns.id)
  date DATE
  spend DECIMAL(10,2)
  impressions INTEGER
  clicks INTEGER
  conversions INTEGER
  revenue DECIMAL(10,2)
  cpc DECIMAL(10,4)  -- cost per click
  cpm DECIMAL(10,4)  -- cost per thousand impressions
  ctr DECIMAL(6,4)   -- click-through rate (0.02 = 2%)
  conversion_rate DECIMAL(6,4)  -- conversions/clicks (0.03 = 3%)
  roas DECIMAL(8,4)  -- return on ad spend (revenue/spend)
  frequency DECIMAL(6,2)  -- average times audience saw ad; NULL for non-Meta channels

TABLE: conversion_touchpoints
  order_id VARCHAR
  customer_id VARCHAR
  touchpoint_sequence INTEGER  -- 1 = first touch, higher = closer to conversion
  channel VARCHAR
  campaign_id VARCHAR
  touchpoint_timestamp TIMESTAMP
  is_conversion_event BOOLEAN  -- TRUE only for the final converting touchpoint

TABLE: customers
  id VARCHAR PRIMARY KEY
  acquisition_channel VARCHAR
  acquisition_date DATE
  total_orders INTEGER
  total_revenue DECIMAL(10,2)
  ltv_tier VARCHAR  -- values: 'high', 'medium', 'low'
  last_order_date DATE
  rfm_segment VARCHAR  -- values: 'champions', 'at_risk', 'lost', 'new'

RULES:
- Write SELECT-only queries. Never INSERT, UPDATE, DELETE, DROP, or ALTER.
- Always include a date filter using the date column.
- Always include LIMIT 500 unless the query is an aggregate (GROUP BY).
- Use COALESCE to handle NULL values in calculations.
- For month-over-month: use DATE_TRUNC('month', date) for grouping.
"""

class SQLResult(BaseModel):
    sql: str = Field(description="The SQL query to execute")
    explanation: str = Field(description="Plain-English explanation of what this query does")
    metric_type: str = Field(description="Primary metric: 'spend', 'roas', 'cac', 'ctr', 'conversion_rate', 'revenue'")


class DataAnalystAgent:
    def __init__(self, llm: AzureChatOpenAI, db_conn_str: str):
        self.llm = llm
        self.db_conn_str = db_conn_str
        self.parser = JsonOutputParser(pydantic_object=SQLResult)

        self.prompt = ChatPromptTemplate.from_messages([
            ("system", f"{SCHEMA_CONTEXT}\n\nGenerate a SQL query to answer the user's marketing analytics question. Return JSON with fields: sql, explanation, metric_type."),
            ("human", "Question: {question}\nDate range: {date_range}")
        ])
        self.chain = self.prompt | self.llm | self.parser

    def _validate_sql(self, sql: str) -> tuple[bool, str]:
        """Ensure generated SQL is safe to execute."""
        sql_upper = sql.upper().strip()

        # Block write operations
        forbidden = ['INSERT', 'UPDATE', 'DELETE', 'DROP', 'ALTER', 'CREATE', 'TRUNCATE', 'EXEC']
        for keyword in forbidden:
            if re.search(r'\b' + keyword + r'\b', sql_upper):
                return False, f"Forbidden keyword: {keyword}"

        # Require SELECT
        if not sql_upper.startswith('SELECT') and 'WITH' not in sql_upper[:10]:
            return False, "Query must start with SELECT or WITH"

        return True, "ok"

    def _sanity_check(self, results: list[dict], metric_type: str) -> list[str]:
        """Flag statistically implausible results."""
        flags = []
        for row in results:
            roas = row.get('roas') or row.get('avg_roas')
            if roas and float(roas) > 50:
                flags.append(f"ROAS > 50 detected ({roas:.1f}x) — verify data integrity")
            cr = row.get('conversion_rate') or row.get('avg_conversion_rate')
            if cr and float(cr) > 0.5:
                flags.append(f"Conversion rate > 50% detected — check touchpoint data")
            cac = row.get('cac')
            if cac and float(cac) < 1:
                flags.append(f"CAC < $1 detected — check spend/conversion data")
        return flags

    async def run(self, question: str, date_range: dict) -> dict:
        """Execute NL→SQL pipeline and return structured results."""
        # Generate SQL
        result = await self.chain.ainvoke({
            "question": question,
            "date_range": f"{date_range['start']} to {date_range['end']}"
        })

        # Validate before execution
        valid, reason = self._validate_sql(result['sql'])
        if not valid:
            return {"error": f"Query validation failed: {reason}", "sql": result['sql']}

        # Execute against read-only connection
        conn = psycopg2.connect(self.db_conn_str)
        cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        try:
            cur.execute(result['sql'])
            rows = [dict(r) for r in cur.fetchall()]
        except Exception as e:
            return {"error": str(e), "sql": result['sql']}
        finally:
            cur.close()
            conn.close()

        flags = self._sanity_check(rows, result['metric_type'])

        return {
            "question": question,
            "sql": result['sql'],
            "explanation": result['explanation'],
            "results": rows,
            "row_count": len(rows),
            "sanity_flags": flags,
            "metric_type": result['metric_type']
        }

# attribution_agent.py
import numpy as np
from collections import defaultdict
from typing import List, Dict
import psycopg2
import psycopg2.extras
from dataclasses import dataclass

@dataclass
class TouchpointJourney:
    order_id: str
    channels: List[str]  # ordered sequence of channels

class AttributionAgent:
    def __init__(self, db_conn_str: str, num_shapley_samples: int = 1000):
        self.db_conn_str = db_conn_str
        self.num_shapley_samples = num_shapley_samples

    def _load_journeys(self, start_date: str, end_date: str) -> List[TouchpointJourney]:
        """Load conversion journeys from touchpoint table."""
        sql = """
        SELECT order_id,
               array_agg(channel ORDER BY touchpoint_sequence) AS channels
        FROM conversion_touchpoints
        WHERE touchpoint_timestamp::date BETWEEN %s AND %s
          AND order_id IN (
              SELECT DISTINCT order_id FROM conversion_touchpoints
              WHERE is_conversion_event = TRUE
                AND touchpoint_timestamp::date BETWEEN %s AND %s
          )
        GROUP BY order_id
        """
        conn = psycopg2.connect(self.db_conn_str)
        cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        cur.execute(sql, (start_date, end_date, start_date, end_date))
        journeys = [
            TouchpointJourney(order_id=r['order_id'], channels=r['channels'])
            for r in cur.fetchall()
        ]
        cur.close()
        conn.close()
        return journeys

    def last_touch(self, journeys: List[TouchpointJourney]) -> Dict[str, float]:
        credits = defaultdict(float)
        for j in journeys:
            if j.channels:
                credits[j.channels[-1]] += 1.0
        total = sum(credits.values())
        return {ch: v / total for ch, v in credits.items()}

    def first_touch(self, journeys: List[TouchpointJourney]) -> Dict[str, float]:
        credits = defaultdict(float)
        for j in journeys:
            if j.channels:
                credits[j.channels[0]] += 1.0
        total = sum(credits.values())
        return {ch: v / total for ch, v in credits.items()}

    def linear(self, journeys: List[TouchpointJourney]) -> Dict[str, float]:
        credits = defaultdict(float)
        for j in journeys:
            if j.channels:
                weight = 1.0 / len(j.channels)
                for ch in j.channels:
                    credits[ch] += weight
        total = sum(credits.values())
        return {ch: v / total for ch, v in credits.items()}

    def time_decay(self, journeys: List[TouchpointJourney],
                   half_life_days: float = 7.0) -> Dict[str, float]:
        """Exponential decay: later touchpoints get more weight."""
        credits = defaultdict(float)
        for j in journeys:
            n = len(j.channels)
            weights = np.array([
                np.exp(np.log(2) * (i - (n - 1)) / half_life_days)
                for i in range(n)
            ])
            weights /= weights.sum()
            for ch, w in zip(j.channels, weights):
                credits[ch] += float(w)
        total = sum(credits.values())
        return {ch: v / total for ch, v in credits.items()}

    def shapley_approximation(
            self, journeys: List[TouchpointJourney]) -> Dict[str, float]:
        """
        Monte Carlo Shapley approximation.
        Exact Shapley is O(2^n) — we sample random orderings instead.
        Conversion probability = 1 if channel is in the coalition, 0 otherwise
        (simplified characteristic function for speed at scale).
        """
        all_channels = list({ch for j in journeys for ch in j.channels})
        shapley_values = defaultdict(float)

        # Build a conversion set per journey for fast lookup
        journey_sets = [set(j.channels) for j in journeys]

        def conversion_prob(coalition: set) -> float:
            """Fraction of journeys where all channels in coalition appeared."""
            if not coalition:
                return 0.0
            return sum(1 for js in journey_sets if coalition.issubset(js)) / len(journey_sets)

        for _ in range(self.num_shapley_samples):
            perm = np.random.permutation(all_channels)
            coalition = set()
            prev_value = 0.0
            for ch in perm:
                coalition.add(ch)
                new_value = conversion_prob(coalition)
                shapley_values[ch] += (new_value - prev_value)
                prev_value = new_value

        # Normalize to percentages
        total = sum(shapley_values.values())
        if total == 0:
            return {ch: 1.0 / len(all_channels) for ch in all_channels}
        return {ch: v / total for ch, v in shapley_values.items()}

    async def run(self, date_range: dict) -> dict:
        journeys = self._load_journeys(date_range['start'], date_range['end'])

        if not journeys:
            return {"error": "No conversion journeys found in date range"}

        models = {
            "last_touch": self.last_touch(journeys),
            "first_touch": self.first_touch(journeys),
            "linear": self.linear(journeys),
            "time_decay": self.time_decay(journeys),
            "shapley": self.shapley_approximation(journeys),
        }

        # Build comparison table
        all_channels = sorted({ch for m in models.values() for ch in m})
        comparison = []
        for ch in all_channels:
            comparison.append({
                "channel": ch,
                "last_touch_pct": round(models["last_touch"].get(ch, 0) * 100, 1),
                "first_touch_pct": round(models["first_touch"].get(ch, 0) * 100, 1),
                "linear_pct": round(models["linear"].get(ch, 0) * 100, 1),
                "time_decay_pct": round(models["time_decay"].get(ch, 0) * 100, 1),
                "shapley_pct": round(models["shapley"].get(ch, 0) * 100, 1),
            })

        # Flag over-credited channels
        flags = []
        for row in comparison:
            gap = row["last_touch_pct"] - row["shapley_pct"]
            if gap > 15:
                flags.append(
                    f"{row['channel']} over-credited by last-touch: "
                    f"+{gap:.1f}pp vs Shapley. Review budget allocation."
                )

        return {
            "journey_count": len(journeys),
            "date_range": date_range,
            "comparison": comparison,
            "attribution_flags": flags,
            "recommendation": (
                "Use Shapley as primary model for budget decisions. "
                "Last-touch is provided for platform reconciliation only."
            )
        }

# campaign_diagnostics_agent.py
from langchain_openai import AzureChatOpenAI
from langchain.prompts import ChatPromptTemplate
from typing import List, Dict, Any
import json

DIAGNOSTIC_THRESHOLDS = {
    "frequency_fatigue": 3.5,
    "ctr_decline_pct": 0.25,
    "cpm_increase_pct": 0.30,
    "conv_rate_drop_pct": 0.20,
    "roas_decline_pct": 0.30,
}

class CampaignDiagnosticsAgent:
    def __init__(self, llm: AzureChatOpenAI):
        self.llm = llm
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a marketing campaign diagnostics specialist.
Given detected signal flags and raw metrics, synthesize a diagnostic report.
For each issue:
1. State the root cause hypothesis
2. Estimate the revenue impact in dollars
3. Recommend one specific action
Be direct. Use the actual numbers from the data. Do not speculate beyond the data."""),
            ("human", "Campaign metrics:\n{metrics}\n\nDetected signals:\n{signals}")
        ])

    def _detect_signals(self, campaign_data: List[Dict]) -> List[Dict]:
        """Apply threshold rules to identify diagnostic signals."""
        signals = []
        for camp in campaign_data:
            name = camp.get('campaign_name', camp.get('campaign_id', 'Unknown'))
            channel = camp.get('channel', '')

            # Frequency fatigue (Meta only)
            freq = camp.get('avg_frequency_recent')
            if freq and channel == 'meta' and float(freq) > DIAGNOSTIC_THRESHOLDS["frequency_fatigue"]:
                signals.append({
                    "campaign": name,
                    "channel": channel,
                    "signal": "frequency_fatigue",
                    "severity": "HIGH",
                    "value": float(freq),
                    "threshold": DIAGNOSTIC_THRESHOLDS["frequency_fatigue"],
                    "context": f"Frequency at {freq:.1f} (threshold: 3.5). "
                               f"Audience has seen this ad too many times."
                })

            # CTR decline
            ctr_recent = camp.get('ctr_last_14d')
            ctr_prior = camp.get('ctr_prior_14d')
            if ctr_recent and ctr_prior and float(ctr_prior) > 0:
                ctr_decline = (float(ctr_prior) - float(ctr_recent)) / float(ctr_prior)
                if ctr_decline > DIAGNOSTIC_THRESHOLDS["ctr_decline_pct"]:
                    signals.append({
                        "campaign": name,
                        "channel": channel,
                        "signal": "ctr_decline",
                        "severity": "HIGH",
                        "value": round(ctr_decline * 100, 1),
                        "context": f"CTR fell {ctr_decline*100:.1f}% in 14 days "
                                   f"({ctr_prior:.3f} → {ctr_recent:.3f})"
                    })

            # Conversion rate drop (with stable CTR — landing page signal)
            cr_recent = camp.get('conv_rate_last_14d')
            cr_prior = camp.get('conv_rate_prior_14d')
            if cr_recent and cr_prior and float(cr_prior) > 0:
                cr_drop = (float(cr_prior) - float(cr_recent)) / float(cr_prior)
                if cr_drop > DIAGNOSTIC_THRESHOLDS["conv_rate_drop_pct"]:
                    # Check if CTR is stable (landing page issue, not creative issue)
                    ctr_stable = (not ctr_recent or not ctr_prior or
                                  abs(float(ctr_recent) - float(ctr_prior)) / float(ctr_prior) < 0.1)
                    signals.append({
                        "campaign": name,
                        "channel": channel,
                        "signal": "conversion_rate_drop",
                        "severity": "HIGH",
                        "value": round(cr_drop * 100, 1),
                        "hypothesis": "landing_page" if ctr_stable else "creative_or_audience",
                        "context": f"Conv rate fell {cr_drop*100:.1f}% "
                                   f"({'CTR stable — likely landing page issue' if ctr_stable else 'CTR also falling — creative issue'})"
                    })

            # ROAS decline
            roas_recent = camp.get('roas_recent')
            roas_prior = camp.get('roas_prior')
            if roas_recent and roas_prior and float(roas_prior) > 0:
                roas_drop = (float(roas_prior) - float(roas_recent)) / float(roas_prior)
                if roas_drop > DIAGNOSTIC_THRESHOLDS["roas_decline_pct"]:
                    signals.append({
                        "campaign": name,
                        "channel": channel,
                        "signal": "roas_decline",
                        "severity": "MEDIUM",
                        "value": round(roas_drop * 100, 1),
                        "context": f"ROAS fell {roas_drop*100:.1f}% "
                                   f"({roas_prior:.2f}x → {roas_recent:.2f}x)"
                    })

        return signals

    async def run(self, analyst_output: dict, date_range: dict) -> dict:
        """Detect signals and synthesize diagnosis."""
        results = analyst_output.get('results', [])
        signals = self._detect_signals(results)

        if not signals:
            return {
                "signals": [],
                "diagnosis": "No diagnostic signals detected above thresholds.",
                "date_range": date_range,
                "severity": "NONE"
            }

        # LLM synthesis
        chain = self.prompt | self.llm
        response = await chain.ainvoke({
            "metrics": json.dumps(results, indent=2, default=str),
            "signals": json.dumps(signals, indent=2)
        })

        return {
            "signals": signals,
            "signal_count": len(signals),
            "high_severity_count": sum(1 for s in signals if s.get('severity') == 'HIGH'),
            "diagnosis": response.content,
            "date_range": date_range
        }

# orchestrator.py
from langgraph.graph import StateGraph, END
from typing import TypedDict, Optional, List
import asyncio
from langchain_openai import AzureChatOpenAI
from langchain.prompts import ChatPromptTemplate

class AnalyticsState(TypedDict):
    session_id: str
    user_question: str
    date_range: dict
    intent: str           # "diagnostic" | "reporting" | "attribution" | "campaign"
    agents_required: List[str]
    data_analyst_output: Optional[dict]
    attribution_output: Optional[dict]
    diagnostics_output: Optional[dict]
    segmentation_output: Optional[dict]
    narrative_output: Optional[str]
    errors: List[str]
    follow_up_questions: List[str]

INTENT_CLASSIFICATION_PROMPT = ChatPromptTemplate.from_messages([
    ("system", """Classify the marketing analytics question into one of these intents:
- diagnostic: questions about why a metric changed (CAC, ROAS, conversion rate)
- reporting: questions requesting a summary or overview
- attribution: questions about channel credit, budget allocation, multi-touch
- campaign: questions about specific campaign performance or issues

Return JSON: {"intent": "...", "agents_required": [...], "date_range_hint": "..."}
Available agents: data_analyst, attribution, diagnostics, segmentation, narrative"""),
    ("human", "{question}")
])

class MarketingAnalyticsOrchestrator:
    def __init__(
        self,
        llm: AzureChatOpenAI,
        data_analyst: 'DataAnalystAgent',
        attribution: 'AttributionAgent',
        diagnostics: 'CampaignDiagnosticsAgent',
        narrative: 'NarrativeAgent',
    ):
        self.llm = llm
        self.data_analyst = data_analyst
        self.attribution = attribution
        self.diagnostics = diagnostics
        self.narrative = narrative
        self.graph = self._build_graph()

    def _build_graph(self):
        workflow = StateGraph(AnalyticsState)

        # Add nodes
        workflow.add_node("classify_intent", self._classify_intent)
        workflow.add_node("run_data_analyst", self._run_data_analyst)
        workflow.add_node("run_attribution", self._run_attribution)
        workflow.add_node("run_diagnostics", self._run_diagnostics)
        workflow.add_node("run_narrative", self._run_narrative)
        workflow.add_node("parallel_analyst_attribution", self._parallel_analyst_attribution)
        workflow.add_node("parallel_analyst_diagnostics", self._parallel_analyst_diagnostics)

        # Entry point
        workflow.set_entry_point("classify_intent")

        # Route from intent classification
        workflow.add_conditional_edges(
            "classify_intent",
            self._route_by_intent,
            {
                "diagnostic": "parallel_analyst_attribution",
                "reporting": "run_data_analyst",
                "attribution": "run_data_analyst",
                "campaign": "parallel_analyst_diagnostics",
            }
        )

        # Diagnostic flow: parallel → diagnostics → narrative
        workflow.add_edge("parallel_analyst_attribution", "run_diagnostics")
        workflow.add_edge("run_diagnostics", "run_narrative")

        # Reporting flow: data_analyst → narrative
        workflow.add_edge("run_data_analyst", "run_narrative")

        # Attribution flow: data_analyst → attribution → narrative
        # (handled in the parallel node for attribution intent)
        workflow.add_edge("run_attribution", "run_narrative")

        # Campaign flow: parallel → narrative
        workflow.add_edge("parallel_analyst_diagnostics", "run_narrative")

        workflow.add_edge("run_narrative", END)

        return workflow.compile()

    def _route_by_intent(self, state: AnalyticsState) -> str:
        return state["intent"]

    async def _classify_intent(self, state: AnalyticsState) -> dict:
        chain = INTENT_CLASSIFICATION_PROMPT | self.llm
        from langchain_core.output_parsers import JsonOutputParser
        parser = JsonOutputParser()
        result = await (INTENT_CLASSIFICATION_PROMPT | self.llm | parser).ainvoke(
            {"question": state["user_question"]}
        )
        return {
            "intent": result.get("intent", "reporting"),
            "agents_required": result.get("agents_required", ["data_analyst", "narrative"])
        }

    async def _parallel_analyst_attribution(self, state: AnalyticsState) -> dict:
        """Run Data Analyst and Attribution in parallel for diagnostic questions."""
        analyst_task = self.data_analyst.run(
            state["user_question"], state["date_range"])
        attribution_task = self.attribution.run(state["date_range"])

        analyst_result, attribution_result = await asyncio.gather(
            analyst_task, attribution_task, return_exceptions=True
        )

        errors = list(state.get("errors", []))
        if isinstance(analyst_result, Exception):
            errors.append(f"DataAnalyst failed: {str(analyst_result)}")
            analyst_result = {}
        if isinstance(attribution_result, Exception):
            errors.append(f"Attribution failed: {str(attribution_result)}")
            attribution_result = {}

        return {
            "data_analyst_output": analyst_result,
            "attribution_output": attribution_result,
            "errors": errors
        }

    async def _parallel_analyst_diagnostics(self, state: AnalyticsState) -> dict:
        """Run Data Analyst and Diagnostics in parallel for campaign questions."""
        analyst_result = await self.data_analyst.run(
            state["user_question"], state["date_range"])
        diag_result = await self.diagnostics.run(analyst_result, state["date_range"])

        return {
            "data_analyst_output": analyst_result,
            "diagnostics_output": diag_result,
        }

    async def _run_data_analyst(self, state: AnalyticsState) -> dict:
        result = await self.data_analyst.run(
            state["user_question"], state["date_range"])
        return {"data_analyst_output": result}

    async def _run_attribution(self, state: AnalyticsState) -> dict:
        result = await self.attribution.run(state["date_range"])
        return {"attribution_output": result}

    async def _run_diagnostics(self, state: AnalyticsState) -> dict:
        result = await self.diagnostics.run(
            state.get("data_analyst_output", {}), state["date_range"])
        return {"diagnostics_output": result}

    async def _run_narrative(self, state: AnalyticsState) -> dict:
        # Narrative agent covered in Part 2
        narrative = await self.narrative.run(state)
        return {"narrative_output": narrative}

    async def answer(self, question: str, date_range: dict,
                     session_id: str = None) -> AnalyticsState:
        import uuid
        initial_state = AnalyticsState(
            session_id=session_id or str(uuid.uuid4()),
            user_question=question,
            date_range=date_range,
            intent="",
            agents_required=[],
            data_analyst_output=None,
            attribution_output=None,
            diagnostics_output=None,
            segmentation_output=None,
            narrative_output=None,
            errors=[],
            follow_up_questions=[]
        )
        return await self.graph.ainvoke(initial_state)