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

llama_stack/providers/remote/vector_io/pgvector/pgvector.py

@@ -37,7 +37,7 @@ from llama_stack_api import (
 )
 from llama_stack_api.internal.kvstore import KVStore
 
-from .config import PGVectorVectorIOConfig
+from .config import PGVectorIndexConfig, PGVectorIndexType, PGVectorVectorIOConfig
 
 log = get_logger(name=__name__, category="vector_io::pgvector")
 
@@ -81,6 +81,26 @@ def upsert_models(conn, keys_models: list[tuple[str, BaseModel]]):
         execute_values(cur, query, values, template="(%s, %s)")
 
 
+def remove_vector_store_metadata(conn: psycopg2.extensions.connection, vector_store_id: str) -> None:
+    """
+    Performs removal of vector store metadata from PGVector metadata_store table when vector store is unregistered
+
+    Args:
+        conn: active PostgreSQL connection
+        vector_store_id: identifier of VectorStore resource
+    """
+    try:
+        with conn.cursor() as cur:
+            cur.execute("DELETE FROM metadata_store WHERE key = %s", (vector_store_id,))
+            if cur.rowcount > 0:
+                log.info(f"Removed metadata for vector store '{vector_store_id}' from PGVector metadata_store table.")
+
+    except Exception as e:
+        raise RuntimeError(
+            f"Error removing metadata from PGVector metadata_store for vector_store: {vector_store_id}"
+        ) from e
+
+
 def load_models(cur, cls):
     cur.execute("SELECT key, data FROM metadata_store")
     rows = cur.fetchall()
@@ -89,22 +109,35 @@ def load_models(cur, cls):
 
 class PGVectorIndex(EmbeddingIndex):
     # reference: https://github.com/pgvector/pgvector?tab=readme-ov-file#querying
+    # Llama Stack supports only search functions that are applied for embeddings with vector type
     PGVECTOR_DISTANCE_METRIC_TO_SEARCH_FUNCTION: dict[str, str] = {
         "L2": "<->",
         "L1": "<+>",
         "COSINE": "<=>",
         "INNER_PRODUCT": "<#>",
-        "HAMMING": "<~>",
-        "JACCARD": "<%>",
     }
 
+    # reference: https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw
+    # Llama Stack supports only index operator classes that are applied for embeddings with vector type
+    PGVECTOR_DISTANCE_METRIC_TO_INDEX_OPERATOR_CLASS: dict[str, str] = {
+        "L2": "vector_l2_ops",
+        "L1": "vector_l1_ops",
+        "COSINE": "vector_cosine_ops",
+        "INNER_PRODUCT": "vector_ip_ops",
+    }
+
+    # pgvector's maximum embedding dimension for HNSW/IVFFlat indexes on column with type vector
+    # references: https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw and https://github.com/pgvector/pgvector?tab=readme-ov-file#ivfflat
+    MAX_EMBEDDING_DIMENSION_FOR_HNSW_AND_IVFFLAT_INDEX = 2000
+
     def __init__(
         self,
         vector_store: VectorStore,
         dimension: int,
         conn: psycopg2.extensions.connection,
+        distance_metric: str,
+        vector_index: PGVectorIndexConfig,
         kvstore: KVStore | None = None,
-        distance_metric: str = "COSINE",
     ):
         self.vector_store = vector_store
         self.dimension = dimension
@@ -112,6 +145,7 @@ class PGVectorIndex(EmbeddingIndex):
         self.kvstore = kvstore
         self.check_distance_metric_availability(distance_metric)
         self.distance_metric = distance_metric
+        self.vector_index = vector_index
         self.table_name = None
 
     async def initialize(self) -> None:
