Skip to content

Commit b0c8002

Browse files
ekzhuvictordibia
andauthored
Use structured output for m1 orchestrator (#6540)
Use structured output when available in m1 orchestrator Co-authored-by: Victor Dibia <victordibia@microsoft.com>
1 parent cd49d71 commit b0c8002

File tree

2 files changed

+33
-1
lines changed

2 files changed

+33
-1
lines changed

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
ORCHESTRATOR_TASK_LEDGER_FULL_PROMPT,
4848
ORCHESTRATOR_TASK_LEDGER_PLAN_PROMPT,
4949
ORCHESTRATOR_TASK_LEDGER_PLAN_UPDATE_PROMPT,
50+
LedgerEntry,
5051
)
5152

5253
trace_logger = logging.getLogger(TRACE_LOGGER_NAME)
@@ -309,7 +310,18 @@ async def _orchestrate_step(self, cancellation_token: CancellationToken) -> None
309310
assert self._max_json_retries > 0
310311
key_error: bool = False
311312
for _ in range(self._max_json_retries):
312-
response = await self._model_client.create(self._get_compatible_context(context), json_output=True)
313+
if self._model_client.model_info.get("structured_output", False):
314+
response = await self._model_client.create(
315+
self._get_compatible_context(context), json_output=LedgerEntry
316+
)
317+
elif self._model_client.model_info.get("json_output", False):
318+
response = await self._model_client.create(
319+
self._get_compatible_context(context), cancellation_token=cancellation_token, json_output=True
320+
)
321+
else:
322+
response = await self._model_client.create(
323+
self._get_compatible_context(context), cancellation_token=cancellation_token
324+
)
313325
ledger_str = response.content
314326
try:
315327
assert isinstance(ledger_str, str)

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_prompts.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from pydantic import BaseModel
2+
13
ORCHESTRATOR_SYSTEM_MESSAGE = ""
24

35

@@ -98,6 +100,24 @@
98100
"""
99101

100102

103+
class LedgerEntryBooleanAnswer(BaseModel):
104+
reason: str
105+
answer: bool
106+
107+
108+
class LedgerEntryStringAnswer(BaseModel):
109+
reason: str
110+
answer: str
111+
112+
113+
class LedgerEntry(BaseModel):
114+
is_request_satisfied: LedgerEntryBooleanAnswer
115+
is_in_loop: LedgerEntryBooleanAnswer
116+
is_progress_being_made: LedgerEntryBooleanAnswer
117+
next_speaker: LedgerEntryStringAnswer
118+
instruction_or_question: LedgerEntryStringAnswer
119+
120+
101121
ORCHESTRATOR_TASK_LEDGER_FACTS_UPDATE_PROMPT = """As a reminder, we are working to solve the following task:
102122
103123
{task}

0 commit comments

Comments
 (0)