feat: support gemini pro system prompt

tmc · May 2, 2024 · 65a1bd5 · 65a1bd5
1 parent 3bb95e1
commit 65a1bd5
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/llms/googleai/googleai.go b/llms/googleai/googleai.go
@@ -23,11 +23,12 @@ var (
 )
 
 const (
-	CITATIONS = "citations"
-	SAFETY    = "safety"
-	RoleModel = "model"
-	RoleUser  = "user"
-	RoleTool  = "tool"
+	CITATIONS  = "citations"
+	SAFETY     = "safety"
+	RoleSystem = "system"
+	RoleModel  = "model"
+	RoleUser   = "user"
+	RoleTool   = "tool"
 )
 
 // Call implements the [llms.Model] interface.
@@ -208,7 +209,7 @@ func convertContent(content llms.MessageContent) (*genai.Content, error) {
 
 	switch content.Role {
 	case llms.ChatMessageTypeSystem:
-		return nil, ErrSystemRoleNotSupported
+		c.Role = RoleSystem
 	case llms.ChatMessageTypeAI:
 		c.Role = RoleModel
 	case llms.ChatMessageTypeHuman:
@@ -258,6 +259,10 @@ func generateFromMessages(ctx context.Context, model *genai.GenerativeModel, mes
 		if err != nil {
 			return nil, err
 		}
+		if mc.Role == RoleSystem {
+			model.SystemInstruction = content
+			continue
+		}
 		history = append(history, content)
 	}