|
| 1 | +"""OpenRouter embeddings models.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +from typing import Any |
| 6 | + |
| 7 | +import openai |
| 8 | +from langchain_core.embeddings import Embeddings |
| 9 | +from langchain_core.utils import from_env, secret_from_env |
| 10 | +from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator |
| 11 | +from typing_extensions import Self |
| 12 | + |
| 13 | +DEFAULT_API_BASE = "https://openrouter.ai/api/v1" |
| 14 | + |
| 15 | + |
| 16 | +class OpenRouterEmbeddings(BaseModel, Embeddings): |
| 17 | + """OpenRouter embedding model integration. |
| 18 | +
|
| 19 | + Setup: |
| 20 | + Install `langchain-openrouter` and set environment variable |
| 21 | + `OPENROUTER_API_KEY`. |
| 22 | +
|
| 23 | + ```bash |
| 24 | + pip install -U langchain-openrouter |
| 25 | + export OPENROUTER_API_KEY="your-api-key" |
| 26 | + ``` |
| 27 | +
|
| 28 | + Key init args — completion params: |
| 29 | + model: |
| 30 | + Name of OpenRouter model to use, e.g. |
| 31 | + `"qwen/qwen3-embedding-8b"`. |
| 32 | + dimensions: |
| 33 | + The number of dimensions the resulting output embeddings |
| 34 | + should have. |
| 35 | +
|
| 36 | + Key init args — client params: |
| 37 | + timeout: |
| 38 | + Timeout for requests. |
| 39 | + max_retries: |
| 40 | + Max number of retries. |
| 41 | + api_key: |
| 42 | + OpenRouter API key. If not passed in will be read from env var |
| 43 | + `OPENROUTER_API_KEY`. |
| 44 | +
|
| 45 | + See full list of supported init args and their descriptions in the |
| 46 | + params section. |
| 47 | +
|
| 48 | + Instantiate: |
| 49 | + ```python |
| 50 | + from langchain_openrouter import OpenRouterEmbeddings |
| 51 | +
|
| 52 | + embeddings = OpenRouterEmbeddings( |
| 53 | + model="qwen/qwen3-embedding-8b", |
| 54 | + # api_key="...", |
| 55 | + # other params... |
| 56 | + ) |
| 57 | + ``` |
| 58 | +
|
| 59 | + Embed single text: |
| 60 | + ```python |
| 61 | + input_text = "The meaning of life is 42" |
| 62 | + vector = embeddings.embed_query(input_text) |
| 63 | + print(vector[:3]) |
| 64 | + ``` |
| 65 | +
|
| 66 | + ```python |
| 67 | + [-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915] |
| 68 | + ``` |
| 69 | +
|
| 70 | + Embed multiple texts: |
| 71 | + ```python |
| 72 | + input_texts = ["Document 1...", "Document 2..."] |
| 73 | + vectors = embeddings.embed_documents(input_texts) |
| 74 | + print(len(vectors)) |
| 75 | + # The first 3 coordinates for the first vector |
| 76 | + print(vectors[0][:3]) |
| 77 | + ``` |
| 78 | +
|
| 79 | + ```python |
| 80 | + 2 |
| 81 | + [-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915] |
| 82 | + ``` |
| 83 | +
|
| 84 | + Async: |
| 85 | + ```python |
| 86 | + vector = await embeddings.aembed_query(input_text) |
| 87 | + print(vector[:3]) |
| 88 | +
|
| 89 | + # multiple: |
| 90 | + # await embeddings.aembed_documents(input_texts) |
| 91 | + ``` |
| 92 | +
|
| 93 | + ```python |
| 94 | + [-0.009100092574954033, 0.005071679595857859, -0.0029193938244134188] |
| 95 | + ``` |
| 96 | + """ |
| 97 | + |
| 98 | + model: str = Field() |
| 99 | + """The name of the model""" |
| 100 | + api_key: SecretStr | None = Field( |
| 101 | + default_factory=secret_from_env("OPENROUTER_API_KEY", default=None), |
| 102 | + ) |
| 103 | + """OpenRouter API key""" |
| 104 | + api_base: str = Field( |
| 105 | + default_factory=from_env("OPENROUTER_API_BASE", default=DEFAULT_API_BASE), |
| 106 | + ) |
| 107 | + """OpenRouter API base URL""" |
| 108 | + dimensions: int | None = None |
| 109 | + """The number of dimensions the resulting output embeddings should have. |
| 110 | +
|
| 111 | + Only supported in text-embedding-3 and later models. |
| 112 | + """ |
| 113 | + |
| 114 | + # Client parameters |
| 115 | + request_timeout: float | None = Field(default=None, alias="timeout") |
| 116 | + """Timeout for requests.""" |
| 117 | + max_retries: int = 3 |
| 118 | + """Maximum number of retries.""" |
| 119 | + default_headers: dict[str, str] | None = None |
| 120 | + """Default headers to include in requests.""" |
| 121 | + default_query: dict[str, str] | None = None |
| 122 | + """Default query parameters to include in requests.""" |
| 123 | + http_client: Any = None |
| 124 | + """Optional httpx.Client to use for making requests.""" |
| 125 | + http_async_client: Any = None |
| 126 | + """Optional httpx.AsyncClient to use for making requests.""" |
| 127 | + |
| 128 | + model_config = ConfigDict(populate_by_name=True) |
| 129 | + |
| 130 | + @property |
| 131 | + def _llm_type(self) -> str: |
| 132 | + """Return type of embeddings.""" |
| 133 | + return "embeddings-openrouter" |
| 134 | + |
| 135 | + @property |
| 136 | + def lc_secrets(self) -> dict[str, str]: |
| 137 | + """A map of constructor argument names to secret ids.""" |
| 138 | + return {"api_key": "OPENROUTER_API_KEY"} |
| 139 | + |
| 140 | + @model_validator(mode="after") |
| 141 | + def validate_environment(self) -> Self: |
| 142 | + """Validate necessary environment vars and client params.""" |
| 143 | + if self.api_base == DEFAULT_API_BASE and not ( |
| 144 | + self.api_key and self.api_key.get_secret_value() |
| 145 | + ): |
| 146 | + msg = "If using default api base, OPENROUTER_API_KEY must be set." |
| 147 | + raise ValueError(msg) |
| 148 | + |
| 149 | + client_params: dict = { |
| 150 | + k: v |
| 151 | + for k, v in { |
| 152 | + "api_key": self.api_key.get_secret_value() if self.api_key else None, |
| 153 | + "base_url": self.api_base, |
| 154 | + "timeout": self.request_timeout, |
| 155 | + "max_retries": self.max_retries, |
| 156 | + "default_headers": self.default_headers, |
| 157 | + "default_query": self.default_query, |
| 158 | + }.items() |
| 159 | + if v is not None |
| 160 | + } |
| 161 | + |
| 162 | + if not (getattr(self, "_root_client", None) or None): |
| 163 | + sync_specific: dict = {"http_client": self.http_client} |
| 164 | + self._root_client = openai.OpenAI(**client_params, **sync_specific) |
| 165 | + self._client = self._root_client.embeddings |
| 166 | + if not (getattr(self, "_root_async_client", None) or None): |
| 167 | + async_specific: dict = {"http_client": self.http_async_client} |
| 168 | + self._root_async_client = openai.AsyncOpenAI( |
| 169 | + **client_params, |
| 170 | + **async_specific, |
| 171 | + ) |
| 172 | + self._async_client = self._root_async_client.embeddings |
| 173 | + return self |
| 174 | + |
| 175 | + def _get_request_payload(self, text_input: str | list[str]) -> dict: |
| 176 | + """Get the request payload for embeddings.""" |
| 177 | + text = ( |
| 178 | + text_input |
| 179 | + if isinstance(text_input, str) |
| 180 | + else text_input[0] |
| 181 | + if text_input |
| 182 | + else "" |
| 183 | + ) |
| 184 | + |
| 185 | + payload: dict[str, Any] = { |
| 186 | + "model": self.model, |
| 187 | + "input": text, |
| 188 | + } |
| 189 | + |
| 190 | + if self.dimensions is not None: |
| 191 | + payload["dimensions"] = self.dimensions |
| 192 | + |
| 193 | + return payload |
| 194 | + |
| 195 | + def _embed_with_payload(self, text: str) -> list[float]: |
| 196 | + """Embed text using request payload.""" |
| 197 | + payload = self._get_request_payload(text) |
| 198 | + response = self._client.create(**payload) |
| 199 | + return response.data[0].embedding |
| 200 | + |
| 201 | + async def _aembed_with_payload(self, text: str) -> list[float]: |
| 202 | + """Embed text using request payload asynchronously.""" |
| 203 | + payload = self._get_request_payload(text) |
| 204 | + response = await self._async_client.create(**payload) |
| 205 | + return response.data[0].embedding |
| 206 | + |
| 207 | + def embed_documents(self, texts: list[str]) -> list[list[float]]: |
| 208 | + """Embed search docs.""" |
| 209 | + if not hasattr(self, "_client") or not self._client: |
| 210 | + msg = "OpenRouter client is not initialized." |
| 211 | + raise ValueError(msg) |
| 212 | + |
| 213 | + # OpenRouter API doesn't support batch embedding, so we need to make |
| 214 | + # individual calls and combine results |
| 215 | + return [self._embed_with_payload(text) for text in texts] |
| 216 | + |
| 217 | + def embed_query(self, text: str) -> list[float]: |
| 218 | + """Embed query text.""" |
| 219 | + return self.embed_documents([text])[0] |
| 220 | + |
| 221 | + async def aembed_documents(self, texts: list[str]) -> list[list[float]]: |
| 222 | + """Embed search docs.""" |
| 223 | + if not hasattr(self, "_async_client") or not self._async_client: |
| 224 | + msg = "OpenRouter async client is not initialized." |
| 225 | + raise ValueError(msg) |
| 226 | + |
| 227 | + # OpenRouter API doesn't support batch embedding, so we need to make |
| 228 | + # individual calls and combine results |
| 229 | + return [await self._aembed_with_payload(text) for text in texts] |
| 230 | + |
| 231 | + async def aembed_query(self, text: str) -> list[float]: |
| 232 | + """Embed query text.""" |
| 233 | + return (await self.aembed_documents([text]))[0] |
0 commit comments