outman/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Support for Vertex AI - load Default Application Configuration (#4641)

Browse Source 
Co-authored-by: miendinh <miendinh@users.noreply.github.com>
Co-authored-by: crazywoola <427733928@qq.com>
This commit is contained in:
miendinh 2024-05-25 12:40:25 +07:00 committed by GitHub
parent 109aabc6f2
commit f804adbff3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 17 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
api/core/model_runtime/model_providers/vertex_ai/llm/llm.py
View File

@ -164,10 +164,13 @@ class VertexAiLargeLanguageModel(LargeLanguageModel):
config_kwargs["stop_sequences"] = stop config_kwargs["stop_sequences"] = stop
service_account_info = json.loads(base64.b64decode(credentials["vertex_service_account_key"])) service_account_info = json.loads(base64.b64decode(credentials["vertex_service_account_key"]))
service_accountSA = service_account.Credentials.from_service_account_info(service_account_info)
project_id = credentials["vertex_project_id"] project_id = credentials["vertex_project_id"]
location = credentials["vertex_location"] location = credentials["vertex_location"]
aiplatform.init(credentials=service_accountSA, project=project_id, location=location) if service_account_info:
service_accountSA = service_account.Credentials.from_service_account_info(service_account_info)
aiplatform.init(credentials=service_accountSA, project=project_id, location=location)
else:
aiplatform.init(project=project_id, location=location)
history = [] history = []
system_instruction = GEMINI_BLOCK_MODE_PROMPT system_instruction = GEMINI_BLOCK_MODE_PROMPT

16
api/core/model_runtime/model_providers/vertex_ai/text_embedding/text_embedding.py
View File

@ -41,15 +41,16 @@ class VertexAiTextEmbeddingModel(_CommonVertexAi, TextEmbeddingModel):
:return: embeddings result :return: embeddings result
""" """
service_account_info = json.loads(base64.b64decode(credentials["vertex_service_account_key"])) service_account_info = json.loads(base64.b64decode(credentials["vertex_service_account_key"]))
service_accountSA = service_account.Credentials.from_service_account_info(service_account_info)
project_id = credentials["vertex_project_id"] project_id = credentials["vertex_project_id"]
location = credentials["vertex_location"] location = credentials["vertex_location"]
aiplatform.init(credentials=service_accountSA, project=project_id, location=location) if service_account_info:
service_accountSA = service_account.Credentials.from_service_account_info(service_account_info)
aiplatform.init(credentials=service_accountSA, project=project_id, location=location)
else:
aiplatform.init(project=project_id, location=location)
client = VertexTextEmbeddingModel.from_pretrained(model) client = VertexTextEmbeddingModel.from_pretrained(model)
embeddings_batch, embedding_used_tokens = self._embedding_invoke( embeddings_batch, embedding_used_tokens = self._embedding_invoke(
client=client, client=client,
texts=texts texts=texts
@ -103,10 +104,13 @@ class VertexAiTextEmbeddingModel(_CommonVertexAi, TextEmbeddingModel):
""" """
try: try:
service_account_info = json.loads(base64.b64decode(credentials["vertex_service_account_key"])) service_account_info = json.loads(base64.b64decode(credentials["vertex_service_account_key"]))
service_accountSA = service_account.Credentials.from_service_account_info(service_account_info)
project_id = credentials["vertex_project_id"] project_id = credentials["vertex_project_id"]
location = credentials["vertex_location"] location = credentials["vertex_location"]
aiplatform.init(credentials=service_accountSA, project=project_id, location=location) if service_account_info:
service_accountSA = service_account.Credentials.from_service_account_info(service_account_info)
aiplatform.init(credentials=service_accountSA, project=project_id, location=location)
else:
aiplatform.init(project=project_id, location=location)
client = VertexTextEmbeddingModel.from_pretrained(model) client = VertexTextEmbeddingModel.from_pretrained(model)

4
api/core/model_runtime/model_providers/vertex_ai/vertex_ai.yaml
View File

@ -36,8 +36,8 @@ provider_credential_schema:
en_US: Enter your Google Cloud Location en_US: Enter your Google Cloud Location
- variable: vertex_service_account_key - variable: vertex_service_account_key
label: label:
en_US: Service Account Key en_US: Service Account Key (Leave blank if you use Application Default Credentials)
type: secret-input type: secret-input
required: true required: false
placeholder: placeholder:
en_US: Enter your Google Cloud Service Account Key in base64 format en_US: Enter your Google Cloud Service Account Key in base64 format