This commit is contained in:
Xi Yan 2025-02-20 19:48:54 -08:00
parent 6d08a935ba
commit 9a07e709ee
2 changed files with 5 additions and 5 deletions

View file

@ -23,7 +23,6 @@ from llama_stack.apis.agents import (
AgentConfig, AgentConfig,
AgentToolGroup, AgentToolGroup,
AgentToolGroupWithArgs, AgentToolGroupWithArgs,
AgentTurnContinueRequest,
AgentTurnCreateRequest, AgentTurnCreateRequest,
AgentTurnResponseEvent, AgentTurnResponseEvent,
AgentTurnResponseEventType, AgentTurnResponseEventType,
@ -34,6 +33,7 @@ from llama_stack.apis.agents import (
AgentTurnResponseTurnAwaitingInputPayload, AgentTurnResponseTurnAwaitingInputPayload,
AgentTurnResponseTurnCompletePayload, AgentTurnResponseTurnCompletePayload,
AgentTurnResponseTurnStartPayload, AgentTurnResponseTurnStartPayload,
AgentTurnResumeRequest,
Attachment, Attachment,
Document, Document,
InferenceStep, InferenceStep,
@ -247,7 +247,7 @@ class ChatAgent(ShieldRunnerMixin):
yield chunk yield chunk
async def continue_turn(self, request: AgentTurnContinueRequest) -> AsyncGenerator: async def continue_turn(self, request: AgentTurnResumeRequest) -> AsyncGenerator:
with tracing.span("continue_turn") as span: with tracing.span("continue_turn") as span:
span.set_attribute("agent_id", self.agent_id) span.set_attribute("agent_id", self.agent_id)
span.set_attribute("session_id", request.session_id) span.set_attribute("session_id", request.session_id)

View file

@ -20,8 +20,8 @@ from llama_stack.apis.agents import (
AgentSessionCreateResponse, AgentSessionCreateResponse,
AgentStepResponse, AgentStepResponse,
AgentToolGroup, AgentToolGroup,
AgentTurnContinueRequest,
AgentTurnCreateRequest, AgentTurnCreateRequest,
AgentTurnResumeRequest,
Document, Document,
Session, Session,
Turn, Turn,
@ -178,7 +178,7 @@ class MetaReferenceAgentsImpl(Agents):
tool_responses: List[ToolResponseMessage], tool_responses: List[ToolResponseMessage],
stream: Optional[bool] = False, stream: Optional[bool] = False,
) -> AsyncGenerator: ) -> AsyncGenerator:
request = AgentTurnContinueRequest( request = AgentTurnResumeRequest(
agent_id=agent_id, agent_id=agent_id,
session_id=session_id, session_id=session_id,
turn_id=turn_id, turn_id=turn_id,
@ -192,7 +192,7 @@ class MetaReferenceAgentsImpl(Agents):
async def _continue_agent_turn_streaming( async def _continue_agent_turn_streaming(
self, self,
request: AgentTurnContinueRequest, request: AgentTurnResumeRequest,
) -> AsyncGenerator: ) -> AsyncGenerator:
agent = await self.get_agent(request.agent_id) agent = await self.get_agent(request.agent_id)
async for event in agent.continue_turn(request): async for event in agent.continue_turn(request):