jaf-py 2.5.10__py3-none-any.whl → 2.5.11__py3-none-any.whl

jaf/a2a/memory/types.py

@@ -16,14 +16,16 @@ from ...memory.types import Failure, Success
 from ..types import A2ATask, TaskState
 
 # Generic types for A2A results
-T = TypeVar('T')
-E = TypeVar('E', bound='A2ATaskError')
+T = TypeVar("T")
+E = TypeVar("E", bound="A2ATaskError")
 
 # A2A Task storage and retrieval types
 
+
 @dataclass(frozen=True)
 class A2ATaskQuery:
     """Query parameters for searching A2A tasks in memory providers"""
+
     task_id: Optional[str] = None
     context_id: Optional[str] = None
     state: Optional[TaskState] = None
@@ -34,9 +36,11 @@ class A2ATaskQuery:
     include_history: bool = True
     include_artifacts: bool = True
 
+
 @dataclass(frozen=True)
 class A2ATaskStorage:
     """Internal storage representation of an A2A task"""
+
     task_id: str
     context_id: str
     state: TaskState
@@ -49,35 +53,32 @@ class A2ATaskStorage:
 
     def __post_init__(self):
         if self.created_at is None:
-            object.__setattr__(self, 'created_at', datetime.now())
+            object.__setattr__(self, "created_at", datetime.now())
         if self.updated_at is None:
-            object.__setattr__(self, 'updated_at', datetime.now())
+            object.__setattr__(self, "updated_at", datetime.now())
+
 
 class A2ATaskProvider(Protocol):
     """
     Protocol defining the interface for A2A task storage providers.
-    
+
     This extends the memory provider pattern for A2A-specific task persistence,
     providing optimized operations for task lifecycle management.
     """
 
     async def store_task(
-        self,
-        task: A2ATask,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> 'A2AResult[None]':
+        self, task: A2ATask, metadata: Optional[Dict[str, Any]] = None
+    ) -> "A2AResult[None]":
         """Store a new A2A task"""
         ...
 
-    async def get_task(self, task_id: str) -> 'A2AResult[Optional[A2ATask]]':
+    async def get_task(self, task_id: str) -> "A2AResult[Optional[A2ATask]]":
         """Retrieve a task by ID"""
         ...
 
     async def update_task(
-        self,
-        task: A2ATask,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> 'A2AResult[None]':
+        self, task: A2ATask, metadata: Optional[Dict[str, Any]] = None
+    ) -> "A2AResult[None]":
         """Update an existing task"""
         ...
 
@@ -86,51 +87,52 @@ class A2ATaskProvider(Protocol):
         task_id: str,
         state: TaskState,
         status_message: Optional[Any] = None,
-        timestamp: Optional[str] = None
-    ) -> 'A2AResult[None]':
+        timestamp: Optional[str] = None,
+    ) -> "A2AResult[None]":
         """Update task status only (optimized for frequent status changes)"""
         ...
 
-    async def find_tasks(self, query: A2ATaskQuery) -> 'A2AResult[List[A2ATask]]':
+    async def find_tasks(self, query: A2ATaskQuery) -> "A2AResult[List[A2ATask]]":
         """Search tasks by query parameters"""
         ...
 
     async def get_tasks_by_context(
-        self,
-        context_id: str,
-        limit: Optional[int] = None
-    ) -> 'A2AResult[List[A2ATask]]':
+        self, context_id: str, limit: Optional[int] = None
+    ) -> "A2AResult[List[A2ATask]]":
         """Get tasks by context ID"""
         ...
 
-    async def delete_task(self, task_id: str) -> 'A2AResult[bool]':
+    async def delete_task(self, task_id: str) -> "A2AResult[bool]":
         """Delete a task and return True if it existed"""
         ...
 
-    async def delete_tasks_by_context(self, context_id: str) -> 'A2AResult[int]':
+    async def delete_tasks_by_context(self, context_id: str) -> "A2AResult[int]":
         """Delete tasks by context ID and return count deleted"""
         ...
 
