Add state functions to manipulate conversation

Author: slimslenderslacks

src/graph.clj

@@ -97,6 +97,10 @@
 
 (declare stream chat-with-tools)
 
+(defn apply-functions [coll] 
+  (fn [state]
+    (reduce (fn [m f] (f m)) state coll)))
+
 (defn sub-graph-node
   "create a sub-graph node that initializes a conversation from the current one,
    creates a new agent graph from the current state and returns the messages to be added 
@@ -109,7 +113,13 @@
              (stream
               ((or construct-graph chat-with-tools) state)
               (->
-               ((or init-state (comp state/construct-initial-state-from-prompts state/add-prompt-ref)) state)
+               ((or 
+                  (and 
+                    init-state 
+                    (if (coll? init-state) 
+                      (apply-functions init-state) 
+                      init-state))
+                  (comp state/construct-initial-state-from-prompts state/add-prompt-ref)) state)
                (update-in [:opts :level] (fnil inc 0)))))]
         ((or next-state state/add-last-message-as-tool-call) state sub-graph-state)))))
 

src/graphs/sql.clj

@@ -83,19 +83,6 @@
       ;; how many times should we try to correct because correct-query will always end up back here 
       :else "correct-query")))
 
-(defn seed-list-tables-conversation [state]
-  (-> state
-      (assoc :finish-reason "tool_calls")
-      (update-in [:functions] (constantly (:tools first-tool-call)))
-      (update-in [:messages] concat (:messages first-tool-call))))
-
-(defn seed-get-schema-conversation [state]
-  ; inherit full conversation
-  ; no prompts
-  ; add the schema tool
-  (-> state
-      (update-in [:functions] (fnil concat []) (:tools model-get-schema))))
-
 (defn seed-correct-query-conversation
   [state]
   ; make one LLM call with the last message (which should be a user query containing the SQL we want to check)
@@ -106,18 +93,27 @@
       (state/construct-initial-state-from-prompts)
       (update-in [:messages] concat [(last (:messages state))])))
 
+(comment
+  [state/messages-reset
+   (state/messages-from-prompt "prompts/sql/query-check.md")
+   (state/messages-take-last 1)])
+
 ;; query-gen has a prompt
 ;; seed-correct-query-conversation has a prompt
 ;; prompts/sql/query-gen.md has a hard-coded db file
 (defn graph [_]
   (graph/construct-graph
    [[["start"                   graph/start]
-     ["list-tables-tool"        (graph/sub-graph-node 
-                                  {:init-state seed-list-tables-conversation
-                                   :construct-graph graph/generate-start-with-tool 
-                                   :next-state (state/take-last-messages 2)})]
+     ["list-tables-tool"        (graph/sub-graph-node
+                                 {:init-state
+                                  [#(assoc % :finish-reason "tool_calls")
+                                   (state/tools-set (:tools first-tool-call))
+                                   (state/messages-append (:messages first-tool-call))]
+                                  :construct-graph graph/generate-start-with-tool
+                                  :next-state (state/take-last-messages 2)})]
      ["model-get-schema"        (graph/sub-graph-node
-                                 {:init-state seed-get-schema-conversation
+                                 {:init-state
+                                  [(state/tools-append (:tools model-get-schema))]
                                   :next-state state/append-new-messages})]
      ["query-gen"               query-gen]
      [:edge                     should-continue]]

src/state.clj

@@ -1,10 +1,11 @@
 (ns state
   (:require
+   [babashka.fs :as fs]
+   [clojure.pprint :refer [pprint]]
    git
    jsonrpc
    prompts
-   tools
-   [clojure.pprint :refer [pprint]]))
+   tools))
 
 (set! *warn-on-reflection* true)
 
@@ -65,6 +66,35 @@
                               (format "failure for prompt configuration:\n %s" (with-out-str (pprint (dissoc opts :pat :jwt))))
                               :exception (str ex)}))))
 
+(defn tools-append [tools]
+  (fn [state]
+    (-> state
+        (update-in [:functions] (fnil concat []) tools))))
+
+(defn tools-set [tools]
+  (fn [state]
+    (-> state
+        (update-in [:functions] (constantly tools)))))
+
+(defn messages-reset [state]
+  (dissoc state :messages))
+
+(defn messages-take-last [n]
+  (fn [state]
+    (-> state
+        (update-in [:messages] (fnil concat []) (take-last n (:messages state))))))
+
+(defn messages-append [coll]
+  (fn [state]
+    (-> state
+        (update-in [:messages] (fnil concat []) coll))))
+
+(defn messages-from-prompt [s]
+  (fn [state]
+    (-> state
+        (update-in [:opts :prompts] (constantly (fs/file s)))
+        (construct-initial-state-from-prompts))))
+
 (defn add-prompt-ref
   [state]
   (let [definition (state/get-function-definition state)