llama-stack 0.4.3__py3-none-any.whl → 0.5.0__py3-none-any.whl

llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from collections.abc import AsyncIterator
 
 from llama_stack.core.datatypes import AccessRule
 from llama_stack.core.storage.kvstore import InmemoryKVStoreImpl, kvstore_impl
@@ -11,21 +12,21 @@ from llama_stack.log import get_logger
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
 from llama_stack_api import (
     Agents,
+    Connectors,
     Conversations,
+    CreateResponseRequest,
+    DeleteResponseRequest,
     Files,
     Inference,
     ListOpenAIResponseInputItem,
     ListOpenAIResponseObject,
+    ListResponseInputItemsRequest,
+    ListResponsesRequest,
     OpenAIDeleteResponseObject,
-    OpenAIResponseInput,
-    OpenAIResponseInputTool,
-    OpenAIResponseInputToolChoice,
     OpenAIResponseObject,
-    OpenAIResponsePrompt,
-    OpenAIResponseText,
-    Order,
+    OpenAIResponseObjectStream,
     Prompts,
-    ResponseGuardrail,
+    RetrieveResponseRequest,
     Safety,
     ToolGroups,
     ToolRuntime,
@@ -50,6 +51,7 @@ class MetaReferenceAgentsImpl(Agents):
         conversations_api: Conversations,
         prompts_api: Prompts,
         files_api: Files,
+        connectors_api: Connectors,
         policy: list[AccessRule],
     ):
         self.config = config
@@ -64,6 +66,7 @@ class MetaReferenceAgentsImpl(Agents):
         self.in_memory_store = InmemoryKVStoreImpl()
         self.openai_responses_impl: OpenAIResponsesImpl | None = None
         self.policy = policy
+        self.connectors_api = connectors_api
 
     async def initialize(self) -> None:
         self.persistence_store = await kvstore_impl(self.config.persistence.agent_state)
@@ -80,6 +83,7 @@ class MetaReferenceAgentsImpl(Agents):
             prompts_api=self.prompts_api,
             files_api=self.files_api,
             vector_stores_config=self.config.vector_stores_config,
+            connectors_api=self.connectors_api,
         )
 
     async def shutdown(self) -> None:
@@ -88,79 +92,72 @@ class MetaReferenceAgentsImpl(Agents):
     # OpenAI responses
     async def get_openai_response(
         self,
-        response_id: str,
+        request: RetrieveResponseRequest,
     ) -> OpenAIResponseObject:
         assert self.openai_responses_impl is not None, "OpenAI responses not initialized"
-        return await self.openai_responses_impl.get_openai_response(response_id)
+        return await self.openai_responses_impl.get_openai_response(request.response_id)
 
     async def create_openai_response(
         self,
-        input: str | list[OpenAIResponseInput],
-        model: str,
-        prompt: OpenAIResponsePrompt | None = None,
-        instructions: str | None = None,
-        parallel_tool_calls: bool | None = True,
-        previous_response_id: str | None = None,
-        conversation: str | None = None,
-        store: bool | None = True,
-        stream: bool | None = False,
-        temperature: float | None = None,
-        text: OpenAIResponseText | None = None,
-        tool_choice: OpenAIResponseInputToolChoice | None = None,
-        tools: list[OpenAIResponseInputTool] | None = None,
-        include: list[str] | None = None,
-        max_infer_iters: int | None = 10,
-        guardrails: list[ResponseGuardrail] | None = None,
-        max_tool_calls: int | None = None,
-        metadata: dict[str, str] | None = None,
-    ) -> OpenAIResponseObject:
+        request: CreateResponseRequest,
+    ) -> OpenAIResponseObject | AsyncIterator[OpenAIResponseObjectStream]:
+        """Create an OpenAI response.
+
+        Returns either a single response object (non-streaming) or an async iterator
+        yielding response stream events (streaming).
+        """
         assert self.openai_responses_impl is not None, "OpenAI responses not initialized"
         result = await self.openai_responses_impl.create_openai_response(
-            input,
-            model,
-            prompt,
-            instructions,
-            previous_response_id,
-            conversation,
-            store,
-            stream,
-            temperature,
-            text,
-            tool_choice,
-            tools,
-            include,
-            max_infer_iters,
-            guardrails,
-            parallel_tool_calls,
-            max_tool_calls,
-            metadata,
+            request.input,
+            request.model,
+            request.prompt,
+            request.instructions,
+            request.previous_response_id,
+            request.conversation,
+            request.store,
+            request.stream,
+            request.temperature,
+            request.text,
+            request.tool_choice,
+            request.tools,
+            request.include,
+            request.max_infer_iters,
+            request.guardrails,
+            request.parallel_tool_calls,
+            request.max_tool_calls,
+            request.max_output_tokens,
+            request.reasoning,
+            request.safety_identifier,
+            request.metadata,
         )
