Merge branch 'feat/vertexai-provider' of https://github.com/yomorun/yomo

into feat/vertexai-provider
yomorun · Dec 14, 2024 · 262f08d · 262f08d
2 parents 90e3d98 + 97ffe98
commit 262f08d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/example/10-ai/zipper.yaml b/example/10-ai/zipper.yaml
@@ -48,3 +48,9 @@ bridge:
       xai:
         api_key:
         model:
+
+      vertexai:
+        project_id:
+        location:
+        model:
+        credentials_file:
diff --git a/pkg/bridge/ai/provider/vertexai/provider.go b/pkg/bridge/ai/provider/vertexai/provider.go
@@ -31,7 +31,7 @@ func NewProvider(projectID, location, model, credentialsFile string) *Provider {
 		log.Fatal("new vertexai client: ", err)
 	}
 	if model == "" {
-		model = "gemini-1.5-pro-latest"
+		model = "gemini-1.5-pro-002"
 	}
 
 	return &Provider{

diff --git a/pkg/bridge/ai/provider/vertexai/response_convert.go b/pkg/bridge/ai/provider/vertexai/response_convert.go
@@ -38,7 +38,7 @@ func convertToResponse(in *genai.GenerateContentResponse, model string) (out ope
 					Index: int(index),
 					Message: openai.ChatCompletionMessage{
 						Content: string(pp),
-						Role:    openai.ChatMessageRoleUser,
+						Role:    openai.ChatMessageRoleAssistant,
 					},
 					FinishReason: toOpenAIFinishReason(candidate.FinishReason),
 				})
@@ -95,7 +95,7 @@ func convertToStreamResponse(id string, in *genai.GenerateContentResponse, model
 					Index: index,
 					Delta: openai.ChatCompletionStreamChoiceDelta{
 						Content: string(pp),
-						Role:    openai.ChatMessageRoleUser,
+						Role:    openai.ChatMessageRoleAssistant,
 					},
 					FinishReason: toOpenAIFinishReason(candidate.FinishReason),
 				})