ag2 0.9.1.post0__py3-none-any.whl → 0.9.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of ag2 might be problematic. Click here for more details.

Files changed (37) hide show
  1. {ag2-0.9.1.post0.dist-info → ag2-0.9.3.dist-info}/METADATA +22 -12
  2. {ag2-0.9.1.post0.dist-info → ag2-0.9.3.dist-info}/RECORD +37 -23
  3. autogen/agentchat/contrib/capabilities/transforms.py +22 -9
  4. autogen/agentchat/conversable_agent.py +37 -34
  5. autogen/agentchat/group/group_utils.py +65 -20
  6. autogen/agentchat/group/handoffs.py +81 -5
  7. autogen/agentchat/group/on_context_condition.py +2 -2
  8. autogen/agentchat/group/patterns/pattern.py +7 -1
  9. autogen/agentchat/groupchat.py +2 -2
  10. autogen/agentchat/realtime/experimental/realtime_swarm.py +12 -4
  11. autogen/agents/experimental/document_agent/document_agent.py +232 -40
  12. autogen/events/agent_events.py +7 -4
  13. autogen/interop/litellm/litellm_config_factory.py +68 -2
  14. autogen/llm_config.py +4 -1
  15. autogen/mcp/__main__.py +78 -0
  16. autogen/mcp/mcp_proxy/__init__.py +19 -0
  17. autogen/mcp/mcp_proxy/fastapi_code_generator_helpers.py +63 -0
  18. autogen/mcp/mcp_proxy/mcp_proxy.py +581 -0
  19. autogen/mcp/mcp_proxy/operation_grouping.py +158 -0
  20. autogen/mcp/mcp_proxy/operation_renaming.py +114 -0
  21. autogen/mcp/mcp_proxy/patch_fastapi_code_generator.py +98 -0
  22. autogen/mcp/mcp_proxy/security.py +400 -0
  23. autogen/mcp/mcp_proxy/security_schema_visitor.py +37 -0
  24. autogen/oai/client.py +11 -2
  25. autogen/oai/gemini.py +20 -3
  26. autogen/oai/gemini_types.py +27 -0
  27. autogen/oai/oai_models/chat_completion.py +1 -1
  28. autogen/tools/experimental/__init__.py +5 -0
  29. autogen/tools/experimental/reliable/__init__.py +10 -0
  30. autogen/tools/experimental/reliable/reliable.py +1316 -0
  31. autogen/version.py +1 -1
  32. templates/client_template/main.jinja2 +69 -0
  33. templates/config_template/config.jinja2 +7 -0
  34. templates/main.jinja2 +61 -0
  35. {ag2-0.9.1.post0.dist-info → ag2-0.9.3.dist-info}/WHEEL +0 -0
  36. {ag2-0.9.1.post0.dist-info → ag2-0.9.3.dist-info}/licenses/LICENSE +0 -0
  37. {ag2-0.9.1.post0.dist-info → ag2-0.9.3.dist-info}/licenses/NOTICE.md +0 -0
@@ -2,7 +2,7 @@
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
- from typing import Optional, Union, overload
5
+ from typing import Union, overload
6
6
 
7
7
  from pydantic import BaseModel, Field
8
8
 
@@ -30,7 +30,7 @@ class Handoffs(BaseModel):
30
30
 
31
31
  context_conditions: list[OnContextCondition] = Field(default_factory=list)
32
32
  llm_conditions: list[OnCondition] = Field(default_factory=list)
33
- after_work: Optional[TransitionTarget] = None
33
+ after_works: list[OnContextCondition] = Field(default_factory=list)
34
34
 
35
35
  def add_context_condition(self, condition: OnContextCondition) -> "Handoffs":
36
36
  """
@@ -102,7 +102,9 @@ class Handoffs(BaseModel):
102
102
 
103
103
  def set_after_work(self, target: TransitionTarget) -> "Handoffs":
104
104
  """
105
- Set the after work target (only one allowed).
105
+ Set the after work target (replaces all after_works with single entry).
106
+
107
+ For backward compatibility, this creates an OnContextCondition with no condition (always true).
106
108
 
