AI-Agent Driven Revenue Investigation: LangGraph Orchestration Meets Medallion Architecture

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]

class SentinelAgent:
    def __init__(self):
        self.tools = [
            RAGTool, MonitoringTool, InfrastructureTool,
            delegate_to_agent, list_available_agents,
            route_task_to_specialist,
        ]
        self.tools_by_name = {t.name: t for t in self.tools}

        # Build the LangGraph StateGraph
        workflow = StateGraph(AgentState)
        workflow.add_node("agent", self._call_model)
        workflow.add_node("tools", self._call_tools)
        workflow.set_entry_point("agent")

        # Conditional edge: if tool_calls exist -> tools, else -> END
        workflow.add_conditional_edges(
            "agent",
            self._should_continue,
            {"continue": "tools", "end": END},
        )
        workflow.add_edge("tools", "agent")
        self.graph = workflow.compile()

def _should_continue(self, state: AgentState):
    """LangGraph conditional edge function.
    If the last message has tool_calls -> continue to tools node.
    Otherwise -> end the graph execution."""
    messages = state["messages"]
    last_message = messages[-1]
    if last_message.tool_calls:
        return "continue"
    return "end"

async def _call_tools(self, state: AgentState):
    """Execute all pending tool calls from the last model response.
    Supports: delegate_to_agent, list_available_agents,
    route_task_to_specialist, RAGTool, MonitoringTool."""
    messages = state["messages"]
    last_message = messages[-1]
    tool_results = []
    for tool_call in last_message.tool_calls:
        tool = self.tools_by_name[tool_call["name"]]
        result = await tool.ainvoke(tool_call["args"])
        tool_results.append(
            ToolMessage(content=str(result), tool_call_id=tool_call["id"])
        )
    return {"messages": tool_results}

# Data Pipeline Agent delegates to Spark for Bronze ingestion
# Tool: etl_pipeline(stage="bronze", source="raw_events")

raw_events = spark.read.json("s3a://raw-events/euroshop/2025-10-*.json.gz")

# Preserve ALL original fields -- Bronze is immutable
bronze_df = raw_events.select(
    col("session_id"),
    col("event_type"),        # page_view | add_to_cart | checkout | purchase
    col("timestamp"),
    col("country"),
    col("device_type"),       # desktop | mobile | tablet
    col("payment_method"),
    col("payment_token"),     # null when gateway fails
    col("order_value"),
    col("page_url"),
    col("user_agent"),
    current_timestamp().alias("_ingested_at"),
    lit("v1").alias("_schema_version"),
)

bronze_df.write \
    .format("delta") \
    .mode("append") \
    .option("mergeSchema", "true") \
    .save("s3a://lakehouse/bronze/events")

# Funnel stage assignment (deterministic, idempotent)
FUNNEL_STAGES = {
    "page_view": 1,
    "add_to_cart": 2,
    "begin_checkout": 3,
    "add_payment_info": 4,
    "purchase": 5,
}

silver_df = bronze_df \
    .withColumn("funnel_stage",
        coalesce(
            *[when(col("event_type") == evt, lit(stage))
              for evt, stage in FUNNEL_STAGES.items()],
            lit(0)
        )
    ) \
    .withColumn("has_payment_token",
        when(col("payment_token").isNotNull(), lit(True))
        .otherwise(lit(False))
    )

# Session-level aggregation: max funnel stage reached
session_funnel = silver_df \
    .groupBy("session_id", "country", "device_type", "date") \
    .agg(
        max("funnel_stage").alias("max_stage"),
        count("*").alias("event_count"),
        max("has_payment_token").alias("payment_completed"),
        first("payment_method").alias("payment_method"),
    )

session_funnel.write \
    .format("delta") \
    .mode("overwrite") \
    .save("s3a://lakehouse/silver/session_funnels")

# Daily conversion metrics by segment
gold_df = session_funnel \
    .groupBy("date", "country", "device_type") \
    .agg(
        count("*").alias("total_sessions"),
        sum(when(col("max_stage") >= 5, 1).otherwise(0))
            .alias("purchases"),
        sum(when(col("max_stage") >= 3, 1).otherwise(0))
            .alias("checkout_starts"),
        avg(when(col("max_stage") >= 5, col("order_value")))
            .alias("avg_order_value"),
        avg(col("payment_completed").cast("int"))
            .alias("payment_success_rate"),
    ) \
    .withColumn("conversion_rate",
        (col("purchases") / col("total_sessions") * 100)
    )

gold_df.write \
    .format("delta") \
    .mode("overwrite") \
    .save("s3a://lakehouse/gold/daily_conversion_metrics")

# >> Data Quality Agent automatically profiled this table
# >> Z-score on conversion_rate: -3.7 (CRITICAL)
# >> Z-score on revenue_total: -3.2 (CRITICAL)
# >> payment_token null rate: 34% vs 2% baseline (ANOMALY)

class DataQualityEngine:
    def detect_anomalies(self, df, threshold=2.5):
        """Z-score anomaly detection across all numeric columns.
        Returns list of {column, z_score, severity, value, baseline}."""
        anomalies = []
        numeric_cols = df.select_dtypes(include=[np.number]).columns

        for col_name in numeric_cols:
            values = df[col_name].dropna()
            if len(values) < 10:
                continue

            mean = values.rolling(30).mean().iloc[-1]
            std = values.rolling(30).std().iloc[-1]

            if std == 0:
                continue

            current = values.iloc[-1]
            z_score = (current - mean) / std

            if abs(z_score) >= threshold:
                anomalies.append({
                    "column": col_name,
                    "z_score": round(z_score, 2),
                    "severity": "critical" if abs(z_score) >= 3.0 else "warning",
                    "current_value": current,
                    "baseline_mean": round(mean, 4),
                    "baseline_std": round(std, 4),
                })

        return anomalies

