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

llama_stack/providers/remote/inference/databricks/databricks.py

@@ -4,13 +4,13 @@
 # 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 Iterable
+from collections.abc import AsyncIterator, Iterable
 
 from databricks.sdk import WorkspaceClient
 
-from llama_stack.apis.inference import OpenAICompletion, OpenAICompletionRequestWithExtraBody
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import OpenAICompletion, OpenAICompletionRequestWithExtraBody
 
 from .config import DatabricksImplConfig
 
@@ -20,6 +20,8 @@ logger = get_logger(name=__name__, category="inference::databricks")
 class DatabricksInferenceAdapter(OpenAIMixin):
     config: DatabricksImplConfig
 
+    provider_data_api_key_field: str = "databricks_api_token"
+
     # source: https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models
     embedding_model_metadata: dict[str, dict[str, int]] = {
         "databricks-gte-large-en": {"embedding_dimension": 1024, "context_length": 8192},
@@ -27,18 +29,26 @@ class DatabricksInferenceAdapter(OpenAIMixin):
     }
 
     def get_base_url(self) -> str:
-        return f"{self.config.url}/serving-endpoints"
+        return str(self.config.base_url)
 
     async def list_provider_model_ids(self) -> Iterable[str]:
+        # Filter out None values from endpoint names
+        api_token = self._get_api_key_from_config_or_provider_data()
+        # WorkspaceClient expects base host without /serving-endpoints suffix
+        base_url_str = str(self.config.base_url)
+        if base_url_str.endswith("/serving-endpoints"):
+            host = base_url_str[:-18]  # Remove '/serving-endpoints'
+        else:
+            host = base_url_str
         return [
-            endpoint.name
+            endpoint.name  # type: ignore[misc]
             for endpoint in WorkspaceClient(
-                host=self.config.url, token=self.get_api_key()
+                host=host, token=api_token
             ).serving_endpoints.list()  # TODO: this is not async
         ]
 
     async def openai_completion(
         self,
         params: OpenAICompletionRequestWithExtraBody,
-    ) -> OpenAICompletion:
+    ) -> OpenAICompletion | AsyncIterator[OpenAICompletion]:
         raise NotImplementedError()

llama_stack/providers/remote/inference/fireworks/config.py

@@ -6,22 +6,22 @@
 
 from typing import Any
 
-from pydantic import Field
+from pydantic import Field, HttpUrl
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 @json_schema_type
 class FireworksImplConfig(RemoteInferenceProviderConfig):
-    url: str = Field(
-        default="https://api.fireworks.ai/inference/v1",
+    base_url: HttpUrl | None = Field(
+        default=HttpUrl("https://api.fireworks.ai/inference/v1"),
         description="The URL for the Fireworks server",
     )
 
     @classmethod
     def sample_run_config(cls, api_key: str = "${env.FIREWORKS_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
-            "url": "https://api.fireworks.ai/inference/v1",
+            "base_url": "https://api.fireworks.ai/inference/v1",
             "api_key": api_key,
         }

llama_stack/providers/remote/inference/fireworks/fireworks.py

@@ -24,4 +24,4 @@ class FireworksInferenceAdapter(OpenAIMixin):
     provider_data_api_key_field: str = "fireworks_api_key"
 
     def get_base_url(self) -> str:
-        return "https://api.fireworks.ai/inference/v1"
+        return str(self.config.base_url)

llama_stack/providers/remote/inference/gemini/config.py

@@ -9,7 +9,7 @@ from typing import Any
 from pydantic import BaseModel, Field
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 class GeminiProviderDataValidator(BaseModel):

llama_stack/providers/remote/inference/gemini/gemini.py

@@ -4,15 +4,15 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from openai import NOT_GIVEN
+from typing import Any
 
-from llama_stack.apis.inference import (
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import (
     OpenAIEmbeddingData,
     OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
 )
-from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import GeminiConfig
 
@@ -37,21 +37,20 @@ class GeminiInferenceAdapter(OpenAIMixin):
         Override embeddings method to handle Gemini's missing usage statistics.
         Gemini's embedding API doesn't return usage information, so we provide default values.
         """
-        # Prepare request parameters
-        request_params = {
+        # Build request params conditionally to avoid NotGiven/Omit type mismatch
+        request_params: dict[str, Any] = {
             "model": await self._get_provider_model_id(params.model),
             "input": params.input,
-            "encoding_format": params.encoding_format if params.encoding_format is not None else NOT_GIVEN,
-            "dimensions": params.dimensions if params.dimensions is not None else NOT_GIVEN,
-            "user": params.user if params.user is not None else NOT_GIVEN,
         }
+        if params.encoding_format is not None:
+            request_params["encoding_format"] = params.encoding_format
+        if params.dimensions is not None:
+            request_params["dimensions"] = params.dimensions
+        if params.user is not None:
+            request_params["user"] = params.user
+        if params.model_extra:
+            request_params["extra_body"] = params.model_extra
 
-        # Add extra_body if present
-        extra_body = params.model_extra
-        if extra_body:
-            request_params["extra_body"] = extra_body
-
-        # Call OpenAI embeddings API with properly typed parameters
         response = await self.client.embeddings.create(**request_params)
 
         data = []

llama_stack/providers/remote/inference/groq/config.py

@@ -6,10 +6,10 @@
 
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, HttpUrl
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 class GroqProviderDataValidator(BaseModel):
@@ -21,14 +21,14 @@ class GroqProviderDataValidator(BaseModel):
 
 @json_schema_type
 class GroqConfig(RemoteInferenceProviderConfig):
-    url: str = Field(
-        default="https://api.groq.com",
+    base_url: HttpUrl | None = Field(
+        default=HttpUrl("https://api.groq.com/openai/v1"),
         description="The URL for the Groq AI server",
     )
 
     @classmethod
     def sample_run_config(cls, api_key: str = "${env.GROQ_API_KEY:=}", **kwargs) -> dict[str, Any]:
         return {
-            "url": "https://api.groq.com",
+            "base_url": "https://api.groq.com/openai/v1",
             "api_key": api_key,
         }

llama_stack/providers/remote/inference/groq/groq.py

@@ -15,4 +15,4 @@ class GroqInferenceAdapter(OpenAIMixin):
     provider_data_api_key_field: str = "groq_api_key"
 
     def get_base_url(self) -> str:
-        return f"{self.config.url}/openai/v1"
+        return str(self.config.base_url)

llama_stack/providers/remote/inference/llama_openai_compat/config.py

@@ -6,10 +6,10 @@
 
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, HttpUrl
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
 
 
 class LlamaProviderDataValidator(BaseModel):
@@ -21,14 +21,14 @@ class LlamaProviderDataValidator(BaseModel):
 
 @json_schema_type
 class LlamaCompatConfig(RemoteInferenceProviderConfig):
-    openai_compat_api_base: str = Field(
-        default="https://api.llama.com/compat/v1/",
+    base_url: HttpUrl | None = Field(
+        default=HttpUrl("https://api.llama.com/compat/v1/"),
         description="The URL for the Llama API server",
     )
 
     @classmethod
     def sample_run_config(cls, api_key: str = "${env.LLAMA_API_KEY}", **kwargs) -> dict[str, Any]:
         return {
-            "openai_compat_api_base": "https://api.llama.com/compat/v1/",
+            "base_url": "https://api.llama.com/compat/v1/",
             "api_key": api_key,
         }

llama_stack/providers/remote/inference/llama_openai_compat/llama.py

@@ -4,15 +4,17 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack.apis.inference.inference import (
+from collections.abc import AsyncIterator
+
+from llama_stack.log import get_logger
+from llama_stack.providers.remote.inference.llama_openai_compat.config import LlamaCompatConfig
+from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import (
     OpenAICompletion,
     OpenAICompletionRequestWithExtraBody,
     OpenAIEmbeddingsRequestWithExtraBody,
     OpenAIEmbeddingsResponse,
 )
-from llama_stack.log import get_logger
-from llama_stack.providers.remote.inference.llama_openai_compat.config import LlamaCompatConfig
-from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 logger = get_logger(name=__name__, category="inference::llama_openai_compat")
 
@@ -31,12 +33,12 @@ class LlamaCompatInferenceAdapter(OpenAIMixin):
 
         :return: The Llama API base URL
         """
-        return self.config.openai_compat_api_base
+        return str(self.config.base_url)
 
     async def openai_completion(
         self,
         params: OpenAICompletionRequestWithExtraBody,
-    ) -> OpenAICompletion:
+    ) -> OpenAICompletion | AsyncIterator[OpenAICompletion]:
         raise NotImplementedError()
 
     async def openai_embeddings(

llama_stack/providers/remote/inference/nvidia/__init__.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack.apis.inference import Inference
+from llama_stack_api import Inference
 
 from .config import NVIDIAConfig
 

llama_stack/providers/remote/inference/nvidia/config.py

@@ -7,10 +7,17 @@
 import os
 from typing import Any
 
-from pydantic import Field
+from pydantic import BaseModel, Field, HttpUrl
 
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
-from llama_stack.schema_utils import json_schema_type
+from llama_stack_api import json_schema_type
+
+
+class NVIDIAProviderDataValidator(BaseModel):
+    nvidia_api_key: str | None = Field(
+        default=None,
+        description="API key for NVIDIA NIM models",
+    )
 
 
 @json_schema_type
@@ -21,6 +28,7 @@ class NVIDIAConfig(RemoteInferenceProviderConfig):
     Attributes:
         url (str): A base url for accessing the NVIDIA NIM, e.g. http://localhost:8000
         api_key (str): The access key for the hosted NIM endpoints
+        rerank_model_to_url (dict[str, str]): Mapping of rerank model identifiers to their API endpoints
 
     There are two ways to access NVIDIA NIMs -
      0. Hosted: Preview APIs hosted at https://integrate.api.nvidia.com
@@ -36,29 +44,31 @@ class NVIDIAConfig(RemoteInferenceProviderConfig):
     URL of your running NVIDIA NIM and do not need to set the api_key.
     """
 
-    url: str = Field(
-        default_factory=lambda: os.getenv("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com"),
+    base_url: HttpUrl | None = Field(
+        default_factory=lambda: os.getenv("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com/v1"),
         description="A base url for accessing the NVIDIA NIM",
     )
     timeout: int = Field(
         default=60,
         description="Timeout for the HTTP requests",
     )
-    append_api_version: bool = Field(
-        default_factory=lambda: os.getenv("NVIDIA_APPEND_API_VERSION", "True").lower() != "false",
-        description="When set to false, the API version will not be appended to the base_url. By default, it is true.",
+    rerank_model_to_url: dict[str, str] = Field(
+        default_factory=lambda: {
+            "nv-rerank-qa-mistral-4b:1": "https://ai.api.nvidia.com/v1/retrieval/nvidia/reranking",
+            "nvidia/nv-rerankqa-mistral-4b-v3": "https://ai.api.nvidia.com/v1/retrieval/nvidia/nv-rerankqa-mistral-4b-v3/reranking",
+            "nvidia/llama-3.2-nv-rerankqa-1b-v2": "https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v2/reranking",
+        },
+        description="Mapping of rerank model identifiers to their API endpoints. ",
     )
 
     @classmethod
     def sample_run_config(
         cls,
-        url: str = "${env.NVIDIA_BASE_URL:=https://integrate.api.nvidia.com}",
+        base_url: HttpUrl | None = "${env.NVIDIA_BASE_URL:=https://integrate.api.nvidia.com/v1}",
         api_key: str = "${env.NVIDIA_API_KEY:=}",
-        append_api_version: bool = "${env.NVIDIA_APPEND_API_VERSION:=True}",
         **kwargs,
     ) -> dict[str, Any]:
         return {
-            "url": url,
+            "base_url": base_url,
             "api_key": api_key,
-            "append_api_version": append_api_version,
         }

llama_stack/providers/remote/inference/nvidia/nvidia.py

@@ -5,8 +5,20 @@
 # the root directory of this source tree.
 
 
+from collections.abc import Iterable
+
+import aiohttp
+
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
+from llama_stack_api import (
+    Model,
+    ModelType,
+    OpenAIChatCompletionContentPartImageParam,
+    OpenAIChatCompletionContentPartTextParam,
+    RerankData,
+    RerankResponse,
+)
 
 from . import NVIDIAConfig
 from .utils import _is_nvidia_hosted
@@ -17,6 +29,8 @@ logger = get_logger(name=__name__, category="inference::nvidia")
 class NVIDIAInferenceAdapter(OpenAIMixin):
     config: NVIDIAConfig
 
+    provider_data_api_key_field: str = "nvidia_api_key"
+
     """
     NVIDIA Inference Adapter for Llama Stack.
     """
@@ -30,11 +44,11 @@ class NVIDIAInferenceAdapter(OpenAIMixin):
     }
 
     async def initialize(self) -> None:
-        logger.info(f"Initializing NVIDIAInferenceAdapter({self.config.url})...")
+        logger.info(f"Initializing NVIDIAInferenceAdapter({self.config.base_url})...")
 
         if _is_nvidia_hosted(self.config):
             if not self.config.auth_credential:
-                raise RuntimeError(
+                logger.error(
                     "API key is required for hosted NVIDIA NIM. Either provide an API key or use a self-hosted NIM."
                 )
 
@@ -58,4 +72,102 @@ class NVIDIAInferenceAdapter(OpenAIMixin):
 
         :return: The NVIDIA API base URL
         """
-        return f"{self.config.url}/v1" if self.config.append_api_version else self.config.url
+        return str(self.config.base_url)
+
+    async def list_provider_model_ids(self) -> Iterable[str]:
+        """
+        Return both dynamic model IDs and statically configured rerank model IDs.
+        """
+        dynamic_ids: Iterable[str] = []
+        try:
+            dynamic_ids = await super().list_provider_model_ids()
+        except Exception:
+            # If the dynamic listing fails, proceed with just configured rerank IDs
+            dynamic_ids = []
+
+        configured_rerank_ids = list(self.config.rerank_model_to_url.keys())
+        return list(dict.fromkeys(list(dynamic_ids) + configured_rerank_ids))  # remove duplicates
+
+    def construct_model_from_identifier(self, identifier: str) -> Model:
+        """
+        Classify rerank models from config; otherwise use the base behavior.
+        """
+        if identifier in self.config.rerank_model_to_url:
+            return Model(
+                provider_id=self.__provider_id__,  # type: ignore[attr-defined]
+                provider_resource_id=identifier,
+                identifier=identifier,
+                model_type=ModelType.rerank,
+            )
+        return super().construct_model_from_identifier(identifier)
+
+    async def rerank(
+        self,
+        model: str,
+        query: str | OpenAIChatCompletionContentPartTextParam | OpenAIChatCompletionContentPartImageParam,
+        items: list[str | OpenAIChatCompletionContentPartTextParam | OpenAIChatCompletionContentPartImageParam],
+        max_num_results: int | None = None,
+    ) -> RerankResponse:
+        provider_model_id = await self._get_provider_model_id(model)
+
+        ranking_url = self.get_base_url()
+
+        if _is_nvidia_hosted(self.config) and provider_model_id in self.config.rerank_model_to_url:
+            ranking_url = self.config.rerank_model_to_url[provider_model_id]
+
+        logger.debug(f"Using rerank endpoint: {ranking_url} for model: {provider_model_id}")
+
+        # Convert query to text format
+        if isinstance(query, str):
+            query_text = query
+        elif isinstance(query, OpenAIChatCompletionContentPartTextParam):
+            query_text = query.text
+        else:
+            raise ValueError("Query must be a string or text content part")
+
+        # Convert items to text format
+        passages = []
+        for item in items:
+            if isinstance(item, str):
+                passages.append({"text": item})
+            elif isinstance(item, OpenAIChatCompletionContentPartTextParam):
+                passages.append({"text": item.text})
+            else:
+                raise ValueError("Items must be strings or text content parts")
+
+        payload = {
+            "model": provider_model_id,
+            "query": {"text": query_text},
+            "passages": passages,
+        }
+
+        headers = {
+            "Authorization": f"Bearer {self.get_api_key()}",
+            "Content-Type": "application/json",
+        }
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(ranking_url, headers=headers, json=payload) as response:
+                    if response.status != 200:
+                        response_text = await response.text()
+                        raise ConnectionError(
+                            f"NVIDIA rerank API request failed with status {response.status}: {response_text}"
+                        )
+
+                    result = await response.json()
+                    rankings = result.get("rankings", [])
+
+                    # Convert to RerankData format
+                    rerank_data = []
+                    for ranking in rankings:
+                        rerank_data.append(RerankData(index=ranking["index"], relevance_score=ranking["logit"]))
+
+                    # Apply max_num_results limit
+                    if max_num_results is not None:
+                        rerank_data = rerank_data[:max_num_results]
+
+                    return RerankResponse(data=rerank_data)
+
+        except aiohttp.ClientError as e:
+            raise ConnectionError(f"Failed to connect to NVIDIA rerank API at {ranking_url}: {e}") from e

llama_stack/providers/remote/inference/nvidia/utils.py

@@ -8,4 +8,4 @@ from . import NVIDIAConfig
 
 
 def _is_nvidia_hosted(config: NVIDIAConfig) -> bool:
-    return "integrate.api.nvidia.com" in config.url
+    return "integrate.api.nvidia.com" in str(config.base_url)

llama_stack/providers/remote/inference/oci/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from llama_stack_api import InferenceProvider
+
+from .config import OCIConfig
+
+
+async def get_adapter_impl(config: OCIConfig, _deps) -> InferenceProvider:
+    from .oci import OCIInferenceAdapter
+
+    adapter = OCIInferenceAdapter(config=config)
+    await adapter.initialize()
+    return adapter

llama_stack/providers/remote/inference/oci/auth.py

@@ -0,0 +1,79 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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 Generator, Mapping
+from typing import Any, override
+
+import httpx
+import oci
+import requests
+from oci.config import DEFAULT_LOCATION, DEFAULT_PROFILE
+
+OciAuthSigner = type[oci.signer.AbstractBaseSigner]
+
+
+class HttpxOciAuth(httpx.Auth):
+    """
+    Custom HTTPX authentication class that implements OCI request signing.
+
+    This class handles the authentication flow for HTTPX requests by signing them
+    using the OCI Signer, which adds the necessary authentication headers for
+    OCI API calls.
+
+    Attributes:
+        signer (oci.signer.Signer): The OCI signer instance used for request signing
+    """
+
+    def __init__(self, signer: OciAuthSigner):
+        self.signer = signer
+
+    @override
+    def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
+        # Read the request content to handle streaming requests properly
+        try:
+            content = request.content
+        except httpx.RequestNotRead:
+            # For streaming requests, we need to read the content first
+            content = request.read()
+
+        req = requests.Request(
+            method=request.method,
+            url=str(request.url),
+            headers=dict(request.headers),
+            data=content,
+        )
+        prepared_request = req.prepare()
+
+        # Sign the request using the OCI Signer
+        self.signer.do_request_sign(prepared_request)  # type: ignore
+
+        # Update the original HTTPX request with the signed headers
+        request.headers.update(prepared_request.headers)
+
+        yield request
+
+
+class OciInstancePrincipalAuth(HttpxOciAuth):
+    def __init__(self, **kwargs: Mapping[str, Any]):
+        self.signer = oci.auth.signers.InstancePrincipalsSecurityTokenSigner(**kwargs)
+
+
+class OciUserPrincipalAuth(HttpxOciAuth):
+    def __init__(self, config_file: str = DEFAULT_LOCATION, profile_name: str = DEFAULT_PROFILE):
+        config = oci.config.from_file(config_file, profile_name)
+        oci.config.validate_config(config)  # type: ignore
+        key_content = ""
+        with open(config["key_file"]) as f:
+            key_content = f.read()
+
+        self.signer = oci.signer.Signer(
+            tenancy=config["tenancy"],
+            user=config["user"],
+            fingerprint=config["fingerprint"],
+            private_key_file_location=config.get("key_file"),
+            pass_phrase="none",  # type: ignore
+            private_key_content=key_content,
+        )

llama_stack/providers/remote/inference/oci/config.py

@@ -0,0 +1,75 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
+from llama_stack_api import json_schema_type
+
+
+class OCIProviderDataValidator(BaseModel):
+    oci_auth_type: str = Field(
+        description="OCI authentication type (must be one of: instance_principal, config_file)",
+    )
+    oci_region: str = Field(
+        description="OCI region (e.g., us-ashburn-1)",
+    )
+    oci_compartment_id: str = Field(
+        description="OCI compartment ID for the Generative AI service",
+    )
+    oci_config_file_path: str | None = Field(
+        default="~/.oci/config",
+        description="OCI config file path (required if oci_auth_type is config_file)",
+    )
+    oci_config_profile: str | None = Field(
+        default="DEFAULT",
+        description="OCI config profile (required if oci_auth_type is config_file)",
+    )
+
+
+@json_schema_type
+class OCIConfig(RemoteInferenceProviderConfig):
+    oci_auth_type: str = Field(
+        description="OCI authentication type (must be one of: instance_principal, config_file)",
+        default_factory=lambda: os.getenv("OCI_AUTH_TYPE", "instance_principal"),
+    )
+    oci_region: str = Field(
+        default_factory=lambda: os.getenv("OCI_REGION", "us-ashburn-1"),
+        description="OCI region (e.g., us-ashburn-1)",
+    )
+    oci_compartment_id: str = Field(
+        default_factory=lambda: os.getenv("OCI_COMPARTMENT_OCID", ""),
+        description="OCI compartment ID for the Generative AI service",
+    )
+    oci_config_file_path: str = Field(
+        default_factory=lambda: os.getenv("OCI_CONFIG_FILE_PATH", "~/.oci/config"),
+        description="OCI config file path (required if oci_auth_type is config_file)",
+    )
+    oci_config_profile: str = Field(
+        default_factory=lambda: os.getenv("OCI_CLI_PROFILE", "DEFAULT"),
+        description="OCI config profile (required if oci_auth_type is config_file)",
+    )
+
+    @classmethod
+    def sample_run_config(
+        cls,
+        oci_auth_type: str = "${env.OCI_AUTH_TYPE:=instance_principal}",
+        oci_config_file_path: str = "${env.OCI_CONFIG_FILE_PATH:=~/.oci/config}",
+        oci_config_profile: str = "${env.OCI_CLI_PROFILE:=DEFAULT}",
+        oci_region: str = "${env.OCI_REGION:=us-ashburn-1}",
+        oci_compartment_id: str = "${env.OCI_COMPARTMENT_OCID:=}",
+        **kwargs,
+    ) -> dict[str, Any]:
+        return {
+            "oci_auth_type": oci_auth_type,
+            "oci_config_file_path": oci_config_file_path,
+            "oci_config_profile": oci_config_profile,
+            "oci_region": oci_region,
+            "oci_compartment_id": oci_compartment_id,
+        }