mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-03 19:57:35 +00:00
c4cb6aa8d9
Prevent sensitive information from being logged in telemetry output by assigning SecretStr type to sensitive fields. API keys, password from KV store are now covered. All providers have been converted. Signed-off-by: Sébastien Han <seb@redhat.com>
37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
# 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.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
|
|
from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
|
|
|
|
from .config import GeminiConfig
|
|
|
|
|
|
class GeminiInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
|
|
embedding_model_metadata = {
|
|
"text-embedding-004": {"embedding_dimension": 768, "context_length": 2048},
|
|
}
|
|
|
|
def __init__(self, config: GeminiConfig) -> None:
|
|
LiteLLMOpenAIMixin.__init__(
|
|
self,
|
|
litellm_provider_name="gemini",
|
|
api_key_from_config=config.api_key.get_secret_value() if config.api_key else None,
|
|
provider_data_api_key_field="gemini_api_key",
|
|
)
|
|
self.config = config
|
|
|
|
get_api_key = LiteLLMOpenAIMixin.get_api_key
|
|
|
|
def get_base_url(self):
|
|
return "https://generativelanguage.googleapis.com/v1beta/openai/"
|
|
|
|
async def initialize(self) -> None:
|
|
await super().initialize()
|
|
|
|
async def shutdown(self) -> None:
|
|
await super().shutdown()
|