Merge branch 'main' into hide-non-openai-inference-apis

llama_stack/providers/utils/inference/model_registry.py

@@ -20,7 +20,7 @@ logger = get_logger(name=__name__, category="providers::utils")
 
 
 class RemoteInferenceProviderConfig(BaseModel):
-    allowed_models: list[str] | None = Field(
+    allowed_models: list[str] | None = Field(  # TODO: make this non-optional and give a list() default
         default=None,
         description="List of models that should be registered with the model registry. If None, all models are allowed.",
     )

llama_stack/providers/utils/inference/openai_mixin.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import base64
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
@@ -26,6 +27,7 @@ from llama_stack.apis.models import ModelType
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 from llama_stack.providers.utils.inference.openai_compat import prepare_openai_completion_params
+from llama_stack.providers.utils.inference.prompt_adapter import localize_image_content
 
 logger = get_logger(name=__name__, category="providers::utils")
 
@@ -51,6 +53,10 @@ class OpenAIMixin(ModelRegistryHelper, ABC):
     # This is useful for providers that do not return a unique id in the response.
     overwrite_completion_id: bool = False
 
+    # Allow subclasses to control whether to download images and convert to base64
+    # for providers that require base64 encoded images instead of URLs.
+    download_images: bool = False
+
     # Embedding model metadata for this provider
     # Can be set by subclasses or instances to provide embedding models
     # Format: {"model_id": {"embedding_dimension": 1536, "context_length": 8192}}
@@ -239,6 +245,24 @@ class OpenAIMixin(ModelRegistryHelper, ABC):
         """
         Direct OpenAI chat completion API call.
         """
+        if self.download_images:
+
+            async def _localize_image_url(m: OpenAIMessageParam) -> OpenAIMessageParam:
+                if isinstance(m.content, list):
+                    for c in m.content:
+                        if c.type == "image_url" and c.image_url and c.image_url.url and "http" in c.image_url.url:
+                            localize_result = await localize_image_content(c.image_url.url)
+                            if localize_result is None:
+                                raise ValueError(
+                                    f"Failed to localize image content from {c.image_url.url[:42]}{'...' if len(c.image_url.url) > 42 else ''}"
+                                )
+                            content, format = localize_result
+                            c.image_url.url = f"data:image/{format};base64,{base64.b64encode(content).decode('utf-8')}"
+                # else it's a string and we don't need to modify it
+                return m
+
+            messages = [await _localize_image_url(m) for m in messages]
+
         resp = await self.client.chat.completions.create(
             **await prepare_openai_completion_params(
                 model=await self._get_provider_model_id(model),

llama_stack/providers/utils/kvstore/config.py

@@ -28,7 +28,7 @@ class CommonConfig(BaseModel):
 
 
 class RedisKVStoreConfig(CommonConfig):
-    type: Literal[KVStoreType.redis.value] = KVStoreType.redis.value
+    type: Literal["redis"] = KVStoreType.redis.value
     host: str = "localhost"
     port: int = 6379
 
@@ -50,7 +50,7 @@ class RedisKVStoreConfig(CommonConfig):
 
 
 class SqliteKVStoreConfig(CommonConfig):
-    type: Literal[KVStoreType.sqlite.value] = KVStoreType.sqlite.value
+    type: Literal["sqlite"] = KVStoreType.sqlite.value
     db_path: str = Field(
         default=(RUNTIME_BASE_DIR / "kvstore.db").as_posix(),
         description="File path for the sqlite database",
@@ -69,7 +69,7 @@ class SqliteKVStoreConfig(CommonConfig):
 
 
 class PostgresKVStoreConfig(CommonConfig):
-    type: Literal[KVStoreType.postgres.value] = KVStoreType.postgres.value
+    type: Literal["postgres"] = KVStoreType.postgres.value
     host: str = "localhost"
     port: int = 5432
     db: str = "llamastack"
@@ -113,11 +113,11 @@ class PostgresKVStoreConfig(CommonConfig):
 
 
 class MongoDBKVStoreConfig(CommonConfig):
-    type: Literal[KVStoreType.mongodb.value] = KVStoreType.mongodb.value
+    type: Literal["mongodb"] = KVStoreType.mongodb.value
     host: str = "localhost"
     port: int = 27017
     db: str = "llamastack"
-    user: str = None
+    user: str | None = None
     password: str | None = None
     collection_name: str = "llamastack_kvstore"
 

llama_stack/providers/utils/kvstore/mongodb/mongodb.py

@@ -7,6 +7,7 @@
 from datetime import datetime
 
 from pymongo import AsyncMongoClient
+from pymongo.asynchronous.collection import AsyncCollection
 
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.kvstore import KVStore
@@ -19,8 +20,13 @@ log = get_logger(name=__name__, category="providers::utils")
 class MongoDBKVStoreImpl(KVStore):
     def __init__(self, config: MongoDBKVStoreConfig):
         self.config = config
-        self.conn = None
-        self.collection = None
+        self.conn: AsyncMongoClient | None = None
+
+    @property
+    def collection(self) -> AsyncCollection:
+        if self.conn is None:
+            raise RuntimeError("MongoDB connection is not initialized")
+        return self.conn[self.config.db][self.config.collection_name]
 
     async def initialize(self) -> None:
         try:
@@ -32,7 +38,6 @@ class MongoDBKVStoreImpl(KVStore):
             }
             conn_creds = {k: v for k, v in conn_creds.items() if v is not None}
             self.conn = AsyncMongoClient(**conn_creds)
-            self.collection = self.conn[self.config.db][self.config.collection_name]
         except Exception as e:
             log.exception("Could not connect to MongoDB database server")
             raise RuntimeError("Could not connect to MongoDB database server") from e

llama_stack/providers/utils/kvstore/sqlite/sqlite.py

@@ -9,9 +9,13 @@ from datetime import datetime
 
 import aiosqlite
 
+from llama_stack.log import get_logger
+
 from ..api import KVStore
 from ..config import SqliteKVStoreConfig
 
+logger = get_logger(name=__name__, category="providers::utils")
+
 
 class SqliteKVStoreImpl(KVStore):
     def __init__(self, config: SqliteKVStoreConfig):
@@ -50,6 +54,9 @@ class SqliteKVStoreImpl(KVStore):
                 if row is None:
                     return None
                 value, expiration = row
+                if not isinstance(value, str):
+                    logger.warning(f"Expected string value for key {key}, got {type(value)}, returning None")
+                    return None
                 return value
 
     async def delete(self, key: str) -> None: