Skip to content

Commit 1edf63c

Browse files
victordibiavictordibia
authored
Merge branch 'main' into declarative_tools_vd
2 parents d3c2e55 + 44b9bff commit 1edf63c

File tree

9 files changed

+368
-32
lines changed

9 files changed

+368
-32
lines changed

‎protos/agent_worker.proto

+25-6
Original file line numberDiff line numberDiff line change
@@ -48,12 +48,12 @@ message Event {
4848
}
4949

5050
message RegisterAgentTypeRequest {
51-
string request_id = 1;
51+
string request_id = 1; // TODO: remove once message based requests are removed
5252
string type = 2;
5353
}
5454

5555
message RegisterAgentTypeResponse {
56-
string request_id = 1;
56+
string request_id = 1; // TODO: remove once message based requests are removed
5757
bool success = 2;
5858
optional string error = 3;
5959
}
@@ -69,27 +69,46 @@ message TypePrefixSubscription {
6969
}
7070

7171
message Subscription {
72+
string id = 1;
7273
oneof subscription {
73-
TypeSubscription typeSubscription = 1;
74-
TypePrefixSubscription typePrefixSubscription = 2;
74+
TypeSubscription typeSubscription = 2;
75+
TypePrefixSubscription typePrefixSubscription = 3;
7576
}
7677
}
7778

7879
message AddSubscriptionRequest {
79-
string request_id = 1;
80+
string request_id = 1; // TODO: remove once message based requests are removed
8081
Subscription subscription = 2;
8182
}
8283

8384
message AddSubscriptionResponse {
84-
string request_id = 1;
85+
string request_id = 1; // TODO: remove once message based requests are removed
8586
bool success = 2;
8687
optional string error = 3;
8788
}
8889

90+
message RemoveSubscriptionRequest {
91+
string id = 1;
92+
}
93+
94+
message RemoveSubscriptionResponse {
95+
bool success = 1;
96+
optional string error = 2;
97+
}
98+
99+
message GetSubscriptionsRequest {}
100+
message GetSubscriptionsResponse {
101+
repeated Subscription subscriptions = 1;
102+
}
103+
89104
service AgentRpc {
90105
rpc OpenChannel (stream Message) returns (stream Message);
91106
rpc GetState(AgentId) returns (GetStateResponse);
92107
rpc SaveState(AgentState) returns (SaveStateResponse);
108+
rpc RegisterAgent(RegisterAgentTypeRequest) returns (RegisterAgentTypeResponse);
109+
rpc AddSubscription(AddSubscriptionRequest) returns (AddSubscriptionResponse);
110+
rpc RemoveSubscription(RemoveSubscriptionRequest) returns (RemoveSubscriptionResponse);
111+
rpc GetSubscriptions(GetSubscriptionsRequest) returns (GetSubscriptionsResponse);
93112
}
94113

95114
message AgentState {

‎python/packages/autogen-core/src/autogen_core/_type_prefix_subscription.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -31,13 +31,13 @@ class TypePrefixSubscription(Subscription):
3131
agent_type (str): Agent type to handle this subscription
3232
"""
3333

34-
def __init__(self, topic_type_prefix: str, agent_type: str | AgentType):
34+
def __init__(self, topic_type_prefix: str, agent_type: str | AgentType, id: str | None = None):
3535
self._topic_type_prefix = topic_type_prefix
3636
if isinstance(agent_type, AgentType):
3737
self._agent_type = agent_type.type
3838
else:
3939
self._agent_type = agent_type
40-
self._id = str(uuid.uuid4())
40+
self._id = id or str(uuid.uuid4())
4141

4242
@property
4343
def id(self) -> str:

‎python/packages/autogen-core/src/autogen_core/_type_subscription.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,13 @@ class TypeSubscription(Subscription):
3030
agent_type (str): Agent type to handle this subscription
3131
"""
3232

33-
def __init__(self, topic_type: str, agent_type: str | AgentType):
33+
def __init__(self, topic_type: str, agent_type: str | AgentType, id: str | None = None):
3434
self._topic_type = topic_type
3535
if isinstance(agent_type, AgentType):
3636
self._agent_type = agent_type.type
3737
else:
3838
self._agent_type = agent_type
39-
self._id = str(uuid.uuid4())
39+
self._id = id or str(uuid.uuid4())
4040

4141
@property
4242
def id(self) -> str:

‎python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime.py

+6-4
Original file line numberDiff line numberDiff line change
@@ -807,25 +807,27 @@ async def add_subscription(self, subscription: Subscription) -> None:
807807
request_id = await self._get_new_request_id()
808808

809809
match subscription:
810-
case TypeSubscription(topic_type=topic_type, agent_type=agent_type):
810+
case TypeSubscription(topic_type=topic_type, agent_type=agent_type, id=id):
811811
message = agent_worker_pb2.Message(
812812
addSubscriptionRequest=agent_worker_pb2.AddSubscriptionRequest(
813813
request_id=request_id,
814814
subscription=agent_worker_pb2.Subscription(
815+
id=id,
815816
typeSubscription=agent_worker_pb2.TypeSubscription(
816817
topic_type=topic_type, agent_type=agent_type
817-
)
818+
),
818819
),
819820
)
820821
)
821-
case TypePrefixSubscription(topic_type_prefix=topic_type_prefix, agent_type=agent_type):
822+
case TypePrefixSubscription(topic_type_prefix=topic_type_prefix, agent_type=agent_type, id=id):
822823
message = agent_worker_pb2.Message(
823824
addSubscriptionRequest=agent_worker_pb2.AddSubscriptionRequest(
824825
request_id=request_id,
825826
subscription=agent_worker_pb2.Subscription(
827+
id=id,
826828
typePrefixSubscription=agent_worker_pb2.TypePrefixSubscription(
827829
topic_type_prefix=topic_type_prefix, agent_type=agent_type
828-
)
830+
),
829831
),
830832
)
831833
)

