"""
AI analysis helper for Data Comparison Agent.
Provides structured output types and analysis functions using Bedrock.
"""
from typing import TypedDict, Annotated, Optional
from config import Config
import json
[docs]
class KeySuggestionResult(TypedDict):
"""Result from AI key suggestion."""
suggested_keys: Annotated[list[str], "Suggested primary key columns"]
confidence: Annotated[float, "Confidence score 0.0-1.0"]
reasoning: Annotated[str, "Explanation for suggestion"]
[docs]
class SchemaMatchResult(TypedDict):
"""Result from schema matching analysis."""
exact_matches: Annotated[list[str], "Columns with exact name match"]
semantic_matches: Annotated[dict[str, str], "Source->target semantic mappings"]
source_only: Annotated[list[str], "Columns only in source"]
target_only: Annotated[list[str], "Columns only in target"]
[docs]
class ComparisonSummaryResult(TypedDict):
"""Final comparison summary from AI analysis."""
summary: Annotated[str, "Human-readable summary"]
key_findings: Annotated[list[str], "List of key findings"]
recommendations: Annotated[list[str], "List of recommendations"]
severity: Annotated[str, "low, medium, high, or critical"]
ai_determined_keys: Annotated[list[str], "Keys AI chose if not user-provided"]
ai_determined_keys_confidence: Annotated[float, "Confidence 0.0-1.0"]
ai_semantic_mappings: Annotated[dict[str, str], "Column mappings"]
ai_mapping_rationale: Annotated[str, "Explanation for mappings"]
[docs]
def get_llm_client(model_id: str = None, temperature: float = 0.1):
"""Get ChatBedrockConverse client for structured output."""
from langchain_aws import ChatBedrockConverse
from botocore.config import Config as BotocoreConfig
config = BotocoreConfig(read_timeout=600)
return ChatBedrockConverse(
model=model_id or Config.CLAUDE_BEDROCK_SONNET,
temperature=temperature,
max_tokens=4096,
config=config,
)
[docs]
def suggest_primary_keys(
source_columns: list[str],
target_columns: list[str],
sample_data: dict,
model_id: str = None,
) -> KeySuggestionResult:
"""
Use AI to suggest primary key columns.
Args:
source_columns: List of source column names
target_columns: List of target column names
sample_data: Sample data with distinct value counts
model_id: Optional model override
Returns:
KeySuggestionResult with suggested keys and confidence
"""
from langchain_core.messages import HumanMessage
llm = get_llm_client(model_id).with_structured_output(KeySuggestionResult)
prompt = f"""Analyze these columns and suggest primary key candidates.
Source columns: {source_columns}
Target columns: {target_columns}
Sample data statistics:
{json.dumps(sample_data, indent=2, default=str)}
Consider:
1. Columns with high cardinality (many distinct values)
2. Columns with no nulls
3. Common naming patterns (id, _id, key, code)
4. Uniqueness within the dataset
Suggest the best primary key column(s) for joining these tables."""
return llm.invoke([HumanMessage(content=prompt)])
[docs]
def analyze_schema_matches(
source_columns: list[str], target_columns: list[str], model_id: str = None
) -> SchemaMatchResult:
"""
Use AI to detect semantic column mappings.
Args:
source_columns: List of source column names
target_columns: List of target column names
model_id: Optional model override
Returns:
SchemaMatchResult with exact and semantic matches
"""
from langchain_core.messages import HumanMessage
llm = get_llm_client(model_id).with_structured_output(SchemaMatchResult)
prompt = f"""Compare these column sets and identify matches.
Source columns: {source_columns}
Target columns: {target_columns}
Identify:
1. Exact matches (same name)
2. Semantic matches (e.g., app_id ≈ application_id, loc ≈ location, locID ≈ location_id)
3. Columns only in source
4. Columns only in target
Be conservative with semantic matches - only suggest high-confidence mappings."""
return llm.invoke([HumanMessage(content=prompt)])
[docs]
def generate_comparison_summary(
schema_diff: dict,
row_counts: dict,
duplicates: dict,
row_level_diff: dict,
custom_prompt: str = None,
key_suggestion: KeySuggestionResult = None,
schema_matches: SchemaMatchResult = None,
model_id: str = None,
) -> ComparisonSummaryResult:
"""
Generate final AI summary of comparison results.
Args:
schema_diff: Schema comparison results
row_counts: Row count comparison
duplicates: Duplicate detection results
row_level_diff: Row-level differences
custom_prompt: Optional additional analysis instructions
key_suggestion: AI key suggestion if used
schema_matches: AI schema matches if detected
model_id: Optional model override
Returns:
ComparisonSummaryResult with findings and recommendations
"""
from langchain_core.messages import HumanMessage
llm = get_llm_client(model_id).with_structured_output(ComparisonSummaryResult)
context = f"""Analyze these data comparison results:
## Schema Comparison
{json.dumps(schema_diff, indent=2, default=str)}
## Row Counts
{json.dumps(row_counts, indent=2, default=str)}
## Duplicates Found
{json.dumps(duplicates, indent=2, default=str)}
## Row-Level Differences (sample)
{json.dumps(row_level_diff, indent=2, default=str)}
"""
if key_suggestion:
context += f"""
## AI-Determined Keys
Keys: {key_suggestion.get('suggested_keys', [])}
Confidence: {key_suggestion.get('confidence', 0)}
Reasoning: {key_suggestion.get('reasoning', '')}
"""
if schema_matches:
context += f"""
## AI-Detected Column Mappings
{json.dumps(schema_matches, indent=2, default=str)}
"""
if custom_prompt:
context += f"""
## Additional Analysis Required
{custom_prompt}
"""
prompt = (
context
+ """
Provide a comprehensive summary including:
1. Overall data quality assessment
2. Key findings (most important issues)
3. Recommendations for data engineers
4. Severity rating (low/medium/high/critical)
Be specific and actionable in your recommendations."""
)
return llm.invoke([HumanMessage(content=prompt)])