-        return result  # type: ignore[no-any-return]
+        return result
 
     async def list_openai_responses(
         self,
-        after: str | None = None,
-        limit: int | None = 50,
-        model: str | None = None,
-        order: Order | None = Order.desc,
+        request: ListResponsesRequest,
     ) -> ListOpenAIResponseObject:
         assert self.openai_responses_impl is not None, "OpenAI responses not initialized"
-        return await self.openai_responses_impl.list_openai_responses(after, limit, model, order)
+        return await self.openai_responses_impl.list_openai_responses(
+            request.after, request.limit, request.model, request.order
+        )
 
     async def list_openai_response_input_items(
         self,
-        response_id: str,
-        after: str | None = None,
-        before: str | None = None,
-        include: list[str] | None = None,
-        limit: int | None = 20,
-        order: Order | None = Order.desc,
+        request: ListResponseInputItemsRequest,
     ) -> ListOpenAIResponseInputItem:
         assert self.openai_responses_impl is not None, "OpenAI responses not initialized"
         return await self.openai_responses_impl.list_openai_response_input_items(
-            response_id, after, before, include, limit, order
+            request.response_id,
+            request.after,
+            request.before,
+            request.include,
+            request.limit,
+            request.order,
         )
 
-    async def delete_openai_response(self, response_id: str) -> OpenAIDeleteResponseObject:
+    async def delete_openai_response(
+        self,
+        request: DeleteResponseRequest,
+    ) -> OpenAIDeleteResponseObject:
         assert self.openai_responses_impl is not None, "OpenAI responses not initialized"
-        return await self.openai_responses_impl.delete_openai_response(response_id)
+        return await self.openai_responses_impl.delete_openai_response(request.response_id)

llama_stack/providers/inline/agents/meta_reference/responses/openai_responses.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import asyncio
 import re
 import time
 import uuid
@@ -19,11 +18,14 @@ from llama_stack.providers.utils.responses.responses_store import (
 )
 from llama_stack.providers.utils.tools.mcp import MCPSessionManager
 from llama_stack_api import (
+    AddItemsRequest,
+    Connectors,
     ConversationItem,
     Conversations,
     Files,
     Inference,
     InvalidConversationIdError,
+    ListItemsRequest,
     ListOpenAIResponseInputItem,
     ListOpenAIResponseObject,
     OpenAIChatCompletionContentPartParam,
@@ -39,6 +41,7 @@ from llama_stack_api import (
     OpenAIResponseObject,
     OpenAIResponseObjectStream,
     OpenAIResponsePrompt,
+    OpenAIResponseReasoning,
     OpenAIResponseText,
     OpenAIResponseTextFormat,
     OpenAISystemMessageParam,
@@ -83,6 +86,7 @@ class OpenAIResponsesImpl:
         conversations_api: Conversations,
         prompts_api: Prompts,
         files_api: Files,
+        connectors_api: Connectors,
         vector_stores_config=None,
     ):
         self.inference_api = inference_api
@@ -100,6 +104,7 @@ class OpenAIResponsesImpl:
         )
         self.prompts_api = prompts_api
         self.files_api = files_api
