Browse Source

Fix a bug in VolcEngine (#3658)

### What problem does this PR solve?

Fix a bug in VolcEngine  #3553

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>
tags/v0.14.1
liuhua 11 months ago
parent
commit
8b35776916
No account linked to committer's email address
2 changed files with 8 additions and 0 deletions
  1. 1
    0
      rag/llm/__init__.py
  2. 7
    0
      rag/llm/embedding_model.py

+ 1
- 0
rag/llm/__init__.py View File

@@ -48,6 +48,7 @@ EmbeddingModel = {
"BaiduYiyan": BaiduYiyanEmbed,
"Voyage AI": VoyageEmbed,
"HuggingFace": HuggingFaceEmbed,
"VolcEngine":VolcEngineEmbed,
}

CvModel = {

+ 7
- 0
rag/llm/embedding_model.py View File

@@ -718,3 +718,10 @@ class HuggingFaceEmbed(Base):
else:
raise Exception(f"Error: {response.status_code} - {response.text}")

class VolcEngineEmbed(OpenAIEmbed):
def __init__(self, key, model_name, base_url="https://ark.cn-beijing.volces.com/api/v3"):
if not base_url:
base_url = "https://ark.cn-beijing.volces.com/api/v3"
ark_api_key = json.loads(key).get('ark_api_key', '')
model_name = json.loads(key).get('ep_id', '') + json.loads(key).get('endpoint_id', '')
super().__init__(ark_api_key,model_name,base_url)

Loading…
Cancel
Save