Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
156 changes: 156 additions & 0 deletions unstract/sdk1/src/unstract/sdk1/adapters/base1.py
Original file line number Diff line number Diff line change
Expand Up @@ -353,6 +353,9 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
_OPENAI_PROVIDER_PREFIX = "openai/"
_CUSTOM_OPENAI_PROVIDER_PREFIX = "custom_openai/"
_OPENAI_REASONING_MODEL_PATTERN = re.compile(r"^(o1|o3|o4|gpt-5)(?:[-/]|$)")
# Keyless gateways still need a non-empty key; the OpenAI SDK rejects a
# null/blank one before any request reaches the endpoint.
_NO_AUTH_API_KEY = "no-auth"


def _is_openai_reasoning_model(model: str) -> bool:
Expand Down Expand Up @@ -494,6 +497,86 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
return f"{_CUSTOM_OPENAI_PROVIDER_PREFIX}{model}"


# Shared validation for branded adapters that reuse the OpenAI-compatible wire
# protocol with a fixed default endpoint.
def _validate_branded_openai_compatible(
adapter_metadata: dict[str, "Any"], default_api_base: str
) -> dict[str, "Any"]:
# Endpoint stays overridable so a provider URL change needs no release.
api_base = adapter_metadata.get("api_base")
if not (isinstance(api_base, str) and api_base.strip()):
api_base = default_api_base
adapter_metadata = {**adapter_metadata, "api_base": api_base}
return OpenAICompatibleLLMParameters.validate(adapter_metadata)


_NVIDIA_BUILD_API_BASE = "https://integrate.api.nvidia.com/v1"
_OPENROUTER_API_BASE = "https://openrouter.ai/api/v1"
_OPENROUTER_PROVIDER_PREFIX = "openrouter/"


class NvidiaBuildLLMParameters(OpenAICompatibleLLMParameters):
"""OpenAI-compatible adapter for NVIDIA's hosted models (build.nvidia.com)."""

# Required str so a directly-constructed instance stays valid.
api_base: str = _NVIDIA_BUILD_API_BASE

@staticmethod
def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]:
return _validate_branded_openai_compatible(
adapter_metadata, _NVIDIA_BUILD_API_BASE
)


class OpenRouterLLMParameters(BaseChatCompletionParameters):
"""Adapter for OpenRouter (openrouter.ai).

Routed through LiteLLM's native `openrouter/` provider so per-token costs
resolve and reasoning params map without provider-specific workarounds.
"""

api_key: str
api_base: str = _OPENROUTER_API_BASE
reasoning_effort: str | None = None

@staticmethod
def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]:
adapter_metadata = dict(adapter_metadata)
api_base = adapter_metadata.get("api_base")
if not (isinstance(api_base, str) and api_base.strip()):
adapter_metadata["api_base"] = _OPENROUTER_API_BASE
adapter_metadata["model"] = OpenRouterLLMParameters.validate_model(
adapter_metadata
)
# Reasoning models reject a non-default temperature; drop it and send
# reasoning_effort only when reasoning is enabled. On a re-validation
# pass `enable_reasoning` is absent (it isn't serialized), so recover the
# state from a surviving reasoning_effort to keep reload idempotent —
# unless the user explicitly opted out with `enable_reasoning: false`.
enable_reasoning = adapter_metadata.get("enable_reasoning", False)
if (
not enable_reasoning
and "enable_reasoning" not in adapter_metadata
and adapter_metadata.get("reasoning_effort") is not None
):
enable_reasoning = True
adapter_metadata.pop("enable_reasoning", None)
if enable_reasoning:
adapter_metadata["temperature"] = None
else:
adapter_metadata.pop("reasoning_effort", None)
return OpenRouterLLMParameters(**adapter_metadata).model_dump()
Comment thread
greptile-apps[bot] marked this conversation as resolved.

@staticmethod
def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
model = str(adapter_metadata.get("model", "")).strip()
if not model:
raise ValueError("model is required for the OpenRouter adapter.")
if model.startswith(_OPENROUTER_PROVIDER_PREFIX):
return model
return f"{_OPENROUTER_PROVIDER_PREFIX}{model}"


class AzureOpenAILLMParameters(BaseChatCompletionParameters):
"""See https://docs.litellm.ai/docs/providers/azure/#completion---using-azure_ad_token-api_base-api_version."""

Expand Down Expand Up @@ -1296,6 +1379,8 @@ class OpenAIEmbeddingParameters(BaseEmbeddingParameters):
api_base: str | None = None
embed_batch_size: int | None = 10
dimensions: int | None = None # For text-embedding-3-* models
# Strict endpoints reject the null LiteLLM sends when this is unset.
encoding_format: str | None = None