-    async def cleanup_expired_tasks(self) -> 'A2AResult[int]':
+    async def cleanup_expired_tasks(self) -> "A2AResult[int]":
         """Clean up expired tasks and return count deleted"""
         ...
 
-    async def get_task_stats(self, context_id: Optional[str] = None) -> 'A2AResult[Dict[str, Any]]':
+    async def get_task_stats(self, context_id: Optional[str] = None) -> "A2AResult[Dict[str, Any]]":
         """Get task statistics"""
         ...
 
-    async def health_check(self) -> 'A2AResult[Dict[str, Any]]':
+    async def health_check(self) -> "A2AResult[Dict[str, Any]]":
         """Check provider health and return status information"""
         ...
 
-    async def close(self) -> 'A2AResult[None]':
+    async def close(self) -> "A2AResult[None]":
         """Close/cleanup the provider"""
         ...
 
+
 # Configuration models for A2A task storage
 
+
 class A2ATaskMemoryConfig(BaseModel):
     """Base configuration for A2A task memory"""
+
     model_config = {"frozen": True}
 
     type: str
@@ -141,15 +143,19 @@ class A2ATaskMemoryConfig(BaseModel):
     enable_history: bool = Field(default=True)  # Store task history
     enable_artifacts: bool = Field(default=True)  # Store task artifacts
 
+
 class A2AInMemoryTaskConfig(A2ATaskMemoryConfig):
     """Configuration for A2A in-memory task provider"""
+
     model_config = {"frozen": True}
 
     type: str = Field(default="memory", frozen=True)
     max_tasks_per_context: int = Field(default=1000)
 
+
 class A2ARedisTaskConfig(A2ATaskMemoryConfig):
     """Configuration for A2A Redis task provider"""
+
     model_config = {"frozen": True}
 
     type: str = Field(default="redis", frozen=True)
@@ -158,8 +164,10 @@ class A2ARedisTaskConfig(A2ATaskMemoryConfig):
     password: Optional[str] = None
     db: int = Field(default=0, ge=0)
 
+
 class A2APostgresTaskConfig(A2ATaskMemoryConfig):
     """Configuration for A2A PostgreSQL task provider"""
+
     model_config = {"frozen": True}
 
     type: str = Field(default="postgres", frozen=True)
@@ -172,14 +180,24 @@ class A2APostgresTaskConfig(A2ATaskMemoryConfig):
     table_name: str = Field(default="a2a_tasks")
     max_connections: int = Field(default=10, ge=1)
 
+
 # Union type for all A2A task provider configurations
 A2ATaskProviderConfig = Union[A2AInMemoryTaskConfig, A2ARedisTaskConfig, A2APostgresTaskConfig]
 
 # Error types specific to A2A task storage
 
+
 class A2ATaskError:
     """Base class for A2A task-related errors"""
-    def __init__(self, message: str, code: str, provider: str, task_id: Optional[str] = None, cause: Optional[Exception] = None):
+
+    def __init__(
+        self,
+        message: str,
+        code: str,
+        provider: str,
+        task_id: Optional[str] = None,
+        cause: Optional[Exception] = None,
+    ):
         self.message = message
         self.code = code
         self.provider = provider
@@ -189,23 +207,44 @@ class A2ATaskError:
     def __eq__(self, other):
         if not isinstance(other, A2ATaskError):
             return False
-        return (self.message == other.message and 
-                self.code == other.code and
-                self.provider == other.provider and
-                self.task_id == other.task_id and
-                self.cause == other.cause)
+        return (
+            self.message == other.message
+            and self.code == other.code
+            and self.provider == other.provider
+            and self.task_id == other.task_id
+            and self.cause == other.cause
+        )
 
     def __repr__(self):
         return f"{self.__class__.__name__}(message={self.message!r}, code={self.code!r}, provider={self.provider!r}, task_id={self.task_id!r}, cause={self.cause!r})"
 
+
 class A2ATaskNotFoundError(A2ATaskError):
     """Error when an A2A task is not found"""
