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

{llama_stack → llama_stack_api}/schema_utils.py

@@ -4,11 +4,9 @@
 # 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 Callable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
-from typing import Any, TypeVar
-
-from .strong_typing.schema import json_schema_type, register_schema  # noqa: F401
+from typing import Any, Literal, TypeVar
 
 
 class ExtraBodyField[T]:
@@ -48,6 +46,98 @@ class ExtraBodyField[T]:
         self.description = description
 
 
+SchemaSource = Literal["json_schema_type", "registered_schema", "dynamic_schema"]
+
+
+@dataclass(frozen=True)
+class SchemaInfo:
+    """Metadata describing a schema entry exposed to OpenAPI generation."""
+
+    name: str
+    type: Any
+    source: SchemaSource
+
+
+_json_schema_types: dict[type, SchemaInfo] = {}
+
+
+def json_schema_type(cls):
+    """
+    Decorator to mark a Pydantic model for top-level component registration.
+
+    Models marked with this decorator will be registered as top-level components
+    in the OpenAPI schema, while unmarked models will be inlined.
+
+    This provides control over schema registration to avoid unnecessary indirection
+    for simple one-off types while keeping complex reusable types as components.
+    """
+    cls._llama_stack_schema_type = True
+    schema_name = getattr(cls, "__name__", f"Anonymous_{id(cls)}")
+    cls._llama_stack_schema_name = schema_name
+    _json_schema_types.setdefault(cls, SchemaInfo(name=schema_name, type=cls, source="json_schema_type"))
+    return cls
+
+
+# Global registries for schemas discoverable by the generator
+_registered_schemas: dict[Any, SchemaInfo] = {}
+_dynamic_schema_types: dict[type, SchemaInfo] = {}
+
+
+def register_schema(schema_type, name: str | None = None):
+    """
+    Register a schema type for top-level component registration.
+
+    This replicates the behavior of strong_typing's register_schema function.
+    It's used for union types and other complex types that should appear as
+    top-level components in the OpenAPI schema.
+
+    Args:
+        schema_type: The type to register (e.g., union types, Annotated types)
+        name: Optional name for the schema in the OpenAPI spec. If not provided,
+              uses the type's __name__ or a generated name.
+    """
+    if name is None:
+        name = getattr(schema_type, "__name__", f"Anonymous_{id(schema_type)}")
+
+    # Store the registration information in a global registry
+    # since union types don't allow setting attributes
+    _registered_schemas[schema_type] = SchemaInfo(name=name, type=schema_type, source="registered_schema")
+
+    return schema_type
+
+
+def get_registered_schema_info(schema_type: Any) -> SchemaInfo | None:
+    """Return the registration metadata for a schema type if present."""
+    return _registered_schemas.get(schema_type)
+
+
+def iter_registered_schema_types() -> Iterable[SchemaInfo]:
+    """Iterate over all explicitly registered schema entries."""
+    return tuple(_registered_schemas.values())
+
+
+def iter_json_schema_types() -> Iterable[type]:
+    """Iterate over all Pydantic models decorated with @json_schema_type."""
+    return tuple(info.type for info in _json_schema_types.values())
+
+
+def iter_dynamic_schema_types() -> Iterable[type]:
+    """Iterate over dynamic models registered at generation time."""
+    return tuple(info.type for info in _dynamic_schema_types.values())
+
+
+def register_dynamic_schema_type(schema_type: type, name: str | None = None) -> type:
+    """Register a dynamic model generated at runtime for schema inclusion."""
+    schema_name = name if name is not None else getattr(schema_type, "__name__", f"Anonymous_{id(schema_type)}")
+    _dynamic_schema_types[schema_type] = SchemaInfo(name=schema_name, type=schema_type, source="dynamic_schema")
+    return schema_type
+
+
+def clear_dynamic_schema_types() -> None:
+    """Clear dynamic schema registrations."""
+    _dynamic_schema_types.clear()
+
+
 @dataclass
 class WebMethod:
     level: str | None = None

{llama_stack/apis/scoring → llama_stack_api}/scoring.py

@@ -8,9 +8,9 @@ from typing import Any, Protocol, runtime_checkable
 
 from pydantic import BaseModel
 
-from llama_stack.apis.scoring_functions import ScoringFn, ScoringFnParams
-from llama_stack.apis.version import LLAMA_STACK_API_V1
-from llama_stack.schema_utils import json_schema_type, webmethod
+from llama_stack_api.schema_utils import json_schema_type, webmethod
+from llama_stack_api.scoring_functions import ScoringFn, ScoringFnParams
+from llama_stack_api.version import LLAMA_STACK_API_V1
 
 # mapping of metric to value
 ScoringResultRow = dict[str, Any]