@staticmethod
def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]:
Expand All @@ -1311,6 +1396,77 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
return model


# custom_openai has no embedding support, so these route through LiteLLM's
# native nvidia_nim provider; its default api_base is pinned for the schema.
_NVIDIA_NIM_PROVIDER_PREFIX = "nvidia_nim/"


class NvidiaBuildEmbeddingParameters(OpenAIEmbeddingParameters):
"""OpenAI-compatible embeddings via NVIDIA's hosted endpoint (build.nvidia.com)."""

# Overridable default endpoint.
api_base: str = _NVIDIA_BUILD_API_BASE

@staticmethod
def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]:
adapter_metadata = dict(adapter_metadata)
# Endpoint stays overridable so a provider URL change needs no release.
api_base = adapter_metadata.get("api_base")
if not (isinstance(api_base, str) and api_base.strip()):
adapter_metadata["api_base"] = _NVIDIA_BUILD_API_BASE
# Strict endpoints reject the null LiteLLM sends; pin a real value.
adapter_metadata.setdefault("encoding_format", "float")
adapter_metadata["model"] = NvidiaBuildEmbeddingParameters.validate_model(
adapter_metadata
)
return OpenAIEmbeddingParameters(**adapter_metadata).model_dump()

@staticmethod
def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
model = str(adapter_metadata.get("model", "")).strip()
if not model:
raise ValueError("model is required for the NVIDIA Build embedding adapter.")
if model.startswith(_NVIDIA_NIM_PROVIDER_PREFIX):
return model
return f"{_NVIDIA_NIM_PROVIDER_PREFIX}{model}"


class OpenAICompatibleEmbeddingParameters(OpenAIEmbeddingParameters):
"""Embeddings for any OpenAI-compatible server (vLLM, self-hosted, etc.).

Routes through the `openai/` provider with a user-supplied `api_base`;
cost stays unresolved since the endpoint is arbitrary.
"""

# Some gateways are keyless; the endpoint is always required.
api_key: str | None = None
api_base: str

@staticmethod
def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]:
adapter_metadata = dict(adapter_metadata)
api_key = adapter_metadata.get("api_key")
if not (isinstance(api_key, str) and api_key.strip()):
adapter_metadata["api_key"] = _NO_AUTH_API_KEY
# Strict endpoints reject the null LiteLLM sends; pin a real value.
adapter_metadata.setdefault("encoding_format", "float")
adapter_metadata["model"] = OpenAICompatibleEmbeddingParameters.validate_model(
adapter_metadata
)
return OpenAICompatibleEmbeddingParameters(**adapter_metadata).model_dump()
Comment thread
coderabbitai[bot] marked this conversation as resolved.

@staticmethod
def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
model = str(adapter_metadata.get("model", "")).strip()
if not model:
raise ValueError(
"model is required for the OpenAI Compatible embedding adapter."
)
if model.startswith(_OPENAI_PROVIDER_PREFIX):
return model
return f"{_OPENAI_PROVIDER_PREFIX}{model}"


class AzureOpenAIEmbeddingParameters(BaseEmbeddingParameters):
"""See https://docs.litellm.ai/docs/providers/azure."""

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,12 @@
from unstract.sdk1.adapters.embedding1.azure_openai import AzureOpenAIEmbeddingAdapter
from unstract.sdk1.adapters.embedding1.bedrock import AWSBedrockEmbeddingAdapter
from unstract.sdk1.adapters.embedding1.gemini import GeminiEmbeddingAdapter
from unstract.sdk1.adapters.embedding1.nvidia_build import NvidiaBuildEmbeddingAdapter
from unstract.sdk1.adapters.embedding1.ollama import OllamaEmbeddingAdapter
from unstract.sdk1.adapters.embedding1.openai import OpenAIEmbeddingAdapter
from unstract.sdk1.adapters.embedding1.openai_compatible import (
OpenAICompatibleEmbeddingAdapter,
)
from unstract.sdk1.adapters.embedding1.vertexai import VertexAIEmbeddingAdapter
from unstract.sdk1.adapters.enums import AdapterTypes

Expand All @@ -18,7 +22,9 @@
"AzureOpenAIEmbeddingAdapter",
"AWSBedrockEmbeddingAdapter",
"GeminiEmbeddingAdapter",
"NvidiaBuildEmbeddingAdapter",
"OpenAIEmbeddingAdapter",
"OpenAICompatibleEmbeddingAdapter",
"VertexAIEmbeddingAdapter",
"OllamaEmbeddingAdapter",
]
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
from typing import Any

from unstract.sdk1.adapters.base1 import BaseAdapter, NvidiaBuildEmbeddingParameters
from unstract.sdk1.adapters.enums import AdapterTypes