@@ -135,6 +169,28 @@ class PGVectorIndex(EmbeddingIndex):
                 """
                 )
 
+                # pgvector's embedding dimensions requirement to create an index for Approximate Nearest Neighbor (ANN) search is up to 2,000 dimensions for column with type vector
+                if self.dimension <= self.MAX_EMBEDDING_DIMENSION_FOR_HNSW_AND_IVFFLAT_INDEX:
+                    if self.vector_index.type == PGVectorIndexType.HNSW:
+                        await self.create_hnsw_vector_index(cur)
+
+                    # Create the index only after the table has some data (https://github.com/pgvector/pgvector?tab=readme-ov-file#ivfflat)
+                    elif (
+                        self.vector_index.type == PGVectorIndexType.IVFFlat
+                        and not await self.check_conflicting_vector_index_exists(cur)
+                    ):
+                        log.info(
+                            f"Creation of {PGVectorIndexType.IVFFlat} vector index in vector_store: {self.vector_store.identifier} was deferred. It will be created when the table has some data."
+                        )
+
+                else:
+                    log.info(
+                        f"Skip creation of {self.vector_index.type} vector index for embedding in PGVector for vector_store: {self.vector_store.identifier}"
+                    )
+                    log.info(
+                        "PGVector requires embedding dimensions are up to 2,000 to successfully create a vector index."
+                    )
+
                 # Create GIN index for full-text search performance
                 cur.execute(
                     f"""
@@ -177,6 +233,13 @@ class PGVectorIndex(EmbeddingIndex):
         with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
             execute_values(cur, query, values, template="(%s, %s, %s::vector, %s, to_tsvector('english', %s))")
 
+            # Create the IVFFlat index only after the table has some data (https://github.com/pgvector/pgvector?tab=readme-ov-file#ivfflat)
+            if (
+                self.vector_index.type == PGVectorIndexType.IVFFlat
+                and self.dimension <= self.MAX_EMBEDDING_DIMENSION_FOR_HNSW_AND_IVFFLAT_INDEX
+            ):
+                await self.create_ivfflat_vector_index(cur)
+
     async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
         """
         Performs vector similarity search using PostgreSQL's search function. Default distance metric is COSINE.
@@ -192,6 +255,14 @@ class PGVectorIndex(EmbeddingIndex):
         pgvector_search_function = self.get_pgvector_search_function()
 
         with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
+            # Specify the number of probes to allow PGVector to use Index Scan using IVFFlat index if it was configured (https://github.com/pgvector/pgvector?tab=readme-ov-file#query-options-1)
+            if self.vector_index.type == PGVectorIndexType.IVFFlat:
+                cur.execute(
+                    f"""
+                    SET ivfflat.probes = {self.vector_index.probes};
+                """
+                )
+
             cur.execute(
                 f"""
             SELECT document, embedding {pgvector_search_function} %s::vector AS distance
@@ -324,6 +395,14 @@ class PGVectorIndex(EmbeddingIndex):
             # Fix: Use proper tuple parameter binding with explicit array cast
             cur.execute(f"DELETE FROM {self.table_name} WHERE id = ANY(%s::text[])", (chunk_ids,))
 
+    def get_pgvector_index_operator_class(self) -> str:
+        """Get the pgvector index operator class for the current distance metric.
+
+        Returns:
+            The operator class name.
+        """
+        return self.PGVECTOR_DISTANCE_METRIC_TO_INDEX_OPERATOR_CLASS[self.distance_metric]
+
     def get_pgvector_search_function(self) -> str:
         return self.PGVECTOR_DISTANCE_METRIC_TO_SEARCH_FUNCTION[self.distance_metric]
 
@@ -343,6 +422,160 @@ class PGVectorIndex(EmbeddingIndex):
                 f"Supported metrics are: {', '.join(supported_metrics)}"
             )
 
