%pip install -U -qqqq databricks-openai databricks-agents mlflow langgraph google-

genai
[Link]()

import os
import sys
import time
import json
from [Link] import WorkspaceClient
from databricks_openai import UCFunctionToolkit, DatabricksFunctionClient
from google import genai

LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class GeminiChat:
"""
Unified Gemini wrapper for LangGraph agents.
Includes:
- safe retry logic
- prompt sanitization
- JSON-safe output
"""

def __init__(self, llm_endpoint: str):

self.llm_endpoint = llm_endpoint
self.workspace_client = WorkspaceClient()
self.model_serving_client: OpenAI = (
self.workspace_client.serving_endpoints.get_open_ai_client()
)
[Link]=llm_endpoint

def invoke(self, prompt: dict):

"""
Safe Gemini invocation with:
- retry logic
- prompt sanitization
"""

# Convert any non-string into JSON-safe string

#if not isinstance(prompt, str):
# try:
# prompt = [Link](prompt)
# except:
# prompt = str(prompt)

# Attempt 3 retries
for attempt in range(3):
try:
response =
self.model_serving_client.[Link](model=self.llm_endpoint,
messages=prompt,response_format={"type": "json_object"})

class Result:
content = [Link][0].[Link]

return Result()

except Exception as e:
 print(f"[Gemini Error] attempt {attempt+1}: {e}")
[Link](1.5)

# FAIL SAFE: return empty JSON

class Result:
content = "{}"
return Result()

from typing import TypedDict, List, Dict, Any, Optional

class WorkflowState(TypedDict, total=False):

"""
Unified state object that flows through the entire LangGraph pipeline.
This stores everything Clean Slate, Watchtower, and Sentinel agents produce.
"""

# ------------------------------
# General / Router Info
# ------------------------------
user_request: Optional[str]
task_type: Optional[str] # "entity_resolution", "data_quality", "governance",
"combined"

# ------------------------------
# Clean Slate (Entity Resolution)
# ------------------------------
schema_profile: Optional[Dict[str, Any]]
candidate_pairs: Optional[List[Dict[str, Any]]]
match_scores: Optional[List[Dict[str, Any]]]
dedup_results: Optional[List[Dict[str, Any]]]
external_lookup_data: Optional[Dict[str, Any]]
feedback_notes: Optional[List[str]]

# ------------------------------
# Watchtower (Data Quality)
# ------------------------------
dq_profile: Optional[Dict[str, Any]]
dq_scores: Optional[Dict[str, float]]
dq_anomalies: Optional[List[Dict[str, Any]]]
remediation_suggestions: Optional[List[str]]

# ------------------------------
# Sentinel (Governance & Compliance)
# ------------------------------
policy_violations: Optional[List[Dict[str, Any]]]
governance_alerts: Optional[List[str]]
audit_log: Optional[List[Dict[str, Any]]]

# ------------------------------
# Internal Agent Communication (A2A)
# ------------------------------
agent_messages: Optional[List[Dict[str, Any]]]

# ------------------------------
# Error Handling
# ------------------------------
errors: Optional[List[str]]
import json, re

LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class DedupAgent:
"""
Clean Slate Agent #3:
Uses LLM logic to identify duplicate entities,
looping until no duplicates remain.
"""

def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

with open("prompts/clean_slate_prompts.md", "r", encoding="utf-8") as f:

self.prompt_file = [Link]()

def _extract_prompt(self):
marker = "## 3. Deduplication / Merge Agent Prompt"
section = self.prompt_file.split(marker)[1]
user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"
return prompt_text.strip()

def run(self, state: WorkflowState):

"""
LangGraph-compatible:
- Accepts only `state`
- Reads all required values from state
"""

records = [Link]("input_records", [])

match_results = [Link]("match_results", [])

# Nothing to dedupe if <2 records

if len(records) < 2:
state["dedup_summary"] = []
state["dedup_done"] = True
return state

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = [Link]("```")[1].replace(
"{{records}}", [Link](records, indent=2)
).replace(
"{{match_results}}", [Link](match_results, indent=2)
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
 ]

response = [Link](prompt)
output = [Link]

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from DedupAgent",
"raw_output": output
}

# Save dedup result

state["dedup_summary"] = parsed

# Check if dedup is complete

# Expecting a boolean in JSON: {"dedup_done": true/false}
state["dedup_done"] = [Link]("dedup_done", True)

# Logging
[Link]("agent_messages", []).append({
"agent": "DedupAgent",
"action": "entity_deduplication",
"info": "Deduplication step completed"
})

return state

class FeedbackAgent:
"""
Clean Slate Agent #4:
Stores SME / automated feedback into workflow state.

In this POC, no human is giving feedback, so we simply

record that the feedback stage was reached.
"""

def __init__(self):
pass

def run(self, state: WorkflowState):

"""
LangGraph-compatible:
- Accepts only state
- Writes feedback info to state
"""

# Initialize feedback container

[Link]("feedback", [])

# Auto feedback placeholder (for POC)

auto_feedback = {
"message": "Auto-feedback placeholder (no SME input in POC).",
"phase": "clean_slate",
"status": "completed"
}

state["feedback"].append(auto_feedback)
 # Log activity
[Link]("agent_messages", []).append({
"agent": "FeedbackAgent",
"action": "store_feedback",
"info": "Feedback stage completed"
})

return state

import json, re

LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class MatchingAgent:
"""
Clean Slate Agent #2:
Performs pairwise record matching using LLM reasoning.
Reads needed values from state instead of function arguments.
"""

def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

with open("prompts/clean_slate_prompts.md", "r", encoding="utf-8") as f:

self.prompt_file = [Link]()

def _extract_prompt(self):
marker = "## 2. Matching Agent Prompt"
section = self.prompt_file.split(marker)[1]
user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"
return prompt_text.strip()

def run(self, state: WorkflowState):

"""
LangGraph-compatible:
- Accepts only `state`
- Reads all required values from state
"""

# Get the dataset from state

records = [Link]("input_records", [])

# If there are fewer than 2 records, skip matching

if len(records) < 2:
state["match_results"] = []
return state

# Prepare prompt
template = self._extract_prompt()

role_text = [Link]("```")[0]
 user_text = [Link]("```")[1].replace(
"{{records}}",
[Link](records, indent=2)
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

# Call Gemini
response = [Link](prompt)
output = [Link]

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from MatchingAgent",
"raw_output": output
}

# Save results
state["match_results"] = parsed

# Log action
[Link]("agent_messages", []).append({
"agent": "MatchingAgent",
"action": "run_matching",
"info": "Record matching complete"
})

return state

import json, re

LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class SchemaProfilerAgent:
"""
Clean Slate Agent #1:
Profiles dataset schema using LLM reasoning.
Reads rows from state["input_records"].
"""

def __init__(self):
# Centralized model loading:
# Only gemini_chat.py controls which model is used.
[Link] = GeminiChat(LLM_ENDPOINT_NAME) # no direct model name inside
the agent

# Load prompt file once

with open("prompts/clean_slate_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()

def _extract_prompt(self):
marker = "## 1. Schema Profiler Agent Prompt"
section = self.prompt_file.split(marker)[1]
user_text = [Link]("```")[1]
 sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"

return prompt_text.strip()

def run(self, state: WorkflowState):

"""
Signature required by LangGraph:
- accepts only state
"""

# Pull input records from state

rows = [Link]("input_records", [])

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = [Link]("```")[1].replace(
"{{records}}",
[Link](rows, indent=2)
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

# Invoke Gemini
response = [Link](prompt)
output = [Link]

# Safe JSON parsing

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from SchemaProfilerAgent",
"raw_output": output
}

# Save into global workflow state

state["schema_profile"] = parsed

# Log activity
[Link]("agent_messages", []).append({
"agent": "SchemaProfilerAgent",
"action": "profile_schema",
"info": "Schema profile completed"
})

return state

import json, re
LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class AlertAgent:
"""
Sentinel Agent #3:
Generates alerts based on governance_report, policy_results,
anomalies, and dq_scores.

Outputs:
state["alerts"]
"""

def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

try:
with open("prompts/sentinel_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/sentinel_prompts.md!")
self.prompt_file = ""

def _extract_prompt(self):
marker = "## 3. Alerting Agent Prompt"

if marker not in self.prompt_file:

print(f"[Warning] Marker not found: {marker}")
return (
"You are the AlertAgent.\n"
"Review the following:\n"
" - policy_results\n"
" - governance_report\n"
" - anomalies\n"
" - dq_scores\n"
"Generate alerts if:\n"
" - Violations exist\n"
" - Severity is high\n"
" - Data issues impact critical systems\n"
"Return STRICT JSON:\n"
"{\n"
" \"alerts\": [...],\n"
" \"priority\": \"high|medium|low\"\n"
"}"
)

section = self.prompt_file.split(marker)[1]

if "```" not in section:

print(f"[Warning] Code block missing for: {marker}")
return "You are the AlertAgent. Generate alerts in JSON."

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"
 return prompt_text.strip()

def run(self, state: WorkflowState):

policies = [Link]("policy_results", {})

governance = [Link]("governance_report", {})
anomalies = [Link]("anomalies", {})
dq_scores = [Link]("dq_scores", {})

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = (
template
.split("```")[1]
.replace("{{policy_results}}", [Link](policies, indent=2))
.replace("{{governance_report}}", [Link](governance, indent=2))
.replace("{{anomalies}}", [Link](anomalies, indent=2))
.replace("{{dq_scores}}", [Link](dq_scores, indent=2))
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

response = [Link](prompt)
output = [Link]

try:
parsed = [Link](output)
except Exception:
parsed = {"error": "Invalid JSON from AlertAgent", "raw_output":
output}

state["alerts"] = parsed

[Link]("agent_messages", []).append({
"agent": "AlertAgent",
"action": "alert_generation",
"info": "Alerts generated"
})

return state

import json, re

LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class AuditSummaryAgent:
"""
Sentinel Agent #4:
Produces an audit summary of the entire pipeline execution.

Reads:
- schema_profile
- match_summary
- dq_profile
 - anomalies
- dq_scores
- remediation_plan
- governance_report
- alerts

Writes:
state["audit_summary"]
"""

def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

try:
with open("prompts/sentinel_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/sentinel_prompts.md!")
self.prompt_file = ""

def _extract_prompt(self):
marker = "## 4. Audit Agent Prompt"

if marker not in self.prompt_file:

print(f"[Warning] Marker not found: {marker}")
return (
"You are the AuditSummaryAgent.\n"
"Summarize the entire pipeline processing lifecycle including:\n"
"schema, matches, anomalies, DQ scores, remediation, governance,
alerts.\n"
"Return STRICT JSON:\n"
"{\n"
" \"summary\": \"...\",\n"
" \"compliance_status\": \"compliant|non_compliant\",\n"
" \"key_findings\": [...],\n"
" \"final_recommendations\": [...]\n"
"}"
)

section = self.prompt_file.split(marker)[1]
if "```" not in section:
print(f"[Warning] Code block missing: {marker}")
return "You are the AuditSummaryAgent. Produce a JSON audit summary."

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"

return prompt_text.strip()

def run(self, state: WorkflowState):

# Gather all important pipeline outputs

collected = {
 "schema_profile": [Link]("schema_profile", {}),
"match_summary": [Link]("match_summary", {}),
"dq_profile": [Link]("dq_profile", {}),
"anomalies": [Link]("anomalies", {}),
"dq_scores": [Link]("dq_scores", {}),
"remediation_plan": [Link]("remediation_plan", {}),
"governance_report": [Link]("governance_report", {}),
"alerts": [Link]("alerts", {})
}

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = [Link]("```")[1].replace(
"{{pipeline_outputs}}", [Link](collected, indent=2)
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

response = [Link](prompt)
output = [Link]

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from AuditSummaryAgent",
"raw_output": output
}

state["audit_summary"] = parsed

[Link]("agent_messages", []).append({
"agent": "AuditSummaryAgent",
"action": "audit_complete",
"info": "Full pipeline audit summary generated"
})

return state

import json, re

LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class GovernanceCheckerAgent:
"""
Sentinel Agent #2:
Performs governance compliance checking by combining:
- policies
- dq_profile
- entity resolution results (er_results)
Produces:
state["governance_report"]
"""

def __init__(self):
 [Link] = GeminiChat(LLM_ENDPOINT_NAME)

# Load prompt file

try:
with open("prompts/sentinel_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/sentinel_prompts.md!")
self.prompt_file = ""

def _extract_prompt(self):
"""
Extract the governance checking prompt section.
Falls back to a safe default if markers are missing.
"""

marker = "## 2. Governance Checker Agent Prompt"

if marker not in self.prompt_file:

print(f"[Warning] Marker not found: {marker}")
return (
"You are the GovernanceCheckerAgent.\n"
"Evaluate whether the dataset complies with governance standards.\
n"
"Inputs:\n"
" - policies: {{policies}}\n"
" - dq_profile: {{dq_profile}}\n"
" - entity resolution results: {{er_results}}\n\n"
"Return STRICT JSON:\n"
"{\n"
" \"governance_status\": \"compliant|non_compliant\",\n"
" \"violations\": [...],\n"
" \"severity\": \"low|medium|high\",\n"
" \"notes\": [...]\n"
"}"
)

section = self.prompt_file.split(marker)[1]

if "```" not in section:

print(f"[Warning] Code block missing after marker: {marker}")
return (
"You are the GovernanceCheckerAgent.\n"
"Evaluate governance compliance. Return JSON only."
)

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"

return prompt_text.strip()

def run(self, state: WorkflowState):

"""
 LangGraph-compatible method.
Reads:
state['policies']
state['dq_profile']
state['er_results'] (from earlier Clean Slate matching)
Writes:
state['governance_report']
"""

policies = [Link]("policies", [])

dq_profile = [Link]("dq_profile", {})
er_results = [Link]("match_summary", {})

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = ([Link]("```")[1]
.replace("{{policies}}", [Link](policies, indent=2))
.replace("{{dq_profile}}", [Link](dq_profile, indent=2))
.replace("{{er_results}}", [Link](er_results, indent=2))
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

# Call Gemini
response = [Link](prompt)
output = [Link]

# Parse JSON safely

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from GovernanceCheckerAgent",
"raw_output": output
}

# Save result
state["governance_report"] = parsed

# Log activity
[Link]("agent_messages", []).append({
"agent": "GovernanceCheckerAgent",
"action": "governance_eval",
"info": "Governance compliance evaluated"
})

return state

import json, re
LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class PolicyAgent:
"""
Sentinel Agent #1:
Checks data against defined governance policies.
 Produces:
state["policy_results"]
"""

def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

try:
with open("prompts/sentinel_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/sentinel_prompts.md!")
self.prompt_file = ""

def _extract_prompt(self):
"""
Extract the policy evaluation prompt from sentinel_prompts.md.
If the marker or code block is missing, fallback to safe default.
"""

marker = "## 1. Policy Agent Prompt"

if marker not in self.prompt_file:

print(f"[Warning] Marker not found: {marker}")
return (
"You are the PolicyAgent.\n"
"Evaluate whether the input data violates any governance rules.\n"
"Policies are inside {{policies}}.\n"
"Dataset is inside {{records}}.\n"
"Return STRICT JSON:\n"
"{\n"
" \"violations\": [...],\n"
" \"compliance\": true/false,\n"
" \"notes\": [...]\n"
"}"
)

section = self.prompt_file.split(marker)[1]

if "```" not in section:

print(f"[Warning] Code block missing after marker: {marker}")
return (
"You are the PolicyAgent.\n"
"Evaluate data against policies. Output JSON only."
)

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"

return prompt_text.strip()

def run(self, state: WorkflowState):

 """
LangGraph-compatible:
Reads:
state["input_records"]
state["policies"] (optional)
Writes:
state["policy_results"]
"""

records = [Link]("input_records", [])

policies = [Link]("policies", [
# Default POC policies (editable)
"No missing required fields.",
"No invalid email formats.",
"Age must be within realistic human range.",
"Duplicates should not appear after Clean Slate stage."
])

# Trim records to avoid overloading model

safe_rows = records[:5]

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = ([Link]("```")[1]
.replace("{{records}}", [Link](safe_rows,
indent=2))
.replace("{{policies}}", [Link](policies,
indent=2))
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

response = [Link](prompt)
output = [Link]

# Parse JSON
try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from PolicyAgent",
"raw_output": output
}

# Save results
state["policy_results"] = parsed

# Logging
[Link]("agent_messages", []).append({
"agent": "PolicyAgent",
"action": "policy_check",
"info": "Policy evaluation completed"
})

return state
import json, re
LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class AnomalyReasonerAgent:
"""
Watchtower Agent #2:
Performs anomaly detection using LLM reasoning.
Accepts only the workflow state and produces:
state["anomalies"]
"""

def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

# Load the prompt file

try:
with open("prompts/watchtower_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/watchtower_prompts.md!")
self.prompt_file = ""

def _extract_prompt(self):
"""
Extract anomaly detection prompt from markdown file.
If markers or code blocks are missing, fallback.
"""
marker = "## 3. Anomaly Reasoner Agent Prompt"

if marker not in self.prompt_file:

print(f"[Warning] Marker not found: {marker}")
return (
"You are the AnomalyReasonerAgent.\n"
"Analyze the dataset for anomalies, outliers, unexpected patterns,
"
"and quality issues. Output VALID JSON."
)

section = self.prompt_file.split(marker)[1]

if "```" not in section:

print(f"[Warning] Code block missing after marker: {marker}")
return (
"You are the AnomalyReasonerAgent.\n"
"Analyze the dataset for anomalies. Output JSON."
)

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"

return prompt_text.strip()
 def run(self, state: WorkflowState):
"""
LangGraph-compatible:
- Reads input records from state
- Generates anomaly summary
"""
records = [Link]("input_records", [])

# Prevent massive prompt sizes

safe_rows = records[:5]

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = [Link]("```")[1].replace(
"{{records}}",
[Link](safe_rows, indent=2)
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

response = [Link](prompt)
output = [Link]

# Attempt to parse JSON

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from AnomalyReasonerAgent",
"raw_output": output
}

# Save anomaly results

state["anomalies"] = parsed

# Log activity
[Link]("agent_messages", []).append({
"agent": "AnomalyReasonerAgent",
"action": "anomaly_detection",
"info": "Anomaly detection completed"
})

return state

import json, re
LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class DQProfileReasonerAgent:
"""
Watchtower Agent #1:
Performs Data Quality (DQ) profiling using LLM reasoning.
Reads records from state["input_records"] and produces:
state["dq_profile"]
Includes robust prompt extraction fallback for missing prompt markers.
"""
 def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

# Load prompt file

try:
with open("prompts/watchtower_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/watchtower_prompts.md !")
self.prompt_file = ""

def _extract_prompt(self):
"""
Extract the DQ Profile Reasoner prompt from the markdown file.
If markers or code blocks are missing, fallback to a safe default.
"""

marker = "## 1. Data Quality Profiler Agent Prompt"

# If the marker isn't found → fallback

if marker not in self.prompt_file:
print(f"\n[Warning] Marker not found in watchtower_prompts.md:
{marker}")
return (
"You are the DQProfileReasonerAgent.\n"
"Analyze the dataset for data-quality issues across dimensions such
as:\n"
"accuracy, completeness, freshness, consistency.\n"
"Output VALID JSON only."
)

# Extract section after marker

section = self.prompt_file.split(marker)[1]

# Ensure code block exists

if "```" not in section:
print(f"\n[Warning] Code block missing after marker: {marker}")
return (
"You are the DQProfileReasonerAgent.\n"
"Analyze the dataset for data-quality issues.\n"
"Output VALID JSON only."
)

# Extract the prompt inside the code block

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"
return prompt_text.strip()

def run(self, state: WorkflowState):

"""
LangGraph-compatible run method.
Accepts ONLY the state.
 Reads input records.
Generates DQ Profile.
"""

records = [Link]("input_records", [])

# Trim records to avoid overload on Gemini

safe_rows = records[:5]

template = self._extract_prompt()

# Build final prompt

role_text = [Link]("```")[0]
user_text = [Link]("```")[1].replace(
"{{records}}",
[Link](safe_rows, indent=2)
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

# Call Gemini through safe wrapper

response = [Link](prompt)
output = [Link]

# Parse JSON safely

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from DQProfileReasonerAgent",
"raw_output": output
}

# Store DQ Profile in state

state["dq_profile"] = parsed

# Log agent activity

[Link]("agent_messages", []).append({
"agent": "DQProfileReasonerAgent",
"action": "dq_profile_generated",
"info": "Data Quality profile created successfully"
})

return state

import json, re
LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class DQScoringAgent:
"""
Watchtower Agent #3:
Combines DQ Profile + Anomalies to compute an overall
DQ score and classification.
Produces state["dq_scores"].
"""
 def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

try:
with open("prompts/watchtower_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/watchtower_prompts.md!")
self.prompt_file = ""

def _extract_prompt(self):
"""
Extracts prompt for DQ scoring.
If missing, fallback.
"""

marker = "## 2. Data Quality Scoring Agent Prompt"

if marker not in self.prompt_file:

print(f"[Warning] Marker not found: {marker}")
return (
"You are the DQScoringAgent.\n"
"Combine the DQ profile and anomaly summary to produce:\n"
"{\n"
" \"overall_score\": <0-100>,\n"
" \"severity\": \"high|medium|low\",\n"
" \"key_issues\": [...],\n"
" \"recommendations\": [...]\n"
"}\n"
"Return VALID JSON only."
)

section = self.prompt_file.split(marker)[1]

if "```" not in section:

print(f"[Warning] Code block missing after marker: {marker}")
return (
"You are the DQScoringAgent.\n"
"Combine the DQ profile and anomalies and output JSON only."
)

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"

return prompt_text.strip()

def run(self, state: WorkflowState):

"""
LangGraph-compatible.
Reads:
state["dq_profile"]
state["anomalies"]
Writes:
 state["dq_scores"]
"""

dq_profile = [Link]("dq_profile", {})

anomalies = [Link]("anomalies", {})

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = ([Link]("```")[1]
.replace("{{dq_profile}}", [Link](dq_profile,
indent=2))
.replace("{{anomalies}}", [Link](anomalies,
indent=2))
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

# Gemini call
response = [Link](prompt)
output = [Link]

# Parse output safely

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from DQScoringAgent",
"raw_output": output
}

# Save results
state["dq_scores"] = parsed

# Logging
[Link]("agent_messages", []).append({
"agent": "DQScoringAgent",
"action": "dq_scoring",
"info": "DQ scoring completed"
})

return state

import json, re
LLM_ENDPOINT_NAME = "databricks-gemini-2-5-flash"

class RemediationAgent:
"""
Watchtower Agent #4:
Suggests remediation actions based on:
- detected anomalies
- computed DQ scores

Produces state["remediation_plan"].
"""
 def __init__(self):
[Link] = GeminiChat(LLM_ENDPOINT_NAME)

try:
with open("prompts/watchtower_prompts.md", "r", encoding="utf-8") as f:
self.prompt_file = [Link]()
except FileNotFoundError:
print("\n[ERROR] Missing prompts/watchtower_prompts.md!")
self.prompt_file = ""

def _extract_prompt(self):
"""
Extract remediation prompt.
Uses fallback if the marker or code block is missing.
"""

marker = "## 4. Remediation Agent Prompt"

if marker not in self.prompt_file:

print(f"[Warning] Marker not found: {marker}")
return (
"You are the RemediationAgent.\n"
"Based on the anomalies and data quality scores, generate a
remediation plan.\n"
"Return STRICT JSON with:\n"
"{\n"
" \"actions\": [...],\n"
" \"priority\": \"high|medium|low\",\n"
" \"notes\": [...]\n"
"}\n"
)

section = self.prompt_file.split(marker)[1]

if "```" not in section:

print(f"[Warning] Code block missing after marker: {marker}")
return (
"You are the RemediationAgent.\n"
"Generate JSON remediation plan using anomalies and dq_scores."
)

user_text = [Link]("```")[1]
sys_text = [Link]("```")[0]
match = [Link](r"\*\*Role:\*\*\s*(.+?)(?=\n\*\*|$)", sys_text,
[Link])
if match:
role_text = [Link](1).strip()

prompt_text = f"{role_text}```{user_text}"

return prompt_text.strip()

def run(self, state: WorkflowState):

"""
LangGraph-compatible run method.
Reads:
state["anomalies"]
state["dq_scores"]
Writes:
 state["remediation_plan"]
"""

anomalies = [Link]("anomalies", {})

dq_scores = [Link]("dq_scores", {})

template = self._extract_prompt()

role_text = [Link]("```")[0]
user_text = ([Link]("```")[1]
.replace("{{anomalies}}", [Link](anomalies, indent=2))
.replace("{{dq_scores}}", [Link](dq_scores, indent=2))
)

prompt = [
{"role": "system", "content": role_text},
{"role": "user", "content": user_text},
]

response = [Link](prompt)
output = [Link]

# Parse JSON safely

try:
parsed = [Link](output)
except Exception:
parsed = {
"error": "Invalid JSON from RemediationAgent",
"raw_output": output
}

# Save output
state["remediation_plan"] = parsed

[Link]("agent_messages", []).append({
"agent": "RemediationAgent",
"action": "remediation_generated",
"info": "Remediation plan created"
})

return state

from typing import Dict, Any

class RouterAgent:
"""
Router Agent:
- Reads user_request from state
- Classifies the task type
- Updates state.task_type accordingly
"""

@staticmethod
def classify_task(state: WorkflowState) -> WorkflowState:
user_request = ([Link]("user_request") or "").lower()

if any(keyword in user_request for keyword in ["entity", "dedup", "match",

"resolution"]):
state["task_type"] = "entity_resolution"
 elif any(keyword in user_request for keyword in ["dq", "quality",
"profile", "anomaly"]):
state["task_type"] = "data_quality"

elif any(keyword in user_request for keyword in ["govern", "policy",

"compliance", "sentinel"]):
state["task_type"] = "governance"

else:
# When unclear, fall back to combined mode
state["task_type"] = "combined"

# Log routing decision into agent messages

message_entry = {
"agent": "RouterAgent",
"decision": state["task_type"],
"info": f"Task classified based on request: {user_request}",
}

if "agent_messages" not in state:

state["agent_messages"] = []

state["agent_messages"].append(message_entry)

return state

from [Link] import StateGraph, END

def build_graph():
graph = StateGraph(WorkflowState)

# ---------------------------------------------------
# Register nodes
# ---------------------------------------------------

# Router
router = RouterAgent()
graph.add_node("router", router.classify_task)

# Clean Slate (Entity Resolution)

schema_profiler = SchemaProfilerAgent()
match_agent = MatchingAgent()
dedup_agent = DedupAgent()
feedback_agent = FeedbackAgent()

graph.add_node("schema_profiler", schema_profiler.run)
graph.add_node("match_agent", match_agent.run)
graph.add_node("dedup_agent", dedup_agent.run)
graph.add_node("feedback_agent", feedback_agent.run)

# Watchtower (Data Quality)

dq_profile_agent = DQProfileReasonerAgent()
anomaly_reasoner = AnomalyReasonerAgent()
dq_scoring_agent = DQScoringAgent()
remediation_agent = RemediationAgent()

graph.add_node("dq_profile_agent", dq_profile_agent.run)
graph.add_node("anomaly_reasoner", anomaly_reasoner.run)
graph.add_node("dq_scoring_agent", dq_scoring_agent.run)
graph.add_node("remediation_agent", remediation_agent.run)

# Sentinel (Governance)
policy_agent = PolicyAgent()
governance_checker = GovernanceCheckerAgent()
alert_agent = AlertAgent()
audit_agent = AuditSummaryAgent()

graph.add_node("policy_agent", policy_agent.run)
graph.add_node("governance_checker", governance_checker.run)
graph.add_node("alert_agent", alert_agent.run)
graph.add_node("audit_agent", audit_agent.run)

# ---------------------------------------------------
# Entry point
# ---------------------------------------------------
graph.set_entry_point("router")

# ---------------------------------------------------
# CONDITIONAL ROUTING (Router → Subsystems)
# ---------------------------------------------------
graph.add_conditional_edges(
"router",
lambda state: state["task_type"],
{
"entity_resolution": "schema_profiler",
"data_quality": "dq_profile_agent",
"governance": "policy_agent",
"combined": "schema_profiler" # full pipeline begins here
}
)

# ---------------------------------------------------
# ENTITY RESOLUTION PATH (with Loop)
# ---------------------------------------------------

# Schema → Matching
graph.add_edge("schema_profiler", "match_agent")

# Matching → Dedup
graph.add_edge("match_agent", "dedup_agent")

# Dedup → loop OR finish

def dedup_loop_condition(state: WorkflowState):
"""
If more record pairs remain, loop back to match_agent.
Otherwise, go to feedback_agent.
"""
remaining = [Link]("remaining_pairs", 0)
return "loop" if remaining > 0 else "finish"

graph.add_conditional_edges(
"dedup_agent",
dedup_loop_condition,
{
"loop": "match_agent",
"finish": "feedback_agent"
}
)

# Combined Mode Logic:

# Feedback → next subsystem (DQ) or → audit directly
def er_exit_condition(state: WorkflowState):
"""
combined mode → continue to DQ
ER only mode → go directly to audit
"""
task = [Link]("task_type", "")
return "combined" if task == "combined" else "audit"

graph.add_conditional_edges(
"feedback_agent",
er_exit_condition,
{
"combined": "dq_profile_agent",
"audit": "audit_agent"
}
)

# ---------------------------------------------------
# DATA QUALITY PATH
# ---------------------------------------------------

graph.add_edge("dq_profile_agent", "anomaly_reasoner")
graph.add_edge("anomaly_reasoner", "dq_scoring_agent")
graph.add_edge("dq_scoring_agent", "remediation_agent")

# remediation → next step depending on mode

def dq_exit_condition(state: WorkflowState):
"""
Combined → continue to governance
DQ only → go to audit
"""
task = [Link]("task_type", "")
return "governance" if task == "combined" else "audit"

graph.add_conditional_edges(
"remediation_agent",
dq_exit_condition,
{
"governance": "policy_agent",
"audit": "audit_agent"
}
)

# ---------------------------------------------------
# GOVERNANCE PATH
# ---------------------------------------------------

graph.add_edge("policy_agent", "governance_checker")
graph.add_edge("governance_checker", "alert_agent")
graph.add_edge("alert_agent", "audit_agent")

# ---------------------------------------------------
# FINAL END
# ---------------------------------------------------
 graph.add_edge("audit_agent", END)

return [Link]()

def run_pipeline(
user_request: str,
records: list = None,
policies: dict = None,
feedback: str = None
):
"""
Runs any of the supported modes:
- entity_resolution
- data_quality
- governance
- combined (ER → DQ → Governance → Audit)
"""

# -----------------------------------
# Build graph
# -----------------------------------
graph = build_graph()

# -----------------------------------
# Initialize workflow state
# -----------------------------------
state = WorkflowState()
state["user_request"] = user_request

# If we get a dataset (for ER or DQ)

if records:
state["input_records"] = records

# Calculate total record pairs for ER loop

total = len(records)
pair_count = (total * (total - 1)) // 2
state["remaining_pairs"] = max(0, pair_count - 1)

# If governance policies provided

if policies:
state["policies"] = policies

# If cleaning feedback provided

if feedback:
state["pending_feedback"] = feedback

# -----------------------------------
# RUN THE GRAPH
# -----------------------------------
final_state = [Link](state)

return final_state

# =================================================
# TEST RUNNER
# =================================================
# -----------------------------------
# TEST DATASET (Records)
# -----------------------------------
#sample_records = [
# {"name": "John Doe", "email": "john@[Link]", "phone": "1234567890",
"age": 32},
# {"name": "Jon Doe", "email": "jon.d@[Link]", "phone": "1234567890",
"age": 32},
# {"name": "Jane Doe", "email": "jane@example", "phone": "", "age": None}
#]
table_data = [Link]("`hackthon-aethermind`.ingestion.table_a")
sample_records = [[Link]() for row in table_data.collect()]
# -----------------------------------
# TEST GOVERNANCE POLICIES
# -----------------------------------
sample_policies = {
"pii_columns": ["name", "email", "phone"],
"required_columns": ["name", "email"],
"rules": [
{"rule": "pii_not_null", "columns": ["name", "email"]},
{"rule": "email_validity", "columns": ["email"]},
{"rule": "phone_validity", "columns": ["phone"]},
{"rule": "non_empty_string", "columns": ["name"]},
{"rule": "completeness_threshold", "threshold": 0.70},
{"rule": "freshness_limit_days", "limit": 60}
]
}

# -----------------------------------
# RUN FULL PIPELINE (ER → DQ → Governance)
# -----------------------------------
result = run_pipeline(
user_request="data_quality", # FULL WORKFLOW
records=sample_records, # Dataset to test
policies=sample_policies # Governance rules
)

# -----------------------------------
# PRINT FINAL STATE
# -----------------------------------
import json
print("\n================ FINAL STATE ================\n")
print([Link](result, indent=2, default=str))