Initial commit

oci_embedding_handler.py

@@ -0,0 +1,153 @@
+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()