From 229f0d5f7c903bf43b45e857e3177f510b896291 Mon Sep 17 00:00:00 2001
From: Zhonglin Han <hanzlfs@users.noreply.github.com>
Date: Tue, 28 Jan 2025 05:05:38 -0800
Subject: [PATCH] Agent response format (#660)

# What does this PR do?

Add response format for agents structured output.

- [ ] Using structured output for agents (interior_design app as an
example) (#issue)
https://github.com/meta-llama/llama-stack-apps/issues/122


## Test Plan
E2E test plan with llama-stack-apps interior_design

Please describe:
 Test ran:

 - provide instructions so it can be reproduced.
 Start your distro:
llama stack run llama_stack/templates/fireworks/run.yaml --env
FIREWORKS_API_KEY=<API_KEY>

Run api test:
```PYTHONPATH=. python examples/interior_design_assistant/api.py localhost 5000 examples/interior_design_assistant/resources/documents/ examples/interior_design_assistant/resources/images/fireplaces```


## Sources
Results:
https://github.com/meta-llama/llama-stack-client-python/pull/72

## Before submitting

- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Ran pre-commit to handle lint / formatting issues.
- [x] Read the [contributor guideline](https://github.com/meta-llama/llama-stack/blob/main/CONTRIBUTING.md),
      Pull Request section?
- [ ] Updated relevant documentation.
- [ ] Wrote necessary unit or integration tests.
---
 llama_stack/apis/agents/agents.py                                | 1 +
 .../providers/inline/agents/meta_reference/agent_instance.py     | 1 +
 2 files changed, 2 insertions(+)

diff --git a/llama_stack/apis/agents/agents.py b/llama_stack/apis/agents/agents.py
index f62d78390..b6d6e6578 100644
--- a/llama_stack/apis/agents/agents.py
+++ b/llama_stack/apis/agents/agents.py
@@ -165,6 +165,7 @@ class AgentConfig(AgentConfigCommon):
     model: str
     instructions: str
     enable_session_persistence: bool
+    response_format: Optional[ResponseFormat] = None
 
 
 class AgentConfigOverridablePerTurn(AgentConfigCommon):
diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
index 32801e514..eca7364d7 100644
--- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
+++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py
@@ -512,6 +512,7 @@ class ChatAgent(ShieldRunnerMixin):
                         if tool_to_group.get(tool.tool_name, None) != RAG_TOOL_GROUP
                     ],
                     tool_prompt_format=self.agent_config.tool_prompt_format,
+                    response_format=self.agent_config.response_format,
                     stream=True,
                     sampling_params=sampling_params,
                 ):