feat(config): support input/output schema by fully-qualified code reference

PiperOrigin-RevId: 785541326
This commit is contained in:
Liang Wu
2025-07-21 12:52:44 -07:00
committed by Copybara-Service
parent 2aab1cf98e
commit dfee06ac06
2 changed files with 33 additions and 0 deletions
@@ -163,6 +163,28 @@
"default": null,
"title": "Disallow Transfer To Peers"
},
"input_schema": {
"anyOf": [
{
"$ref": "#/$defs/CodeConfig"
},
{
"type": "null"
}
],
"default": null
},
"output_schema": {
"anyOf": [
{
"$ref": "#/$defs/CodeConfig"
},
{
"type": "null"
}
],
"default": null
},
"output_key": {
"anyOf": [
{
+11
View File
@@ -563,6 +563,7 @@ class LlmAgent(BaseAgent):
config_abs_path: str,
) -> LlmAgent:
from .config_agent_utils import resolve_callbacks
from .config_agent_utils import resolve_code_reference
agent = super().from_config(config, config_abs_path)
if config.model:
@@ -575,6 +576,10 @@ class LlmAgent(BaseAgent):
agent.disallow_transfer_to_peers = config.disallow_transfer_to_peers
if config.include_contents != 'default':
agent.include_contents = config.include_contents
if config.input_schema:
agent.input_schema = resolve_code_reference(config.input_schema)
if config.output_schema:
agent.output_schema = resolve_code_reference(config.output_schema)
if config.output_key:
agent.output_key = config.output_key
if config.tools:
@@ -619,6 +624,12 @@ class LlmAgentConfig(BaseAgentConfig):
disallow_transfer_to_peers: Optional[bool] = None
"""Optional. LlmAgent.disallow_transfer_to_peers."""
input_schema: Optional[CodeConfig] = None
"""Optional. LlmAgent.input_schema."""
output_schema: Optional[CodeConfig] = None
"""Optional. LlmAgent.output_schema."""
output_key: Optional[str] = None
"""Optional. LlmAgent.output_key."""