You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: agent generate config err (#1305)
* fix: agent generate config err * fix: resovle comment --------- Co-authored-by: Hangfei Lin <hangfei@google.com> Co-authored-by: genquan9 <49327371+genquan9@users.noreply.github.com>
This commit is contained in:
@@ -23,6 +23,7 @@ from typing import cast
|
||||
from typing import Dict
|
||||
from typing import Generator
|
||||
from typing import Iterable
|
||||
from typing import List
|
||||
from typing import Literal
|
||||
from typing import Optional
|
||||
from typing import Tuple
|
||||
@@ -481,16 +482,22 @@ def _message_to_generate_content_response(
|
||||
|
||||
def _get_completion_inputs(
|
||||
llm_request: LlmRequest,
|
||||
) -> tuple[Iterable[Message], Iterable[dict]]:
|
||||
"""Converts an LlmRequest to litellm inputs.
|
||||
) -> Tuple[
|
||||
List[Message],
|
||||
Optional[List[Dict]],
|
||||
Optional[types.SchemaUnion],
|
||||
Optional[Dict],
|
||||
]:
|
||||
"""Converts an LlmRequest to litellm inputs and extracts generation params.
|
||||
|
||||
Args:
|
||||
llm_request: The LlmRequest to convert.
|
||||
|
||||
Returns:
|
||||
The litellm inputs (message list, tool dictionary and response format).
|
||||
The litellm inputs (message list, tool dictionary, response format, and generation params).
|
||||
"""
|
||||
messages = []
|
||||
# 1. Construct messages
|
||||
messages: List[Message] = []
|
||||
for content in llm_request.contents or []:
|
||||
message_param_or_list = _content_to_message_param(content)
|
||||
if isinstance(message_param_or_list, list):
|
||||
@@ -507,7 +514,8 @@ def _get_completion_inputs(
|
||||
),
|
||||
)
|
||||
|
||||
tools = None
|
||||
# 2. Convert tool declarations
|
||||
tools: Optional[List[Dict]] = None
|
||||
if (
|
||||
llm_request.config
|
||||
and llm_request.config.tools
|
||||
@@ -518,12 +526,39 @@ def _get_completion_inputs(
|
||||
for tool in llm_request.config.tools[0].function_declarations
|
||||
]
|
||||
|
||||
response_format = None
|
||||
# 3. Handle response format
|
||||
response_format: Optional[types.SchemaUnion] = (
|
||||
llm_request.config.response_schema if llm_request.config else None
|
||||
)
|
||||
|
||||
if llm_request.config.response_schema:
|
||||
response_format = llm_request.config.response_schema
|
||||
# 4. Extract generation parameters
|
||||
generation_params: Optional[Dict] = None
|
||||
if llm_request.config:
|
||||
config_dict = llm_request.config.model_dump(exclude_none=True)
|
||||
# Generate LiteLlm parameters here,
|
||||
# Following https://docs.litellm.ai/docs/completion/input.
|
||||
generation_params = {}
|
||||
param_mapping = {
|
||||
"max_output_tokens": "max_completion_tokens",
|
||||
"stop_sequences": "stop",
|
||||
}
|
||||
for key in (
|
||||
"temperature",
|
||||
"max_output_tokens",
|
||||
"top_p",
|
||||
"top_k",
|
||||
"stop_sequences",
|
||||
"presence_penalty",
|
||||
"frequency_penalty",
|
||||
):
|
||||
if key in config_dict:
|
||||
mapped_key = param_mapping.get(key, key)
|
||||
generation_params[mapped_key] = config_dict[key]
|
||||
|
||||
return messages, tools, response_format
|
||||
if not generation_params:
|
||||
generation_params = None
|
||||
|
||||
return messages, tools, response_format, generation_params
|
||||
|
||||
|
||||
def _build_function_declaration_log(
|
||||
@@ -660,7 +695,9 @@ class LiteLlm(BaseLlm):
|
||||
self._maybe_append_user_content(llm_request)
|
||||
logger.debug(_build_request_log(llm_request))
|
||||
|
||||
messages, tools, response_format = _get_completion_inputs(llm_request)
|
||||
messages, tools, response_format, generation_params = (
|
||||
_get_completion_inputs(llm_request)
|
||||
)
|
||||
|
||||
completion_args = {
|
||||
"model": self.model,
|
||||
@@ -668,7 +705,13 @@ class LiteLlm(BaseLlm):
|
||||
"tools": tools,
|
||||
"response_format": response_format,
|
||||
}
|
||||
completion_args.update(self._additional_args)
|
||||
|
||||
# Merge additional arguments and generation parameters safely
|
||||
if hasattr(self, "_additional_args") and self._additional_args:
|
||||
completion_args.update(self._additional_args)
|
||||
|
||||
if generation_params:
|
||||
completion_args.update(generation_params)
|
||||
|
||||
if stream:
|
||||
text = ""
|
||||
|
||||
@@ -13,7 +13,6 @@
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import Mock
|
||||
|
||||
@@ -1430,3 +1429,35 @@ async def test_generate_content_async_non_compliant_multiple_function_calls(
|
||||
assert final_response.content.parts[1].function_call.name == "function_2"
|
||||
assert final_response.content.parts[1].function_call.id == "1"
|
||||
assert final_response.content.parts[1].function_call.args == {"arg": "value2"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
def test_get_completion_inputs_generation_params():
|
||||
# Test that generation_params are extracted and mapped correctly
|
||||
req = LlmRequest(
|
||||
contents=[
|
||||
types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
|
||||
],
|
||||
config=types.GenerateContentConfig(
|
||||
temperature=0.33,
|
||||
max_output_tokens=123,
|
||||
top_p=0.88,
|
||||
top_k=7,
|
||||
stop_sequences=["foo", "bar"],
|
||||
presence_penalty=0.1,
|
||||
frequency_penalty=0.2,
|
||||
),
|
||||
)
|
||||
from google.adk.models.lite_llm import _get_completion_inputs
|
||||
|
||||
_, _, _, generation_params = _get_completion_inputs(req)
|
||||
assert generation_params["temperature"] == 0.33
|
||||
assert generation_params["max_completion_tokens"] == 123
|
||||
assert generation_params["top_p"] == 0.88
|
||||
assert generation_params["top_k"] == 7
|
||||
assert generation_params["stop"] == ["foo", "bar"]
|
||||
assert generation_params["presence_penalty"] == 0.1
|
||||
assert generation_params["frequency_penalty"] == 0.2
|
||||
# Should not include max_output_tokens
|
||||
assert "max_output_tokens" not in generation_params
|
||||
assert "stop_sequences" not in generation_params
|
||||
|
||||
Reference in New Issue
Block a user