+    async def create_hnsw_vector_index(self, cur: cursor) -> None:
+        """Create PGVector HNSW vector index for Approximate Nearest Neighbor (ANN) search
+
+        Args:
+            cur: PostgreSQL cursor
+
+        Raises:
+            RuntimeError: If the error occurred when creating vector index in PGVector
+        """
+
+        # prevents from creating index for the table that already has conflicting index (HNSW or IVFFlat)
+        if await self.check_conflicting_vector_index_exists(cur):
+            return
+
+        try:
+            index_operator_class = self.get_pgvector_index_operator_class()
+
+            # Create HNSW (Hierarchical Navigable Small Worlds) index on embedding column to allow efficient and performant vector search in pgvector
+            # HNSW finds the approximate nearest neighbors by only calculating distance metric for vectors it visits during graph traversal instead of processing all vectors
+            cur.execute(
+                f"""
+                CREATE INDEX IF NOT EXISTS {self.table_name}_hnsw_idx
+                ON {self.table_name} USING hnsw(embedding {index_operator_class}) WITH (m = {self.vector_index.m}, ef_construction = {self.vector_index.ef_construction});
+            """
+            )
+            log.info(
+                f"{PGVectorIndexType.HNSW} vector index was created with parameters m = {self.vector_index.m}, ef_construction = {self.vector_index.ef_construction} for vector_store: {self.vector_store.identifier}."
+            )
+
+        except psycopg2.Error as e:
+            raise RuntimeError(
+                f"Failed to create {PGVectorIndexType.HNSW} vector index for vector_store: {self.vector_store.identifier}: {e}"
+            ) from e
+
+    async def create_ivfflat_vector_index(self, cur: cursor) -> None:
+        """Create PGVector IVFFlat vector index for Approximate Nearest Neighbor (ANN) search
+
+        Args:
+            cur: PostgreSQL cursor
+
+        Raises:
+            RuntimeError: If the error occurred when creating vector index in PGVector
+        """
+
+        # prevents from creating index for the table that already has conflicting index (HNSW or IVFFlat)
+        if await self.check_conflicting_vector_index_exists(cur):
+            return
+
+        # don't create index too early as it decreases a performance (https://github.com/pgvector/pgvector?tab=readme-ov-file#ivfflat)
+        # create IVFFLAT index only if vector store has rows >= lists * 1000
+        if await self.fetch_number_of_records(cur) < self.vector_index.lists * 1000:
+            log.info(
+                f"IVFFlat index wasn't created for vector_store {self.vector_store.identifier} because table doesn't have enough records."
+            )
+            return
+
+        try:
+            index_operator_class = self.get_pgvector_index_operator_class()
+
+            # Create Inverted File with Flat Compression (IVFFlat) index on embedding column to allow efficient and performant vector search in pgvector
+            # IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector
+            # Index should be created only after the table has some data (https://github.com/pgvector/pgvector?tab=readme-ov-file#ivfflat)
+            cur.execute(
+                f"""
+                CREATE INDEX IF NOT EXISTS {self.table_name}_ivfflat_idx
+                ON {self.table_name} USING ivfflat(embedding {index_operator_class}) WITH (lists = {self.vector_index.lists});
+            """
+            )
+            log.info(
+                f"{PGVectorIndexType.IVFFlat} vector index was created with parameter lists = {self.vector_index.lists} for vector_store: {self.vector_store.identifier}."
+            )
+
+        except psycopg2.Error as e:
+            raise RuntimeError(
+                f"Failed to create {PGVectorIndexType.IVFFlat} vector index for vector_store: {self.vector_store.identifier}: {e}"
+            ) from e
+
+    async def check_conflicting_vector_index_exists(self, cur: cursor) -> bool:
+        """Check if vector index of any type has already been created for the table to prevent the conflict
+
+        Args:
+            cur: PostgreSQL cursor
+
+        Returns:
+            True if exists, otherwise False
+
+        Raises:
+            RuntimeError: If the error occurred when checking vector index exists in PGVector
+        """
+        try:
+            log.info(
+                f"Checking vector_store: {self.vector_store.identifier} for conflicting vector index in PGVector..."
+            )
+            cur.execute(
+                """
+                SELECT indexname FROM pg_indexes
+                WHERE (indexname LIKE %s OR indexname LIKE %s) AND tablename = %s;
+                """,
+                (
+                    "%hnsw%",
+                    "%ivfflat%",
+                    self.table_name,
+                ),
+            )
+            result = cur.fetchone()
+
+            if result:
+                log.warning(
+                    f"Conflicting vector index {result[0]} already exists in vector_store: {self.vector_store.identifier}"
+                )
+                log.warning(
+                    f"vector_store: {self.vector_store.identifier} will continue to use vector index {result[0]} to preserve performance."
+                )
+                return True
+
+            log.info(f"vector_store: {self.vector_store.identifier} currently doesn't have conflicting vector index")
+            log.info(f"Proceeding with creation of vector index for {self.vector_store.identifier}")
+            return False
+
+        except psycopg2.Error as e:
+            raise RuntimeError(f"Failed to check if vector index exists in PGVector: {e}") from e
+
+    async def fetch_number_of_records(self, cur: cursor) -> int:
+        """Returns number of records in a vector store
+
+        Args:
+            cur: PostgreSQL cursor
+
+        Returns:
+            number of records in a vector store
+
+        Raises:
+            RuntimeError: If the error occurred when fetching a number of records in a vector store in PGVector
+        """
+        try:
+            log.info(f"Fetching number of records in vector_store: {self.vector_store.identifier}...")
+            cur.execute(
+                f"""
+                SELECT COUNT(DISTINCT id)
+                FROM {self.table_name};
+                """
+            )
+            result = cur.fetchone()
+
+            if result:
+                log.info(f"vector_store: {self.vector_store.identifier} has {result[0]} records.")
+                return result[0]
+
+            log.info(f"vector_store: {self.vector_store.identifier} currently doesn't have any records.")
+            return 0
+
+        except psycopg2.Error as e:
+            raise RuntimeError(f"Failed to check if vector store has records in PGVector: {e}") from e
+
 
 class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtocolPrivate):
     def __init__(
@@ -401,6 +634,8 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
                 dimension=vector_store.embedding_dimension,
                 conn=self.conn,
                 kvstore=self.kvstore,
+                distance_metric=self.config.distance_metric,
+                vector_index=self.config.vector_index,
             )
             await pgvector_index.initialize()
             index = VectorStoreWithIndex(vector_store, index=pgvector_index, inference_api=self.inference_api)
