Create simple conversation state manipulators

* need to create operators for composing different langraph nodes

Author: slimslenderslacks

prompts/sql/query-check.md

@@ -14,7 +14,7 @@ tools:
     container:
       image: vonwig/sqlite:latest
       command:
-        - "{{database}}"
+        - "./Chinook.db"
         - "{{sql}}"
 tool_choice: required
 ---

src/docker/main.clj

@@ -14,6 +14,7 @@
    jsonrpc
    [logging :refer [warn]]
    prompts
+   state
    user-loop)
   (:gen-class))
 
@@ -176,7 +177,7 @@
                    (user-loop/start-jsonrpc-loop
                     (user-loop/create-step
                      (fn [state]
-                       (let [m (graph/construct-initial-state-from-prompts
+                       (let [m (state/construct-initial-state-from-prompts
                                  (assoc state :opts
                                         (-> (with-options opts (rest args))
                                             (assoc :thread-id thread-id))))]

src/graph.clj

@@ -2,7 +2,6 @@
   (:require
    [babashka.fs :as fs]
    [clojure.core.async :as async]
-   [clojure.pprint :refer [pprint]]
    [clojure.core.match :refer [match]]
    git
    jsonrpc
@@ -72,7 +71,6 @@
   [state]
   (run-llm (:messages state) (dissoc (:metadata state) :agent) (:functions state) (:opts state)))
 
-;; TODO does the LangGraph Tool Node always search for the a tool_call
 (defn tool
   "execute the tool_calls from the last AI message in the conversation"
   [state]
@@ -97,63 +95,8 @@
   [_]
   (async/go {}))
 
-(defn construct-initial-state-from-prompts [{{:keys [prompts] :as opts} :opts :as state}]
-  (try
-    (-> state
-        (merge
-         {:metadata (prompts/collect-metadata prompts)
-          :functions (prompts/collect-functions prompts)
-          :opts (merge opts {:level (or (:level opts) 0)})})
-        (update
-         :messages
-         (fnil concat [])
-         (when (not (seq (:messages state)))
-           (let [new-prompts (prompts/get-prompts opts)]
-             (jsonrpc/notify :prompts {:messages new-prompts})
-             new-prompts))))
-    (catch Throwable ex
-      (jsonrpc/notify :error {:content
-                              (format "failure for prompt configuration:\n %s" (with-out-str (pprint (dissoc opts :pat :jwt))))
-                              :exception (str ex)}))))
-
-; tool_calls are maps with an id and a function with arguments an name
-; look up the full tool definition using the name
-
-(defn add-prompt-ref [state]
-  (let [definition (state/get-function-definition state)
-        arg-context (let [raw-args (-> state :messages last :tool_calls first :function :arguments)]
-                      (tools/arg-context raw-args))]
-    (-> state
-        (dissoc :messages)
-        (update-in [:opts :level] (fnil inc 0))
-        (update-in [:opts :prompts] (constantly (git/prompt-file (-> definition :function :ref))))
-        (update-in [:opts :parameters] (constantly arg-context)))))
-
-(comment
-  ;; TODO move this into the thingy
-  (add-prompt-ref {:messages [{:tool_calls [{:function {:name "sql_db_list_tables"
-                                                        :arguments "{\"arg\": 1}"}}]}]
-                   :functions [{:function {:name "sql_db_list_tables"
-                                           :description "List all tables in the database"
-                                           :parameters {:type "object"
-                                                        :properties
-                                                        {:database {:type "string" :description "the database to query"}}}
-                                           :ref "github:docker/labs-ai-tools-for-devs?path=prompts/curl/README.md"}}]}))
-
 (declare stream chat-with-tools)
 
-(defn add-last-message-as-tool-call
-  [state sub-graph-state]
-  {:messages [(-> sub-graph-state
-                  :messages
-                  last
-                  (state/add-tool-call-id (-> state :messages last :tool_calls first :id)))]})
-
-(defn append-new-messages
-  [state sub-graph-state]
-  {:messages (->> (:messages sub-graph-state)
-                  (filter (complement (fn [m] (some #(= m %) (:messages state))))))})
-
 (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 
@@ -165,8 +108,10 @@
             (async/<!
              (stream
               ((or construct-graph chat-with-tools) state)
-              ((or init-state (comp construct-initial-state-from-prompts add-prompt-ref)) state)))]
-        ((or next-state add-last-message-as-tool-call) state sub-graph-state)))))
+              (->
+               ((or 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)))))
 
 ; =====================================================
 ; edge functions takes state and returns next node
@@ -280,7 +225,7 @@
     ["tools-query"]]
    [["end" end]]])
 
-(defn one-tool-call [_]
+(defn generate-one-tool-call [_]
   (-> {}
       (add-node "start" start)
       (add-node "completion" completion)
@@ -302,22 +247,12 @@
     ["completion"]]
    [["end" end]]])
 
-(comment
-  (alter-var-root #'jsonrpc/notify (fn [_] (partial jsonrpc/-println {:debug true})))
-  (let [x {:prompts (fs/file "/Users/slim/docker/labs-ai-tools-for-devs/prompts/curl/README.md")
-           :platform "darwin"
-           :user "jimclark106"
-           :thread-id "thread"
-           :host-dir "/Users/slim"
-           :stream true}]
-    (state/summarize (async/<!! (stream (chat-with-tools x) x)))))
-
-(comment
-  (def x {:stream true,
-          :host-dir "/Users/slim/docker/labs-make-runbook",
-          :prompts "/Users/slim/docker/labs-ai-tools-for-devs/prompts/hub/default.md"
-          :platform "darwin", :user "jimclark106",
-          :thread-id "3e61ffe7-840e-4177-b84a-f6f7db58b24d"})
-  (state/summarize
-   (async/<!! (stream (chat-with-tools x) x))))
+; requires finish-reason to be set to tool_calls
+(def start-with-tool
+  [[["start" start]
+    [:edge tool-or-end]]
+   [["sub-graph" (sub-graph-node {})]
+    ["end" end]]
+   [["tool" (tool-node {})]
+    ["end"]]])
 

src/graphs/sql.clj

@@ -3,7 +3,8 @@
    [babashka.fs :as fs]
    [clojure.core.async :as async]
    [clojure.string :as string]
-   [graph]))
+   [graph]
+   state))
 
 (def first-tool-call
   {:messages [{:role "assistant"
@@ -36,10 +37,6 @@
                        {:image "vonwig/sqlite:latest"
                         :command ["{{database}}" ".schema {{table}}"]}}}]})
 
-(defn list-tables-inject-tool [_]
-  (async/go
-    first-tool-call))
-
 (defn failed-tool-call-message [tool-call-name]
   ;; this is awful - binds the should-continue edge to the format of this string
   (format
@@ -56,7 +53,7 @@
              (dissoc :functions)
              (update-in [:opts :level] (fnil inc 0))
              (update-in [:opts :prompts] (constantly (fs/file "prompts/sql/query-gen.md")))
-             (graph/construct-initial-state-from-prompts)
+             (state/construct-initial-state-from-prompts)
              (update-in [:messages] concat (:messages state)))
           {:keys [messages _finish-reason]} (async/<! (graph/run-llm
                                                        (:messages x)
@@ -81,15 +78,22 @@
   (let [last-message (last messages)]
     (cond
       (contains? last-message :tool_calls) "end"
+      ;; prevent inifinite loops of errors
       (string/starts-with? (:content last-message) "Error:") "query-gen"
+      ;; 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 [:opts :level] (fnil inc 0))
       (update-in [:functions] (fnil concat []) (:tools model-get-schema))))
 
 (defn seed-correct-query-conversation
@@ -98,25 +102,29 @@
   ; add the last message to the conversation
   (-> state
       (dissoc :messages)
-      (update-in [:opts :level] (fnil inc 0))
       (update-in [:opts :prompts] (constantly (fs/file "prompts/sql/query-check.md")))
-      (graph/construct-initial-state-from-prompts)
+      (state/construct-initial-state-from-prompts)
       (update-in [:messages] concat [(last (:messages state))])))
 
+;; 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-inject-tool" list-tables-inject-tool]
-     ["list-tables-tool" (graph/tool-node {})]
-     ["model-get-schema" (graph/sub-graph-node
-                          {:init-state seed-get-schema-conversation
-                           :next-state graph/append-new-messages})]
-     ["query-gen" query-gen]
-     [:edge should-continue]]
-    [["correct-query" (graph/sub-graph-node
-                       {:init-state seed-correct-query-conversation
-                        :construct-graph graph/one-tool-call
-                        :next-state graph/append-new-messages})]
+   [[["start"                   graph/start]
+     ["list-tables-tool"        (graph/sub-graph-node 
+                                  {:init-state seed-list-tables-conversation
+                                   :construct-graph (fn [_] (graph/construct-graph graph/start-with-tool))
+                                   :next-state state/take-last-two-messages})]
+     ["model-get-schema"        (graph/sub-graph-node
+                                 {:init-state seed-get-schema-conversation
+                                  :next-state state/append-new-messages})]
+     ["query-gen"               query-gen]
+     [:edge                     should-continue]]
+    [["correct-query"           (graph/sub-graph-node
+                                 {:init-state seed-correct-query-conversation
+                                  :construct-graph graph/generate-one-tool-call
+                                  :next-state state/append-new-messages})]
      ["query-gen"]]
-    [["end" graph/end]]]))
+    [["end"                     graph/end]]]))
 

