-
Notifications
You must be signed in to change notification settings - Fork 407
Expand file tree
/
Copy pathmodels.py
More file actions
194 lines (150 loc) · 6.21 KB
/
models.py
File metadata and controls
194 lines (150 loc) · 6.21 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
from datetime import datetime
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from typing_extensions import override
else:
def override(func): # noqa: ANN001, ANN201
"""Override decorator."""
return func
from a2a.types.a2a_pb2 import Artifact, Message, TaskStatus
try:
from sqlalchemy import JSON, DateTime, Index, LargeBinary, String
from sqlalchemy.orm import (
DeclarativeBase,
Mapped,
declared_attr,
mapped_column,
)
except ImportError as e:
raise ImportError(
'Database models require SQLAlchemy. '
'Install with one of: '
"'pip install a2a-sdk[postgresql]', "
"'pip install a2a-sdk[mysql]', "
"'pip install a2a-sdk[sqlite]', "
"or 'pip install a2a-sdk[sql]'"
) from e
# Base class for all database models
class Base(DeclarativeBase):
"""Base class for declarative models in A2A SDK."""
# TaskMixin that can be used with any table name
class TaskMixin:
"""Mixin providing standard task columns with proper type handling."""
id: Mapped[str] = mapped_column(String(36), primary_key=True, index=True)
context_id: Mapped[str] = mapped_column(String(36), nullable=False)
kind: Mapped[str] = mapped_column(
String(16), nullable=False, default='task'
)
owner: Mapped[str] = mapped_column(String(255), nullable=True)
last_updated: Mapped[datetime | None] = mapped_column(
DateTime, nullable=True
)
status: Mapped[TaskStatus] = mapped_column(JSON, nullable=False)
artifacts: Mapped[list[Artifact] | None] = mapped_column(
JSON, nullable=True
)
history: Mapped[list[Message] | None] = mapped_column(JSON, nullable=True)
protocol_version: Mapped[str | None] = mapped_column(
String(16), nullable=True
)
# Using declared_attr to avoid conflict with Pydantic's metadata
@declared_attr
@classmethod
def task_metadata(cls) -> Mapped[dict[str, Any] | None]:
"""Define the 'metadata' column, avoiding name conflicts with Pydantic."""
return mapped_column(JSON, nullable=True, name='metadata')
@override
def __repr__(self) -> str:
"""Return a string representation of the task."""
return (
f'<{self.__class__.__name__}(id="{self.id}", '
f'context_id="{self.context_id}", status="{self.status}")>'
)
@declared_attr.directive
@classmethod
def __table_args__(cls) -> tuple[Any, ...]:
"""Define a composite index (owner, last_updated) for each table that uses the mixin."""
tablename = getattr(cls, '__tablename__', 'tasks')
return (
Index(
f'idx_{tablename}_owner_last_updated', 'owner', 'last_updated'
),
)
def create_task_model(
table_name: str = 'tasks', base: type[DeclarativeBase] = Base
) -> type:
"""Create a TaskModel class with a configurable table name.
Args:
table_name: Name of the database table. Defaults to 'tasks'.
base: Base declarative class to use. Defaults to the SDK's Base class.
Returns:
TaskModel class with the specified table name.
Example:
.. code-block:: python
# Create a task model with default table name
TaskModel = create_task_model()
# Create a task model with custom table name
CustomTaskModel = create_task_model('my_tasks')
# Use with a custom base
from myapp.database import Base as MyBase
TaskModel = create_task_model('tasks', MyBase)
"""
class TaskModel(TaskMixin, base): # type: ignore
__tablename__ = table_name
@override
def __repr__(self) -> str:
"""Return a string representation of the task."""
return (
f'<TaskModel[{table_name}](id="{self.id}", '
f'context_id="{self.context_id}", status="{self.status}")>'
)
# Set a dynamic name for better debugging
TaskModel.__name__ = f'TaskModel_{table_name}'
TaskModel.__qualname__ = f'TaskModel_{table_name}'
return TaskModel
# Default TaskModel for backward compatibility
class TaskModel(TaskMixin, Base):
"""Default task model with standard table name."""
__tablename__ = 'tasks'
# PushNotificationConfigMixin that can be used with any table name
class PushNotificationConfigMixin:
"""Mixin providing standard push notification config columns."""
task_id: Mapped[str] = mapped_column(String(36), primary_key=True)
config_id: Mapped[str] = mapped_column(String(255), primary_key=True)
config_data: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
owner: Mapped[str] = mapped_column(String(255), nullable=True, index=True)
protocol_version: Mapped[str | None] = mapped_column(
String(16), nullable=True
)
@override
def __repr__(self) -> str:
"""Return a string representation of the push notification config."""
return (
f'<{self.__class__.__name__}(task_id="{self.task_id}", '
f'config_id="{self.config_id}")>'
)
def create_push_notification_config_model(
table_name: str = 'push_notification_configs',
base: type[DeclarativeBase] = Base,
) -> type:
"""Create a PushNotificationConfigModel class with a configurable table name."""
class PushNotificationConfigModel(PushNotificationConfigMixin, base): # type: ignore
__tablename__ = table_name
@override
def __repr__(self) -> str:
"""Return a string representation of the push notification config."""
return (
f'<PushNotificationConfigModel[{table_name}]('
f'task_id="{self.task_id}", config_id="{self.config_id}")>'
)
PushNotificationConfigModel.__name__ = (
f'PushNotificationConfigModel_{table_name}'
)
PushNotificationConfigModel.__qualname__ = (
f'PushNotificationConfigModel_{table_name}'
)
return PushNotificationConfigModel
# Default PushNotificationConfigModel for backward compatibility
class PushNotificationConfigModel(PushNotificationConfigMixin, Base):
"""Default push notification config model with standard table name."""
__tablename__ = 'push_notification_configs'