modify for runtime

microsoft · Feb 11, 2025 · 7ccf6b0 · 7ccf6b0
1 parent c3e69df
commit 7ccf6b0
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 4 deletions.
diff --git a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py
@@ -13,6 +13,7 @@
 
 from ._constants import GRPC_IMPORT_ERROR_STR
 from ._utils import subscription_from_proto
+from .event_store.base import EventStore
 from .event_store.redis import RedisEventStore
 from .event_store.memory import MemoryEventStore
 
@@ -43,7 +44,7 @@ async def get_client_id_or_abort(context: grpc.aio.ServicerContext[Any, Any]) ->
 
     return client_id  # type: ignore
 
-def get_send_queue(client_id):
+def get_send_queue(client_id) -> EventStore:
     redis_url = os.getenv("REDIS_URL")
     if redis_url:
         import redis
@@ -61,7 +62,7 @@ class ChannelConnection(ABC, Generic[SendT, ReceiveT]):
     def __init__(self, request_iterator: AsyncIterator[ReceiveT], client_id: str) -> None:
         self._request_iterator = request_iterator
         self._client_id = client_id
-        self._send_queue: asyncio.Queue[agent_worker_pb2.Message] = get_send_queue(client_id)
+        self._send_queue: EventStore[agent_worker_pb2.Message] = get_send_queue(client_id)
         self._receiving_task = asyncio.create_task(self._receive_messages(client_id, request_iterator))
 
     async def _receive_messages(self, client_id: ClientConnectionId, request_iterator: AsyncIterator[ReceiveT]) -> None:

diff --git a/...en_ext/runtimes/grpc/event_store/store.py → ...gen_ext/runtimes/grpc/event_store/base.py b/...en_ext/runtimes/grpc/event_store/store.py → ...gen_ext/runtimes/grpc/event_store/base.py
diff --git a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/event_store/memory.py b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/event_store/memory.py
@@ -1,5 +1,5 @@
 import asyncio
-from .store import EventStore
+from .base import EventStore
 
 class MemoryEventStore(EventStore):
     def __init__(self, client_id, **kwags):

diff --git a/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/event_store/redis.py b/python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/event_store/redis.py
@@ -1,5 +1,5 @@
 import redis
-from .store import EventStore
+from .base import EventStore
 
 class RedisEventStore(EventStore):
     def __init__(self, redis_instance: redis.Redis, client_id):