‎python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/_worker_runtime_host_servicer.py

+40-1
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,9 @@ async def _process_add_subscription_request(
227227
add_subscription_req.subscription.typeSubscription
228228
)
229229
subscription = TypeSubscription(
230-
topic_type=type_subscription_msg.topic_type, agent_type=type_subscription_msg.agent_type
230+
topic_type=type_subscription_msg.topic_type,
231+
agent_type=type_subscription_msg.agent_type,
232+
id=add_subscription_req.subscription.id,
231233
)
232234

233235
case "typePrefixSubscription":
@@ -237,6 +239,7 @@ async def _process_add_subscription_request(
237239
subscription = TypePrefixSubscription(
238240
topic_type_prefix=type_prefix_subscription_msg.topic_type_prefix,
239241
agent_type=type_prefix_subscription_msg.agent_type,
242+
id=add_subscription_req.subscription.id,
240243
)
241244
case None:
242245
logger.warning("Received empty subscription message")
@@ -260,6 +263,42 @@ async def _process_add_subscription_request(
260263
)
261264
)
262265

266+
def RegisterAgent( # type: ignore
267+
self,
268+
request: agent_worker_pb2.RegisterAgentTypeRequest,
269+
context: grpc.aio.ServicerContext[
270+
agent_worker_pb2.RegisterAgentTypeRequest, agent_worker_pb2.RegisterAgentTypeResponse
271+
],
272+
) -> agent_worker_pb2.RegisterAgentTypeResponse:
273+
raise NotImplementedError("Method not implemented.")
274+
275+
def AddSubscription( # type: ignore
276+
self,
277+
request: agent_worker_pb2.AddSubscriptionRequest,
278+
context: grpc.aio.ServicerContext[
279+
agent_worker_pb2.AddSubscriptionRequest, agent_worker_pb2.AddSubscriptionResponse
280+
],
281+
) -> agent_worker_pb2.AddSubscriptionResponse:
282+
raise NotImplementedError("Method not implemented.")
283+
284+
def RemoveSubscription( # type: ignore
285+
self,
286+
request: agent_worker_pb2.RemoveSubscriptionRequest,
287+
context: grpc.aio.ServicerContext[
288+
agent_worker_pb2.RemoveSubscriptionRequest, agent_worker_pb2.RemoveSubscriptionResponse
289+
],
290+
) -> agent_worker_pb2.RemoveSubscriptionResponse:
291+
raise NotImplementedError("Method not implemented.")
292+
293+
def GetSubscriptions( # type: ignore
294+
self,
295+
request: agent_worker_pb2.GetSubscriptionsRequest,
296+
context: grpc.aio.ServicerContext[
297+
agent_worker_pb2.GetSubscriptionsRequest, agent_worker_pb2.GetSubscriptionsResponse
298+
],
299+
) -> agent_worker_pb2.GetSubscriptionsResponse:
300+
raise NotImplementedError("Method not implemented.")
301+
263302
async def GetState( # type: ignore
264303
self,
265304
request: agent_worker_pb2.AgentId,

‎python/packages/autogen-ext/src/autogen_ext/runtimes/grpc/protos/agent_worker_pb2.py

+24-16
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)