Просмотр исходного кода

fix: nvidia special embedding model payload close #11193 (#11239)

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
yihong 4 месяцев назад
Родитель
Сommit
239bf97b47

+ 5 - 1
api/core/model_runtime/model_providers/openai_api_compatible/text_embedding/text_embedding.py

@@ -139,13 +139,17 @@ class OAICompatEmbeddingModel(_CommonOaiApiCompat, TextEmbeddingModel):
             if api_key:
                 headers["Authorization"] = f"Bearer {api_key}"
 
-            endpoint_url = credentials.get("endpoint_url")
+            endpoint_url = credentials.get("endpoint_url", "")
             if not endpoint_url.endswith("/"):
                 endpoint_url += "/"
 
             endpoint_url = urljoin(endpoint_url, "embeddings")
 
             payload = {"input": "ping", "model": model}
+            # For nvidia models, the "input_type":"query" need in the payload
+            # more to check issue #11193 or NvidiaTextEmbeddingModel
+            if model.startswith("nvidia/"):
+                payload["input_type"] = "query"
 
             response = requests.post(url=endpoint_url, headers=headers, data=json.dumps(payload), timeout=(10, 300))