@app.post("/analyze/quality")
async def analyze_quality(request: QualityRequest):
    """Full quality profile: null rates, z-scores, distribution stats.
    Called by Sentinel via delegate_to_agent tool."""
    engine = DataQualityEngine()

    # Connect to PostgreSQL for historical baselines
    async with get_db_connection() as conn:
        df = await conn.fetch_dataframe(request.table_name)

    profile = engine.profile_dataframe(df)
    anomalies = engine.detect_anomalies(df, threshold=request.threshold)

    # Persist checkpoint for audit trail
    await conn.execute(
        "INSERT INTO quality_checkpoints (table_name, profile, anomalies, ts) "
        "VALUES ($1, $2, $3, NOW())",
        request.table_name, json.dumps(profile), json.dumps(anomalies),
    )

    return {
        "table": request.table_name,
        "columns_profiled": len(profile),
        "anomalies": anomalies,
        "profile_summary": profile,
    }

@dataclass
class AgentServerConfig:
    name: str
    version: str = "1.0.0"
    description: str = ""
    default_model: str = "qwen2.5:7b"
    capabilities: list = field(default_factory=list)
    tools: list = field(default_factory=list)
    coordinator_url: str = "http://localhost:8080"

def create_agent_server(config: AgentServerConfig) -> FastAPI:
    """Factory function that creates a standardized agent server.
    Every agent in the fleet uses this template."""
    app = FastAPI(
        title=f"{config.name} Agent",
        version=config.version,
        description=config.description,
    )

    # Standard endpoints
    @app.get("/health")
    async def health():
        return {"status": "healthy", "agent": config.name,
                "version": config.version, "uptime": get_uptime()}

    @app.post("/query")
    async def query(request: QueryRequest):
        return await process_query(request, config)

    @app.post("/task")
    async def handle_task(request: TaskRequest):
        return await execute_task(request, config)

    @app.post("/delegate")
    async def delegate(request: DelegateRequest):
        coordinator = AgentCoordinator(config.coordinator_url)
        return await coordinator.delegate(request)

    @app.get("/capabilities")
    async def capabilities():
        return {"name": config.name, "tools": config.tools,
                "capabilities": config.capabilities}

    return app

class DataPipelineAgent(BaseAgent):
    """Extends BaseAgent with ETL-specific tools.
    Registered with Sentinel for delegation via /capabilities."""

    def __init__(self):
        super().__init__(
            name="data_pipeline_agent",
            description="ETL orchestration, data lineage, pipeline health",
            tools=[
                etl_pipeline,        # Trigger Bronze/Silver/Gold transforms
                data_quality_check,  # Pre/post transform validation
                backup_data,         # Delta Lake snapshot backup to MinIO
                data_lineage,        # Track column-level lineage graph
                query_knowledge_base # RAG over pipeline documentation
            ],
            model="qwen2.5:7b",
        )

    async def handle_investigation(self, task):
        """Called by Sentinel via delegate_to_agent tool.
        Executes ETL replay for specific date range and segments."""
        result = await self.run_tool("etl_pipeline", {
            "stage": task.stage,
            "date_range": task.date_range,
            "segments": task.segments,
        })
        quality = await self.run_tool("data_quality_check", {
            "table": f"lakehouse/{task.stage}/*",
        })
        return {"etl_result": result, "quality": quality}

@tool
async def delegate_to_agent(agent_name: str, task: str, params: dict = {}):
    """Delegate a task to a specialist agent via HTTP.
    The Sentinel uses this tool to coordinate the investigation.

    Args:
        agent_name: Target agent (e.g., 'data_quality_agent')
        task: Task description for the agent
        params: Additional parameters for the task
    """
    agent_registry = {
        "data_quality_agent": "http://dq-agent:8016",
        "data_pipeline_agent": "http://pipeline-agent:8017",
        "bi_analytics_agent": "http://bi-agent:8084",
        "infrastructure_agent": "http://infra-agent:8083",
    }

    url = agent_registry.get(agent_name)
    if not url:
        return {"error": f"Unknown agent: {agent_name}"}

    async with httpx.AsyncClient(timeout=30) as client:
        response = await client.post(
            f"{url}/task",
            json={"task": task, "params": params,
                   "source": "sentinel_agent"},
        )
        return response.json()

services:
  spark-master:
    image: bitnami/spark:3.5.3
    ports: ["8080:8080", "7077:7077"]
    environment:
      SPARK_MODE: master

  spark-worker-1:
    image: bitnami/spark:3.5.3
    environment:
      SPARK_MODE: worker
      SPARK_MASTER_URL: spark://spark-master:7077

  minio:
    image: minio/minio
    ports: ["9000:9000", "9001:9001"]
    command: server /data --console-address ":9001"

  airflow:
    image: apache/airflow:2.9.3
    ports: ["8081:8080"]

  jupyter:
    image: jupyter/pyspark-notebook
    ports: ["8888:8888"]

  postgres:
    image: postgres:16
    ports: ["5432:5432"]

  # AI Agent containers
  sentinel-agent:
    build: ./agents/sentinel
    ports: ["8080:8080"]
    depends_on: [postgres]

  data-quality-agent:
    build: ./agents/data-quality
    ports: ["8016:8016"]
    depends_on: [postgres]

  data-pipeline-agent:
    build: ./agents/data-pipeline
    ports: ["8017:8017"]