+        self.connectors_api = connectors_api
 
     async def _prepend_previous_response(
         self,
@@ -150,7 +155,9 @@ class OpenAIResponsesImpl:
 
             tool_context.recover_tools_from_previous_response(previous_response)
         elif conversation is not None:
-            conversation_items = await self.conversations_api.list_items(conversation, order="asc")
+            conversation_items = await self.conversations_api.list_items(
+                ListItemsRequest(conversation_id=conversation, order="asc")
+            )
 
             # Use stored messages as source of truth (like previous_response.messages)
             stored_messages = await self.responses_store.get_conversation_messages(conversation)
@@ -324,6 +331,125 @@ class OpenAIResponsesImpl:
             messages=messages,
         )
 
+    def _prepare_input_items_for_storage(
+        self,
+        input: str | list[OpenAIResponseInput],
+    ) -> list[OpenAIResponseInput]:
+        """Prepare input items for storage, adding IDs where needed.
+
+        This method is called once at the start of streaming to prepare input items
+        that will be reused across multiple persistence calls during streaming.
+        """
+        new_input_id = f"msg_{uuid.uuid4()}"
+        input_items_data: list[OpenAIResponseInput] = []
+
+        if isinstance(input, str):
+            input_content = OpenAIResponseInputMessageContentText(text=input)
+            input_content_item = OpenAIResponseMessage(
+                role="user",
+                content=[input_content],
+                id=new_input_id,
+            )
+            input_items_data = [input_content_item]
+        else:
+            for input_item in input:
+                if isinstance(input_item, OpenAIResponseMessage):
+                    input_item_dict = input_item.model_dump()
+                    if "id" not in input_item_dict:
+                        input_item_dict["id"] = new_input_id
+                    input_items_data.append(OpenAIResponseMessage(**input_item_dict))
+                else:
+                    input_items_data.append(input_item)
+
+        return input_items_data
+
+    async def _persist_streaming_state(
+        self,
+        stream_chunk: OpenAIResponseObjectStream,
+        orchestrator,
+        input_items: list[OpenAIResponseInput],
+        output_items: list,
+    ) -> None:
+        """Persist response state at significant streaming events.
+
+        This enables clients to poll GET /v1/responses/{response_id} during streaming
+        to see in-progress turn state instead of empty results.
+
+        Persistence occurs at:
+        - response.in_progress: Initial INSERT with empty output
+        - response.output_item.done: UPDATE with accumulated output items
+        - response.completed/response.incomplete: Final UPDATE with complete state
+        - response.failed: UPDATE with error state
+
+        :param stream_chunk: The current streaming event.
+        :param orchestrator: The streaming orchestrator (for snapshotting response).
+        :param input_items: Pre-prepared input items for storage.
+        :param output_items: Accumulated output items so far.
+        """
+        try:
+            match stream_chunk.type:
+                case "response.in_progress":
+                    # Initial persistence when response starts
+                    in_progress_response = stream_chunk.response
+                    await self.responses_store.upsert_response_object(
+                        response_object=in_progress_response,
+                        input=input_items,
+                        messages=[],
+                    )
+
+                case "response.output_item.done":
+                    # Incremental update when an output item completes (tool call, message)
+                    current_snapshot = orchestrator._snapshot_response(
+                        status="in_progress",
+                        outputs=output_items,
+                    )
+                    # Get current messages (filter out system messages)
+                    messages_to_store = list(
+                        filter(
+                            lambda x: not isinstance(x, OpenAISystemMessageParam),
+                            orchestrator.final_messages or orchestrator.ctx.messages,
+                        )
+                    )
+                    await self.responses_store.upsert_response_object(
+                        response_object=current_snapshot,
+                        input=input_items,
+                        messages=messages_to_store,
+                    )
+
+                case "response.completed" | "response.incomplete":
+                    # Final persistence when response finishes
+                    final_response = stream_chunk.response
+                    messages_to_store = list(
+                        filter(
+                            lambda x: not isinstance(x, OpenAISystemMessageParam),
+                            orchestrator.final_messages,
+                        )
+                    )
+                    await self.responses_store.upsert_response_object(
+                        response_object=final_response,
+                        input=input_items,
+                        messages=messages_to_store,
+                    )
+
+                case "response.failed":
+                    # Persist failed state so GET shows error
+                    failed_response = stream_chunk.response
+                    # Preserve any accumulated non-system messages for failed responses
+                    messages_to_store = list(
+                        filter(
+                            lambda x: not isinstance(x, OpenAISystemMessageParam),
+                            orchestrator.final_messages or orchestrator.ctx.messages,
+                        )
+                    )
+                    await self.responses_store.upsert_response_object(
+                        response_object=failed_response,
+                        input=input_items,
+                        messages=messages_to_store,
+                    )
+        except Exception as e:
+            # Best-effort persistence: log error but don't fail the stream
+            logger.warning(f"Failed to persist streaming state for {stream_chunk.type}: {e}")
+
     async def create_openai_response(
         self,
         input: str | list[OpenAIResponseInput],
@@ -343,6 +469,9 @@ class OpenAIResponsesImpl:
         guardrails: list[str | ResponseGuardrailSpec] | None = None,
         parallel_tool_calls: bool | None = None,
         max_tool_calls: int | None = None,
+        reasoning: OpenAIResponseReasoning | None = None,
+        max_output_tokens: int | None = None,
+        safety_identifier: str | None = None,
         metadata: dict[str, str] | None = None,
     ):
         stream = bool(stream)
