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

jaf/a2a/memory/providers/postgres.py

@@ -30,7 +30,7 @@ from ..types import (
 
 # SQL queries for A2A task operations
 SQL_QUERIES = {
-    'CREATE_TABLE': '''
+    "CREATE_TABLE": """
         CREATE TABLE IF NOT EXISTS {table_name} (
             task_id VARCHAR(255) PRIMARY KEY,
             context_id VARCHAR(255) NOT NULL,
@@ -47,150 +47,158 @@ SQL_QUERIES = {
         CREATE INDEX IF NOT EXISTS idx_{table_name}_state ON {table_name} (state);
         CREATE INDEX IF NOT EXISTS idx_{table_name}_created_at ON {table_name} (created_at);
         CREATE INDEX IF NOT EXISTS idx_{table_name}_expires_at ON {table_name} (expires_at) WHERE expires_at IS NOT NULL;
-    ''',
-    'INSERT_TASK': '''
+    """,
+    "INSERT_TASK": """
         INSERT INTO {table_name} (
             task_id, context_id, state, task_data, status_message, 
             created_at, updated_at, expires_at, metadata
         ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
-    ''',
-    'SELECT_TASK': '''
+    """,
+    "SELECT_TASK": """
         SELECT task_id, context_id, state, task_data, status_message, 
                created_at, updated_at, expires_at, metadata
         FROM {table_name} 
         WHERE task_id = $1 
           AND (expires_at IS NULL OR expires_at > NOW())
-    ''',
-    'UPDATE_TASK': '''
+    """,
+    "UPDATE_TASK": """
         UPDATE {table_name} 
         SET state = $2, task_data = $3, status_message = $4, 
             updated_at = $5, metadata = $6
         WHERE task_id = $1
           AND (expires_at IS NULL OR expires_at > NOW())
-    ''',
-    'DELETE_TASK': '''
+    """,
+    "DELETE_TASK": """
         DELETE FROM {table_name} 
         WHERE task_id = $1
-    ''',
-    'DELETE_TASKS_BY_CONTEXT': '''
+    """,
+    "DELETE_TASKS_BY_CONTEXT": """
         DELETE FROM {table_name} 
         WHERE context_id = $1
-    ''',
-    'CLEANUP_EXPIRED': '''
+    """,
+    "CLEANUP_EXPIRED": """
         DELETE FROM {table_name} 
         WHERE expires_at IS NOT NULL AND expires_at <= NOW()
-    ''',
-    'COUNT_TASKS': '''
+    """,
+    "COUNT_TASKS": """
         SELECT COUNT(*) as total 
         FROM {table_name} 
         WHERE (expires_at IS NULL OR expires_at > NOW())
-    ''',
-    'STATS_BY_STATE': '''
+    """,
+    "STATS_BY_STATE": """
         SELECT state, COUNT(*) as count 
         FROM {table_name} 
         WHERE (expires_at IS NULL OR expires_at > NOW())
           AND ($1::text IS NULL OR context_id = $1)
         GROUP BY state
-    ''',
-    'DATE_RANGE': '''
+    """,
+    "DATE_RANGE": """
         SELECT MIN(created_at) as oldest, MAX(created_at) as newest
         FROM {table_name} 
         WHERE (expires_at IS NULL OR expires_at > NOW())
           AND ($1::text IS NULL OR context_id = $1)
-    '''
+    """,
 }
 
+
 async def create_a2a_postgres_task_provider(
-    config: A2APostgresTaskConfig,
-    postgres_client: Any
+    config: A2APostgresTaskConfig, postgres_client: Any
 ) -> A2AResult[A2ATaskProvider]:
     """
     Create a PostgreSQL-based A2A task provider
-    
+
     Args:
         config: Configuration for the PostgreSQL provider
         postgres_client: PostgreSQL client instance (asyncpg connection/pool)
-        
+
     Returns:
         A2AResult containing the task provider or an error
     """
     try:
-        table_name = config.table_name or 'a2a_tasks'
+        table_name = config.table_name or "a2a_tasks"
 
         # Initialize database schema
-        create_table_sql = SQL_QUERIES['CREATE_TABLE'].format(table_name=table_name)
+        create_table_sql = SQL_QUERIES["CREATE_TABLE"].format(table_name=table_name)
         await postgres_client.execute(create_table_sql)
 
         def row_to_serialized_task(row: Any) -> A2ATaskSerialized:
             """Convert database row to serialized task"""
             return A2ATaskSerialized(
-                task_id=row['task_id'],
-                context_id=row['context_id'],
-                state=row['state'],
-                task_data=row['task_data'] if isinstance(row['task_data'], str) else json.dumps(row['task_data']),
-                status_message=row['status_message'] if isinstance(row['status_message'], str) else json.dumps(row['status_message']) if row['status_message'] else None,
-                created_at=row['created_at'].isoformat(),
-                updated_at=row['updated_at'].isoformat(),
-                metadata=row['metadata'] if isinstance(row['metadata'], str) else json.dumps(row['metadata']) if row['metadata'] else None
+                task_id=row["task_id"],
+                context_id=row["context_id"],
+                state=row["state"],
+                task_data=row["task_data"]
+                if isinstance(row["task_data"], str)
+                else json.dumps(row["task_data"]),
+                status_message=row["status_message"]
+                if isinstance(row["status_message"], str)
+                else json.dumps(row["status_message"])
+                if row["status_message"]
+                else None,
+                created_at=row["created_at"].isoformat(),
+                updated_at=row["updated_at"].isoformat(),
+                metadata=row["metadata"]
+                if isinstance(row["metadata"], str)
+                else json.dumps(row["metadata"])
+                if row["metadata"]
+                else None,
             )
 
         def build_where_clause(query: A2ATaskQuery) -> tuple[str, list]:
             """Build WHERE clause for queries"""
-            conditions = ['(expires_at IS NULL OR expires_at > NOW())']
+            conditions = ["(expires_at IS NULL OR expires_at > NOW())"]
             params = []
             param_index = 1
 
             if query.task_id:
-                conditions.append(f'task_id = ${param_index}')
+                conditions.append(f"task_id = ${param_index}")
                 params.append(query.task_id)
                 param_index += 1
 
             if query.context_id:
-                conditions.append(f'context_id = ${param_index}')
+                conditions.append(f"context_id = ${param_index}")
                 params.append(query.context_id)
                 param_index += 1
 
             if query.state:
-                conditions.append(f'state = ${param_index}')
+                conditions.append(f"state = ${param_index}")
                 params.append(query.state.value)
                 param_index += 1
 
             if query.since:
                 # Use metadata created_at if available, otherwise use database created_at
-                conditions.append(f'''(
+                conditions.append(f"""(
                     CASE 
                         WHEN metadata ? 'created_at' THEN 
                             (metadata->>'created_at')::timestamp with time zone >= ${param_index}
                         ELSE 
                             created_at >= ${param_index}
                     END
-                )''')
+                )""")
                 params.append(query.since)
                 param_index += 1
 
             if query.until:
                 # Use metadata created_at if available, otherwise use database created_at
-                conditions.append(f'''(
+                conditions.append(f"""(
                     CASE 
                         WHEN metadata ? 'created_at' THEN 
                             (metadata->>'created_at')::timestamp with time zone <= ${param_index}
                         ELSE 
                             created_at <= ${param_index}
                     END
-                )''')
+                )""")
                 params.append(query.until)
                 param_index += 1
 
-            where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ''
+            where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
             return where_clause, params
 
         class PostgresA2ATaskProvider:
             """PostgreSQL implementation of A2ATaskProvider"""
 
             async def store_task(
-                self,
-                task: A2ATask,
-                metadata: Optional[Dict[str, Any]] = None
+                self, task: A2ATask, metadata: Optional[Dict[str, Any]] = None
             ) -> A2AResult[None]:
                 """Store a new A2A task in PostgreSQL"""
                 try:
@@ -205,7 +213,7 @@ async def create_a2a_postgres_task_provider(
                         return serialize_result
 
                     serialized = serialize_result.data
-                    query = SQL_QUERIES['INSERT_TASK'].format(table_name=table_name)
+                    query = SQL_QUERIES["INSERT_TASK"].format(table_name=table_name)
 
                     await postgres_client.execute(
                         query,
@@ -216,21 +224,21 @@ async def create_a2a_postgres_task_provider(
                         serialized.status_message,
                         datetime.fromisoformat(serialized.created_at),
                         datetime.fromisoformat(serialized.updated_at),
-                        metadata.get('expires_at') if metadata else None,
-                        json.dumps(metadata) if metadata else None
+                        metadata.get("expires_at") if metadata else None,
+                        json.dumps(metadata) if metadata else None,
                     )
 
                     return create_a2a_success(None)
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('store', 'postgres', task.id, error)
+                        create_a2a_task_storage_error("store", "postgres", task.id, error)
                     )
 
             async def get_task(self, task_id: str) -> A2AResult[Optional[A2ATask]]:
                 """Retrieve a task by ID from PostgreSQL"""
                 try:
-                    query = SQL_QUERIES['SELECT_TASK'].format(table_name=table_name)
+                    query = SQL_QUERIES["SELECT_TASK"].format(table_name=table_name)
                     row = await postgres_client.fetchrow(query, task_id)
 
                     if not row:
@@ -246,13 +254,11 @@ async def create_a2a_postgres_task_provider(
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('get', 'postgres', task_id, error)
+                        create_a2a_task_storage_error("get", "postgres", task_id, error)
                     )
 
             async def update_task(
-                self,
-                task: A2ATask,
-                metadata: Optional[Dict[str, Any]] = None
+                self, task: A2ATask, metadata: Optional[Dict[str, Any]] = None
             ) -> A2AResult[None]:
                 """Update an existing task in PostgreSQL"""
                 try:
@@ -263,7 +269,7 @@ async def create_a2a_postgres_task_provider(
 
                     if not existing_result.data:
                         return create_a2a_failure(
-                            create_a2a_task_not_found_error(task.id, 'postgres')
+                            create_a2a_task_not_found_error(task.id, "postgres")
                         )
 
                     # Validate and sanitize task
@@ -272,9 +278,13 @@ async def create_a2a_postgres_task_provider(
                         return sanitize_result
 
                     # Get existing metadata
-                    existing_query = SQL_QUERIES['SELECT_TASK'].format(table_name=table_name)
+                    existing_query = SQL_QUERIES["SELECT_TASK"].format(table_name=table_name)
                     existing_data = await postgres_client.fetchrow(existing_query, task.id)
-                    existing_metadata = existing_data['metadata'] if existing_data and existing_data['metadata'] else {}
+                    existing_metadata = (
+                        existing_data["metadata"]
+                        if existing_data and existing_data["metadata"]
+                        else {}
+                    )
                     if isinstance(existing_metadata, str):
                         existing_metadata = json.loads(existing_metadata)
 
@@ -286,7 +296,7 @@ async def create_a2a_postgres_task_provider(
                         return serialize_result
 
                     serialized = serialize_result.data
-                    query = SQL_QUERIES['UPDATE_TASK'].format(table_name=table_name)
+                    query = SQL_QUERIES["UPDATE_TASK"].format(table_name=table_name)
 
                     result = await postgres_client.execute(
                         query,
@@ -295,20 +305,20 @@ async def create_a2a_postgres_task_provider(
                         serialized.task_data,
                         serialized.status_message,
                         datetime.fromisoformat(serialized.updated_at),
-                        json.dumps(merged_metadata)
+                        json.dumps(merged_metadata),
                     )
 
                     # Check if any rows were affected
-                    if result == 'UPDATE 0':
+                    if result == "UPDATE 0":
                         return create_a2a_failure(
-                            create_a2a_task_not_found_error(task.id, 'postgres')
+                            create_a2a_task_not_found_error(task.id, "postgres")
                         )
 
                     return create_a2a_success(None)
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('update', 'postgres', task.id, error)
+                        create_a2a_task_storage_error("update", "postgres", task.id, error)
                     )
 
             async def update_task_status(
@@ -316,7 +326,7 @@ async def create_a2a_postgres_task_provider(
                 task_id: str,
                 state: TaskState,
                 status_message: Optional[Any] = None,
-                timestamp: Optional[str] = None
+                timestamp: Optional[str] = None,
             ) -> A2AResult[None]:
                 """Update task status only"""
                 try:
@@ -327,7 +337,7 @@ async def create_a2a_postgres_task_provider(
 
                     if not existing_result.data:
                         return create_a2a_failure(
-                            create_a2a_task_not_found_error(task_id, 'postgres')
+                            create_a2a_task_not_found_error(task_id, "postgres")
                         )
 
                     task = existing_result.data
@@ -339,22 +349,22 @@ async def create_a2a_postgres_task_provider(
 
                     # Update task status
                     from ...types import A2ATaskStatus
+
                     updated_status = A2ATaskStatus(
                         state=state,
                         message=status_message or task.status.message,
-                        timestamp=timestamp or datetime.now().isoformat()
+                        timestamp=timestamp or datetime.now().isoformat(),
                     )
 
-                    updated_task = task.model_copy(update={
-                        'status': updated_status,
-                        'history': updated_history
-                    })
+                    updated_task = task.model_copy(
+                        update={"status": updated_status, "history": updated_history}
+                    )
 
                     return await self.update_task(updated_task)
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('update-status', 'postgres', task_id, error)
+                        create_a2a_task_storage_error("update-status", "postgres", task_id, error)
                     )
 
             async def find_tasks(self, query: A2ATaskQuery) -> A2AResult[List[A2ATask]]:
@@ -362,19 +372,19 @@ async def create_a2a_postgres_task_provider(
                 try:
                     where_clause, params = build_where_clause(query)
 
-                    sql = f'''
+                    sql = f"""
                         SELECT task_id, context_id, state, task_data, status_message, 
                                created_at, updated_at, expires_at, metadata
                         FROM {table_name} 
                         {where_clause}
                         ORDER BY created_at DESC
-                    '''
+                    """
 
                     # Add pagination
                     if query.limit:
-                        sql += f' LIMIT {query.limit}'
+                        sql += f" LIMIT {query.limit}"
                     if query.offset:
-                        sql += f' OFFSET {query.offset}'
+                        sql += f" OFFSET {query.offset}"
 
                     rows = await postgres_client.fetch(sql, *params)
                     tasks: List[A2ATask] = []
@@ -390,13 +400,11 @@ async def create_a2a_postgres_task_provider(
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('find', 'postgres', None, error)
+                        create_a2a_task_storage_error("find", "postgres", None, error)
                     )
 
             async def get_tasks_by_context(
-                self,
-                context_id: str,
-                limit: Optional[int] = None
+                self, context_id: str, limit: Optional[int] = None
             ) -> A2AResult[List[A2ATask]]:
                 """Get tasks by context ID"""
                 return await self.find_tasks(A2ATaskQuery(context_id=context_id, limit=limit))
@@ -404,87 +412,86 @@ async def create_a2a_postgres_task_provider(
             async def delete_task(self, task_id: str) -> A2AResult[bool]:
                 """Delete a task and return True if it existed"""
                 try:
-                    query = SQL_QUERIES['DELETE_TASK'].format(table_name=table_name)
+                    query = SQL_QUERIES["DELETE_TASK"].format(table_name=table_name)
                     result = await postgres_client.execute(query, task_id)
 
                     # Check if any rows were affected
-                    rows_affected = int(result.split()[-1]) if 'DELETE' in result else 0
+                    rows_affected = int(result.split()[-1]) if "DELETE" in result else 0
                     return create_a2a_success(rows_affected > 0)
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('delete', 'postgres', task_id, error)
+                        create_a2a_task_storage_error("delete", "postgres", task_id, error)
                     )
 
             async def delete_tasks_by_context(self, context_id: str) -> A2AResult[int]:
                 """Delete tasks by context ID and return count deleted"""
                 try:
-                    query = SQL_QUERIES['DELETE_TASKS_BY_CONTEXT'].format(table_name=table_name)
+                    query = SQL_QUERIES["DELETE_TASKS_BY_CONTEXT"].format(table_name=table_name)
                     result = await postgres_client.execute(query, context_id)
 
                     # Extract count from result
-                    rows_affected = int(result.split()[-1]) if 'DELETE' in result else 0
+                    rows_affected = int(result.split()[-1]) if "DELETE" in result else 0
                     return create_a2a_success(rows_affected)
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('delete-by-context', 'postgres', None, error)
+                        create_a2a_task_storage_error("delete-by-context", "postgres", None, error)
                     )
 
             async def cleanup_expired_tasks(self) -> A2AResult[int]:
                 """Clean up expired tasks and return count deleted"""
                 try:
-                    query = SQL_QUERIES['CLEANUP_EXPIRED'].format(table_name=table_name)
+                    query = SQL_QUERIES["CLEANUP_EXPIRED"].format(table_name=table_name)
                     result = await postgres_client.execute(query)
 
                     # Extract count from result
-                    rows_affected = int(result.split()[-1]) if 'DELETE' in result else 0
+                    rows_affected = int(result.split()[-1]) if "DELETE" in result else 0
                     return create_a2a_success(rows_affected)
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('cleanup', 'postgres', None, error)
+                        create_a2a_task_storage_error("cleanup", "postgres", None, error)
                     )
 
             async def get_task_stats(
-                self,
-                context_id: Optional[str] = None
+                self, context_id: Optional[str] = None
             ) -> A2AResult[Dict[str, Any]]:
                 """Get task statistics"""
                 try:
                     tasks_by_state = {state.value: 0 for state in TaskState}
 
                     # Get state counts
-                    state_query = SQL_QUERIES['STATS_BY_STATE'].format(table_name=table_name)
+                    state_query = SQL_QUERIES["STATS_BY_STATE"].format(table_name=table_name)
                     state_rows = await postgres_client.fetch(state_query, context_id)
 
                     total_tasks = 0
                     for row in state_rows:
-                        state = row['state']
-                        count = int(row['count'])
+                        state = row["state"]
+                        count = int(row["count"])
                         if state in tasks_by_state:
                             tasks_by_state[state] = count
                         total_tasks += count
 
                     # Get date range
-                    date_query = SQL_QUERIES['DATE_RANGE'].format(table_name=table_name)
+                    date_query = SQL_QUERIES["DATE_RANGE"].format(table_name=table_name)
                     date_row = await postgres_client.fetchrow(date_query, context_id)
 
                     oldest_task = None
                     newest_task = None
 
-                    if date_row and date_row['oldest']:
-                        oldest_task = date_row['oldest']
-                        newest_task = date_row['newest']
+                    if date_row and date_row["oldest"]:
+                        oldest_task = date_row["oldest"]
+                        newest_task = date_row["newest"]
 
                     # Build stats dict with individual state counts as top-level keys
                     stats = {
-                        'total_tasks': total_tasks,
-                        'tasks_by_state': tasks_by_state,
-                        'oldest_task': oldest_task,
-                        'newest_task': newest_task
+                        "total_tasks": total_tasks,
+                        "tasks_by_state": tasks_by_state,
+                        "oldest_task": oldest_task,
+                        "newest_task": newest_task,
                     }
-                    
+
                     # Also add individual state counts for backwards compatibility
                     for state in TaskState:
                         stats[state.value] = tasks_by_state[state.value]
@@ -493,7 +500,7 @@ async def create_a2a_postgres_task_provider(
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('stats', 'postgres', None, error)
+                        create_a2a_task_storage_error("stats", "postgres", None, error)
                     )
 
             async def health_check(self) -> A2AResult[Dict[str, Any]]:
@@ -502,20 +509,15 @@ async def create_a2a_postgres_task_provider(
                     start_time = datetime.now()
 
                     # Simple query to check database connectivity
-                    await postgres_client.fetchval('SELECT 1')
+                    await postgres_client.fetchval("SELECT 1")
 
                     latency_ms = (datetime.now() - start_time).total_seconds() * 1000
-                    return create_a2a_success({
-                        'healthy': True,
-                        'provider': 'postgres',
-                        'latency_ms': latency_ms
-                    })
+                    return create_a2a_success(
+                        {"healthy": True, "provider": "postgres", "latency_ms": latency_ms}
+                    )
 
                 except Exception as error:
-                    return create_a2a_success({
-                        'healthy': False,
-                        'error': str(error)
-                    })
+                    return create_a2a_success({"healthy": False, "error": str(error)})
 
             async def close(self) -> A2AResult[None]:
                 """Close/cleanup the provider"""
@@ -526,12 +528,12 @@ async def create_a2a_postgres_task_provider(
 
                 except Exception as error:
                     return create_a2a_failure(
-                        create_a2a_task_storage_error('close', 'postgres', None, error)
+                        create_a2a_task_storage_error("close", "postgres", None, error)
                     )
 
         return create_a2a_success(PostgresA2ATaskProvider())
 
     except Exception as error:
         return create_a2a_failure(
-            create_a2a_task_storage_error('create-postgres-provider', 'postgres', None, error)
+            create_a2a_task_storage_error("create-postgres-provider", "postgres", None, error)
         )