@@ -427,7 +662,12 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
 
         # Create and cache the PGVector index table for the vector DB
         pgvector_index = PGVectorIndex(
-            vector_store=vector_store, dimension=vector_store.embedding_dimension, conn=self.conn, kvstore=self.kvstore
+            vector_store=vector_store,
+            dimension=vector_store.embedding_dimension,
+            conn=self.conn,
+            kvstore=self.kvstore,
+            distance_metric=self.config.distance_metric,
+            vector_index=self.config.vector_index,
         )
         await pgvector_index.initialize()
         index = VectorStoreWithIndex(vector_store, index=pgvector_index, inference_api=self.inference_api)
@@ -444,6 +684,9 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
             raise RuntimeError("KVStore not initialized. Call initialize() before unregistering vector stores.")
         await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_store_id}")
 
+        # Delete vector store metadata from PGVector metadata_store table
+        remove_vector_store_metadata(self.conn, vector_store_id)
+
     async def insert_chunks(
         self, vector_store_id: str, chunks: list[EmbeddedChunk], ttl_seconds: int | None = None
     ) -> None:
@@ -470,7 +713,13 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
             raise VectorStoreNotFoundError(vector_store_id)
 
         vector_store = VectorStore.model_validate_json(vector_store_data)
-        index = PGVectorIndex(vector_store, vector_store.embedding_dimension, self.conn)
+        index = PGVectorIndex(
+            vector_store,
+            vector_store.embedding_dimension,
+            self.conn,
+            distance_metric=self.config.distance_metric,
+            vector_index=self.config.vector_index,
+        )
         await index.initialize()
         self.cache[vector_store_id] = VectorStoreWithIndex(vector_store, index, self.inference_api)
         return self.cache[vector_store_id]

llama_stack/providers/remote/vector_io/qdrant/qdrant.py

@@ -16,6 +16,7 @@ from qdrant_client.models import PointStruct
 from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
+from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
 from llama_stack.providers.utils.vector_io.vector_utils import load_embedded_chunk_with_backward_compat
@@ -80,11 +81,16 @@ class QdrantIndex(EmbeddingIndex):
         points = []
         for chunk in chunks:
             chunk_id = chunk.chunk_id