107
109
  Args:
108
110
  target: The after work TransitionTarget to set
@@ -113,7 +115,81 @@ class Handoffs(BaseModel):
113
115
  if not isinstance(target, TransitionTarget):
114
116
  raise TypeError(f"Expected a TransitionTarget instance, got {type(target).__name__}")
115
117
 
116
- self.after_work = target
118
+ # Create OnContextCondition with no condition (always true)
119
+ after_work_condition = OnContextCondition(target=target, condition=None)
120
+ self.after_works = [after_work_condition]
121
+ return self
122
+
123
+ def add_after_work(self, condition: OnContextCondition) -> "Handoffs":
124
+ """
125
+ Add a single after-work condition.
126
+
127
+ If the condition has condition=None, it will replace any existing
128
+ condition=None entry and be placed at the end.
129
+
130
+ Args:
131
+ condition: The OnContextCondition to add
132
+
133
+ Returns:
134
+ Self for method chaining
135
+ """
136
+ if not isinstance(condition, OnContextCondition):
137
+ raise TypeError(f"Expected an OnContextCondition instance, got {type(condition).__name__}")
138
+
139
+ if condition.condition is None:
140
+ # Remove any existing condition=None entries
141
+ self.after_works = [c for c in self.after_works if c.condition is not None]
142
+ # Add the new one at the end
143
+ self.after_works.append(condition)
144
+ else:
145
+ # For regular conditions, check if we need to move condition=None to the end
146
+ none_conditions = [c for c in self.after_works if c.condition is None]
147
+ if none_conditions:
148
+ # Remove the None condition temporarily
149
+ self.after_works = [c for c in self.after_works if c.condition is not None]
150
+ # Add the new regular condition
151
+ self.after_works.append(condition)
152
+ # Re-add the None condition at the end
153
+ self.after_works.append(none_conditions[0])
154
+ else:
155
+ # No None condition exists, just append
156
+ self.after_works.append(condition)
157
+
158
+ return self
159
+
160
+ def add_after_works(self, conditions: list[OnContextCondition]) -> "Handoffs":
161
+ """
162
+ Add multiple after-work conditions.
163
+
164
+ Special handling for condition=None entries:
165
+ - Only one condition=None entry is allowed (the fallback)
166
+ - It will always be placed at the end of the list
167
+ - If multiple condition=None entries are provided, only the last one is kept
168
+
169
+ Args:
170
+ conditions: List of OnContextConditions to add
171
+
172
+ Returns:
173
+ Self for method chaining
174
+ """
175
+ # Validate that it is a list of OnContextConditions
176
+ if not all(isinstance(condition, OnContextCondition) for condition in conditions):
177
+ raise TypeError("All conditions must be of type OnContextCondition")
178
+
179
+ # Separate conditions with None and without None
180
+ none_conditions = [c for c in conditions if c.condition is None]
181
+ regular_conditions = [c for c in conditions if c.condition is not None]
182
+
183
+ # Remove any existing condition=None entries
184
+ self.after_works = [c for c in self.after_works if c.condition is not None]
185
+
186
+ # Add regular conditions
187
+ self.after_works.extend(regular_conditions)
188
+
189
+ # Add at most one None condition at the end
190
+ if none_conditions:
191
+ self.after_works.append(none_conditions[-1]) # Use the last one if multiple provided
192
+
117
193
  return self
118
194
 
119
195
  @overload
@@ -186,7 +262,7 @@ class Handoffs(BaseModel):
186
262
  """
187
263
  self.context_conditions.clear()
188
264
  self.llm_conditions.clear()
189
- self.after_work = None
265
+ self.after_works.clear()
190
266
  return self
191
267
 
192
268
  def get_llm_conditions_by_target_type(self, target_type: type) -> list[OnCondition]:
@@ -24,12 +24,12 @@ class OnContextCondition(BaseModel): # noqa: N801
24
24
 
25
25
  Args:
26
26
  target (TransitionTarget): The transition (essentially an agent) to hand off to.
27
- condition (ContextCondition): The context variable based condition for transitioning to the target agent.
27
+ condition (Optional[ContextCondition]): The context variable based condition for transitioning to the target agent. If None, the condition always evaluates to True.
28
28
  available (AvailableCondition): Optional condition to determine if this OnCondition is included for the LLM to evaluate based on context variables using classes like StringAvailableCondition and ContextExpressionAvailableCondition.
29
29
  """
