AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-12 06:38:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix a bug in VolcEngine (#3658)

Browse Source 
### 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>
This commit is contained in:
liuhua 2024-11-27 09:30:49 +08:00 committed by GitHub
parent b6f3f15f0b
commit 8b35776916
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
rag/llm/__init__.py
View File

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

7
rag/llm/embedding_model.py
View File

@ -718,3 +718,10 @@ class HuggingFaceEmbed(Base):
else: else:
raise Exception(f"Error: {response.status_code} - {response.text}") 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)