fix(go): fixed googleai/vertexai system prompt handling (#732)

firebase · Aug 1, 2024 · a9af53c · a9af53c
1 parent cb36ca9
commit a9af53c
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 6 deletions.
diff --git a/go/plugins/googleai/googleai.go b/go/plugins/googleai/googleai.go
@@ -232,7 +232,10 @@ func generate(
  input *ai.GenerateRequest,
  cb func(context.Context, *ai.GenerateResponseChunk) error,
 ) (*ai.GenerateResponse, error) {
- gm := newModel(client, model, input)
+ gm, err := newModel(client, model, input)
+ if err != nil {
+ return nil, err
+ }
  cs, err := startChat(gm, input)
  if err != nil {
  return nil, err
@@ -300,7 +303,7 @@ func generate(
  return r, nil
 }
 
-func newModel(client *genai.Client, model string, input *ai.GenerateRequest) *genai.GenerativeModel {
+func newModel(client *genai.Client, model string, input *ai.GenerateRequest) (*genai.GenerativeModel, error) {
  gm := client.GenerativeModel(model)
  gm.SetCandidateCount(int32(input.Candidates))
  if c, ok := input.Config.(*ai.GenerationCommonConfig); ok && c != nil {
@@ -320,7 +323,21 @@ func newModel(client *genai.Client, model string, input *ai.GenerateRequest) *ge
  gm.SetTopP(float32(c.TopP))
  }
  }
- return gm
+ for _, m := range input.Messages {
+ systemParts, err := convertParts(m.Content)
+ if err != nil {
+ return nil, err
+
+ }
+ // system prompts go into GenerativeModel.SystemInstruction field.
+ if m.Role == ai.RoleSystem {
+ gm.SystemInstruction = &genai.Content{
+ Parts: systemParts,
+ Role: string(m.Role),
+ }
+ }
+ }
+ return gm, nil
 }
 
 // startChat starts a chat session and configures it with the input messages.
@@ -332,6 +349,12 @@ func startChat(gm *genai.GenerativeModel, input *ai.GenerateRequest) (*genai.Cha
  for len(messages) > 1 {
  m := messages[0]
  messages = messages[1:]
+
+ // skip system prompt message, it's handled separately.
+ if m.Role == ai.RoleSystem {
+ continue
+ }
+
  parts, err := convertParts(m.Content)
  if err != nil {
  return nil, err

diff --git a/go/plugins/vertexai/vertexai.go b/go/plugins/vertexai/vertexai.go
@@ -239,7 +239,10 @@ func generate(
  input *ai.GenerateRequest,
  cb func(context.Context, *ai.GenerateResponseChunk) error,
 ) (*ai.GenerateResponse, error) {
- gm := newModel(client, model, input)
+ gm, err := newModel(client, model, input)
+ if err != nil {
+ return nil, err
+ }
  cs, err := startChat(gm, input)
  if err != nil {
  return nil, err
@@ -307,7 +310,7 @@ func generate(
  return r, nil
 }
 
-func newModel(client *genai.Client, model string, input *ai.GenerateRequest) *genai.GenerativeModel {
+func newModel(client *genai.Client, model string, input *ai.GenerateRequest) (*genai.GenerativeModel, error) {
  gm := client.GenerativeModel(model)
  gm.SetCandidateCount(int32(input.Candidates))
  if c, ok := input.Config.(*ai.GenerationCommonConfig); ok && c != nil {
@@ -327,7 +330,21 @@ func newModel(client *genai.Client, model string, input *ai.GenerateRequest) *ge
  gm.SetTopP(float32(c.TopP))
  }
  }
- return gm
+ for _, m := range input.Messages {
+ systemParts, err := convertParts(m.Content)
+ if err != nil {
+ return nil, err
+
+ }
+ // system prompts go into GenerativeModel.SystemInstruction field.
+ if m.Role == ai.RoleSystem {
+ gm.SystemInstruction = &genai.Content{
+ Parts: systemParts,
+ Role: string(m.Role),
+ }
+ }
+ }
+ return gm, nil
 }
 
 // startChat starts a chat session and configures it with the input messages.
@@ -339,6 +356,12 @@ func startChat(gm *genai.GenerativeModel, input *ai.GenerateRequest) (*genai.Cha
  for len(messages) > 1 {
  m := messages[0]
  messages = messages[1:]
+
+ // skip system prompt message, it's handled separately.
+ if m.Role == ai.RoleSystem {
+ continue
+ }
+
  parts, err := convertParts(m.Content)
  if err != nil {
  return nil, err