From fc7d4bdf985a6906945454c59d4cd9d52906601d Mon Sep 17 00:00:00 2001
From: Lynn <lynn_inf@hotmail.com>
Date: Wed, 22 Jul 2026 13:17:52 +0800
Subject: [PATCH] Feat: openrouter embedding (#17213)

---
 rag/llm/embedding_model.py | 39 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py
index ed5418ad91..3d246e1b8e 100644
--- a/rag/llm/embedding_model.py
+++ b/rag/llm/embedding_model.py
@@ -1323,3 +1323,42 @@ class NewAPIEmbed(OpenAIEmbed):
             raise ValueError("url cannot be None")
         self.client = OpenAI(api_key=key, base_url=base_url)
         self.model_name = model_name.split("___")[0]
+
+
+class OpenRouterEmbed(Base):
+    _FACTORY_NAME = "OpenRouter"
+
+    def __init__(self, key, model_name, base_url="https://openrouter.ai/api/v1", **kwargs):
+        if not base_url:
+            base_url = "https://openrouter.ai/api/v1"
+        self.base_url = ensure_v1(base_url)
+        try:
+            payload = json.loads(key)
+        except (JSONDecodeError, TypeError):
+            api_key = key
+            provider_order = ""
+        else:
+            if isinstance(payload, dict):
+                api_key = payload.get("api_key", "")
+                provider_order = payload.get("provider_order", "")
+            else:
+                api_key = key
+                provider_order = ""
+        self.client = OpenAI(api_key=api_key, base_url=self.base_url)
+        self.model_name = model_name
+        self.provider_order = provider_order
+
+    def _call(self, batch):
+        extra_body = {"drop_params": True}
+        if self.provider_order:
+            order = [s.strip() for s in self.provider_order.split(",") if s.strip()]
+            extra_body["provider"] = {"order": order, "allow_fallbacks": False}
+        res = self.client.embeddings.create(input=batch, model=self.model_name, encoding_format="float", extra_body=extra_body)
+        return [d.embedding for d in _sorted_by_index(res.data)], total_token_count_from_response(res)
+
+    def encode(self, texts: list):
+        return self._batched_encode(texts, self._call, batch_size=16, truncate_to=8191)
+
+    def encode_queries(self, text):
+        vectors, token_count = self._batched_encode([text], self._call, batch_size=16, truncate_to=8191)
+        return vectors[0], token_count