30
30
 
31
31
  target: TransitionTarget
32
- condition: ContextCondition
32
+ condition: Optional[ContextCondition] = None
33
33
  available: Optional[AvailableCondition] = None
34
34
 
35
35
  def has_target_type(self, target_type: type) -> bool:
@@ -152,7 +152,13 @@ class Pattern(ABC):
152
152
  manager = create_group_manager(groupchat, self.group_manager_args, self.agents, self.group_after_work)
153
153
 
154
154
  # Point all agent's context variables to this function's context_variables
155
- setup_context_variables(tool_executor, self.agents, manager, self.context_variables)
155
+ setup_context_variables(
156
+ tool_execution=tool_executor,
157
+ agents=self.agents,
158
+ manager=manager,
159
+ user_agent=self.user_agent,
160
+ context_variables=self.context_variables,
161
+ )
156
162
 
157
163
  # Link all agents with the GroupChatManager to allow access to the group chat
158
164
  link_agents_to_group_manager(groupchat.agents, manager)
@@ -1489,10 +1489,10 @@ class GroupChatManager(ConversableAgent):
1489
1489
  for agent in self._groupchat.agents:
1490
1490
  if agent.name == message["name"]:
1491
1491
  # An agent`s message is sent to the Group Chat Manager
1492
- agent.a_send(message, self, request_reply=False, silent=True)
1492
+ await agent.a_send(message, self, request_reply=False, silent=True)
1493
1493
  else:
1494
1494
  # Otherwise, messages are sent from the Group Chat Manager to the agent
1495
- self.a_send(message, agent, request_reply=False, silent=True)
1495
+ await self.a_send(message, agent, request_reply=False, silent=True)
1496
1496
 
1497
1497
  # Add previous message to the new groupchat, if it's an admin message the name may not match so add the message directly
1498
1498
  if message_speaker_agent:
