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

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
This commit is contained in:
yihong 2024-12-02 10:25:15 +08:00 committed by GitHub
parent 858db2f239
commit 239bf97b47
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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))