|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | +from pydantic import AliasChoices |
| 17 | +from pydantic import ConfigDict |
| 18 | +from pydantic import Field |
| 19 | + |
| 20 | +from nat.builder.builder import Builder |
| 21 | +from nat.builder.llm import LLMProviderInfo |
| 22 | +from nat.cli.register_workflow import register_llm_provider |
| 23 | +from nat.data_models.llm import LLMBaseConfig |
| 24 | +from nat.data_models.optimizable import OptimizableField |
| 25 | +from nat.data_models.optimizable import OptimizableMixin |
| 26 | +from nat.data_models.optimizable import SearchSpace |
| 27 | +from nat.data_models.retry_mixin import RetryMixin |
| 28 | +from nat.data_models.ssl_verification_mixin import SSLVerificationMixin |
| 29 | +from nat.data_models.thinking_mixin import ThinkingMixin |
| 30 | + |
| 31 | +class OCIModelConfig(LLMBaseConfig, RetryMixin, OptimizableMixin, ThinkingMixin, SSLVerificationMixin, name="oci"): |
| 32 | + """OCI Generative AI LLM provider.""" |
| 33 | + |
| 34 | + model_config = ConfigDict(protected_namespaces=(), extra="allow") |
| 35 | + |
| 36 | + endpoint: str | None = Field( |
| 37 | + default=None, |
| 38 | + validation_alias=AliasChoices("endpoint", "service_endpoint", "base_url"), |
| 39 | + description="OCI Generative AI service endpoint URL.", |
| 40 | + ) |
| 41 | + compartment_id: str | None = Field(default=None, description="OCI compartment OCID for Generative AI requests.") |
| 42 | + auth_type: str = Field(default="API_KEY", |
| 43 | + description="OCI SDK authentication type: API_KEY, SECURITY_TOKEN, INSTANCE_PRINCIPAL, " |
| 44 | + "or RESOURCE_PRINCIPAL.") |
| 45 | + auth_profile: str = Field(default="DEFAULT", |
| 46 | + description="OCI config profile to use for API_KEY or SECURITY_TOKEN auth.") |
| 47 | + auth_file_location: str = Field(default="~/.oci/config", |
| 48 | + description="Path to the OCI config file used for SDK authentication.") |
| 49 | + model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"), |
| 50 | + serialization_alias="model", |
| 51 | + description="The OCI Generative AI model ID.") |
| 52 | + provider: str | None = Field(default=None, |
| 53 | + description="Optional OCI provider override such as cohere, google, meta, or openai.") |
| 54 | + context_size: int | None = Field( |
| 55 | + default=1024, |
| 56 | + gt=0, |
| 57 | + description="The maximum number of tokens available for input.", |
| 58 | + ) |
| 59 | + seed: int | None = Field(default=None, description="Random seed to set for generation.") |
| 60 | + max_retries: int = Field(default=10, description="The max number of retries for the request.") |
| 61 | + max_tokens: int | None = Field(default=None, gt=0, description="Maximum number of output tokens.") |
| 62 | + temperature: float | None = OptimizableField( |
| 63 | + default=None, |
| 64 | + ge=0.0, |
| 65 | + description="Sampling temperature to control randomness in the output.", |
| 66 | + space=SearchSpace(high=0.9, low=0.1, step=0.2)) |
| 67 | + top_p: float | None = OptimizableField(default=None, |
| 68 | + ge=0.0, |
| 69 | + le=1.0, |
| 70 | + description="Top-p for distribution sampling.", |
| 71 | + space=SearchSpace(high=1.0, low=0.5, step=0.1)) |
| 72 | + request_timeout: float | None = Field(default=None, gt=0.0, description="HTTP request timeout in seconds.") |
| 73 | + |
| 74 | + |
| 75 | +@register_llm_provider(config_type=OCIModelConfig) |
| 76 | +async def oci_llm(config: OCIModelConfig, _builder: Builder): |
| 77 | + |
| 78 | + yield LLMProviderInfo(config=config, description="An OCI Generative AI model for use with an LLM client.") |
0 commit comments