DESCRIPTION = (
"Adapter for NVIDIA's OpenAI-compatible embedding models (build.nvidia.com). "
"Supply a model name and your NVIDIA API key; the endpoint is preconfigured."
)


class NvidiaBuildEmbeddingAdapter(NvidiaBuildEmbeddingParameters, BaseAdapter):
@staticmethod
def get_id() -> str:
return "nvidiabuild|2afcdf59-5323-4086-97fb-d9a0432e7795"

@staticmethod
def get_metadata() -> dict[str, Any]:
return {
"name": "NVIDIA Build",
"version": "1.0.0",
"adapter": NvidiaBuildEmbeddingAdapter,
"description": DESCRIPTION,
"is_active": True,
}

@staticmethod
def get_name() -> str:
return "NVIDIA Build"

@staticmethod
def get_description() -> str:
return DESCRIPTION

@staticmethod
def get_provider() -> str:
return "nvidia_build"

@staticmethod
def get_icon() -> str:
return "/icons/adapter-icons/NvidiaBuild.png"

@staticmethod
def get_adapter_type() -> AdapterTypes:
return AdapterTypes.EMBEDDING
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
from typing import Any

from unstract.sdk1.adapters.base1 import (
BaseAdapter,
OpenAICompatibleEmbeddingParameters,
)
from unstract.sdk1.adapters.enums import AdapterTypes

DESCRIPTION = (
"Embedding adapter for servers that implement the OpenAI Embeddings API "
"(vLLM, self-hosted gateways, and third-party providers). "
"Use OpenAI for the official OpenAI service."
)


class OpenAICompatibleEmbeddingAdapter(OpenAICompatibleEmbeddingParameters, BaseAdapter):
@staticmethod
def get_id() -> str:
return "openaicompatible|65573de7-2ea5-4631-bb49-492717972455"

@staticmethod
def get_metadata() -> dict[str, Any]:
return {
"name": "OpenAI Compatible",
"version": "1.0.0",
"adapter": OpenAICompatibleEmbeddingAdapter,
"description": DESCRIPTION,
"is_active": True,
}

@staticmethod
def get_name() -> str:
return "OpenAI Compatible"

@staticmethod
def get_description() -> str:
return DESCRIPTION

@staticmethod
def get_provider() -> str:
return "custom_openai"

@staticmethod
def get_icon() -> str:
return "/icons/adapter-icons/OpenAICompatible.png"

@staticmethod
def get_adapter_type() -> AdapterTypes:
return AdapterTypes.EMBEDDING
Original file line number Diff line number Diff line change
Expand Up @@ -54,13 +54,6 @@
"title": "Dimensions",
"description": "Output embedding dimensions. Only supported by text-embedding-3-* models. Leave empty for default dimensions."
},
"embed_batch_size": {
"type": "number",
"minimum": 0,
"multipleOf": 1,
"title": "Embedding Batch Size",
"default": 5
},
"max_retries": {
"type": "number",
"minimum": 0,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
{
"title": "OpenAI Compatible Embedding",
"type": "object",
"required": [
"adapter_name",
"model",
"api_base"
],
"properties": {
"adapter_name": {
"type": "string",
"title": "Name",
"default": "",
"description": "Provide a unique name for this adapter instance. Example: compatible-emb-1"
},
"model": {
"type": "string",
"title": "Model",
"description": "The embedding model name expected by your OpenAI-compatible endpoint. Examples: text-embedding-3-small, BAAI/bge-m3, nomic-embed-text"
},
"api_key": {
"type": [
"string",
"null"
],
"title": "API Key",
"format": "password",
"description": "API key for your OpenAI-compatible endpoint. Leave empty if the endpoint does not require one."
},
"api_base": {
"type": "string",
"format": "uri",
"title": "API Base",
"description": "Base URL for the OpenAI-compatible embeddings endpoint. Examples: https://gateway.example.com/v1, https://llm.example.net/openai/v1"
},
"max_retries": {
"type": "number",
"minimum": 0,
"multipleOf": 1,
"title": "Max Retries",
"default": 3,
"description": "The maximum number of times to retry a request if it fails."
},
"timeout": {
"type": "number",
"minimum": 0,
"multipleOf": 1,
"title": "Timeout",
"default": 240,
"description": "Timeout in seconds"
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,6 @@
"default": "",
"format": "password"
},
"embed_batch_size": {
"type": "number",
"minimum": 1,
"multipleOf": 1,
"title": "Embed Batch Size",
"description": "Number of texts to embed in a single batch. Leave empty to use the system default."
},
"timeout": {
"type": "number",
"minimum": 0,
Expand Down
Loading
Loading