dbx-unifiedchat

Architecture

Multi-Agent System for Cross-Domain Queries with Databricks Genie.

System Overview

This system enables intelligent querying across multiple data domains (patients, medications, diagnoses, treatments, etc.) using a multi-agent architecture built with LangGraph.

High-Level Architecture

User Query
    ↓
SupervisorAgent (orchestrates all agents)
    ↓
ThinkingPlanningAgent (analyzes & plans)
    ├── Vector Search (semantic retrieval)
    └── Genie Space Metadata
    ↓
    ├─ Single Space → GenieAgent
    ├─ Multiple Spaces (No Join) → Multiple GenieAgents → Verbal Merge
    └─ Multiple Spaces (Join) → Fast/Genie Route → SQLSynthesis → SQLExecution
    ↓
Response with reasoning

Architecture Diagrams

Visual representations of the system architecture:

All diagrams are in the architecture/ directory in multiple formats (SVG, PNG, PDF, Mermaid).

Core Components

1. SupervisorAgent

Role: Central orchestrator Responsibilities:

2. ThinkingPlanningAgent

Role: Query analysis and planning Responsibilities:

Tools:

3. GenieAgent(s)

Role: Query individual Genie spaces Responsibilities:

Multiple instances: One per Genie space being queried

4. SQLSynthesisAgent

Role: Combines SQL queries across tables/spaces Responsibilities:

Variants:

5. SQLExecutionAgent

Role: Executes synthesized SQL Responsibilities:

6. ClarificationAgent

Role: Handles ambiguous queries Responsibilities:

7. SummarizeAgent

Role: Final response formatting Responsibilities:

Data Flow

Scenario 1: Single Genie Space Query

User: "Show me patient demographics"
   ↓
ThinkingPlanningAgent
   ├─ Vector Search: Finds "patient_demographics" space
   └─ Decision: Single space query
   ↓
GenieAgent (patient_demographics space)
   └─ Query Genie space directly
   ↓
SummarizeAgent
   └─ Format response
   ↓
User: Response with patient demographics

Scenario 2: Multi-Space Query (No Join)

User: "Show me patients and their medications"
   ↓
ThinkingPlanningAgent
   ├─ Vector Search: Finds "patients" AND "medications" spaces
   └─ Decision: Multiple spaces, no join needed
   ↓
GenieAgent (patients) + GenieAgent (medications)
   └─ Query both spaces in parallel
   ↓
SummarizeAgent
   └─ Verbal merge: "Here are patients... and their medications..."
   ↓
User: Combined response

Scenario 3: Multi-Space Query (With Join)

User: "Show me patients with high blood pressure AND their medications"
   ↓
ThinkingPlanningAgent
   ├─ Vector Search: Finds "patients" AND "medications" spaces
   └─ Decision: Multiple spaces, JOIN required
   ↓
SQLSynthesisAgent
   ├─ Get table schemas
   ├─ Get sample data
   └─ Generate JOIN query
   ↓
SQLExecutionAgent
   └─ Execute SQL on SQL Warehouse
   ↓
SummarizeAgent
   └─ Format results with reasoning
   ↓
User: Synthesized response

Scenario 4: Ambiguous Query

User: "Show me data"
   ↓
ThinkingPlanningAgent
   └─ Decision: Query too ambiguous
   ↓
ClarificationAgent
   └─ "What type of data are you interested in?"
   ↓
User: "Patient data"
   ↓
ThinkingPlanningAgent
   └─ Continue with Scenario 1

State Management

Short-Term Memory (CheckpointSaver)

Uses Lakebase PostgreSQL for conversation state:

Long-Term Memory (DatabricksStore)

Uses Lakebase with semantic search:

State Schema

class AgentState(TypedDict):
    messages: list  # Conversation history
    relevant_spaces: list  # Genie spaces for current query
    sql_query: Optional[str]  # Generated SQL
    sql_results: Optional[dict]  # Execution results
    final_response: Optional[str]  # Final answer
    # ... more fields

Technology Stack

Core Framework

Databricks Services

Models

Deployment Architecture

┌─────────────────────────────────────────────────────┐
│              Databricks Model Serving                │
│  ┌───────────────────────────────────────────────┐  │
│  │   Agent Container (Auto-scaled)              │  │
│  │   ├─ agent.py (MLflow wrapper)               │  │
│  │   ├─ src/multi_agent/ (packaged code)        │  │
│  │   └─ prod_config.yaml (runtime config)       │  │
│  └───────────────────────────────────────────────┘  │
│                     ↓                                │
│  ┌───────────────────────────────────────────────┐  │
│  │   Databricks Services                         │  │
│  │   ├─ Genie Spaces (data querying)            │  │
│  │   ├─ Vector Search (semantic retrieval)      │  │
│  │   ├─ SQL Warehouse (query execution)         │  │
│  │   ├─ Lakebase (state management)             │  │
│  │   ├─ Unity Catalog (metadata)                │  │
│  │   └─ LLM Endpoints (Claude models)           │  │
│  └───────────────────────────────────────────────┘  │
└─────────────────────────────────────────────────────┘

Security

Authentication

Authorization

Data Privacy

Performance

Latency

Optimization

Design Decisions

Why Multi-Agent?

Why LangGraph?

Why Modular Code?

Scalability

The system scales in multiple dimensions:

  1. Genie Spaces: Add more spaces by updating config
  2. Agents: Add new agents by extending graph
  3. LLM Models: Swap models per agent for cost/performance balance
  4. Workload: Model Serving auto-scales based on traffic

Extensibility

Adding New Agent

  1. Create agent in src/multi_agent/agents/new_agent.py
  2. Register in src/multi_agent/core/graph.py
  3. Add routing logic if needed
  4. Test locally and in Databricks
  5. Redeploy

Adding New Tool

  1. Create tool in src/multi_agent/tools/new_tool.py
  2. Register with appropriate agent
  3. Test tool independently
  4. Integrate with agent workflow

Supporting New Data Source

  1. Add Genie space ID to config
  2. Run ETL to enrich metadata
  3. Rebuild vector search index
  4. Agents automatically discover via vector search

See Also

Want to understand the code? See ../src/multi_agent/README.md for code structure guide! 💡

This site is open source. Improve this page.