+            content_text = interleaved_content_as_str(chunk.content)
             points.append(
                 PointStruct(
                     id=convert_id(chunk_id),
-                    vector=chunk.embedding,  # Already a list[float]
-                    payload={"chunk_content": chunk.model_dump()} | {CHUNK_ID_KEY: chunk_id},
+                    vector=chunk.embedding,
+                    payload={
+                        "chunk_content": chunk.model_dump(),
+                        "content_text": content_text,
+                        CHUNK_ID_KEY: chunk_id,
+                    },
                 )
             )
 
@@ -144,32 +150,32 @@ class QdrantIndex(EmbeddingIndex):
             QueryChunksResponse with chunks and scores matching the keyword query
         """
         try:
-            results = (
-                await self.client.query_points(
-                    collection_name=self.collection_name,
-                    query_filter=models.Filter(
-                        must=[
-                            models.FieldCondition(
-                                key="chunk_content.content", match=models.MatchText(text=query_string)
-                            )
-                        ]
-                    ),
-                    limit=k,
-                    with_payload=True,
-                    with_vectors=False,
-                    score_threshold=score_threshold,
-                )
-            ).points
+            # Use scroll for keyword-only search since query_points requires a query vector
+            # Scroll allows filtering without a query vector
+            query_words = query_string.lower().split()
+            if not query_words:
+                return QueryChunksResponse(chunks=[], scores=[])
+            scroll_result = await self.client.scroll(
+                collection_name=self.collection_name,
+                scroll_filter=models.Filter(
+                    should=[
+                        models.FieldCondition(key="content_text", match=models.MatchText(text=word))
+                        for word in query_words
+                    ]
+                ),
+                limit=k,
+                with_payload=True,
+                with_vectors=False,
+            )
+            results = scroll_result[0]
         except Exception as e:
             log.error(f"Error querying keyword search in Qdrant collection {self.collection_name}: {e}")
             raise
 
         chunks, scores = [], []
         for point in results:
-            if not isinstance(point, models.ScoredPoint):
-                raise RuntimeError(f"Expected ScoredPoint from Qdrant query, got {type(point).__name__}")
             if point.payload is None:
-                raise RuntimeError("Qdrant query returned point with no payload")
+                raise RuntimeError("Qdrant scroll returned point with no payload")
 
             try:
                 chunk = load_embedded_chunk_with_backward_compat(point.payload["chunk_content"])
@@ -182,8 +188,13 @@ class QdrantIndex(EmbeddingIndex):
                 )
                 continue
 
+            # For keyword search, use a fixed score of 1.0 since we're not doing vector similarity
+            score = 1.0
+            if score < score_threshold:
+                continue
+
             chunks.append(chunk)
-            scores.append(point.score)
+            scores.append(score)
 
         return QueryChunksResponse(chunks=chunks, scores=scores)
 
@@ -214,22 +225,35 @@ class QdrantIndex(EmbeddingIndex):
             QueryChunksResponse with filtered vector search results
         """
         try:
-            results = (
-                await self.client.query_points(
-                    collection_name=self.collection_name,
-                    query=embedding.tolist(),
-                    query_filter=models.Filter(
-                        must=[
-                            models.FieldCondition(
-                                key="chunk_content.content", match=models.MatchText(text=query_string)
-                            )
-                        ]
-                    ),
-                    limit=k,
-                    with_payload=True,
-                    score_threshold=score_threshold,
-                )
-            ).points
+            query_words = query_string.lower().split()
+            if not query_words:
+                # If no words, just do vector search without keyword filter
+                results = (
+                    await self.client.query_points(
+                        collection_name=self.collection_name,
+                        query=embedding.tolist(),
+                        limit=k,
+                        with_payload=True,
+                        score_threshold=score_threshold,
+                    )
+                ).points
+            else:
+                # Use should to match any of the query words
+                results = (
+                    await self.client.query_points(
+                        collection_name=self.collection_name,
+                        query=embedding.tolist(),
+                        query_filter=models.Filter(
+                            should=[
+                                models.FieldCondition(key="content_text", match=models.MatchText(text=word))
+                                for word in query_words
+                            ]
+                        ),
+                        limit=k,
+                        with_payload=True,
+                        score_threshold=score_threshold,
+                    )
+                ).points
         except Exception as e:
             log.error(f"Error querying hybrid search in Qdrant collection {self.collection_name}: {e}")
             raise

