From 6839423331d6aeaf9801639b2fd72f5404cdc971 Mon Sep 17 00:00:00 2001
From: macie <macie@users.noreply.github.com>
Date: Tue, 2 Jan 2024 22:42:22 +0100
Subject: [PATCH] feat: Read system prompt from config file

---
 cmd/boludo/config.go      | 5 +++++
 cmd/boludo/config_test.go | 9 +++++++++
 2 files changed, 14 insertions(+)

diff --git a/cmd/boludo/config.go b/cmd/boludo/config.go
index 67a028a..6e45c57 100644
--- a/cmd/boludo/config.go
+++ b/cmd/boludo/config.go
@@ -191,6 +191,7 @@ type ConfigFile map[string]ModelSpec
 // ModelSpec represents a model specification in the configuration file.
 type ModelSpec struct {
 	Model         string
+	SystemPrompt  string
 	InitialPrompt string
 	Format        string
 	Creativity    float32
@@ -229,6 +230,7 @@ func (c *ConfigFile) UnmarshalTOML(data interface{}) error {
 	for configId := range definedConfigs {
 		defaultSpec := ModelSpec{
 			Model:         "",
+			SystemPrompt:  "",
 			InitialPrompt: "",
 			Format:        "",
 			Creativity:    llama.DefaultOptions.Temp,
@@ -244,6 +246,8 @@ func (c *ConfigFile) UnmarshalTOML(data interface{}) error {
 				defaultSpec.Cutoff = (float32)(v.(float64))
 			case "format":
 				defaultSpec.Format = v.(string)
+			case "system-prompt":
+				defaultSpec.SystemPrompt = v.(string)
 			case "initial-prompt":
 				defaultSpec.InitialPrompt = v.(string)
 			}
@@ -274,6 +278,7 @@ func (c *ConfigFile) Prompt(configId string) llama.Prompt {
 	if spec, ok := (*c)[configId]; ok {
 		return llama.Prompt{
 			Format: spec.Format,
+			System: spec.SystemPrompt,
 		}
 	}
 	return llama.Prompt{}
diff --git a/cmd/boludo/config_test.go b/cmd/boludo/config_test.go
index d2edda0..e62e119 100644
--- a/cmd/boludo/config_test.go
+++ b/cmd/boludo/config_test.go
@@ -119,6 +119,15 @@ func TestParseFile(t *testing.T) {
 				Cutoff:     0.0,
 			},
 		}},
+		{"[assistant]\nmodel = \"model.gguf\"\ninitial-prompt = 'Reword:'\nsystem-prompt = 'You are an assistant.'", ConfigFile{
+			"assistant": ModelSpec{
+				Model:         "model.gguf",
+				Format:        "",
+				InitialPrompt: "Reword:",
+				SystemPrompt:  "You are an assistant.",
+				Creativity:    1.0,
+			},
+		}},
 	}
 	for _, tc := range testcases {
 		tc := tc