src/state.clj

@@ -1,4 +1,10 @@
-(ns state)
+(ns state 
+  (:require
+   git
+   jsonrpc
+   prompts
+   tools
+   [clojure.pprint :refer [pprint]]))
 
 (set! *warn-on-reflection* true)
 
@@ -40,3 +46,51 @@
 (def prompt-tool? (comp prompt? get-function-definition))
 
 (defn add-tool-call-id [m id] (assoc m :role "tool" :tool_call_id id))
+
+(defn construct-initial-state-from-prompts [{{:keys [prompts] :as opts} :opts :as state}]
+  (try
+    (-> state
+        (merge
+         {:metadata (prompts/collect-metadata prompts)
+          :functions (prompts/collect-functions prompts)})
+        (update
+         :messages
+         (fnil concat [])
+         (when (not (seq (:messages state)))
+           (let [new-prompts (prompts/get-prompts opts)]
+             (jsonrpc/notify :prompts {:messages new-prompts})
+             new-prompts))))
+    (catch Throwable ex
+      (jsonrpc/notify :error {:content
+                              (format "failure for prompt configuration:\n %s" (with-out-str (pprint (dissoc opts :pat :jwt))))
+                              :exception (str ex)}))))
+
+(defn add-prompt-ref
+  [state]
+  (let [definition (state/get-function-definition state)
+        arg-context (let [raw-args (-> state :messages last :tool_calls first :function :arguments)]
+                      (tools/arg-context raw-args))]
+    (-> state
+        (dissoc :messages)
+        (update-in [:opts :prompts] (constantly (git/prompt-file (-> definition :function :ref))))
+        (update-in [:opts :parameters] (constantly arg-context)))))
+
+
+(defn add-last-message-as-tool-call
+  [state sub-graph-state]
+  {:messages [(-> sub-graph-state
+                  :messages
+                  last
+                  (state/add-tool-call-id (-> state :messages last :tool_calls first :id)))]})
+
+(defn append-new-messages
+  [state sub-graph-state]
+  {:messages (->> (:messages sub-graph-state)
+                  (filter (complement (fn [m] (some #(= m %) (:messages state))))))})
+
+(defn take-last-two-messages
+  [_ sub-graph-state]
+  {:messages (->> (:messages sub-graph-state)
+                  (take-last 2))})
+
+

src/tools.clj

@@ -32,7 +32,6 @@
 
 (defn arg-context [json-arg-string]
   (merge
-   ;; TODO raw is a bad name when merging
    {:raw (if json-arg-string
            json-arg-string
            "{}")}

test/graph_t.clj

@@ -1,8 +1,46 @@
 (ns graph_t
-  (:require [clojure.test :as t]
-            [graph]
-            [graphs.sql]))
+  (:require
+   [babashka.fs :as fs]
+   [clojure.core.async :as async]
+   [clojure.test :as t]
+   [graph]
+   [graphs.sql]
+   jsonrpc
+   state))
 
 (t/deftest
   (graph/construct-graph graph/chat-with-tools-representation)
-  (graph/construct-graph graphs.sql/graph-data))
+  (graph/construct-graph graphs.sql/graph))
+
+(comment
+  ;; requires network to test
+  (state/add-prompt-ref
+   {:messages [{:tool_calls [{:function {:name "sql_db_list_tables"
+                                         :arguments "{\"arg\": 1}"}}]}]
+    :functions [{:function {:name "sql_db_list_tables"
+                            :description "List all tables in the database"
+                            :parameters {:type "object"
+                                         :properties
+                                         {:database {:type "string" :description "the database to query"}}}
+                            :ref "github:docker/labs-ai-tools-for-devs?path=prompts/curl/README.md"}}]}))
+
+(comment
+  (alter-var-root #'jsonrpc/notify (fn [_] (partial jsonrpc/-println {:debug true})))
+  (let [x {:prompts (fs/file "/Users/slim/docker/labs-ai-tools-for-devs/prompts/curl/README.md")
+           :platform "darwin"
+           :user "jimclark106"
+           :thread-id "thread"
+           :host-dir "/Users/slim"
+           :stream true}]
+    (state/summarize (async/<!! (graph/stream (graph/chat-with-tools x) x)))))
+
+(comment
+  (def x {:stream true,
+          :host-dir "/Users/slim/docker/labs-make-runbook",
+          :prompts "/Users/slim/docker/labs-ai-tools-for-devs/prompts/hub/default.md"
+          :platform "darwin", :user "jimclark106",
+          :thread-id "3e61ffe7-840e-4177-b84a-f6f7db58b24d"})
+  (state/summarize
+   (async/<!! (graph/stream (graph/chat-with-tools x) x))))
+
+