Merge pull request #194 from esnible/update-contact-extractor

bug: Modify contact extractor example for a2a-sdk refactor

Author: esnible

a2a/a2a_contact_extractor/__main__.py

@@ -10,7 +10,11 @@
 import httpx
 from a2a.server.apps import A2AStarletteApplication
 from a2a.server.request_handlers import DefaultRequestHandler
-from a2a.server.tasks import InMemoryPushNotifier, InMemoryTaskStore
+from a2a.server.tasks import (
+    BasePushNotificationSender,
+    InMemoryPushNotificationConfigStore,
+    InMemoryTaskStore,
+)
 from a2a.types import AgentCapabilities, AgentCard, AgentSkill
 from dotenv import load_dotenv
 from pydantic import BaseModel, EmailStr, Field
@@ -55,10 +59,16 @@ def main(host, port, result_type, instructions):
         exit(1)
     agent = ExtractorAgent(instructions=instructions, result_type=result_type)
     httpx_client = httpx.AsyncClient()
+    push_notification_config_store = InMemoryPushNotificationConfigStore()
+    push_notification_sender = BasePushNotificationSender(
+        httpx_client, config_store=push_notification_config_store
+    )
+
     request_handler = DefaultRequestHandler(
         agent_executor=ExtractorAgentExecutor(agent=agent),
         task_store=InMemoryTaskStore(),
-        push_notifier=InMemoryPushNotifier(httpx_client),
+        push_config_store=push_notification_config_store,
+        push_sender=push_notification_sender,
     )
     server = A2AStarletteApplication(
         agent_card=get_agent_card(host, port), http_handler=request_handler

a2a/a2a_contact_extractor/agent_executor.py

@@ -36,9 +36,9 @@ async def execute(
         task = context.current_task
         if not task:
             task = new_task(context.message)
-            event_queue.enqueue_event(task)
+            await event_queue.enqueue_event(task)
 
-        async for item in self.agent.stream(query, task.contextId):
+        async for item in self.agent.stream(query, task.context_id):
             is_task_complete = item["is_task_complete"]
             require_user_input = item["require_user_input"]
             # content = item["content"]
@@ -48,7 +48,7 @@ async def execute(
                 f"Stream item received: complete={is_task_complete}, require_input={require_user_input}, content_len={len(content)}"
             )
 
-            agent_outcome = await self.agent.invoke(query, task.contextId)
+            agent_outcome = await self.agent.invoke(query, task.context_id)
             is_task_complete = agent_outcome["is_task_complete"]
             require_user_input = not is_task_complete
             content = agent_outcome.get("text_parts", [])
@@ -77,20 +77,20 @@ async def execute(
                             state=TaskState.input_required,
                             message=new_agent_text_message(
                                 content,
-                                task.contextId,
+                                task.context_id,
                                 task.id,
                             ),
                         ),
                         final=True,
-                        contextId=task.contextId,
+                        contextId=task.context_id,
                         taskId=task.id,
                     )
                 )
             elif is_task_complete:
                 event_queue.enqueue_event(
                     TaskArtifactUpdateEvent(
                         append=False,
-                        contextId=task.contextId,
+                        contextId=task.context_id,
                         taskId=task.id,
                         lastChunk=True,
                         artifact=artifact,
@@ -100,7 +100,7 @@ async def execute(
                     TaskStatusUpdateEvent(
                         status=TaskStatus(state=TaskState.completed),
                         final=True,
-                        contextId=task.contextId,
+                        contextId=task.context_id,
                         taskId=task.id,
                     )
                 )
@@ -111,12 +111,12 @@ async def execute(
                             state=TaskState.working,
                             message=new_agent_text_message(
                                 "Analyzing your text...",
-                                task.contextId,
+                                task.context_id,
                                 task.id,
                             ),
                         ),
                         final=False,
-                        contextId=task.contextId,
+                        contextId=task.context_id,
                         taskId=task.id,
                     )
                 )