@@ -380,9 +509,6 @@ class OpenAIResponsesImpl:
             if not conversation.startswith("conv_"):
                 raise InvalidConversationIdError(conversation)
 
-        if max_tool_calls is not None and max_tool_calls < 1:
-            raise ValueError(f"Invalid {max_tool_calls=}; should be >= 1")
-
         stream_gen = self._create_streaming_response(
             input=input,
             conversation=conversation,
@@ -399,6 +525,9 @@ class OpenAIResponsesImpl:
             guardrail_ids=guardrail_ids,
             parallel_tool_calls=parallel_tool_calls,
             max_tool_calls=max_tool_calls,
+            reasoning=reasoning,
+            max_output_tokens=max_output_tokens,
+            safety_identifier=safety_identifier,
             metadata=metadata,
             include=include,
         )
@@ -454,6 +583,9 @@ class OpenAIResponsesImpl:
         guardrail_ids: list[str] | None = None,
         parallel_tool_calls: bool | None = True,
         max_tool_calls: int | None = None,
+        reasoning: OpenAIResponseReasoning | None = None,
+        max_output_tokens: int | None = None,
+        safety_identifier: str | None = None,
         metadata: dict[str, str] | None = None,
         include: list[ResponseItemInclude] | None = None,
     ) -> AsyncIterator[OpenAIResponseObjectStream]:
@@ -493,42 +625,45 @@ class OpenAIResponsesImpl:
 
         # Create a per-request MCP session manager for session reuse (fix for #4452)
         # This avoids redundant tools/list calls when making multiple MCP tool invocations
-        mcp_session_manager = MCPSessionManager()
-
-        # Create a per-request ToolExecutor with the session manager
-        request_tool_executor = ToolExecutor(
-            tool_groups_api=self.tool_groups_api,
-            tool_runtime_api=self.tool_runtime_api,
-            vector_io_api=self.vector_io_api,
-            vector_stores_config=self.tool_executor.vector_stores_config,
-            mcp_session_manager=mcp_session_manager,
-        )
+        async with MCPSessionManager() as mcp_session_manager:
+            request_tool_executor = ToolExecutor(
+                tool_groups_api=self.tool_groups_api,
+                tool_runtime_api=self.tool_runtime_api,
+                vector_io_api=self.vector_io_api,
+                vector_stores_config=self.tool_executor.vector_stores_config,
+                mcp_session_manager=mcp_session_manager,
+            )
 
-        orchestrator = StreamingResponseOrchestrator(
-            inference_api=self.inference_api,
-            ctx=ctx,
-            response_id=response_id,
-            created_at=created_at,
-            prompt=prompt,
-            text=text,
-            max_infer_iters=max_infer_iters,
-            parallel_tool_calls=parallel_tool_calls,
-            tool_executor=request_tool_executor,
-            safety_api=self.safety_api,
-            guardrail_ids=guardrail_ids,
-            instructions=instructions,
-            max_tool_calls=max_tool_calls,
-            metadata=metadata,
-            include=include,
-        )
+            orchestrator = StreamingResponseOrchestrator(
+                inference_api=self.inference_api,
+                ctx=ctx,
+                response_id=response_id,
+                created_at=created_at,
+                prompt=prompt,
+                text=text,
+                max_infer_iters=max_infer_iters,
+                parallel_tool_calls=parallel_tool_calls,
+                tool_executor=request_tool_executor,
+                safety_api=self.safety_api,
+                connectors_api=self.connectors_api,
+                guardrail_ids=guardrail_ids,
+                instructions=instructions,
+                max_tool_calls=max_tool_calls,
+                reasoning=reasoning,
+                max_output_tokens=max_output_tokens,
+                safety_identifier=safety_identifier,
+                metadata=metadata,
+                include=include,
+                store=store,
+            )
 
-        # Stream the response
-        final_response = None
-        failed_response = None
+            final_response = None
+            failed_response = None
+
+            output_items: list[ConversationItem] = []
+
+            input_items_for_storage = self._prepare_input_items_for_storage(all_input)
 
-        # Type as ConversationItem to avoid list invariance issues
-        output_items: list[ConversationItem] = []
-        try:
             async for stream_chunk in orchestrator.create_response():
                 match stream_chunk.type:
                     case "response.completed" | "response.incomplete":
@@ -541,6 +676,16 @@ class OpenAIResponsesImpl:
                     case _:
                         pass  # Other event types
 
+                # Incremental persistence: persist on significant state changes
+                # This enables clients to poll GET /v1/responses/{response_id} during streaming
+                if store:
+                    await self._persist_streaming_state(
+                        stream_chunk=stream_chunk,
+                        orchestrator=orchestrator,
+                        input_items=input_items_for_storage,
+                        output_items=output_items,
+                    )
+
                 # Store and sync before yielding terminal events
                 # This ensures the storage/syncing happens even if the consumer breaks after receiving the event
                 if (
@@ -548,32 +693,14 @@ class OpenAIResponsesImpl:
                     and final_response
                     and failed_response is None
                 ):
-                    messages_to_store = list(
-                        filter(lambda x: not isinstance(x, OpenAISystemMessageParam), orchestrator.final_messages)
-                    )
-                    if store:
-                        # TODO: we really should work off of output_items instead of "final_messages"
-                        await self._store_response(
-                            response=final_response,
-                            input=all_input,
-                            messages=messages_to_store,
-                        )
-
                     if conversation:
+                        messages_to_store = list(
+                            filter(lambda x: not isinstance(x, OpenAISystemMessageParam), orchestrator.final_messages)
+                        )
                         await self._sync_response_to_conversation(conversation, input, output_items)
                         await self.responses_store.store_conversation_messages(conversation, messages_to_store)
 
                 yield stream_chunk
-        finally:
-            # Clean up MCP sessions at the end of the request (fix for #4452)
-            # Use shield() to prevent cancellation from interrupting cleanup and leaking resources
-            # Wrap in try/except as cleanup errors should not mask the original response
-            try:
-                await asyncio.shield(mcp_session_manager.close_all())
-            except BaseException as e:
-                # Debug level - cleanup errors are expected in streaming scenarios where
-                # anyio cancel scopes may be in a different task context
-                logger.debug(f"Error during MCP session cleanup: {e}")
 
     async def delete_openai_response(self, response_id: str) -> OpenAIDeleteResponseObject:
         return await self.responses_store.delete_response_object(response_id)
@@ -596,4 +723,4 @@ class OpenAIResponsesImpl:
 
         adapter = TypeAdapter(list[ConversationItem])
         validated_items = adapter.validate_python(conversation_items)
-        await self.conversations_api.add_items(conversation_id, validated_items)
+        await self.conversations_api.add_items(conversation_id, AddItemsRequest(items=validated_items))