Skip to content

Commit 34b4de0

Browse files
wyf7107copybara-github
authored andcommitted
fix(cli): Serialize LiteLlm graph models safely
The graph serializer currently passes model fields through when their value is not a simple collection or nested agent. For LiteLlm that leaves the runtime LiteLLMClient object inside the serialized graph payload, so JSON encoding the graph fails. This patch serializes BaseLlm values using their model name string value.model to ensure JSON safety and match the web UI frontend's expected string schema. Merge #5956 closes #5949 Co-authored-by: Yifan Wang <wanyif@google.com> PiperOrigin-RevId: 927453412
1 parent 09003fe commit 34b4de0

2 files changed

Lines changed: 17 additions & 0 deletions

File tree

src/google/adk/cli/utils/graph_serialization.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
logger = logging.getLogger("google_adk." + __name__)
2323

2424
from ...agents.base_agent import BaseAgent
25+
from ...models.base_llm import BaseLlm
2526
from ...tools.base_toolset import BaseToolset
2627

2728
# Node type mapping for cleaner lookup
@@ -227,6 +228,8 @@ def serialize_agent(agent: BaseAgent) -> dict[str, Any]:
227228
# Handle nested agents
228229
if isinstance(value, BaseAgent):
229230
agent_dict[field_name] = serialize_agent(value)
231+
elif isinstance(value, BaseLlm):
232+
agent_dict[field_name] = value.model
230233
# Handle simple types and collections
231234
elif isinstance(value, (str, int, float, bool, list, dict)):
232235
agent_dict[field_name] = value

tests/unittests/cli/utils/test_graph_serialization.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,11 @@
1313
# limitations under the License.
1414

1515
"""Tests for graph_serialization edge handling with routing maps."""
16+
import json
1617

18+
from google.adk.agents import LlmAgent
1719
from google.adk.cli.utils.graph_serialization import serialize_agent
20+
from google.adk.models.lite_llm import LiteLlm
1821
from google.adk.tools.base_toolset import BaseToolset
1922
from google.adk.workflow import START
2023
from google.adk.workflow import Workflow
@@ -126,3 +129,14 @@ def __init__(self):
126129
assert len(result['tools']) == 1
127130
assert result['tools'][0]['name'] == 'MockToolset'
128131
assert result['tools'][0]['type'] == 'tool'
132+
133+
def test_serialize_agent_with_litellm_model_is_json_safe() -> None:
134+
agent = LlmAgent(
135+
name='repro',
136+
model=LiteLlm(model='ollama_chat/llama3'),
137+
)
138+
139+
result = serialize_agent(agent)
140+
141+
assert result['model'] == 'ollama_chat/llama3'
142+
json.dumps(result)

0 commit comments

Comments
 (0)