""" Agent result models for mapping JSON responses from lineage agents. This module contains classes for representing the structured results returned by lineage analysis agents in a type-safe manner. """ from typing import Dict, Any, List, Optional class SchemaField: """Schema field configuration for agent results""" def __init__(self, name: str, type: str, description: str): self.name = name self.type = type self.description = description @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'SchemaField': """Create SchemaField from dictionary""" return cls( name=data.get('name', ''), type=data.get('type', ''), description=data.get('description', '') ) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { 'name': self.name, 'type': self.type, 'description': self.description } class Schema: """Schema configuration for agent results""" def __init__(self, fields: List[SchemaField]): self.fields = fields @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'Schema': """Create Schema from dictionary""" fields = [SchemaField.from_dict(field) for field in data.get('fields', [])] return cls(fields=fields) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { 'fields': [field.to_dict() for field in self.fields] } class Transformation: """Transformation configuration for column lineage""" def __init__(self, type: str, subtype: str, description: str, masking: bool = False): self.type = type self.subtype = subtype self.description = description self.masking = masking @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'Transformation': """Create Transformation from dictionary""" return cls( type=data.get('type', ''), subtype=data.get('subtype', ''), description=data.get('description', ''), masking=data.get('masking', False) ) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { 'type': self.type, 'subtype': self.subtype, 'description': self.description, 'masking': self.masking } class InputField: """Input field configuration for column lineage""" def __init__(self, namespace: str, name: str, field: str, transformations: List[Transformation]): self.namespace = namespace self.name = name self.field = field self.transformations = transformations @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'InputField': """Create InputField from dictionary""" transformations = [Transformation.from_dict(t) for t in data.get('transformations', [])] return cls( namespace=data.get('namespace', ''), name=data.get('name', ''), field=data.get('field', ''), transformations=transformations ) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { 'namespace': self.namespace, 'name': self.name, 'field': self.field, 'transformations': [t.to_dict() for t in self.transformations] } class ColumnLineageField: """Column lineage field configuration""" def __init__(self, input_fields: List[InputField]): self.input_fields = input_fields @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'ColumnLineageField': """Create ColumnLineageField from dictionary""" input_fields = [InputField.from_dict(field) for field in data.get('inputFields', [])] return cls(input_fields=input_fields) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { 'inputFields': [field.to_dict() for field in self.input_fields] } class ColumnLineage: """Column lineage configuration""" def __init__(self, fields: Dict[str, ColumnLineageField]): self.fields = fields @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'ColumnLineage': """Create ColumnLineage from dictionary""" fields = { field_name: ColumnLineageField.from_dict(field_data) for field_name, field_data in data.get('fields', {}).items() } return cls(fields=fields) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { 'fields': { field_name: field_data.to_dict() for field_name, field_data in self.fields.items() } } class InputFacets: """Input facets configuration for agent results""" def __init__(self, schema: Optional[Schema] = None): self.schema = schema @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'InputFacets': """Create InputFacets from dictionary""" schema = Schema.from_dict(data.get('schema', {})) if data.get('schema') else None return cls(schema=schema) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" result = {} if self.schema: result['schema'] = self.schema.to_dict() return result class Input: """Input configuration for agent results""" def __init__(self, namespace: str, name: str, facets: Optional[InputFacets] = None): self.namespace = namespace self.name = name self.facets = facets @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'Input': """Create Input from dictionary""" facets = InputFacets.from_dict(data.get('facets', {})) if data.get('facets') else None return cls( namespace=data.get('namespace', ''), name=data.get('name', ''), facets=facets ) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" result = { 'namespace': self.namespace, 'name': self.name } if self.facets: result['facets'] = self.facets.to_dict() return result class OutputFacets: """Output facets configuration for agent results""" def __init__(self, column_lineage: Optional[ColumnLineage] = None): self.column_lineage = column_lineage @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'OutputFacets': """Create OutputFacets from dictionary""" column_lineage = ColumnLineage.from_dict(data.get('columnLineage', {})) if data.get('columnLineage') else None return cls(column_lineage=column_lineage) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" result = {} if self.column_lineage: result['columnLineage'] = self.column_lineage.to_dict() return result class Output: """Output configuration for agent results""" def __init__(self, namespace: str, name: str, facets: Optional[OutputFacets] = None): self.namespace = namespace self.name = name self.facets = facets @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'Output': """Create Output from dictionary""" facets = OutputFacets.from_dict(data.get('facets', {})) if data.get('facets') else None return cls( namespace=data.get('namespace', ''), name=data.get('name', ''), facets=facets ) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" result = { 'namespace': self.namespace, 'name': self.name } if self.facets: result['facets'] = self.facets.to_dict() return result class AgentResult: """Main result class for agent lineage analysis""" def __init__(self, inputs: List[Input], outputs: List[Output]): self.inputs = inputs self.outputs = outputs @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'AgentResult': """Create AgentResult from dictionary""" inputs = [Input.from_dict(input_data) for input_data in data.get('inputs', [])] outputs = [Output.from_dict(output_data) for output_data in data.get('outputs', [])] return cls(inputs=inputs, outputs=outputs) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { 'inputs': [input_obj.to_dict() for input_obj in self.inputs], 'outputs': [output_obj.to_dict() for output_obj in self.outputs] } def __str__(self) -> str: """String representation""" return f"AgentResult(inputs={len(self.inputs)}, outputs={len(self.outputs)})" def __repr__(self) -> str: """Detailed string representation""" return f"AgentResult(inputs={self.inputs}, outputs={self.outputs})"