-    def __init__(self, message: str, code: str, provider: str, task_id: str, cause: Optional[Exception] = None):
+
+    def __init__(
+        self,
+        message: str,
+        code: str,
+        provider: str,
+        task_id: str,
+        cause: Optional[Exception] = None,
+    ):
         super().__init__(message, code, provider, task_id, cause)
 
+
 class A2ATaskStorageError(A2ATaskError):
     """Error for A2A task storage operation failures"""
-    def __init__(self, message: str, code: str, provider: str, operation: str, task_id: Optional[str] = None, cause: Optional[Exception] = None):
+
+    def __init__(
+        self,
+        message: str,
+        code: str,
+        provider: str,
+        operation: str,
+        task_id: Optional[str] = None,
+        cause: Optional[Exception] = None,
+    ):
         super().__init__(message, code, provider, task_id, cause)
         self.operation = operation
 
@@ -217,6 +256,7 @@ class A2ATaskStorageError(A2ATaskError):
     def __repr__(self):
         return f"{self.__class__.__name__}(message={self.message!r}, code={self.code!r}, provider={self.provider!r}, operation={self.operation!r}, task_id={self.task_id!r}, cause={self.cause!r})"
 
+
 # Union of all possible A2A task errors
 A2ATaskErrorUnion = Union[A2ATaskError, A2ATaskNotFoundError, A2ATaskStorageError]
 
@@ -225,39 +265,30 @@ A2AResult = Union[Success[T], Failure[A2ATaskErrorUnion]]
 
 # Error factory functions
 
+
 def create_a2a_task_error(
     message: str,
     code: str,
     provider: str,
     task_id: Optional[str] = None,
-    cause: Optional[Exception] = None
+    cause: Optional[Exception] = None,
 ) -> A2ATaskError:
     """Create an A2A task error"""
-    return A2ATaskError(
-        message=message,
-        code=code,
-        provider=provider,
-        task_id=task_id,
-        cause=cause
-    )
+    return A2ATaskError(message=message, code=code, provider=provider, task_id=task_id, cause=cause)
+
 
-def create_a2a_task_not_found_error(
-    task_id: str,
-    provider: str
-) -> A2ATaskNotFoundError:
+def create_a2a_task_not_found_error(task_id: str, provider: str) -> A2ATaskNotFoundError:
     """Create an A2A task not found error"""
     return A2ATaskNotFoundError(
         message=f"A2A task {task_id} not found",
         code="TASK_NOT_FOUND",
         provider=provider,
-        task_id=task_id
+        task_id=task_id,
     )
 
+
 def create_a2a_task_storage_error(
-    operation: str,
-    provider: str,
-    task_id: Optional[str] = None,
-    cause: Optional[Exception] = None
+    operation: str, provider: str, task_id: Optional[str] = None, cause: Optional[Exception] = None
 ) -> A2ATaskStorageError:
     """Create an A2A task storage error"""
     message = f"Failed to {operation} A2A task"
@@ -271,29 +302,36 @@ def create_a2a_task_storage_error(
         provider=provider,
         operation=operation,
         task_id=task_id,
-        cause=cause
+        cause=cause,
     )
 
+
 # Error checking functions
 
+
 def is_a2a_task_error(error: Any) -> bool:
     """Check if error is an A2A task error"""
     return isinstance(error, (A2ATaskError, A2ATaskNotFoundError, A2ATaskStorageError))
 
+
 def is_a2a_task_not_found_error(error: Any) -> bool:
     """Check if error is an A2A task not found error"""
     return isinstance(error, A2ATaskNotFoundError)
 
+
 def is_a2a_task_storage_error(error: Any) -> bool:
     """Check if error is an A2A task storage error"""
     return isinstance(error, A2ATaskStorageError)
 
+
 # A2A-specific Result factory functions
 
+
 def create_a2a_success(data: T) -> Success[T]:
     """Create an A2A success result"""
     return Success(data=data)
 
+
 def create_a2a_failure(error: A2ATaskErrorUnion) -> Failure[A2ATaskErrorUnion]:
     """Create an A2A failure result"""
     return Failure(error=error)