@@ -104,7 +104,7 @@ def parse_oai_message(message: Union[dict[str, Any], str], role: str, adressee:
104
104
  return oai_message
105
105
 
106
106
 
107
- class SwarmableAgent:
107
+ class SwarmableAgent(Agent):
108
108
  """A class for an agent that can participate in a swarm chat."""
109
109
 
110
110
  def __init__(
@@ -239,7 +239,7 @@ class SwarmableAgent:
239
239
  sender: Optional["Agent"] = None,
240
240
  **kwargs: Any,
241
241
  ) -> Union[str, dict[str, Any], None]:
242
- raise NotImplementedError
242
+ return self.generate_reply(messages=messages, sender=sender, **kwargs)
243
243
 
244
244
  async def a_receive(
245
245
  self,
@@ -247,7 +247,7 @@ class SwarmableAgent:
247
247
  sender: "Agent",
248
248
  request_reply: Optional[bool] = None,
249
249
  ) -> None:
250
- raise NotImplementedError
250
+ self.receive(message, sender, request_reply)
251
251
 
252
252
  async def a_send(
253
253
  self,
@@ -255,7 +255,7 @@ class SwarmableAgent:
255
255
  recipient: "Agent",
256
256
  request_reply: Optional[bool] = None,
257
257
  ) -> None:
258
- raise NotImplementedError
258
+ self.send(message, recipient, request_reply)
259
259
 
260
260
  @property
261
261
  def chat_messages(self) -> dict[Agent, list[dict[str, Any]]]:
@@ -293,6 +293,14 @@ class SwarmableAgent:
293
293
  def _raise_exception_on_async_reply_functions(self) -> None:
294
294
  pass
295
295
 
296
+ def set_ui_tools(self, tools: Optional[list] = None) -> None:
297
+ """Set UI tools for the agent."""
298
+ pass
299
+
300
+ def unset_ui_tools(self) -> None:
301
+ """Unset UI tools for the agent."""
302
+ pass
303
+
296
304
  @staticmethod
297
305
  def _last_msg_as_summary(sender: Agent, recipient: Agent, summary_args: Optional[dict[str, Any]]) -> str:
298
306
  """Get a chat summary from the last message of the recipient."""
@@ -42,7 +42,8 @@ TASK_MANAGER_NAME = "TaskManagerAgent"
42
42
  TASK_MANAGER_SYSTEM_MESSAGE = """
43
43
  You are a task manager agent. You have 2 priorities:
44
44
  1. You initiate the tasks which updates the context variables based on the task decisions (DocumentTask) from the DocumentTriageAgent.
45
- If the DocumentTriageAgent has suggested any ingestions or queries, call initiate_tasks to record them.
45
+ ALWAYS call initiate_tasks first when you receive a message from the DocumentTriageAgent, even if you think there are no new tasks.
46
+ This ensures that any new ingestions or queries from the triage agent are properly recorded.
46
47
  Put all ingestion and query tasks into the one tool call.
47
48
  i.e. output
48
49
  {
@@ -75,7 +76,7 @@ TASK_MANAGER_SYSTEM_MESSAGE = """
75
76
  Transfer to the summary agent if all ingestion and query tasks are done.
76
77
  """
77
78
 
78
- DEFAULT_ERROR_SWARM_MESSAGE: str = """
79
+ DEFAULT_ERROR_GROUP_CHAT_MESSAGE: str = """
79
80
  Document Agent failed to perform task.
80
81
  """
81
82
 
@@ -147,7 +148,7 @@ class DocAgent(ConversableAgent):
147
148
  """
148
149
  The DocAgent is responsible for ingest and querying documents.
149
150
 
150
- Internally, it generates a group of swarm agents to solve tasks.
151
+ Internally, it generates a group chat with a set of agents to ingest, query, and summarize.
151
152
  """
152
153
 
153
154
  def __init__(
@@ -196,7 +197,7 @@ class DocAgent(ConversableAgent):
196
197
  llm_config=llm_config,
197
198
  human_input_mode="NEVER",
198
199
  )
199
- self.register_reply([ConversableAgent, None], self.generate_inner_swarm_reply, position=0)
200
+ self.register_reply([ConversableAgent, None], self.generate_inner_group_chat_reply, position=0)
200
201
 
201
202
  self.context_variables: ContextVariables = ContextVariables(
202
203
  data={
@@ -210,7 +211,15 @@ class DocAgent(ConversableAgent):
210
211
  self._triage_agent = DocumentTriageAgent(llm_config=llm_config)
211
212
 
212
213
  def create_error_agent_prompt(agent: ConversableAgent, messages: list[dict[str, Any]]) -> str:
213
- """Create the error agent prompt, primarily used to update ingested documents for ending"""
214
+ """Create the error agent prompt, primarily used to update ingested documents for ending.
215
+
216
+ Args:
217
+ agent: The conversable agent requesting the prompt
218
+ messages: List of conversation messages
219
+
220
+ Returns:
221
+ str: The error manager system message
222
+ """
214
223
  update_ingested_documents()
215
224
 
216
225
  return ERROR_MANAGER_SYSTEM_MESSAGE
@@ -223,7 +232,11 @@ class DocAgent(ConversableAgent):
223
232
  )
224
233
 
225
234
  def update_ingested_documents() -> None:
226
- """Updates the list of ingested documents, persisted so we can keep a list over multiple replies"""
235
+ """Updates the list of ingested documents, persisted so we can keep a list over multiple replies.
236
+
237
+ This function updates self.documents_ingested with any new documents that have been ingested
238
+ by the triage agent, ensuring persistence across multiple DocAgent interactions.
239
+ """
227
240
  agent_documents_ingested = self._triage_agent.context_variables.get("DocumentsIngested", [])
228
241
  # Update self.documents_ingested with any new documents ingested
229
242
  for doc in agent_documents_ingested: # type: ignore[union-attr]
@@ -234,21 +247,162 @@ class DocAgent(ConversableAgent):
234
247
  ingestions: Annotated[list[Ingest], Field(description="List of documents, files, and URLs to ingest")]
235
248
  queries: Annotated[list[Query], Field(description="List of queries to run")]
236
249
 
250
+ def _deduplicate_ingestions(
251
+ new_ingestions: list[Ingest], existing_ingestions: list[Ingest], documents_ingested: list[str]
252
+ ) -> tuple[list[Ingest], list[str]]:
253
+ """Deduplicate ingestions against existing pending and already ingested documents.
254
+
255
+ Args:
256
+ new_ingestions: List of new ingestion requests to process
257
+ existing_ingestions: List of ingestions already pending
258
+ documents_ingested: List of document paths already ingested
259
+
260
+ Returns:
261
+ tuple: (new_unique_ingestions, ignored_duplicate_paths)
262
+ """
263
+ unique_ingestions = []
264
+ ignored_paths = []
265
+
266
+ for ingestion in new_ingestions:
267
+ ingestion_path = ingestion.path_or_url
268
+ # Check if already in pending ingestions
269
+ already_pending = any(existing.path_or_url == ingestion_path for existing in existing_ingestions)
270
+ # Check if already ingested
271
+ already_ingested = ingestion_path in documents_ingested
272
+
273
+ if already_pending or already_ingested:
274
+ ignored_paths.append(ingestion_path)
275
+ else:
276
+ unique_ingestions.append(ingestion)
277
+
278
+ return unique_ingestions, ignored_paths
279
+
280
+ def _deduplicate_queries(
281
+ new_queries: list[Query], existing_queries: list[Query]
282
+ ) -> tuple[list[Query], list[str]]:
283
+ """Deduplicate queries against existing pending queries.
284
+
285
+ Args:
286
+ new_queries: List of new query requests to process
287
+ existing_queries: List of queries already pending
288
+
289
+ Returns:
290
+ tuple: (new_unique_queries, ignored_duplicate_query_texts)
291
+ """
292
+ unique_queries = []
293
+ ignored_query_texts = []
294
+
295
+ for query in new_queries:
296
+ query_text = query.query
297
+ # Check if query already exists in pending queries
298
+ already_pending = any(existing.query == query_text for existing in existing_queries)
299
+
300
+ if already_pending:
301
+ ignored_query_texts.append(query_text)
302
+ else:
303
+ unique_queries.append(query)
304
+
305
+ return unique_queries, ignored_query_texts
306
+
307
+ def _build_response_message(
308
+ added_ingestions: int, ignored_ingestions: list[str], added_queries: int, ignored_queries: list[str]
309
+ ) -> str:
310
+ """Build a descriptive response message about what was added/ignored.
311
+
312
+ Args:
313
+ added_ingestions: Number of unique ingestions added
314
+ ignored_ingestions: List of duplicate ingestion paths ignored
315
+ added_queries: Number of unique queries added
316
+ ignored_queries: List of duplicate query texts ignored
317
+
318
+ Returns:
319
+ str: Formatted message describing the results
320
+ """
321
+ messages = []
322
+
323
+ if added_ingestions > 0:
324
+ messages.append(f"Added {added_ingestions} new document(s) for ingestion")
325
+
326
+ if ignored_ingestions:
327
+ messages.append(
328
+ f"Ignored {len(ignored_ingestions)} duplicate document(s): {', '.join(ignored_ingestions)}"
329
+ )
330
+
331
+ if added_queries > 0:
332
+ messages.append(f"Added {added_queries} new query/queries")
333
+
334
+ if ignored_queries:
335
+ messages.append(f"Ignored {len(ignored_queries)} duplicate query/queries: {', '.join(ignored_queries)}")
336
+
337
+ if messages:
338
+ return "; ".join(messages)
339
+ else:
340
+ return "All requested tasks were duplicates and ignored"
341
+
237
342
  def initiate_tasks(
238
343
  task_init_info: Annotated[TaskInitInfo, "Documents, Files, URLs to ingest and the queries to run"],
239
344
  context_variables: Annotated[ContextVariables, "Context variables"],
240
345
  ) -> ReplyResult:
241
- """Add documents to ingest and queries to answer when received."""
346
+ """Add documents to ingest and queries to answer when received.
347
+
348
+ Args:
349
+ task_init_info: Information about documents to ingest and queries to run
350
+ context_variables: The current context variables containing task state
351
+
352
+ Returns:
353
+ ReplyResult: Contains response message, updated context, and target agent
354
+ """
242
355
  ingestions = task_init_info.ingestions
243
356
  queries = task_init_info.queries
244
357
 
245
358
  if "TaskInitiated" in context_variables:
246
- return ReplyResult(message="Task already initiated", context_variables=context_variables)
247
- context_variables["DocumentsToIngest"] = ingestions
248
- context_variables["QueriesToRun"] = [query for query in queries]
249
- context_variables["TaskInitiated"] = True
359
+ # Handle follow-up tasks with deduplication
360
+ added_ingestions_count = 0
361
+ ignored_ingestions = []
362
+ added_queries_count = 0
363
+ ignored_queries = []
364
+
365
+ if ingestions:
366
+ existing_ingestions: list[Ingest] = context_variables.get("DocumentsToIngest", []) # type: ignore[assignment]
367
+ documents_ingested: list[str] = context_variables.get("DocumentsIngested", []) # type: ignore[assignment]
368
+
369
+ unique_ingestions, ignored_ingestion_paths = _deduplicate_ingestions(
370
+ ingestions, existing_ingestions, documents_ingested
371
+ )
372
+
373
+ if unique_ingestions:
374
+ context_variables["DocumentsToIngest"] = existing_ingestions + unique_ingestions
375
+ added_ingestions_count = len(unique_ingestions)
376
+
377
+ ignored_ingestions = ignored_ingestion_paths
378
+
379
+ if queries:
380
+ existing_queries: list[Query] = context_variables.get("QueriesToRun", []) # type: ignore[assignment]
381
+
382
+ unique_queries, ignored_query_texts = _deduplicate_queries(queries, existing_queries)
383
+
384
+ if unique_queries:
385
+ context_variables["QueriesToRun"] = existing_queries + unique_queries
386
+ added_queries_count = len(unique_queries)
387
+
388
+ ignored_queries = ignored_query_texts
389
+
390
+ if not ingestions and not queries:
391
+ return ReplyResult(message="No new tasks to initiate", context_variables=context_variables)
392
+
393
+ response_message = _build_response_message(
394
+ added_ingestions_count, ignored_ingestions, added_queries_count, ignored_queries
395
+ )
396
+
397
+ else:
398
+ # First time initialization - no deduplication needed
399
+ context_variables["DocumentsToIngest"] = ingestions
400
+ context_variables["QueriesToRun"] = [query for query in queries]
401
+ context_variables["TaskInitiated"] = True
402
+ response_message = "Updated context variables with task decisions"
403
+
250
404
  return ReplyResult(
251
- message="Updated context variables with task decisions",
405
+ message=response_message,
252
406
  context_variables=context_variables,
253
407
  target=AgentNameTarget(agent_name=TASK_MANAGER_NAME),
254
408
  )
@@ -271,7 +425,14 @@ class DocAgent(ConversableAgent):
271
425
  )
272
426
 
273
427
  def execute_rag_query(context_variables: ContextVariables) -> ReplyResult: # type: ignore[type-arg]
274
- """Execute outstanding RAG queries, call the tool once for each outstanding query. Call this tool with no arguments."""
428
+ """Execute outstanding RAG queries, call the tool once for each outstanding query. Call this tool with no arguments.
429
+
430
+ Args:
431
+ context_variables: The current context variables containing queries to run
432
+
433
+ Returns:
434
+ ReplyResult: Contains query answer, updated context, and target agent
435
+ """
275
436
  if len(context_variables["QueriesToRun"]) == 0:
276
437
  return ReplyResult(
277
438
  target=AgentNameTarget(agent_name=TASK_MANAGER_NAME),
@@ -303,6 +464,9 @@ class DocAgent(ConversableAgent):
303
464
  context_variables["QueriesToRun"].pop(0)
304
465
  context_variables["CompletedTaskCount"] += 1
305
466
  context_variables["QueryResults"].append({"query": query, "answer": answer, "citations": txt_citations})
467
+
468
+ # Query completed
469
+
306
470
  return ReplyResult(message=answer, context_variables=context_variables)
307
471
  except Exception as e:
308
472
  return ReplyResult(
@@ -322,9 +486,17 @@ class DocAgent(ConversableAgent):
322
486
  functions=[execute_rag_query],
323
487
  )
324
488
 
325
- # Summary agent prompt will include the results of the ingestions and swarms
489
+ # Summary agent prompt will include the results of the ingestions and queries
326
490
  def create_summary_agent_prompt(agent: ConversableAgent, messages: list[dict[str, Any]]) -> str:
327
- """Create the summary agent prompt and updates ingested documents"""
491
+ """Create the summary agent prompt and updates ingested documents.
492
+
493
+ Args:
494
+ agent: The conversable agent requesting the prompt
495
+ messages: List of conversation messages
496
+
497
+ Returns:
498
+ str: The summary agent system message with context information
499
+ """
328
500
  update_ingested_documents()
329
501
 
330
502
  documents_to_ingest: list[Ingest] = cast(list[Ingest], agent.context_variables.get("DocumentsToIngest", []))
@@ -368,14 +540,7 @@ class DocAgent(ConversableAgent):
368
540
  expression=ContextExpression(expression="len(${QueriesToRun}) > 0")
369
541
  ),
370
542
  ),