llama_stack/providers/utils/bedrock/client.py

@@ -49,9 +49,9 @@ def create_bedrock_client(config: BedrockBaseConfig, service_name: str = "bedroc
         boto3_config = Config(**config_args)
 
         session_args = {
-            "aws_access_key_id": config.aws_access_key_id,
-            "aws_secret_access_key": config.aws_secret_access_key,
-            "aws_session_token": config.aws_session_token,
+            "aws_access_key_id": config.aws_access_key_id.get_secret_value(),
+            "aws_secret_access_key": config.aws_secret_access_key.get_secret_value(),
+            "aws_session_token": config.aws_session_token.get_secret_value() if config.aws_session_token else None,
             "region_name": config.region_name,
             "profile_name": config.profile_name,
             "session_ttl": config.session_ttl,

llama_stack/providers/utils/bedrock/config.py

@@ -6,23 +6,23 @@
 
 import os
 
-from pydantic import Field
+from pydantic import Field, SecretStr
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
 
 
 class BedrockBaseConfig(RemoteInferenceProviderConfig):
     auth_credential: None = Field(default=None, exclude=True)
-    aws_access_key_id: str | None = Field(
-        default_factory=lambda: os.getenv("AWS_ACCESS_KEY_ID"),
+    aws_access_key_id: SecretStr | None = Field(
+        default_factory=lambda: SecretStr(val) if (val := os.getenv("AWS_ACCESS_KEY_ID")) else None,
         description="The AWS access key to use. Default use environment variable: AWS_ACCESS_KEY_ID",
     )
-    aws_secret_access_key: str | None = Field(
-        default_factory=lambda: os.getenv("AWS_SECRET_ACCESS_KEY"),
+    aws_secret_access_key: SecretStr | None = Field(
+        default_factory=lambda: SecretStr(val) if (val := os.getenv("AWS_SECRET_ACCESS_KEY")) else None,
         description="The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY",
     )
-    aws_session_token: str | None = Field(
-        default_factory=lambda: os.getenv("AWS_SESSION_TOKEN"),
+    aws_session_token: SecretStr | None = Field(
+        default_factory=lambda: SecretStr(val) if (val := os.getenv("AWS_SESSION_TOKEN")) else None,
         description="The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN",
     )
     region_name: str | None = Field(

llama_stack/providers/utils/inference/__init__.py

@@ -5,30 +5,5 @@
 # the root directory of this source tree.
 
 from llama_stack.models.llama.sku_list import all_registered_models
-from llama_stack.models.llama.sku_types import *  # noqa: F403
-
-
-def is_supported_safety_model(model: Model) -> bool:
-    if model.quantization_format != CheckpointQuantizationFormat.bf16:
-        return False
-
-    model_id = model.core_model_id
-    return model_id in [
-        CoreModelId.llama_guard_3_8b,
-        CoreModelId.llama_guard_3_1b,
-        CoreModelId.llama_guard_3_11b_vision,
-    ]
-
-
-def supported_inference_models() -> list[Model]:
-    return [
-        m
-        for m in all_registered_models()
-        if (
-            m.model_family in {ModelFamily.llama3_1, ModelFamily.llama3_2, ModelFamily.llama3_3, ModelFamily.llama4}
-            or is_supported_safety_model(m)
-        )
-    ]
-
 
 ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR = {m.huggingface_repo: m.descriptor() for m in all_registered_models()}

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -23,6 +23,7 @@ from llama_stack_api import (
     OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
+    validate_embeddings_input_is_text,
 )
 
 EMBEDDING_MODELS: dict[str, "SentenceTransformer"] = {}
@@ -41,6 +42,9 @@ class SentenceTransformerEmbeddingMixin:
         self,
         params: OpenAIEmbeddingsRequestWithExtraBody,
     ) -> OpenAIEmbeddingsResponse:
+        # Validate that input contains only text, not token arrays
+        validate_embeddings_input_is_text(params)
+
         # Convert input to list format if it's a single string
         input_list = [params.input] if isinstance(params.input, str) else params.input
         if not input_list: