llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

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

@@ -0,0 +1,383 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import uuid
+from collections.abc import AsyncGenerator
+from datetime import UTC, datetime
+
+from llama_stack.apis.agents import (
+    Agent,
+    AgentConfig,
+    AgentCreateResponse,
+    Agents,
+    AgentSessionCreateResponse,
+    AgentStepResponse,
+    AgentToolGroup,
+    AgentTurnCreateRequest,
+    AgentTurnResumeRequest,
+    Document,
+    ListOpenAIResponseInputItem,
+    ListOpenAIResponseObject,
+    OpenAIResponseInput,
+    OpenAIResponseInputTool,
+    OpenAIResponseObject,
+    Order,
+    Session,
+    Turn,
+)
+from llama_stack.apis.agents.agents import ResponseGuardrail
+from llama_stack.apis.agents.openai_responses import OpenAIResponseText
+from llama_stack.apis.common.responses import PaginatedResponse
+from llama_stack.apis.conversations import Conversations
+from llama_stack.apis.inference import (
+    Inference,
+    ToolConfig,
+    ToolResponse,
+    ToolResponseMessage,
+    UserMessage,
+)
+from llama_stack.apis.safety import Safety
+from llama_stack.apis.tools import ToolGroups, ToolRuntime
+from llama_stack.apis.vector_io import VectorIO
+from llama_stack.core.datatypes import AccessRule
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
+from llama_stack.providers.utils.pagination import paginate_records
+from llama_stack.providers.utils.responses.responses_store import ResponsesStore
+
+from .agent_instance import ChatAgent
+from .config import MetaReferenceAgentsImplConfig
+from .persistence import AgentInfo
+from .responses.openai_responses import OpenAIResponsesImpl
+
+logger = get_logger(name=__name__, category="agents::meta_reference")
+
+
+class MetaReferenceAgentsImpl(Agents):
+    def __init__(
+        self,
+        config: MetaReferenceAgentsImplConfig,
+        inference_api: Inference,
+        vector_io_api: VectorIO,
+        safety_api: Safety,
+        tool_runtime_api: ToolRuntime,
+        tool_groups_api: ToolGroups,
+        conversations_api: Conversations,
+        policy: list[AccessRule],
+        telemetry_enabled: bool = False,
+    ):
+        self.config = config
+        self.inference_api = inference_api
+        self.vector_io_api = vector_io_api
+        self.safety_api = safety_api
+        self.tool_runtime_api = tool_runtime_api
+        self.tool_groups_api = tool_groups_api
+        self.conversations_api = conversations_api
+        self.telemetry_enabled = telemetry_enabled
+
+        self.in_memory_store = InmemoryKVStoreImpl()
+        self.openai_responses_impl: OpenAIResponsesImpl | None = None
+        self.policy = policy
+
+    async def initialize(self) -> None:
+        self.persistence_store = await kvstore_impl(self.config.persistence.agent_state)
+        self.responses_store = ResponsesStore(self.config.persistence.responses, self.policy)
+        await self.responses_store.initialize()
+        self.openai_responses_impl = OpenAIResponsesImpl(
+            inference_api=self.inference_api,
+            tool_groups_api=self.tool_groups_api,
+            tool_runtime_api=self.tool_runtime_api,
+            responses_store=self.responses_store,
+            vector_io_api=self.vector_io_api,
+            safety_api=self.safety_api,
+            conversations_api=self.conversations_api,
+        )
+
+    async def create_agent(
+        self,
+        agent_config: AgentConfig,
+    ) -> AgentCreateResponse:
+        agent_id = str(uuid.uuid4())
+        created_at = datetime.now(UTC)
+
+        agent_info = AgentInfo(
+            **agent_config.model_dump(),
+            created_at=created_at,
+        )
+
+        # Store the agent info
+        await self.persistence_store.set(
+            key=f"agent:{agent_id}",
+            value=agent_info.model_dump_json(),
+        )
+
+        return AgentCreateResponse(
+            agent_id=agent_id,
+        )
+
+    async def _get_agent_impl(self, agent_id: str) -> ChatAgent:
+        agent_info_json = await self.persistence_store.get(
+            key=f"agent:{agent_id}",
+        )
+        if not agent_info_json:
+            raise ValueError(f"Could not find agent info for {agent_id}")
+
+        try:
+            agent_info = AgentInfo.model_validate_json(agent_info_json)
+        except Exception as e:
+            raise ValueError(f"Could not validate agent info for {agent_id}") from e
+
+        return ChatAgent(
+            agent_id=agent_id,
+            agent_config=agent_info,
+            inference_api=self.inference_api,
+            safety_api=self.safety_api,
+            vector_io_api=self.vector_io_api,
+            tool_runtime_api=self.tool_runtime_api,
+            tool_groups_api=self.tool_groups_api,
+            persistence_store=(
+                self.persistence_store if agent_info.enable_session_persistence else self.in_memory_store
+            ),
+            created_at=agent_info.created_at,
+            policy=self.policy,
+            telemetry_enabled=self.telemetry_enabled,
+        )
+
+    async def create_agent_session(
+        self,
+        agent_id: str,
+        session_name: str,
+    ) -> AgentSessionCreateResponse:
+        agent = await self._get_agent_impl(agent_id)
+
+        session_id = await agent.create_session(session_name)
+        return AgentSessionCreateResponse(
+            session_id=session_id,
+        )
+
+    async def create_agent_turn(
+        self,
+        agent_id: str,
+        session_id: str,
+        messages: list[UserMessage | ToolResponseMessage],
+        toolgroups: list[AgentToolGroup] | None = None,
+        documents: list[Document] | None = None,
+        stream: bool | None = False,
+        tool_config: ToolConfig | None = None,
+    ) -> AsyncGenerator:
+        request = AgentTurnCreateRequest(
+            agent_id=agent_id,
+            session_id=session_id,
+            messages=messages,
+            stream=True,
+            toolgroups=toolgroups,
+            documents=documents,
+            tool_config=tool_config,
+        )
+        if stream:
+            return self._create_agent_turn_streaming(request)
+        else:
+            raise NotImplementedError("Non-streaming agent turns not yet implemented")
+
+    async def _create_agent_turn_streaming(
+        self,
+        request: AgentTurnCreateRequest,
+    ) -> AsyncGenerator:
+        agent = await self._get_agent_impl(request.agent_id)
+        async for event in agent.create_and_execute_turn(request):
+            yield event
+
+    async def resume_agent_turn(
+        self,
+        agent_id: str,
+        session_id: str,
+        turn_id: str,
+        tool_responses: list[ToolResponse],
+        stream: bool | None = False,
+    ) -> AsyncGenerator:
+        request = AgentTurnResumeRequest(
+            agent_id=agent_id,
+            session_id=session_id,
+            turn_id=turn_id,
+            tool_responses=tool_responses,
+            stream=stream,
+        )
+        if stream:
+            return self._continue_agent_turn_streaming(request)
+        else:
+            raise NotImplementedError("Non-streaming agent turns not yet implemented")
+
+    async def _continue_agent_turn_streaming(
+        self,
+        request: AgentTurnResumeRequest,
+    ) -> AsyncGenerator:
+        agent = await self._get_agent_impl(request.agent_id)
+        async for event in agent.resume_turn(request):
+            yield event
+
+    async def get_agents_turn(self, agent_id: str, session_id: str, turn_id: str) -> Turn:
+        agent = await self._get_agent_impl(agent_id)
+        turn = await agent.storage.get_session_turn(session_id, turn_id)
+        return turn
+
+    async def get_agents_step(self, agent_id: str, session_id: str, turn_id: str, step_id: str) -> AgentStepResponse:
+        turn = await self.get_agents_turn(agent_id, session_id, turn_id)
+        for step in turn.steps:
+            if step.step_id == step_id:
+                return AgentStepResponse(step=step)
+        raise ValueError(f"Provided step_id {step_id} could not be found")
+
+    async def get_agents_session(
+        self,
+        agent_id: str,
+        session_id: str,
+        turn_ids: list[str] | None = None,
+    ) -> Session:
+        agent = await self._get_agent_impl(agent_id)
+
+        session_info = await agent.storage.get_session_info(session_id)
+        turns = await agent.storage.get_session_turns(session_id)
+        if turn_ids:
+            turns = [turn for turn in turns if turn.turn_id in turn_ids]
+        return Session(
+            session_name=session_info.session_name,
+            session_id=session_id,
+            turns=turns,
+            started_at=session_info.started_at,
+        )
+
+    async def delete_agents_session(self, agent_id: str, session_id: str) -> None:
+        agent = await self._get_agent_impl(agent_id)
+
+        # Delete turns first, then the session
+        await agent.storage.delete_session_turns(session_id)
+        await agent.storage.delete_session(session_id)
+
+    async def delete_agent(self, agent_id: str) -> None:
+        # First get all sessions for this agent
+        agent = await self._get_agent_impl(agent_id)
+        sessions = await agent.storage.list_sessions()
+
+        # Delete all sessions
+        for session in sessions:
+            await self.delete_agents_session(agent_id, session.session_id)
+
+        # Finally delete the agent itself
+        await self.persistence_store.delete(f"agent:{agent_id}")
+
+    async def list_agents(self, start_index: int | None = None, limit: int | None = None) -> PaginatedResponse:
+        agent_keys = await self.persistence_store.keys_in_range("agent:", "agent:\xff")
+        agent_list: list[Agent] = []
+        for agent_key in agent_keys:
+            agent_id = agent_key.split(":")[1]
+
+            # Get the agent info using the key
+            agent_info_json = await self.persistence_store.get(agent_key)
+            if not agent_info_json:
+                logger.error(f"Could not find agent info for key {agent_key}")
+                continue
+
+            try:
+                agent_info = AgentInfo.model_validate_json(agent_info_json)
+                agent_list.append(
+                    Agent(
+                        agent_id=agent_id,
+                        agent_config=agent_info,
+                        created_at=agent_info.created_at,
+                    )
+                )
+            except Exception as e:
+                logger.error(f"Error parsing agent info for {agent_id}: {e}")
+                continue
+
+        # Convert Agent objects to dictionaries
+        agent_dicts = [agent.model_dump() for agent in agent_list]
+        return paginate_records(agent_dicts, start_index, limit)
+
+    async def get_agent(self, agent_id: str) -> Agent:
+        chat_agent = await self._get_agent_impl(agent_id)
+        agent = Agent(
+            agent_id=agent_id,
+            agent_config=chat_agent.agent_config,
+            created_at=chat_agent.created_at,
+        )
+        return agent
+
+    async def list_agent_sessions(
+        self, agent_id: str, start_index: int | None = None, limit: int | None = None
+    ) -> PaginatedResponse:
+        agent = await self._get_agent_impl(agent_id)
+        sessions = await agent.storage.list_sessions()
+        # Convert Session objects to dictionaries
+        session_dicts = [session.model_dump() for session in sessions]
+        return paginate_records(session_dicts, start_index, limit)
+
+    async def shutdown(self) -> None:
+        pass
+
+    # OpenAI responses
+    async def get_openai_response(
+        self,
+        response_id: str,
+    ) -> OpenAIResponseObject:
+        return await self.openai_responses_impl.get_openai_response(response_id)
+
+    async def create_openai_response(
+        self,
+        input: str | list[OpenAIResponseInput],
+        model: str,
+        instructions: str | None = None,
+        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,
+        tools: list[OpenAIResponseInputTool] | None = None,
+        include: list[str] | None = None,
+        max_infer_iters: int | None = 10,
+        guardrails: list[ResponseGuardrail] | None = None,
+    ) -> OpenAIResponseObject:
+        return await self.openai_responses_impl.create_openai_response(
+            input,
+            model,
+            instructions,
+            previous_response_id,
+            conversation,
+            store,
+            stream,
+            temperature,
+            text,
+            tools,
+            include,
+            max_infer_iters,
+            guardrails,
+        )
+
+    async def list_openai_responses(
+        self,
+        after: str | None = None,
+        limit: int | None = 50,
+        model: str | None = None,
+        order: Order | None = Order.desc,
+    ) -> ListOpenAIResponseObject:
+        return await self.openai_responses_impl.list_openai_responses(after, limit, model, 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,
+    ) -> ListOpenAIResponseInputItem:
+        return await self.openai_responses_impl.list_openai_response_input_items(
+            response_id, after, before, include, limit, order
+        )
+
+    async def delete_openai_response(self, response_id: str) -> None:
+        return await self.openai_responses_impl.delete_openai_response(response_id)

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

@@ -0,0 +1,37 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+from pydantic import BaseModel
+
+from llama_stack.core.storage.datatypes import KVStoreReference, ResponsesStoreReference
+
+
+class AgentPersistenceConfig(BaseModel):
+    """Nested persistence configuration for agents."""
+
+    agent_state: KVStoreReference
+    responses: ResponsesStoreReference
+
+
+class MetaReferenceAgentsImplConfig(BaseModel):
+    persistence: AgentPersistenceConfig
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str) -> dict[str, Any]:
+        return {
+            "persistence": {
+                "agent_state": KVStoreReference(
+                    backend="kv_default",
+                    namespace="agents",
+                ).model_dump(exclude_none=True),
+                "responses": ResponsesStoreReference(
+                    backend="sql_default",
+                    table_name="responses",
+                ).model_dump(exclude_none=True),
+            }
+        }

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

@@ -0,0 +1,228 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+import uuid
+from datetime import UTC, datetime
+
+from llama_stack.apis.agents import AgentConfig, Session, ToolExecutionStep, Turn
+from llama_stack.apis.common.errors import SessionNotFoundError
+from llama_stack.core.access_control.access_control import AccessDeniedError, is_action_allowed
+from llama_stack.core.access_control.datatypes import AccessRule
+from llama_stack.core.datatypes import User
+from llama_stack.core.request_headers import get_authenticated_user
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore import KVStore
+
+log = get_logger(name=__name__, category="agents::meta_reference")
+
+
+class AgentSessionInfo(Session):
+    # TODO: is this used anywhere?
+    vector_db_id: str | None = None
+    started_at: datetime
+    owner: User | None = None
+    identifier: str | None = None
+    type: str = "session"
+
+
+class AgentInfo(AgentConfig):
+    created_at: datetime
+
+
+class AgentPersistence:
+    def __init__(self, agent_id: str, kvstore: KVStore, policy: list[AccessRule]):
+        self.agent_id = agent_id
+        self.kvstore = kvstore
+        self.policy = policy
+
+    async def create_session(self, name: str) -> str:
+        session_id = str(uuid.uuid4())
+
+        # Get current user's auth attributes for new sessions
+        user = get_authenticated_user()
+
+        session_info = AgentSessionInfo(
+            session_id=session_id,
+            session_name=name,
+            started_at=datetime.now(UTC),
+            owner=user,
+            turns=[],
+            identifier=name,  # should this be qualified in any way?
+        )
+        if not is_action_allowed(self.policy, "create", session_info, user):
+            raise AccessDeniedError("create", session_info, user)
+
+        await self.kvstore.set(
+            key=f"session:{self.agent_id}:{session_id}",
+            value=session_info.model_dump_json(),
+        )
+        return session_id
+
+    async def get_session_info(self, session_id: str) -> AgentSessionInfo:
+        value = await self.kvstore.get(
+            key=f"session:{self.agent_id}:{session_id}",
+        )
+        if not value:
+            raise SessionNotFoundError(session_id)
+
+        session_info = AgentSessionInfo(**json.loads(value))
+
+        # Check access to session
+        if not self._check_session_access(session_info):
+            return None
+
+        return session_info
+
+    def _check_session_access(self, session_info: AgentSessionInfo) -> bool:
+        """Check if current user has access to the session."""
+        # Handle backward compatibility for old sessions without access control
+        if not hasattr(session_info, "access_attributes") and not hasattr(session_info, "owner"):
+            return True
+
+        return is_action_allowed(self.policy, "read", session_info, get_authenticated_user())
+
+    async def get_session_if_accessible(self, session_id: str) -> AgentSessionInfo | None:
+        """Get session info if the user has access to it. For internal use by sub-session methods."""
+        session_info = await self.get_session_info(session_id)
+        if not session_info:
+            return None
+
+        return session_info
+
+    async def add_vector_db_to_session(self, session_id: str, vector_db_id: str):
+        session_info = await self.get_session_if_accessible(session_id)
+        if session_info is None:
+            raise SessionNotFoundError(session_id)
+
+        session_info.vector_db_id = vector_db_id
+        await self.kvstore.set(
+            key=f"session:{self.agent_id}:{session_id}",
+            value=session_info.model_dump_json(),
+        )
+
+    async def add_turn_to_session(self, session_id: str, turn: Turn):
+        if not await self.get_session_if_accessible(session_id):
+            raise SessionNotFoundError(session_id)
+
+        await self.kvstore.set(
+            key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}",
+            value=turn.model_dump_json(),
+        )
+
+    async def get_session_turns(self, session_id: str) -> list[Turn]:
+        if not await self.get_session_if_accessible(session_id):
+            raise SessionNotFoundError(session_id)
+
+        values = await self.kvstore.values_in_range(
+            start_key=f"session:{self.agent_id}:{session_id}:",
+            end_key=f"session:{self.agent_id}:{session_id}:\xff\xff\xff\xff",
+        )
+        turns = []
+        for value in values:
+            try:
+                turn = Turn(**json.loads(value))
+                turns.append(turn)
+            except Exception as e:
+                log.error(f"Error parsing turn: {e}")
+                continue
+
+        # The kvstore does not guarantee order, so we sort by started_at
+        # to ensure consistent ordering of turns.
+        turns.sort(key=lambda t: t.started_at)
+
+        return turns
+
+    async def get_session_turn(self, session_id: str, turn_id: str) -> Turn | None:
+        if not await self.get_session_if_accessible(session_id):
+            raise SessionNotFoundError(session_id)
+
+        value = await self.kvstore.get(
+            key=f"session:{self.agent_id}:{session_id}:{turn_id}",
+        )
+        if not value:
+            return None
+        return Turn(**json.loads(value))
+
+    async def set_in_progress_tool_call_step(self, session_id: str, turn_id: str, step: ToolExecutionStep):
+        if not await self.get_session_if_accessible(session_id):
+            raise SessionNotFoundError(session_id)
+
+        await self.kvstore.set(
+            key=f"in_progress_tool_call_step:{self.agent_id}:{session_id}:{turn_id}",
+            value=step.model_dump_json(),
+        )
+
+    async def get_in_progress_tool_call_step(self, session_id: str, turn_id: str) -> ToolExecutionStep | None:
+        if not await self.get_session_if_accessible(session_id):
+            return None
+
+        value = await self.kvstore.get(
+            key=f"in_progress_tool_call_step:{self.agent_id}:{session_id}:{turn_id}",
+        )
+        return ToolExecutionStep(**json.loads(value)) if value else None
+
+    async def set_num_infer_iters_in_turn(self, session_id: str, turn_id: str, num_infer_iters: int):
+        if not await self.get_session_if_accessible(session_id):
+            raise SessionNotFoundError(session_id)
+
+        await self.kvstore.set(
+            key=f"num_infer_iters_in_turn:{self.agent_id}:{session_id}:{turn_id}",
+            value=str(num_infer_iters),
+        )
+
+    async def get_num_infer_iters_in_turn(self, session_id: str, turn_id: str) -> int | None:
+        if not await self.get_session_if_accessible(session_id):
+            return None
+
+        value = await self.kvstore.get(
+            key=f"num_infer_iters_in_turn:{self.agent_id}:{session_id}:{turn_id}",
+        )
+        return int(value) if value else None
+
+    async def list_sessions(self) -> list[Session]:
+        values = await self.kvstore.values_in_range(
+            start_key=f"session:{self.agent_id}:",
+            end_key=f"session:{self.agent_id}:\xff\xff\xff\xff",
+        )
+        sessions = []
+        for value in values:
+            try:
+                data = json.loads(value)
+                if "turn_id" in data:
+                    continue
+
+                session_info = Session(**data)
+                sessions.append(session_info)
+            except Exception as e:
+                log.error(f"Error parsing session info: {e}")
+                continue
+        return sessions
+
+    async def delete_session_turns(self, session_id: str) -> None:
+        """Delete all turns and their associated data for a session.
+
+        Args:
+            session_id: The ID of the session whose turns should be deleted.
+        """
+        turns = await self.get_session_turns(session_id)
+        for turn in turns:
+            await self.kvstore.delete(key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}")
+
+    async def delete_session(self, session_id: str) -> None:
+        """Delete a session and all its associated turns.
+
+        Args:
+            session_id: The ID of the session to delete.
+
+        Raises:
+            ValueError: If the session does not exist.
+        """
+        session_info = await self.get_session_info(session_id)
+        if session_info is None:
+            raise SessionNotFoundError(session_id)
+
+        await self.kvstore.delete(key=f"session:{self.agent_id}:{session_id}")

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

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.