From 3b361d1a3d4ba9207ae9aeda6e3d8302adced2eb Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@localai.io>
Date: Wed, 10 Apr 2024 23:49:13 +0200
Subject: [PATCH] Add prompt blocks

---
 agent/agent.go            | 13 +++++++++++++
 agent/options.go          | 17 +++++++++++++++++
 example/webui/create.html |  4 ++--
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index d9aa4da..cd46c90 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -199,7 +199,19 @@ func (a *Agent) consumeJob(job *Job, role string) {
 	//if job.Image != "" {
 	// TODO: Use llava to explain the image content
 	//}
+	// Add custom prompts
+	for _, prompt := range a.options.prompts {
+		message := prompt.Render(a)
+		if !Messages(a.currentConversation).Exist(a.options.systemPrompt) {
+			a.currentConversation = append([]openai.ChatCompletionMessage{
+				{
+					Role:    prompt.Role(),
+					Content: message,
+				}}, a.currentConversation...)
+		}
+	}
 
+	// TODO: move to a Promptblock?
 	if a.options.systemPrompt != "" {
 		if !Messages(a.currentConversation).Exist(a.options.systemPrompt) {
 			a.currentConversation = append([]openai.ChatCompletionMessage{
@@ -210,6 +222,7 @@ func (a *Agent) consumeJob(job *Job, role string) {
 		}
 	}
 
+	// TODO: move to a promptblock?
 	// RAG
 	if memory {
 		// Walk conversation from bottom to top, and find the first message of the user
diff --git a/agent/options.go b/agent/options.go
index e5f7cfe..d606bd9 100644
--- a/agent/options.go
+++ b/agent/options.go
@@ -30,6 +30,8 @@ type options struct {
 	kbResults                                         int
 	ragdb                                             RAGDB
 
+	prompts []PromptBlock
+
 	systemPrompt string
 
 	// callbacks
@@ -37,6 +39,11 @@ type options struct {
 	resultCallback    func(ActionState)
 }
 
+type PromptBlock interface {
+	Render(a *Agent) string
+	Role() string
+}
+
 func defaultOptions() *options {
 	return &options{
 		periodicRuns: 15 * time.Minute,
@@ -140,6 +147,16 @@ func WithCharacterFile(path string) Option {
 	}
 }
 
+// WithPrompts adds additional block prompts to the agent
+// to be rendered internally in the conversation
+// when processing the conversation to the LLM
+func WithPrompts(prompts ...PromptBlock) Option {
+	return func(o *options) error {
+		o.prompts = prompts
+		return nil
+	}
+}
+
 func WithLLMAPIKey(key string) Option {
 	return func(o *options) error {
 		o.LLMAPI.APIKey = key
diff --git a/example/webui/create.html b/example/webui/create.html
index 29434ce..a748aeb 100644
--- a/example/webui/create.html
+++ b/example/webui/create.html
@@ -41,7 +41,7 @@
                         </div>
                         <div class="mb-4">
                             <label for="connectorConfig${newConnectorIndex}" class="block text-lg font-medium text-gray-400">Connector Config (JSON)</label>
-                            <textarea id="connectorConfig${newConnectorIndex}" name="connectors[${newConnectorIndex}].config" class="mt-1 focus:ring-blue-500 focus:border-blue-500 block w-full shadow-sm sm:text-sm border-gray-300 rounded-md bg-gray-700 text-white" placeholder='{"token":"sk-mg3.."}'></textarea>
+                            <textarea id="connectorConfig${newConnectorIndex}" name="connectors[${newConnectorIndex}].config" class="mt-1 focus:ring-blue-500 focus:border-blue-500 block w-full shadow-sm sm:text-sm border-gray-300 rounded-md bg-gray-700 text-white" placeholder='{"token":"sk-mg3.."}'>{}</textarea>
                         </div>
                     </div>
                 `;
@@ -67,7 +67,7 @@
                     `+actions+`</select>
                     <div class="mb-4">
                         <label for="actionsConfig${ii}" class="block text-lg font-medium text-gray-400">Action Config (JSON)</label>
-                        <textarea id="actionsConfig${ii}" name="actions[${ii}].config" class="mt-1 focus:ring-blue-500 focus:border-blue-500 block w-full shadow-sm sm:text-sm border-gray-300 rounded-md bg-gray-700 text-white" placeholder='{"results":"5"}'></textarea>
+                        <textarea id="actionsConfig${ii}" name="actions[${ii}].config" class="mt-1 focus:ring-blue-500 focus:border-blue-500 block w-full shadow-sm sm:text-sm border-gray-300 rounded-md bg-gray-700 text-white" placeholder='{"results":"5"}'>{}</textarea>
                     </div>
                 </div>
                 `;