{llama_stack/apis/scoring_functions → llama_stack_api}/scoring_functions.py

@@ -16,10 +16,10 @@ from typing import (
 
 from pydantic import BaseModel, Field
 
-from llama_stack.apis.common.type_system import ParamType
-from llama_stack.apis.resource import Resource, ResourceType
-from llama_stack.apis.version import LLAMA_STACK_API_V1
-from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
+from llama_stack_api.common.type_system import ParamType
+from llama_stack_api.resource import Resource, ResourceType
+from llama_stack_api.schema_utils import json_schema_type, register_schema, webmethod
+from llama_stack_api.version import LLAMA_STACK_API_V1
 
 
 # Perhaps more structure can be imposed on these functions. Maybe they could be associated
@@ -155,6 +155,7 @@ class ScoringFnInput(CommonScoringFnFields, BaseModel):
     provider_scoring_fn_id: str | None = None
 
 
+@json_schema_type
 class ListScoringFunctionsResponse(BaseModel):
     data: list[ScoringFn]
 
@@ -178,7 +179,7 @@ class ScoringFunctions(Protocol):
         """
         ...
 
-    @webmethod(route="/scoring-functions", method="POST", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/scoring-functions", method="POST", level=LLAMA_STACK_API_V1, deprecated=True)
     async def register_scoring_function(
         self,
         scoring_fn_id: str,
@@ -199,7 +200,9 @@ class ScoringFunctions(Protocol):
         """
         ...
 
-    @webmethod(route="/scoring-functions/{scoring_fn_id:path}", method="DELETE", level=LLAMA_STACK_API_V1)
+    @webmethod(
+        route="/scoring-functions/{scoring_fn_id:path}", method="DELETE", level=LLAMA_STACK_API_V1, deprecated=True
+    )
     async def unregister_scoring_function(self, scoring_fn_id: str) -> None:
         """Unregister a scoring function.
 

{llama_stack/apis/shields → llama_stack_api}/shields.py

@@ -8,10 +8,9 @@ from typing import Any, Literal, Protocol, runtime_checkable
 
 from pydantic import BaseModel
 
-from llama_stack.apis.resource import Resource, ResourceType
-from llama_stack.apis.version import LLAMA_STACK_API_V1
-from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
-from llama_stack.schema_utils import json_schema_type, webmethod
+from llama_stack_api.resource import Resource, ResourceType
+from llama_stack_api.schema_utils import json_schema_type, webmethod
+from llama_stack_api.version import LLAMA_STACK_API_V1
 
 
 class CommonShieldFields(BaseModel):
@@ -43,12 +42,12 @@ class ShieldInput(CommonShieldFields):
     provider_shield_id: str | None = None
 
 
+@json_schema_type
 class ListShieldsResponse(BaseModel):
     data: list[Shield]
 
 
 @runtime_checkable
-@trace_protocol
 class Shields(Protocol):
     @webmethod(route="/shields", method="GET", level=LLAMA_STACK_API_V1)
     async def list_shields(self) -> ListShieldsResponse:
@@ -67,7 +66,7 @@ class Shields(Protocol):
         """
         ...
 
-    @webmethod(route="/shields", method="POST", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/shields", method="POST", level=LLAMA_STACK_API_V1, deprecated=True)
     async def register_shield(
         self,
         shield_id: str,
@@ -85,7 +84,7 @@ class Shields(Protocol):
         """
         ...
 
-    @webmethod(route="/shields/{identifier:path}", method="DELETE", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/shields/{identifier:path}", method="DELETE", level=LLAMA_STACK_API_V1, deprecated=True)
     async def unregister_shield(self, identifier: str) -> None:
         """Unregister a shield.
 

{llama_stack/apis/tools → llama_stack_api}/tools.py

@@ -10,13 +10,10 @@ from typing import Any, Literal, Protocol
 from pydantic import BaseModel
 from typing_extensions import runtime_checkable
 
-from llama_stack.apis.common.content_types import URL, InterleavedContent
-from llama_stack.apis.resource import Resource, ResourceType
-from llama_stack.apis.version import LLAMA_STACK_API_V1
-from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol
-from llama_stack.schema_utils import json_schema_type, webmethod
-
-from .rag_tool import RAGToolRuntime
+from llama_stack_api.common.content_types import URL, InterleavedContent
+from llama_stack_api.resource import Resource, ResourceType
+from llama_stack_api.schema_utils import json_schema_type, webmethod
+from llama_stack_api.version import LLAMA_STACK_API_V1
 
 
 @json_schema_type
@@ -90,6 +87,7 @@ class ToolStore(Protocol):
     async def get_tool_group(self, toolgroup_id: str) -> ToolGroup: ...
 
 
+@json_schema_type
 class ListToolGroupsResponse(BaseModel):
     """Response containing a list of tool groups.
 
@@ -99,6 +97,7 @@ class ListToolGroupsResponse(BaseModel):
     data: list[ToolGroup]
 
 
+@json_schema_type
 class ListToolDefsResponse(BaseModel):
     """Response containing a list of tool definitions.
 
@@ -109,9 +108,8 @@ class ListToolDefsResponse(BaseModel):
 
 
 @runtime_checkable
-@trace_protocol
 class ToolGroups(Protocol):
-    @webmethod(route="/toolgroups", method="POST", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/toolgroups", method="POST", level=LLAMA_STACK_API_V1, deprecated=True)
     async def register_tool_group(
         self,
         toolgroup_id: str,
@@ -128,7 +126,7 @@ class ToolGroups(Protocol):
         """
         ...
 
-    @webmethod(route="/toolgroups/{toolgroup_id:path}", method="GET", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/toolgroups/{toolgroup_id:path}", method="GET", level=LLAMA_STACK_API_V1, deprecated=True)
     async def get_tool_group(
         self,
         toolgroup_id: str,
@@ -140,7 +138,7 @@ class ToolGroups(Protocol):
         """
         ...
 
-    @webmethod(route="/toolgroups", method="GET", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/toolgroups", method="GET", level=LLAMA_STACK_API_V1, deprecated=True)
     async def list_tool_groups(self) -> ListToolGroupsResponse:
         """List tool groups with optional provider.
 
@@ -148,7 +146,7 @@ class ToolGroups(Protocol):
         """
         ...
 
-    @webmethod(route="/tools", method="GET", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/tools", method="GET", level=LLAMA_STACK_API_V1, deprecated=True)
     async def list_tools(self, toolgroup_id: str | None = None) -> ListToolDefsResponse:
         """List tools with optional tool group.
 
@@ -157,7 +155,7 @@ class ToolGroups(Protocol):
         """
         ...
 
-    @webmethod(route="/tools/{tool_name:path}", method="GET", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/tools/{tool_name:path}", method="GET", level=LLAMA_STACK_API_V1, deprecated=True)
     async def get_tool(
         self,
         tool_name: str,
@@ -169,7 +167,7 @@ class ToolGroups(Protocol):
         """
         ...
 
-    @webmethod(route="/toolgroups/{toolgroup_id:path}", method="DELETE", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/toolgroups/{toolgroup_id:path}", method="DELETE", level=LLAMA_STACK_API_V1, deprecated=True)
     async def unregister_toolgroup(
         self,
         toolgroup_id: str,
@@ -191,31 +189,38 @@ class SpecialToolGroup(Enum):
 
 
 @runtime_checkable
-@trace_protocol
 class ToolRuntime(Protocol):
     tool_store: ToolStore | None = None
 
-    rag_tool: RAGToolRuntime | None = None
-
     # TODO: This needs to be renamed once OPEN API generator name conflict issue is fixed.
-    @webmethod(route="/tool-runtime/list-tools", method="GET", level=LLAMA_STACK_API_V1)
+    @webmethod(route="/tool-runtime/list-tools", method="GET", level=LLAMA_STACK_API_V1, deprecated=True)
     async def list_runtime_tools(
-        self, tool_group_id: str | None = None, mcp_endpoint: URL | None = None
+        self,
+        tool_group_id: str | None = None,
+        mcp_endpoint: URL | None = None,
+        authorization: str | None = None,
     ) -> ListToolDefsResponse:
         """List all tools in the runtime.
 
         :param tool_group_id: The ID of the tool group to list tools for.
         :param mcp_endpoint: The MCP endpoint to use for the tool group.
+        :param authorization: (Optional) OAuth access token for authenticating with the MCP server.
         :returns: A ListToolDefsResponse.
         """
         ...
 
-    @webmethod(route="/tool-runtime/invoke", method="POST", level=LLAMA_STACK_API_V1)
-    async def invoke_tool(self, tool_name: str, kwargs: dict[str, Any]) -> ToolInvocationResult:
+    @webmethod(route="/tool-runtime/invoke", method="POST", level=LLAMA_STACK_API_V1, deprecated=True)
+    async def invoke_tool(
+        self,
+        tool_name: str,
+        kwargs: dict[str, Any],
+        authorization: str | None = None,
+    ) -> ToolInvocationResult:
         """Run a tool with the given arguments.
 
         :param tool_name: The name of the tool to invoke.
         :param kwargs: A dictionary of arguments to pass to the tool.
+        :param authorization: (Optional) OAuth access token for authenticating with the MCP server.
         :returns: A ToolInvocationResult.
         """
         ...