371
- OnContextCondition( # Go to Summary agent if no documents or queries left to run and we have query results
372
- target=AgentTarget(agent=self._summary_agent),
373
- condition=ExpressionContextCondition(
374
- expression=ContextExpression(
375
- expression="len(${DocumentsToIngest}) == 0 and len(${QueriesToRun}) == 0 and len(${QueryResults}) > 0"
376
- )
377
- ),
378
- ),
543
+ # Removed automatic context condition - let task manager decide when to summarize
379
544
  OnCondition(
380
545
  target=AgentTarget(agent=self._summary_agent),
381
546
  condition=StringLLMCondition(
@@ -396,28 +561,45 @@ class DocAgent(ConversableAgent):
396
561
  # The Error Agent always terminates the DocumentAgent
397
562
  self._error_agent.handoffs.set_after_work(target=TerminateTarget())
398
563
 
399
- self.register_reply([Agent, None], DocAgent.generate_inner_swarm_reply)
564
+ self.register_reply([Agent, None], DocAgent.generate_inner_group_chat_reply)
400
565
 
401
566
  self.documents_ingested: list[str] = []
567
+ self._group_chat_context_variables: Optional[ContextVariables] = None
402
568
 
403
- def generate_inner_swarm_reply(
569
+ def generate_inner_group_chat_reply(
404
570
  self,
405
571
  messages: Optional[Union[list[dict[str, Any]], str]] = None,
406
572
  sender: Optional[Agent] = None,
407
573
  config: Optional[OpenAIWrapper] = None,
408
574
  ) -> tuple[bool, Optional[Union[str, dict[str, Any]]]]:
409
- """Reply function that generates the inner swarm reply for the DocAgent."""
410
- context_variables: ContextVariables = ContextVariables(
411
- data={
412
- "CompletedTaskCount": 0,
413
- "DocumentsToIngest": [],
414
- "DocumentsIngested": self.documents_ingested,
415
- "QueriesToRun": [],
416
- "QueryResults": [],
417
- }
418
- )
575
+ """Reply function that generates the inner group chat reply for the DocAgent.
576
+
577
+ Args:
578
+ messages: Input messages to process
579
+ sender: The agent that sent the message
580
+ config: OpenAI wrapper configuration
581
+
582
+ Returns:
583
+ tuple: (should_terminate, reply_message)
584
+ """
585
+ # Use existing context_variables if available, otherwise create new ones
586
+ if hasattr(self, "_group_chat_context_variables") and self._group_chat_context_variables is not None:
587
+ context_variables = self._group_chat_context_variables
588
+ # Reset for the new run
589
+ context_variables["DocumentsToIngest"] = [] # type: ignore[index]
590
+ else:
591
+ context_variables = ContextVariables(
592
+ data={
593
+ "CompletedTaskCount": 0,
594
+ "DocumentsToIngest": [],
595
+ "DocumentsIngested": self.documents_ingested,
596
+ "QueriesToRun": [],
597
+ "QueryResults": [],
598
+ }
599
+ )
600
+ self._group_chat_context_variables = context_variables
419
601
 
420
- swarm_agents = [
602
+ group_chat_agents = [
421
603
  self._triage_agent,
422
604
  self._task_manager_agent,
423
605
  self._data_ingestion_agent,
@@ -428,7 +610,7 @@ class DocAgent(ConversableAgent):
428
610
 
429
611
  agent_pattern = DefaultPattern(
430
612
  initial_agent=self._triage_agent,
431
- agents=swarm_agents,
613
+ agents=group_chat_agents,
432
614
  context_variables=context_variables,
433
615
  group_after_work=TerminateTarget(),
434
616
  )
@@ -441,13 +623,23 @@ class DocAgent(ConversableAgent):
441
623
  # If we finish with the error agent, we return their message which contains the error
442
624
  return True, chat_result.summary
443
625
  if last_speaker != self._summary_agent:
444
- # If the swarm finished but not with the summary agent, we assume something has gone wrong with the flow
445
- return True, DEFAULT_ERROR_SWARM_MESSAGE
626
+ # If the group chat finished but not with the summary agent, we assume something has gone wrong with the flow
627
+ return True, DEFAULT_ERROR_GROUP_CHAT_MESSAGE
446
628
 
447
629
  return True, chat_result.summary
448
630
 
449
631
  def _get_document_input_message(self, messages: Optional[Union[list[dict[str, Any]], str]]) -> str: # type: ignore[type-arg]
450
- """Gets and validates the input message(s) for the document agent."""
632
+ """Gets and validates the input message(s) for the document agent.
633
+
634
+ Args:
635
+ messages: Input messages as string or list of message dictionaries
636
+
637
+ Returns:
638
+ str: The extracted message content
639
+
640
+ Raises:
641
+ NotImplementedError: If messages format is invalid
642
+ """
451
643
  if isinstance(messages, str):
452
644
  return messages
453
645
  elif (
@@ -763,9 +763,10 @@ class ExecuteFunctionEvent(BaseEvent):
763
763
  class ExecutedFunctionEvent(BaseEvent):
764
764
  func_name: str
765
765
  call_id: Optional[str] = None
766
- arguments: dict[str, Any]
767
- content: str
766
+ arguments: Optional[dict[str, Any]]
767
+ content: Any
768
768
  recipient: str
769
+ is_exec_success: bool = True
769
770
 
770
771
  def __init__(
771
772
  self,
@@ -773,9 +774,10 @@ class ExecutedFunctionEvent(BaseEvent):
773
774
  uuid: Optional[UUID] = None,
774
775
  func_name: str,
775
776
  call_id: Optional[str] = None,
776
- arguments: dict[str, Any],
777
- content: str,
777
+ arguments: Optional[dict[str, Any]],
778
+ content: Any,
778
779
  recipient: Union["Agent", str],
780
+ is_exec_success: bool = True,
779
781
  ):
780
782
  super().__init__(
781
783
  uuid=uuid,
@@ -785,6 +787,7 @@ class ExecutedFunctionEvent(BaseEvent):
785
787
  content=content,
786
788
  recipient=recipient.name if hasattr(recipient, "name") else recipient,
787
789
  )
790
+ self.is_exec_success = is_exec_success
788
791
 
789
792
  def print(self, f: Optional[Callable[..., Any]] = None) -> None:
790
793
  f = f or print