client sdk test fixes

This commit is contained in:
Dinesh Yeduguru 2025-01-07 16:59:01 -08:00
parent c3865faf37
commit efe3189728
3 changed files with 21 additions and 25 deletions

View file

@ -402,7 +402,6 @@ class ChatAgent(ShieldRunnerMixin):
# if the session has a memory bank id, let the memory tool use it
if session_info.memory_bank_id:
query_args["memory_bank_id"] = session_info.memory_bank_id
serialized_args = tracing.serialize_value(query_args)
yield AgentTurnResponseStreamChunk(
event=AgentTurnResponseEvent(
payload=AgentTurnResponseStepProgressPayload(
@ -412,8 +411,8 @@ class ChatAgent(ShieldRunnerMixin):
parse_status=ToolCallParseStatus.success,
content=ToolCall(
call_id="",
tool_name="memory",
arguments=serialized_args,
tool_name=MEMORY_QUERY_TOOL,
arguments={},
),
),
)
@ -435,14 +434,14 @@ class ChatAgent(ShieldRunnerMixin):
tool_calls=[
ToolCall(
call_id="",
tool_name="memory",
arguments=serialized_args,
tool_name=MEMORY_QUERY_TOOL,
arguments={},
)
],
tool_responses=[
ToolResponse(
call_id="",
tool_name="memory",
tool_name=MEMORY_QUERY_TOOL,
content=result.content,
)
],
@ -456,7 +455,7 @@ class ChatAgent(ShieldRunnerMixin):
span.set_attribute("output", result.content)
span.set_attribute("error_code", result.error_code)
span.set_attribute("error_message", result.error_message)
span.set_attribute("tool_name", "memory")
span.set_attribute("tool_name", MEMORY_QUERY_TOOL)
if result.error_code == 0:
last_message = input_messages[-1]
last_message.context = result.content

View file

@ -56,7 +56,7 @@ class MemoryToolRuntimeImpl(ToolsProtocolPrivate, ToolRuntime):
) -> List[ToolDef]:
return [
ToolDef(
name="memory",
name="query_memory",
description="Retrieve context from memory",
parameters=[
ToolParameter(