11from datetime import datetime
2- from typing import TYPE_CHECKING , Any , Generic , TypeVar
2+ from typing import TYPE_CHECKING , Any
33
44
55if TYPE_CHECKING :
@@ -11,24 +11,14 @@ def override(func): # noqa: ANN001, ANN201
1111 return func
1212
1313
14- from google .protobuf .json_format import MessageToDict , ParseDict
15- from google .protobuf .message import Message as ProtoMessage
16- from pydantic import BaseModel
17-
18- from a2a .types .a2a_pb2 import Artifact , Message , TaskStatus
19-
20-
2114try :
22- from sqlalchemy import JSON , DateTime , Dialect , Index , LargeBinary , String
15+ from sqlalchemy import JSON , DateTime , Index , LargeBinary , String
2316 from sqlalchemy .orm import (
2417 DeclarativeBase ,
2518 Mapped ,
2619 declared_attr ,
2720 mapped_column ,
2821 )
29- from sqlalchemy .types import (
30- TypeDecorator ,
31- )
3222except ImportError as e :
3323 raise ImportError (
3424 'Database models require SQLAlchemy. '
@@ -40,101 +30,6 @@ def override(func): # noqa: ANN001, ANN201
4030 ) from e
4131
4232
43- T = TypeVar ('T' )
44-
45-
46- class PydanticType (TypeDecorator [T ], Generic [T ]):
47- """SQLAlchemy type that handles Pydantic model and Protobuf message serialization."""
48-
49- impl = JSON
50- cache_ok = True
51-
52- def __init__ (self , pydantic_type : type [T ], ** kwargs : dict [str , Any ]):
53- """Initialize the PydanticType.
54-
55- Args:
56- pydantic_type: The Pydantic model or Protobuf message type to handle.
57- **kwargs: Additional arguments for TypeDecorator.
58- """
59- self .pydantic_type = pydantic_type
60- super ().__init__ (** kwargs )
61-
62- def process_bind_param (
63- self , value : T | None , dialect : Dialect
64- ) -> dict [str , Any ] | None :
65- """Convert Pydantic model or Protobuf message to a JSON-serializable dictionary for the database."""
66- if value is None :
67- return None
68- if isinstance (value , ProtoMessage ):
69- return MessageToDict (value , preserving_proto_field_name = False )
70- if isinstance (value , BaseModel ):
71- return value .model_dump (mode = 'json' )
72- return value # type: ignore[return-value]
73-
74- def process_result_value (
75- self , value : dict [str , Any ] | None , dialect : Dialect
76- ) -> T | None :
77- """Convert a JSON-like dictionary from the database back to a Pydantic model or Protobuf message."""
78- if value is None :
79- return None
80- # Check if it's a protobuf message class
81- if isinstance (self .pydantic_type , type ) and issubclass (
82- self .pydantic_type , ProtoMessage
83- ):
84- return ParseDict (value , self .pydantic_type ()) # type: ignore[return-value]
85- # Assume it's a Pydantic model
86- return self .pydantic_type .model_validate (value ) # type: ignore[attr-defined]
87-
88-
89- class PydanticListType (TypeDecorator , Generic [T ]):
90- """SQLAlchemy type that handles lists of Pydantic models or Protobuf messages."""
91-
92- impl = JSON
93- cache_ok = True
94-
95- def __init__ (self , pydantic_type : type [T ], ** kwargs : dict [str , Any ]):
96- """Initialize the PydanticListType.
97-
98- Args:
99- pydantic_type: The Pydantic model or Protobuf message type for items in the list.
100- **kwargs: Additional arguments for TypeDecorator.
101- """
102- self .pydantic_type = pydantic_type
103- super ().__init__ (** kwargs )
104-
105- def process_bind_param (
106- self , value : list [T ] | None , dialect : Dialect
107- ) -> list [dict [str , Any ]] | None :
108- """Convert a list of Pydantic models or Protobuf messages to a JSON-serializable list for the DB."""
109- if value is None :
110- return None
111- result : list [dict [str , Any ]] = []
112- for item in value :
113- if isinstance (item , ProtoMessage ):
114- result .append (
115- MessageToDict (item , preserving_proto_field_name = False )
116- )
117- elif isinstance (item , BaseModel ):
118- result .append (item .model_dump (mode = 'json' ))
119- else :
120- result .append (item ) # type: ignore[arg-type]
121- return result
122-
123- def process_result_value (
124- self , value : list [dict [str , Any ]] | None , dialect : Dialect
125- ) -> list [T ] | None :
126- """Convert a JSON-like list from the DB back to a list of Pydantic models or Protobuf messages."""
127- if value is None :
128- return None
129- # Check if it's a protobuf message class
130- if isinstance (self .pydantic_type , type ) and issubclass (
131- self .pydantic_type , ProtoMessage
132- ):
133- return [ParseDict (item , self .pydantic_type ()) for item in value ] # type: ignore[misc]
134- # Assume it's a Pydantic model
135- return [self .pydantic_type .model_validate (item ) for item in value ] # type: ignore[attr-defined]
136-
137-
13833# Base class for all database models
13934class Base (DeclarativeBase ):
14035 """Base class for declarative models in A2A SDK."""
@@ -153,14 +48,12 @@ class TaskMixin:
15348 last_updated : Mapped [datetime | None ] = mapped_column (
15449 DateTime , nullable = True
15550 )
156-
157- # Properly typed Pydantic fields with automatic serialization
158- status : Mapped [TaskStatus ] = mapped_column (PydanticType (TaskStatus ))
159- artifacts : Mapped [list [Artifact ] | None ] = mapped_column (
160- PydanticListType (Artifact ), nullable = True
51+ status : Mapped [dict [str , Any ] | None ] = mapped_column (JSON , nullable = True )
52+ artifacts : Mapped [list [dict [str , Any ]] | None ] = mapped_column (
53+ JSON , nullable = True
16154 )
162- history : Mapped [list [Message ] | None ] = mapped_column (
163- PydanticListType ( Message ) , nullable = True
55+ history : Mapped [list [dict [ str , Any ] ] | None ] = mapped_column (
56+ JSON , nullable = True
16457 )
16558 protocol_version : Mapped [str | None ] = mapped_column (
16659 String (16 ), nullable = True
0 commit comments