llama-stack 0.4.4__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -7,7 +7,6 @@
 import re
 import uuid
 from string import Template
-from typing import Any
 
 from llama_stack.core.datatypes import Api
 from llama_stack.log import get_logger
@@ -17,6 +16,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 from llama_stack_api import (
+    GetShieldRequest,
     ImageContentItem,
     Inference,
     ModerationObject,
@@ -24,6 +24,8 @@ from llama_stack_api import (
     OpenAIChatCompletionRequestWithExtraBody,
     OpenAIMessageParam,
     OpenAIUserMessageParam,
+    RunModerationRequest,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -161,17 +163,12 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
         # The routing table handles the removal from the registry
         pass
 
-    async def run_shield(
-        self,
-        shield_id: str,
-        messages: list[OpenAIMessageParam],
-        params: dict[str, Any] = None,
-    ) -> RunShieldResponse:
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Unknown shield {shield_id}")
+            raise ValueError(f"Unknown shield {request.shield_id}")
 
-        messages = messages.copy()
+        messages = request.messages.copy()
         # some shields like llama-guard require the first message to be a user message
         # since this might be a tool call, first role might not be user
         if len(messages) > 0 and messages[0].role != "user":
@@ -200,30 +197,30 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
 
         return await impl.run(messages)
 
-    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
-        if model is None:
+    async def run_moderation(self, request: RunModerationRequest) -> ModerationObject:
+        if request.model is None:
             raise ValueError("Llama Guard moderation requires a model identifier.")
 
-        if isinstance(input, list):
-            messages = input.copy()
+        if isinstance(request.input, list):
+            messages = request.input.copy()
         else:
-            messages = [input]
+            messages = [request.input]
 
         # convert to user messages format with role
         messages = [OpenAIUserMessageParam(content=m) for m in messages]
 
         # Determine safety categories based on the model type
         # For known Llama Guard models, use specific categories
-        if model in LLAMA_GUARD_MODEL_IDS:
+        if request.model in LLAMA_GUARD_MODEL_IDS:
             # Use the mapped model for categories but the original model_id for inference
-            mapped_model = LLAMA_GUARD_MODEL_IDS[model]
+            mapped_model = LLAMA_GUARD_MODEL_IDS[request.model]
             safety_categories = MODEL_TO_SAFETY_CATEGORIES_MAP.get(mapped_model, DEFAULT_LG_V3_SAFETY_CATEGORIES)
         else:
             # For unknown models, use default Llama Guard 3 8B categories
             safety_categories = DEFAULT_LG_V3_SAFETY_CATEGORIES + [CAT_CODE_INTERPRETER_ABUSE]
 
         impl = LlamaGuardShield(
-            model=model,
+            model=request.model,
             inference_api=self.inference_api,
             excluded_categories=self.config.excluded_categories,
             safety_categories=safety_categories,

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

@@ -4,17 +4,19 @@
 # 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
-
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 from llama_stack.core.utils.model_utils import model_local_dir
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
+from llama_stack.providers.utils.inference.prompt_adapter import (
+    interleaved_content_as_str,
+)
+from llama_stack.providers.utils.safety import ShieldToModerationMixin
 from llama_stack_api import (
-    ModerationObject,
+    GetShieldRequest,
     OpenAIMessageParam,
+    RunShieldRequest,
     RunShieldResponse,
     Safety,
     SafetyViolation,
@@ -31,7 +33,7 @@ log = get_logger(name=__name__, category="safety")
 PROMPT_GUARD_MODEL = "Prompt-Guard-86M"
 
 
-class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
+class PromptGuardSafetyImpl(ShieldToModerationMixin, Safety, ShieldsProtocolPrivate):
     shield_store: ShieldStore
 
     def __init__(self, config: PromptGuardConfig, _deps) -> None:
@@ -51,20 +53,12 @@ class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
     async def unregister_shield(self, identifier: str) -> None:
         pass
 
-    async def run_shield(
-        self,
-        shield_id: str,
-        messages: list[OpenAIMessageParam],
-        params: dict[str, Any],
-    ) -> RunShieldResponse:
-        shield = await self.shield_store.get_shield(shield_id)
+    async def run_shield(self, request: RunShieldRequest) -> RunShieldResponse:
+        shield = await self.shield_store.get_shield(GetShieldRequest(identifier=request.shield_id))
         if not shield:
-            raise ValueError(f"Unknown shield {shield_id}")
-
-        return await self.shield.run(messages)
+            raise ValueError(f"Unknown shield {request.shield_id}")
 
-    async def run_moderation(self, input: str | list[str], model: str | None = None) -> ModerationObject:
-        raise NotImplementedError("run_moderation is not implemented for Prompt Guard")
+        return await self.shield.run(request.messages)
 
 
 class PromptGuardShield:

llama_stack/providers/inline/scoring/basic/scoring.py

@@ -3,16 +3,17 @@
 #
 # 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 llama_stack_api import (
     DatasetIO,
     Datasets,
+    IterRowsRequest,
+    ScoreBatchRequest,
     ScoreBatchResponse,
+    ScoreRequest,
     ScoreResponse,
     Scoring,
     ScoringFn,
-    ScoringFnParams,
     ScoringFunctionsProtocolPrivate,
     ScoringResult,
 )
@@ -75,19 +76,15 @@ class BasicScoringImpl(
 
     async def score_batch(
         self,
-        dataset_id: str,
-        scoring_functions: dict[str, ScoringFnParams | None] = None,
-        save_results_dataset: bool = False,
+        request: ScoreBatchRequest,
     ) -> ScoreBatchResponse:
-        all_rows = await self.datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=-1,
-        )
-        res = await self.score(
+        all_rows = await self.datasetio_api.iterrows(IterRowsRequest(dataset_id=request.dataset_id, limit=-1))
+        score_request = ScoreRequest(
             input_rows=all_rows.data,
-            scoring_functions=scoring_functions,
+            scoring_functions=request.scoring_functions,
         )
-        if save_results_dataset:
+        res = await self.score(score_request)
+        if request.save_results_dataset:
             # TODO: persist and register dataset on to server for reading
             # self.datasets_api.register_dataset()
             raise NotImplementedError("Save results dataset not implemented yet")
@@ -98,16 +95,15 @@ class BasicScoringImpl(
 
     async def score(
         self,
-        input_rows: list[dict[str, Any]],
-        scoring_functions: dict[str, ScoringFnParams | None] = None,
+        request: ScoreRequest,
     ) -> ScoreResponse:
         res = {}
-        for scoring_fn_id in scoring_functions.keys():
+        for scoring_fn_id in request.scoring_functions.keys():
             if scoring_fn_id not in self.scoring_fn_id_impls:
                 raise ValueError(f"Scoring function {scoring_fn_id} is not supported.")
             scoring_fn = self.scoring_fn_id_impls[scoring_fn_id]
-            scoring_fn_params = scoring_functions.get(scoring_fn_id, None)
-            score_results = await scoring_fn.score(input_rows, scoring_fn_id, scoring_fn_params)
+            scoring_fn_params = request.scoring_functions.get(scoring_fn_id, None)
+            score_results = await scoring_fn.score(request.input_rows, scoring_fn_id, scoring_fn_params)
             agg_results = await scoring_fn.aggregate(score_results, scoring_fn_id, scoring_fn_params)
             res[scoring_fn_id] = ScoringResult(
                 score_rows=score_results,

llama_stack/providers/inline/scoring/braintrust/braintrust.py

@@ -29,11 +29,13 @@ from llama_stack.providers.utils.scoring.aggregation_utils import aggregate_metr
 from llama_stack_api import (
     DatasetIO,
     Datasets,
+    IterRowsRequest,
+    ScoreBatchRequest,
     ScoreBatchResponse,
+    ScoreRequest,
     ScoreResponse,
     Scoring,
     ScoringFn,
-    ScoringFnParams,
     ScoringFunctionsProtocolPrivate,
     ScoringResult,
     ScoringResultRow,
@@ -158,18 +160,17 @@ class BraintrustScoringImpl(
 
     async def score_batch(
         self,
-        dataset_id: str,
-        scoring_functions: dict[str, ScoringFnParams | None],
-        save_results_dataset: bool = False,
+        request: ScoreBatchRequest,
     ) -> ScoreBatchResponse:
         await self.set_api_key()
 
-        all_rows = await self.datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=-1,
+        all_rows = await self.datasetio_api.iterrows(IterRowsRequest(dataset_id=request.dataset_id, limit=-1))
+        score_request = ScoreRequest(
+            input_rows=all_rows.data,
+            scoring_functions=request.scoring_functions,
         )
-        res = await self.score(input_rows=all_rows.data, scoring_functions=scoring_functions)
-        if save_results_dataset:
+        res = await self.score(score_request)
+        if request.save_results_dataset:
             # TODO: persist and register dataset on to server for reading
             # self.datasets_api.register_dataset()
             raise NotImplementedError("Save results dataset not implemented yet")
@@ -198,21 +199,20 @@ class BraintrustScoringImpl(
 
     async def score(
         self,
-        input_rows: list[dict[str, Any]],
-        scoring_functions: dict[str, ScoringFnParams | None],
+        request: ScoreRequest,
     ) -> ScoreResponse:
         await self.set_api_key()
         res = {}
-        for scoring_fn_id in scoring_functions:
+        for scoring_fn_id in request.scoring_functions:
             if scoring_fn_id not in self.supported_fn_defs_registry:
                 raise ValueError(f"Scoring function {scoring_fn_id} is not supported.")
 
-            score_results = [await self.score_row(input_row, scoring_fn_id) for input_row in input_rows]
+            score_results = [await self.score_row(input_row, scoring_fn_id) for input_row in request.input_rows]
             aggregation_functions = self.supported_fn_defs_registry[scoring_fn_id].params.aggregation_functions
 
             # override scoring_fn params if provided
-            if scoring_functions[scoring_fn_id] is not None:
-                override_params = scoring_functions[scoring_fn_id]
+            if request.scoring_functions[scoring_fn_id] is not None:
+                override_params = request.scoring_functions[scoring_fn_id]
                 if override_params.aggregation_functions:
                     aggregation_functions = override_params.aggregation_functions
 

llama_stack/providers/inline/scoring/llm_as_judge/scoring.py

@@ -3,17 +3,18 @@
 #
 # 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 llama_stack_api import (
     DatasetIO,
     Datasets,
     Inference,
+    IterRowsRequest,
+    ScoreBatchRequest,
     ScoreBatchResponse,
+    ScoreRequest,
     ScoreResponse,
     Scoring,
     ScoringFn,
-    ScoringFnParams,
     ScoringFunctionsProtocolPrivate,
     ScoringResult,
 )
@@ -64,19 +65,15 @@ class LlmAsJudgeScoringImpl(
 
     async def score_batch(
         self,
-        dataset_id: str,
-        scoring_functions: dict[str, ScoringFnParams | None] = None,
-        save_results_dataset: bool = False,
+        request: ScoreBatchRequest,
     ) -> ScoreBatchResponse:
-        all_rows = await self.datasetio_api.iterrows(
-            dataset_id=dataset_id,
-            limit=-1,
-        )
-        res = await self.score(
+        all_rows = await self.datasetio_api.iterrows(IterRowsRequest(dataset_id=request.dataset_id, limit=-1))
+        score_request = ScoreRequest(
             input_rows=all_rows.data,
-            scoring_functions=scoring_functions,
+            scoring_functions=request.scoring_functions,
         )
-        if save_results_dataset:
+        res = await self.score(score_request)
+        if request.save_results_dataset:
             # TODO: persist and register dataset on to server for reading
             # self.datasets_api.register_dataset()
             raise NotImplementedError("Save results dataset not implemented yet")
@@ -87,14 +84,13 @@ class LlmAsJudgeScoringImpl(
 
     async def score(
         self,
-        input_rows: list[dict[str, Any]],
-        scoring_functions: dict[str, ScoringFnParams | None] = None,
+        request: ScoreRequest,
     ) -> ScoreResponse:
         res = {}
-        for scoring_fn_id in scoring_functions.keys():
+        for scoring_fn_id in request.scoring_functions.keys():
             scoring_fn = self.llm_as_judge_fn
-            scoring_fn_params = scoring_functions.get(scoring_fn_id, None)
-            score_results = await scoring_fn.score(input_rows, scoring_fn_id, scoring_fn_params)
+            scoring_fn_params = request.scoring_functions.get(scoring_fn_id, None)
+            score_results = await scoring_fn.score(request.input_rows, scoring_fn_id, scoring_fn_params)
             agg_results = await scoring_fn.aggregate(score_results, scoring_fn_id, scoring_fn_params)
             res[scoring_fn_id] = ScoringResult(
                 score_rows=score_results,

llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -59,7 +59,7 @@ def serialize_vector(vector: list[float]) -> bytes:
     return struct.pack(f"{len(vector)}f", *vector)
 
 
-def _create_sqlite_connection(db_path):
+def _create_sqlite_connection(db_path: str):
     """Create a SQLite connection with sqlite_vec extension loaded."""
     connection = sqlite3.connect(db_path)
     connection.enable_load_extension(True)

llama_stack/providers/registry/agents.py

@@ -37,6 +37,7 @@ def available_providers() -> list[ProviderSpec]:
                 Api.conversations,
                 Api.prompts,
                 Api.files,
+                Api.connectors,
             ],
             optional_api_dependencies=[
                 Api.safety,

llama_stack/providers/registry/inference.py

@@ -28,14 +28,6 @@ META_REFERENCE_DEPS = [
 
 def available_providers() -> list[ProviderSpec]:
     return [
-        InlineProviderSpec(
-            api=Api.inference,
-            provider_type="inline::meta-reference",
-            pip_packages=META_REFERENCE_DEPS,
-            module="llama_stack.providers.inline.inference.meta_reference",
-            config_class="llama_stack.providers.inline.inference.meta_reference.MetaReferenceInferenceConfig",
-            description="Meta's reference implementation of inference with support for various model formats and optimization techniques.",
-        ),
         InlineProviderSpec(
             api=Api.inference,
             provider_type="inline::sentence-transformers",
@@ -223,7 +215,7 @@ def available_providers() -> list[ProviderSpec]:
 
 Configuration:
 - Set VERTEX_AI_PROJECT environment variable (required)
-- Set VERTEX_AI_LOCATION environment variable (optional, defaults to us-central1)
+- Set VERTEX_AI_LOCATION environment variable (optional, defaults to global)
 - Use Google Cloud Application Default Credentials or service account key
 
 Authentication Setup:

llama_stack/providers/registry/vector_io.py

@@ -419,6 +419,7 @@ There are three implementations of search for PGVectoIndex available:
   - Semantic understanding - finds documents similar in meaning even if they don't share keywords
   - Works with high-dimensional vector embeddings (typically 768, 1024, or higher dimensions)
   - Best for: Finding conceptually related content, handling synonyms, cross-language search
+  - By default, Llama Stack creates a HNSW (Hierarchical Navigable Small Worlds) index on a column "embedding" in a vector store table enabling production-ready, performant and scalable vector search for large datasets out of the box.
 
 2. Keyword Search
 - How it works:
@@ -448,6 +449,7 @@ There are three implementations of search for PGVectoIndex available:
   - Best for: General-purpose search where you want both precision and recall
 
 4. Database Schema
+
 The PGVector implementation stores data optimized for all three search types:
 CREATE TABLE vector_store_xxx (
     id TEXT PRIMARY KEY,
@@ -457,9 +459,6 @@ CREATE TABLE vector_store_xxx (
     tokenized_content TSVECTOR          -- For keyword search
 );
 
--- Indexes for performance
-CREATE INDEX content_gin_idx ON table USING GIN(tokenized_content);  -- Keyword search
--- Vector index created automatically by pgvector
 
 ## Usage
 
@@ -469,32 +468,55 @@ To use PGVector in your Llama Stack project, follow these steps:
 2. Configure your Llama Stack project to use pgvector. (e.g. remote::pgvector).
 3. Start storing and querying vectors.
 
-## This is an example how you can set up your environment for using PGVector
+## This is an example how you can set up your environment for using PGVector (you can use either Podman or Docker)
 
-1. Export env vars:
+1. Export PGVector environment variables:
 ```bash
-export ENABLE_PGVECTOR=true
+export PGVECTOR_DB=testvectordb
 export PGVECTOR_HOST=localhost
 export PGVECTOR_PORT=5432
-export PGVECTOR_DB=llamastack
-export PGVECTOR_USER=llamastack
-export PGVECTOR_PASSWORD=llamastack
+export PGVECTOR_USER=user
+export PGVECTOR_PASSWORD=password
 ```
 
-2. Create DB:
+2. Pull pgvector image with that tag you want:
+
+Via Podman:
 ```bash
-psql -h localhost -U postgres -c "CREATE ROLE llamastack LOGIN PASSWORD 'llamastack';"
-psql -h localhost -U postgres -c "CREATE DATABASE llamastack OWNER llamastack;"
-psql -h localhost -U llamastack -d llamastack -c "CREATE EXTENSION IF NOT EXISTS vector;"
+podman pull pgvector/pgvector:0.8.1-pg18-trixie
 ```
 
-## Installation
+Via Docker:
+```bash
+docker pull pgvector/pgvector:0.8.1-pg18-trixie
+```
+
+3. Run container with PGVector:
 
-You can install PGVector using docker:
+Via Podman
+```bash
+podman run -d \
+  --name pgvector \
+  -e POSTGRES_PASSWORD=password \
+  -e POSTGRES_USER=user \
+  -e POSTGRES_DB=testvectordb \
+  -p 5432:5432 \
+  -v pgvector_data:/var/lib/postgresql \
+  pgvector/pgvector:0.8.1-pg18-trixie
+```
 
+Via Docker
 ```bash
-docker pull pgvector/pgvector:pg17
+docker run -d \
+  --name pgvector \
+  -e POSTGRES_PASSWORD=password \
+  -e POSTGRES_USER=user \
+  -e POSTGRES_DB=testvectordb \
+  -p 5432:5432 \
+  -v pgvector_data:/var/lib/postgresql \
+  pgvector/pgvector:0.8.1-pg18-trixie
 ```
+
 ## Documentation
 See [PGVector's documentation](https://github.com/pgvector/pgvector) for more details about PGVector in general.
 """,
@@ -823,6 +845,104 @@ For more details on TLS configuration, refer to the [TLS setup guide](https://mi
             optional_api_dependencies=[Api.files, Api.models],
             description="""
 Please refer to the remote provider documentation.
+""",
+        ),
+        RemoteProviderSpec(
+            api=Api.vector_io,
+            adapter_type="elasticsearch",
+            provider_type="remote::elasticsearch",
+            pip_packages=["elasticsearch>=8.16.0,<9.0.0"] + DEFAULT_VECTOR_IO_DEPS,
+            module="llama_stack.providers.remote.vector_io.elasticsearch",
+            config_class="llama_stack.providers.remote.vector_io.elasticsearch.ElasticsearchVectorIOConfig",
+            api_dependencies=[Api.inference],
+            optional_api_dependencies=[Api.files, Api.models],
+            description="""
+[Elasticsearch](https://www.elastic.co/) is a vector database provider for Llama Stack.
+It allows you to store and query vectors directly within an Elasticsearch database.
+That means you're not limited to storing vectors in memory or in a separate service.
+
+## Features
+Elasticsearch supports:
+- Store embeddings and their metadata
+- Vector search
+- Full-text search
+- Fuzzy search
+- Hybrid search
+- Document storage
+- Metadata filtering
+- Inference service
+- Machine Learning integrations
+
+## Usage
+
+To use Elasticsearch in your Llama Stack project, follow these steps:
+
+1. Install the necessary dependencies.
+2. Configure your Llama Stack project to use Elasticsearch.
+3. Start storing and querying vectors.
+
+## Installation
+
+You can test Elasticsearch locally by running this script in the terminal:
+
+```bash
+curl -fsSL https://elastic.co/start-local | sh
+```
+
+Or you can [start a free trial](https://www.elastic.co/cloud/cloud-trial-overview?utm_campaign=llama-stack-integration) on Elastic Cloud.
+For more information on how to deploy Elasticsearch, see the [official documentation](https://www.elastic.co/docs/deploy-manage/deploy).
+
+## Documentation
+See [Elasticsearch's documentation](https://www.elastic.co/docs/solutions/search) for more details about Elasticsearch in general.
+""",
+        ),
+        RemoteProviderSpec(
+            api=Api.vector_io,
+            adapter_type="oci",
+            provider_type="remote::oci",
+            pip_packages=["oracledb", "numpy"] + DEFAULT_VECTOR_IO_DEPS,
+            module="llama_stack.providers.remote.vector_io.oci",
+            config_class="llama_stack.providers.remote.vector_io.oci.OCI26aiVectorIOConfig",
+            api_dependencies=[Api.inference],
+            optional_api_dependencies=[Api.files, Api.models],
+            description="""
+[Oracle 26ai](https://docs.oracle.com/en/database/oracle/oracle-database/26/index.html)
+is a remote vector database provider for Llama Stack. It allows you to store and query vectors directly
+in an Oracle 26ai database.
+## Features
+- Easy to use
+- Fully integrated with Llama Stack
+- Supports vector search, keyword search, and hybrid search
+## Usage
+To use Oracle 26ai in your Llama Stack project, follow these steps:
+1. Install the necessary dependencies.
+2. Configure your Llama Stack project to use Oracle 26ai.
+3. Start storing and querying vectors.
+## Installation
+You can install the Oracle 26ai client using pip:
+```bash
+pip install oracledb
+```
+## Configuration
+```yaml
+vector_io:
+- provider_id: oci
+  provider_type: remote::oci
+  config:
+    conn_str: "${env.OCI26AI_CONNECTION_STRING}"
+    user: "${env.OCI26AI_USER}"
+    password: "${env.OCI26AI_PASSWORD}"
+    tnsnames_loc: "${env.OCI26AI_TNSNAMES_LOC}"
+    ewallet_pem_loc: "${env.OCI26AI_EWALLET_PEM_LOC}"
+    ewallet_password: "${env.OCI26AI_EWALLET_PWD}"
+    vector_datatype: "${env.OCI26AI_VECTOR_DATATYPE:=FLOAT32}"
+    persistence:
+      namespace: vector_id::oci26ai
+      backend: kv_default
+```
+## Documentation
+See the [Oracle 26ai documentation](https://docs.oracle.com/en/database/oracle/oracle-database/26/index.html)
+for more details about Oracle 26ai in general.
 """,
         ),
     ]