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

llama_stack/core/datatypes.py

@@ -11,7 +11,7 @@ from urllib.parse import urlparse
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
-from llama_stack.core.access_control.datatypes import AccessRule
+from llama_stack.core.access_control.datatypes import AccessRule, RouteAccessRule
 from llama_stack.core.storage.datatypes import (
     KVStoreReference,
     StorageBackendType,
@@ -329,13 +329,17 @@ AuthProviderConfig = Annotated[
 class AuthenticationConfig(BaseModel):
     """Top-level authentication configuration."""
 
-    provider_config: AuthProviderConfig = Field(
-        ...,
-        description="Authentication provider configuration",
+    provider_config: AuthProviderConfig | None = Field(
+        default=None,
+        description="Authentication provider configuration (optional if only using route_policy)",
+    )
+    route_policy: list[RouteAccessRule] = Field(
+        default=[],
+        description="Rules for determining access to API routes (infrastructure-level)",
     )
     access_policy: list[AccessRule] = Field(
         default=[],
-        description="Rules for determining access to resources",
+        description="Rules for determining access to resources (data-level)",
     )
 
 
@@ -348,6 +352,7 @@ class QualifiedModel(BaseModel):
 
     provider_id: str
     model_id: str
+    embedding_dimensions: int | None = None
 
 
 class RewriteQueryParams(BaseModel):
@@ -654,7 +659,6 @@ class RegisteredResources(BaseModel):
     scoring_fns: list[ScoringFnInput] = Field(default_factory=list)
     benchmarks: list[BenchmarkInput] = Field(default_factory=list)
     tool_groups: list[ToolGroupInput] = Field(default_factory=list)
-    connectors: list[ConnectorInput] = Field(default_factory=list)
 
 
 class ServerConfig(BaseModel):
@@ -703,11 +707,20 @@ class ServerConfig(BaseModel):
 class StackConfig(BaseModel):
     version: int = LLAMA_STACK_RUN_CONFIG_VERSION
 
-    image_name: str = Field(
-        ...,
+    distro_name: str | None = Field(
+        default=None,
         description="""
 Reference to the distribution this package refers to. For unregistered (adhoc) packages,
 this could be just a hash
+""",
+    )
+    image_name: str | None = Field(
+        default=None,
+        deprecated=True,
+        description="""
+DEPRECATED: Use 'distro_name' instead. This field is maintained for backward compatibility.
+Reference to the distribution this package refers to. For unregistered (adhoc) packages,
+this could be just a hash
 """,
     )
     container_image: str | None = Field(
@@ -763,6 +776,11 @@ can be instantiated multiple times (with different configs) if necessary.
         description="Configuration for default moderations model",
     )
 
+    connectors: list[ConnectorInput] = Field(
+        default_factory=list,
+        description="List of connectors to register at stack startup",
+    )
+
     @field_validator("external_providers_dir")
     @classmethod
     def validate_external_providers_dir(cls, v):
@@ -772,6 +790,34 @@ can be instantiated multiple times (with different configs) if necessary.
             return Path(v)
         return v
 
+    @model_validator(mode="after")
+    def validate_distro_name_migration(self) -> "StackConfig":
+        """Handle migration from image_name to distro_name."""
+        import warnings
+
+        if self.distro_name is None and self.image_name is None:
+            raise ValueError("Either 'distro_name' or 'image_name' must be provided")
+
+        if self.image_name is not None and self.distro_name is None:
+            # Migrate from image_name to distro_name
+            warnings.warn(
+                "The 'image_name' field is deprecated. Please use 'distro_name' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            self.distro_name = self.image_name
+        elif self.image_name is not None and self.distro_name is not None:
+            # Both provided - warn and prefer distro_name
+            warnings.warn(
+                "Both 'image_name' and 'distro_name' were provided. "
+                "The 'image_name' field is deprecated and will be ignored. "
+                "Please use only 'distro_name' in your configuration.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        return self
+
     @model_validator(mode="after")
     def validate_server_stores(self) -> "StackConfig":
         backend_map = self.storage.backends

llama_stack/core/library_client.py

@@ -20,7 +20,7 @@ import httpx
 import yaml
 from fastapi import Response as FastAPIResponse
 
-from llama_stack.core.utils.type_inspection import is_unwrapped_body_param
+from llama_stack.core.utils.type_inspection import is_body_param, is_unwrapped_body_param
 
 try:
     from llama_stack_client import (
@@ -504,11 +504,30 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
         # Prepare body for the function call (handles both Pydantic and traditional params)
         body = self._convert_body(func, body)
 
+        result = await func(**body)
+        content_type = "application/json"
+        if isinstance(result, FastAPIResponse):
+            content_type = result.media_type or content_type
+
         async def gen():
-            async for chunk in await func(**body):
-                data = json.dumps(convert_pydantic_to_json_value(chunk))
-                sse_event = f"data: {data}\n\n"
-                yield sse_event.encode("utf-8")
+            # Handle FastAPI StreamingResponse (returned by router endpoints)
+            # Extract the async generator from the StreamingResponse body
+            from fastapi.responses import StreamingResponse
+
+            if isinstance(result, StreamingResponse):
+                # StreamingResponse.body_iterator is the async generator
+                async for chunk in result.body_iterator:
+                    # Chunk is already SSE-formatted string from sse_generator, encode to bytes
+                    if isinstance(chunk, str):
+                        yield chunk.encode("utf-8")
+                    else:
+                        yield chunk
+            else:
+                # Direct async generator from implementation
+                async for chunk in result:
+                    data = json.dumps(convert_pydantic_to_json_value(chunk))
+                    sse_event = f"data: {data}\n\n"
+                    yield sse_event.encode("utf-8")
 
         wrapped_gen = preserve_contexts_async_generator(gen(), [PROVIDER_DATA_VAR])
 
@@ -516,7 +535,7 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
             status_code=httpx.codes.OK,
             content=wrapped_gen,
             headers={
-                "Content-Type": "application/json",
+                "Content-Type": content_type,
             },
             request=httpx.Request(
                 method=options.method,
@@ -549,10 +568,26 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
         sig = inspect.signature(func)
         params_list = [p for p in sig.parameters.values() if p.name != "self"]
 
+        # Resolve string annotations (from `from __future__ import annotations`) to actual types
+        try:
+            type_hints = typing.get_type_hints(func, include_extras=True)
+        except NameError as e:
+            # Forward reference could not be resolved - fall back to raw annotations
+            logger.debug(f"Could not resolve type hints for {func.__name__}: {e}")
+            type_hints = {}
+        except Exception as e:
+            # Unexpected error - log and fall back
+            logger.warning(f"Failed to resolve type hints for {func.__name__}: {e}")
+            type_hints = {}
+
+        # Helper to get the resolved type for a parameter
+        def get_param_type(param: inspect.Parameter) -> Any:
+            return type_hints.get(param.name, param.annotation)
+
         # Flatten if there's a single unwrapped body parameter (BaseModel or Annotated[BaseModel, Body(embed=False)])
         if len(params_list) == 1:
             param = params_list[0]
-            param_type = param.annotation
+            param_type = get_param_type(param)
             if is_unwrapped_body_param(param_type):
                 base_type = get_args(param_type)[0]
                 return {param.name: base_type(**body)}
@@ -563,16 +598,22 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
         # Check if there's an unwrapped body parameter among multiple parameters
         # (e.g., path param + body param like: vector_store_id: str, params: Annotated[Model, Body(...)])
         unwrapped_body_param = None
+        unwrapped_body_param_type = None
+        body_param = None
         for param in params_list:
-            if is_unwrapped_body_param(param.annotation):
+            param_type = get_param_type(param)
+            if is_unwrapped_body_param(param_type):
                 unwrapped_body_param = param
+                unwrapped_body_param_type = param_type
                 break
+            if body_param is None and is_body_param(param_type):
+                body_param = param
 
         # Check for parameters with Depends() annotation (FastAPI router endpoints)
         # These need special handling: construct the request model from body
         depends_param = None
         for param in params_list:
-            param_type = param.annotation
+            param_type = get_param_type(param)
             if get_origin(param_type) is typing.Annotated:
                 args = get_args(param_type)
                 if len(args) > 1:
@@ -595,11 +636,12 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
                 if param_name in exclude_params:
                     converted_body[param_name] = value
                 else:
-                    converted_body[param_name] = convert_to_pydantic(param.annotation, value)
+                    resolved_type = get_param_type(param)
+                    converted_body[param_name] = convert_to_pydantic(resolved_type, value)
 
         # Handle Depends parameter: construct request model from body
         if depends_param and depends_param.name not in converted_body:
-            param_type = depends_param.annotation
+            param_type = get_param_type(depends_param)
             if get_origin(param_type) is typing.Annotated:
                 base_type = get_args(param_type)[0]
                 # Handle Union types (e.g., SomeRequestModel | None) - extract the non-None type
@@ -619,10 +661,15 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
                     converted_body[depends_param.name] = base_type(**body)
 
         # handle unwrapped body parameter after processing all named parameters
-        if unwrapped_body_param:
-            base_type = get_args(unwrapped_body_param.annotation)[0]
+        if unwrapped_body_param and unwrapped_body_param_type:
+            base_type = get_args(unwrapped_body_param_type)[0]
             # extract only keys not already used by other params
             remaining_keys = {k: v for k, v in body.items() if k not in converted_body}
             converted_body[unwrapped_body_param.name] = base_type(**remaining_keys)
+        elif body_param and body_param.name not in converted_body:
+            body_param_type = get_param_type(body_param)
+            base_type = get_args(body_param_type)[0]
+            remaining_keys = {k: v for k, v in body.items() if k not in converted_body}
+            converted_body[body_param.name] = base_type(**remaining_keys)
 
         return converted_body

llama_stack/core/prompts/prompts.py

@@ -11,7 +11,17 @@ from pydantic import BaseModel
 
 from llama_stack.core.datatypes import StackConfig
 from llama_stack.core.storage.kvstore import KVStore, kvstore_impl
-from llama_stack_api import ListPromptsResponse, Prompt, Prompts
+from llama_stack_api import (
+    CreatePromptRequest,
+    DeletePromptRequest,
+    GetPromptRequest,
+    ListPromptsResponse,
+    ListPromptVersionsRequest,
+    Prompt,
+    Prompts,
+    SetDefaultVersionRequest,
+    UpdatePromptRequest,
+)
 
 
 class PromptServiceConfig(BaseModel):
@@ -114,26 +124,23 @@ class PromptServiceImpl(Prompts):
         prompts.sort(key=lambda p: p.prompt_id or "", reverse=True)
         return ListPromptsResponse(data=prompts)
 
-    async def get_prompt(self, prompt_id: str, version: int | None = None) -> Prompt:
+    async def get_prompt(self, request: GetPromptRequest) -> Prompt:
         """Get a prompt by its identifier and optional version."""
-        key = await self._get_prompt_key(prompt_id, version)
+        key = await self._get_prompt_key(request.prompt_id, request.version)
         data = await self.kvstore.get(key)
         if data is None:
-            raise ValueError(f"Prompt {prompt_id}:{version if version else 'default'} not found")
+            raise ValueError(
+                f"Prompt {request.prompt_id}:{request.version if request.version else 'default'} not found"
+            )
         return self._deserialize_prompt(data)
 
-    async def create_prompt(
-        self,
-        prompt: str,
-        variables: list[str] | None = None,
-    ) -> Prompt:
+    async def create_prompt(self, request: CreatePromptRequest) -> Prompt:
         """Create a new prompt."""
-        if variables is None:
-            variables = []
+        variables = request.variables if request.variables is not None else []
 
         prompt_obj = Prompt(
             prompt_id=Prompt.generate_prompt_id(),
-            prompt=prompt,
+            prompt=request.prompt,
             version=1,
             variables=variables,
         )
@@ -147,55 +154,49 @@ class PromptServiceImpl(Prompts):
 
         return prompt_obj
 
-    async def update_prompt(
-        self,
-        prompt_id: str,
-        prompt: str,
-        version: int,
-        variables: list[str] | None = None,
-        set_as_default: bool = True,
-    ) -> Prompt:
+    async def update_prompt(self, request: UpdatePromptRequest) -> Prompt:
         """Update an existing prompt (increments version)."""
-        if version < 1:
+        if request.version < 1:
             raise ValueError("Version must be >= 1")
-        if variables is None:
-            variables = []
+        variables = request.variables if request.variables is not None else []
 
-        prompt_versions = await self.list_prompt_versions(prompt_id)
+        prompt_versions = await self.list_prompt_versions(ListPromptVersionsRequest(prompt_id=request.prompt_id))
         latest_prompt = max(prompt_versions.data, key=lambda x: int(x.version))
 
-        if version and latest_prompt.version != version:
+        if request.version and latest_prompt.version != request.version:
             raise ValueError(
-                f"'{version}' is not the latest prompt version for prompt_id='{prompt_id}'. Use the latest version '{latest_prompt.version}' in request."
+                f"'{request.version}' is not the latest prompt version for prompt_id='{request.prompt_id}'. Use the latest version '{latest_prompt.version}' in request."
             )
 
-        current_version = latest_prompt.version if version is None else version
+        current_version = latest_prompt.version if request.version is None else request.version
         new_version = current_version + 1
 
-        updated_prompt = Prompt(prompt_id=prompt_id, prompt=prompt, version=new_version, variables=variables)
+        updated_prompt = Prompt(
+            prompt_id=request.prompt_id, prompt=request.prompt, version=new_version, variables=variables
+        )
 
-        version_key = self._get_version_key(prompt_id, str(new_version))
+        version_key = self._get_version_key(request.prompt_id, str(new_version))
         data = self._serialize_prompt(updated_prompt)
         await self.kvstore.set(version_key, data)
 
-        if set_as_default:
-            await self.set_default_version(prompt_id, new_version)
+        if request.set_as_default:
+            await self.set_default_version(SetDefaultVersionRequest(prompt_id=request.prompt_id, version=new_version))
 
         return updated_prompt
 
-    async def delete_prompt(self, prompt_id: str) -> None:
+    async def delete_prompt(self, request: DeletePromptRequest) -> None:
         """Delete a prompt and all its versions."""
-        await self.get_prompt(prompt_id)
+        await self.get_prompt(GetPromptRequest(prompt_id=request.prompt_id))
 
-        prefix = f"prompts:v1:{prompt_id}:"
+        prefix = f"prompts:v1:{request.prompt_id}:"
         keys = await self.kvstore.keys_in_range(prefix, prefix + "\xff")
 
         for key in keys:
             await self.kvstore.delete(key)
 
-    async def list_prompt_versions(self, prompt_id: str) -> ListPromptsResponse:
+    async def list_prompt_versions(self, request: ListPromptVersionsRequest) -> ListPromptsResponse:
         """List all versions of a specific prompt."""
-        prefix = f"prompts:v1:{prompt_id}:"
+        prefix = f"prompts:v1:{request.prompt_id}:"
         keys = await self.kvstore.keys_in_range(prefix, prefix + "\xff")
 
         default_version = None
@@ -211,7 +212,7 @@ class PromptServiceImpl(Prompts):
                     prompts.append(prompt_obj)
 
         if not prompts:
-            raise ValueError(f"Prompt {prompt_id} not found")
+            raise ValueError(f"Prompt {request.prompt_id} not found")
 
         for prompt in prompts:
             prompt.is_default = str(prompt.version) == default_version
@@ -219,15 +220,15 @@ class PromptServiceImpl(Prompts):
         prompts.sort(key=lambda x: x.version)
         return ListPromptsResponse(data=prompts)
 
-    async def set_default_version(self, prompt_id: str, version: int) -> Prompt:
+    async def set_default_version(self, request: SetDefaultVersionRequest) -> Prompt:
         """Set which version of a prompt should be the default, If not set. the default is the latest."""
-        version_key = self._get_version_key(prompt_id, str(version))
+        version_key = self._get_version_key(request.prompt_id, str(request.version))
         data = await self.kvstore.get(version_key)
         if data is None:
-            raise ValueError(f"Prompt {prompt_id} version {version} not found")
+            raise ValueError(f"Prompt {request.prompt_id} version {request.version} not found")
 
-        default_key = self._get_default_key(prompt_id)
-        await self.kvstore.set(default_key, str(version))
+        default_key = self._get_default_key(request.prompt_id)
+        await self.kvstore.set(default_key, str(request.version))
 
         return self._deserialize_prompt(data)
 

llama_stack/core/routers/datasets.py

@@ -7,7 +7,15 @@
 from typing import Any
 
 from llama_stack.log import get_logger
-from llama_stack_api import DatasetIO, DatasetPurpose, DataSource, PaginatedResponse, RoutingTable
+from llama_stack_api import (
+    AppendRowsRequest,
+    DatasetIO,
+    DatasetPurpose,
+    DataSource,
+    IterRowsRequest,
+    PaginatedResponse,
+    RoutingTable,
+)
 
 logger = get_logger(name=__name__, category="core::routers")
 
@@ -45,26 +53,21 @@ class DatasetIORouter(DatasetIO):
             dataset_id=dataset_id,
         )
 
-    async def iterrows(
-        self,
-        dataset_id: str,
-        start_index: int | None = None,
-        limit: int | None = None,
-    ) -> PaginatedResponse:
+    async def iterrows(self, request: IterRowsRequest) -> PaginatedResponse:
         logger.debug(
-            f"DatasetIORouter.iterrows: {dataset_id}, {start_index=} {limit=}",
+            f"DatasetIORouter.iterrows: {request.dataset_id}, start_index={request.start_index} limit={request.limit}",
         )
-        provider = await self.routing_table.get_provider_impl(dataset_id)
+        provider = await self.routing_table.get_provider_impl(request.dataset_id)
         return await provider.iterrows(
-            dataset_id=dataset_id,
-            start_index=start_index,
-            limit=limit,
+            dataset_id=request.dataset_id,
+            start_index=request.start_index,
+            limit=request.limit,
         )
 
-    async def append_rows(self, dataset_id: str, rows: list[dict[str, Any]]) -> None:
-        logger.debug(f"DatasetIORouter.append_rows: {dataset_id}, {len(rows)} rows")
-        provider = await self.routing_table.get_provider_impl(dataset_id)
+    async def append_rows(self, request: AppendRowsRequest) -> None:
+        logger.debug(f"DatasetIORouter.append_rows: {request.dataset_id}, {len(request.rows)} rows")
+        provider = await self.routing_table.get_provider_impl(request.dataset_id)
         return await provider.append_rows(
-            dataset_id=dataset_id,
-            rows=rows,
+            dataset_id=request.dataset_id,
+            rows=request.rows,
         )