154 lines
5.4 KiB
Python
154 lines
5.4 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import oci
|
|
from litellm import CustomLLM
|
|
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
|
|
|
|
|
|
class OCIEmbeddingLLM(CustomLLM):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._client = None
|
|
self._client_signature = None
|
|
|
|
@staticmethod
|
|
def _pick_value(optional_params: dict, param_name: str, env_name: str, default: Optional[str] = None):
|
|
value = optional_params.get(param_name)
|
|
if value is None:
|
|
value = os.getenv(env_name, default)
|
|
if isinstance(value, str):
|
|
value = value.strip()
|
|
return value or default
|
|
|
|
def _build_oci_config(self, optional_params: dict) -> dict:
|
|
# Build OCI SDK config from LiteLLM params first, then fallback to env vars.
|
|
oci_config = {
|
|
"user": self._pick_value(optional_params, "oci_user", "OCI_USER"),
|
|
"fingerprint": self._pick_value(optional_params, "oci_fingerprint", "OCI_FINGERPRINT"),
|
|
"tenancy": self._pick_value(optional_params, "oci_tenancy", "OCI_TENANCY"),
|
|
"region": self._pick_value(optional_params, "oci_region", "OCI_REGION"),
|
|
"key_file": self._pick_value(optional_params, "oci_key_file", "OCI_KEY_FILE"),
|
|
"pass_phrase": self._pick_value(optional_params, "oci_pass_phrase", "OCI_PASS_PHRASE"),
|
|
}
|
|
|
|
if oci_config.get("key_file"):
|
|
oci_config["key_file"] = str(Path(oci_config["key_file"]).expanduser())
|
|
|
|
# Remove optional empty values so OCI SDK handles defaults cleanly.
|
|
if not oci_config.get("pass_phrase"):
|
|
oci_config.pop("pass_phrase", None)
|
|
|
|
missing = [
|
|
key
|
|
for key in ("user", "fingerprint", "tenancy", "region", "key_file")
|
|
if not oci_config.get(key)
|
|
]
|
|
if missing:
|
|
raise ValueError(
|
|
"Missing OCI config values for embedding provider: " + ", ".join(missing)
|
|
)
|
|
|
|
return oci_config
|
|
|
|
def _get_client(self, optional_params: dict):
|
|
"""Lazy init OCI client and recreate only when OCI auth/region changes."""
|
|
config = self._build_oci_config(optional_params)
|
|
signature = tuple(sorted(config.items()))
|
|
|
|
if self._client is None or self._client_signature != signature:
|
|
self._client = oci.generative_ai_inference.GenerativeAiInferenceClient(config)
|
|
self._client_signature = signature
|
|
|
|
return self._client
|
|
|
|
def _build_serving_mode(self, model_id: str, optional_params: dict):
|
|
serving_mode = self._pick_value(
|
|
optional_params,
|
|
"oci_serving_mode",
|
|
"OCI_SERVING_MODE",
|
|
default="ON_DEMAND",
|
|
).upper()
|
|
|
|
if serving_mode == "DEDICATED":
|
|
endpoint_id = self._pick_value(optional_params, "oci_endpoint_id", "OCI_ENDPOINT_ID")
|
|
if not endpoint_id:
|
|
raise ValueError(
|
|
"oci_endpoint_id/OCI_ENDPOINT_ID is required when oci_serving_mode is DEDICATED"
|
|
)
|
|
return oci.generative_ai_inference.models.DedicatedServingMode(endpoint_id=endpoint_id)
|
|
|
|
return oci.generative_ai_inference.models.OnDemandServingMode(model_id=model_id)
|
|
|
|
def embedding(
|
|
self,
|
|
model: str,
|
|
input: list,
|
|
model_response: EmbeddingResponse,
|
|
optional_params: dict,
|
|
encoding=None,
|
|
api_key: Optional[str] = None,
|
|
**kwargs,
|
|
) -> EmbeddingResponse:
|
|
del encoding, api_key, kwargs # not used by OCI embedding endpoint
|
|
|
|
optional_params = optional_params or {}
|
|
|
|
# model arrives as "oci-embed/cohere.embed-multilingual-v3.0"
|
|
model_id = model.split("/", 1)[-1]
|
|
|
|
compartment_id = self._pick_value(
|
|
optional_params, "oci_compartment_id", "OCI_COMPARTMENT_ID"
|
|
)
|
|
if not compartment_id:
|
|
raise ValueError("oci_compartment_id/OCI_COMPARTMENT_ID is required")
|
|
|
|
input_type = self._pick_value(
|
|
optional_params,
|
|
"input_type",
|
|
"OCI_INPUT_TYPE",
|
|
default="SEARCH_DOCUMENT",
|
|
).upper()
|
|
|
|
truncate = self._pick_value(
|
|
optional_params,
|
|
"truncate",
|
|
"OCI_EMBED_TRUNCATE",
|
|
default="NONE",
|
|
).upper()
|
|
|
|
inputs = input if isinstance(input, list) else [input]
|
|
|
|
detail = oci.generative_ai_inference.models.EmbedTextDetails(
|
|
inputs=inputs,
|
|
serving_mode=self._build_serving_mode(model_id, optional_params),
|
|
compartment_id=compartment_id,
|
|
input_type=input_type,
|
|
truncate=truncate,
|
|
)
|
|
|
|
resp = self._get_client(optional_params).embed_text(detail)
|
|
|
|
model_response.model = model_id
|
|
model_response.data = [
|
|
Embedding(object="embedding", index=i, embedding=emb)
|
|
for i, emb in enumerate(resp.data.embeddings)
|
|
]
|
|
model_response.usage = Usage(
|
|
prompt_tokens=resp.data.usage.prompt_tokens if resp.data.usage else 0,
|
|
completion_tokens=0,
|
|
total_tokens=resp.data.usage.total_tokens if resp.data.usage else 0,
|
|
)
|
|
|
|
return model_response
|
|
|
|
async def aembedding(self, *args, **kwargs) -> EmbeddingResponse:
|
|
# LiteLLM proxy is async; OCI SDK call is sync.
|
|
return self.embedding(*args, **kwargs)
|
|
|
|
|
|
oci_embedding_llm = OCIEmbeddingLLM()
|