add flag allow_turn_resume

This commit is contained in:
Xi Yan 2025-02-20 21:49:06 -08:00
parent 702e74da8e
commit e2bfd165d2
5 changed files with 12 additions and 1 deletions

View file

@ -4287,6 +4287,9 @@
}, },
"tool_config": { "tool_config": {
"$ref": "#/components/schemas/ToolConfig" "$ref": "#/components/schemas/ToolConfig"
},
"allow_turn_resume": {
"type": "boolean"
} }
}, },
"additionalProperties": false, "additionalProperties": false,

View file

@ -2779,6 +2779,8 @@ components:
$ref: '#/components/schemas/AgentTool' $ref: '#/components/schemas/AgentTool'
tool_config: tool_config:
$ref: '#/components/schemas/ToolConfig' $ref: '#/components/schemas/ToolConfig'
allow_turn_resume:
type: boolean
additionalProperties: false additionalProperties: false
required: required:
- messages - messages

View file

@ -296,6 +296,9 @@ class AgentTurnCreateRequest(AgentConfigOverridablePerTurn):
stream: Optional[bool] = False stream: Optional[bool] = False
tool_config: Optional[ToolConfig] = None tool_config: Optional[ToolConfig] = None
# TODO (xiyan): used for backward compatibility, update for 0.1.5
allow_turn_resume: Optional[bool] = False
@json_schema_type @json_schema_type
class AgentTurnResumeRequest(BaseModel): class AgentTurnResumeRequest(BaseModel):
@ -352,6 +355,7 @@ class Agents(Protocol):
documents: Optional[List[Document]] = None, documents: Optional[List[Document]] = None,
toolgroups: Optional[List[AgentToolGroup]] = None, toolgroups: Optional[List[AgentToolGroup]] = None,
tool_config: Optional[ToolConfig] = None, tool_config: Optional[ToolConfig] = None,
allow_turn_resume: Optional[bool] = False,
) -> Union[Turn, AsyncIterator[AgentTurnResponseStreamChunk]]: ... ) -> Union[Turn, AsyncIterator[AgentTurnResponseStreamChunk]]: ...
@webmethod( @webmethod(

View file

@ -227,7 +227,7 @@ class ChatAgent(ShieldRunnerMixin):
) )
await self.storage.add_turn_to_session(request.session_id, turn) await self.storage.add_turn_to_session(request.session_id, turn)
if output_message.tool_calls: if output_message.tool_calls and request.allow_turn_resume:
chunk = AgentTurnResponseStreamChunk( chunk = AgentTurnResponseStreamChunk(
event=AgentTurnResponseEvent( event=AgentTurnResponseEvent(
payload=AgentTurnResponseTurnAwaitingInputPayload( payload=AgentTurnResponseTurnAwaitingInputPayload(

View file

@ -146,6 +146,7 @@ class MetaReferenceAgentsImpl(Agents):
documents: Optional[List[Document]] = None, documents: Optional[List[Document]] = None,
stream: Optional[bool] = False, stream: Optional[bool] = False,
tool_config: Optional[ToolConfig] = None, tool_config: Optional[ToolConfig] = None,
allow_turn_resume: Optional[bool] = False,
) -> AsyncGenerator: ) -> AsyncGenerator:
request = AgentTurnCreateRequest( request = AgentTurnCreateRequest(
agent_id=agent_id, agent_id=agent_id,
@ -155,6 +156,7 @@ class MetaReferenceAgentsImpl(Agents):
toolgroups=toolgroups, toolgroups=toolgroups,
documents=documents, documents=documents,
tool_config=tool_config, tool_config=tool_config,
allow_turn_resume=allow_turn_resume,
) )
if stream: if stream:
return self._create_agent_turn_streaming(request) return self._create_agent_turn_streaming(request)