add possiblity to use different instruct sets

rwth-acis · Feb 27, 2024 · f0dbb99 · f0dbb99
1 parent 392e0bf
commit f0dbb99
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/ollama/ollama.go b/ollama/ollama.go
@@ -15,12 +15,18 @@ import (
 type ResponseData struct {
 	Response string `json:"response"`
 	PromptID string `json:"promptID"`
+	Instruct string `json:"instruct"`
 }
 
 func GenerateResponse(prompt map[string]interface{}) (ResponseData, error) {
 	url := os.Getenv("OLLAMA_URL") + "/api/generate"
 
-	instruct, err := redis.GetSetMember("default")
+	set, ok := prompt["instructType"].(string)
+	if !ok {
+		set = "default"
+	}
+
+	instruct, err := redis.GetSetMember(set)
 	if err != nil {
 		return ResponseData{}, errors.New("no data available")
 	}
@@ -103,6 +109,7 @@ func GenerateResponse(prompt map[string]interface{}) (ResponseData, error) {
 	responseData := ResponseData{
 		Response: response,
 		PromptID: PromptID,
+		Instruct: instruct,
 	}
 
 	return responseData, nil