Sql agent graph

Author: slimslenderslacks

prompts/sql/Chinook.db

@@ -4,4 +4,5 @@ agent: sql
 
 # prompt user
 
-find all artists in database ./Chinook.db
+find all artists in the database ./Chinook.db
+

prompts/sql/prompt.md

@@ -1,18 +1,22 @@
 ---
 tools:
-  - name: sqlite3
+  - name: sql_db_query_tool
     description: execute the DB query
     parameters:
       type: object
       properties:
+        database:
+          type: string
+          description: the database to query
         sql:
           type: string
           description: the sql statement to run
     container:
       image: vonwig/sqlite:latest
       command:
-        - "./Chinook.db"
+        - "{{database}}"
         - "{{sql}}"
+tool_choice: required
 ---
 
 # prompt system
@@ -32,6 +36,3 @@ If there are any of the above mistakes, rewrite the query. If there are no mista
 
 You will call the appropriate tool to execute the query after running this check.
 
-# prompt user
-
-SELECT * FROM Artist LIMIT 10;

prompts/sql/query-check.md

@@ -1,4 +1,13 @@
 ---
+tools:
+  - name: SubmitFinalAnswer
+    description: Submit the final answer to the user based on the query results
+    parameters:
+      type: object
+      properties:
+        final_answer:
+          type: string
+          description: The final answer to the user
 ---
 
 # prompt system

prompts/sql/query-gen.md

@@ -66,15 +66,14 @@
     c))
 
 (defn completion
-  "get the next llm completion
-     uses the current conversation messages and outputs
-     outputs the next set of messages and finish-reason to be added to the conversation"
+  "generate the next AI message
+     passes the whole converation to the AI model"
   [state]
-  (run-llm (:messages state) (:metadata state) (:functions state) (:opts 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
-  "make docker container tool calls"
+  "execute the tool_calls from the last AI message in the conversation"
   [state]
   (let [calls (-> (:messages state) last :tool_calls)]
     (async/go
@@ -88,7 +87,9 @@
                     calls)
                    (async/reduce conj []))))})))
 
-(defn tool-node [_]
+(defn tool-node
+  "add a tool node that will run tool_calls from the last AI message in the conversation"
+  [_]
   tool)
 
 (defn tools-query
@@ -128,6 +129,7 @@
         (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"
@@ -139,18 +141,31 @@
 
 (declare stream chat-with-tools)
 
-(defn sub-graph-node [{:keys [init-state construct-graph]}]
+(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 
+   to the parent conversation"
+  [{:keys [init-state construct-graph next-state]}]
   (fn [state]
     (async/go
       (let [sub-graph-state
             (async/<!
              (stream
               ((or construct-graph chat-with-tools) state)
               ((or init-state (comp construct-initial-state-from-prompts add-prompt-ref)) state)))]
-        {:messages [(-> sub-graph-state
-                        :messages
-                        last
-                        (state/add-tool-call-id (-> state :messages last :tool_calls first :id)))]}))))
+        ((or next-state add-last-message-as-tool-call) state sub-graph-state)))))
 
 ; =====================================================
 ; edge functions takes state and returns next node
@@ -194,6 +209,7 @@
     [state m
      node "start"]
      (jsonrpc/notify :message {:debug (format "\n-> entering %s\n\n" node)})
+     #_(jsonrpc/notify :message {:debug (with-out-str (pprint (state/summarize (dissoc state :opts))))})
      ;; TODO handling bad graphs with missing nodes
      (let [enter-node (get-in graph [:nodes node])
            new-state (state-reducer state (async/<! (enter-node state)))]
@@ -222,6 +238,18 @@
       (add-edge "tools-query" "completion")
       (add-conditional-edges "completion" tool-or-end)))
 
+(defn one-tool-call [_]
+  (-> {}
+      (add-node "start" start)
+      (add-node "completion" completion)
+      (add-node "tool" (tool-node nil))
+      (add-node "end" end)
+      (add-node "sub-graph" (sub-graph-node nil))
+      (add-edge "start" "completion")
+      (add-edge "sub-graph" "end")
+      (add-edge "tool" "end")
+      (add-conditional-edges "completion" tool-or-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")

src/graph.clj

@@ -1,25 +1,10 @@
 (ns graphs.sql
   (:require
+   [babashka.fs :as fs]
    [clojure.core.async :as async]
    [clojure.string :as string]
    [graph]))
 
-(def db-query-tool-call
-  {:messages [{:content ""
-               :tool_calls [{:name "sql_db_query_tool"
-                             :arguments "{}"
-                             :id "tool_abc123"}]}]
-   :tools [{:name "sql_db_query_tool"
-            :description "List all tables in the database"
-            :parameters
-            {:type "object"
-             :properties
-             {:database {:type "string" :description "the database to query"}
-              :query {:type "string" :description "the sql statement to run"}}}
-            :container
-            {:image "vonwig/sqlite:latest"
-             :command ["{{database}}" "{{query}}"]}}]})
-
 (def first-tool-call
   {:messages [{:role "assistant"
                :content ""
@@ -55,30 +40,48 @@
   (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
+   "Error: The wrong tool was called: %s. Please fix your mistakes. Remember to only call SubmitFinalAnswer to submit the final answer. Generated queries should be outputted WITHOUT a tool call."
+   tool-call-name))
+
 (defn query-gen
   "Assistant+Tool Node: has it's own prompt but also adds checks for proper final answers"
-  [_]
-  ;; note that we should not bind tools here because the only acceptable output is a proper final answer
-  ;; so this sub-graph needs to remove all tools
-  )
-
-;; TODO pull out the query check system
-(defn correct-query
-  "Assistant Node: double check the query 
-      - should generate a tool call to execute the query"
-  [_])
-
-(defn execute-query
-  "Tool Node: runs db query"
-  [_])
+  [state]
+  (async/go
+    (let [x (->
+             state
+             (dissoc :messages)
+             (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)
+             (update-in [:messages] concat (:messages state)))
+          {:keys [messages _finish-reason]} (async/<! (graph/run-llm
+                                                       (:messages x)
+                                                       (dissoc (:metadata x) :agent)
+                                                       (:functions x)
+                                                       (:opts x)))]
+
+      ; check for bad tool_calls and create failed Tool messages for them 
+      {:messages
+       (concat
+        messages
+        (->> (:tool_calls (last messages))
+             (filter (complement #(= "SubmitFinalAnswer" (-> % :function :name))))
+             (map (fn [{:keys [id] :as tc}]
+                    {:role "tool"
+                     :content (failed-tool-call-message (-> tc :function :name))
+                     :tool_call_id id}))))})))
 
 (defn should-continue
   "end, correct-query, or query-gen"
   [{:keys [messages]}]
   (let [last-message (last messages)]
     (cond
       (contains? last-message :tool_calls) "end"
-      (string/starts-with? last-message "Error:") "query-gen"
+      (string/starts-with? (:content last-message) "Error:") "query-gen"
       :else "correct-query")))
 
 (defn seed-get-schema-conversation [state]
@@ -90,6 +93,18 @@
       (update-in [:opts :parameters] (constantly {:database "./Chinook.db"}))
       (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)
+  ; 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")))
+      (update-in [:opts :parameters] (constantly {:database "./Chinook.db"}))
+      (graph/construct-initial-state-from-prompts)
+      (update-in [:messages] concat [(last (:messages state))])))
+
 (defn graph [_]
   (-> {}
       (graph/add-node "start" graph/start)
@@ -99,22 +114,18 @@
       (graph/add-node "list-tables-tool" (graph/tool-node nil))
       (graph/add-edge "list-tables-inject-tool" "list-tables-tool")
 
-      (graph/add-node "model-get-schema" (graph/sub-graph-node {:init-state seed-get-schema-conversation}))       ; assistant
+      ; TODO replace the conversation state, don't append
+      (graph/add-node "model-get-schema" (graph/sub-graph-node {:init-state seed-get-schema-conversation
+                                                                :next-state graph/append-new-messages}))       ; assistant
       (graph/add-edge "list-tables-tool" "model-get-schema")
 
-      (graph/add-node "end" graph/end)
-      (graph/add-edge "model-get-schema" "end")
-      ;(graph/add-node "query-gen" query-gen)                     ; assistant - might just end if it generates the right response
-                                                                 ;;           - might just loop back to query-gen if there's an error
-                                                                 ;;           - otherwise switch to correct-query
-
-      ;(graph/add-node "correct-query" correct-query)             ; assistant 
-      ;(graph/add-node "execute-query" execute-query)             ; tool
+      (graph/add-node "query-gen" query-gen)
+      (graph/add-edge "model-get-schema" "query-gen")
 
-      ;(graph/add-edge "list-table-sub-graph" "model-get-schema")
-      ;(graph/add-edge "model-get-schema" "query-gen")
-      ;(graph/add-conditional-edges "query-gen" should-continue)
+      (graph/add-node "end" graph/end)
+      (graph/add-node "correct-query" (graph/sub-graph-node {:init-state seed-correct-query-conversation
+                                                             :construct-graph graph/one-tool-call
+                                                             :next-state graph/append-new-messages}))
+      (graph/add-conditional-edges "query-gen" should-continue)
 
-      ;(graph/add-edge "correct-query" "execute-query")
-      ;(graph/add-edge "execute-query" "query-gen")
-      ))
+      (graph/add-edge "correct-query" "query-gen")))

src/graphs/sql.clj

@@ -16,7 +16,10 @@
 
 (defn summarize [state]
   (-> state
-      (update :messages (each summarize-content summarize-tool-calls))))
+      (update :messages (each 
+                          ;summarize-content 
+                          ;summarize-tool-calls
+                          ))))
 
 (defn prompt? [m]
   (= "prompt" (-> m :function :type)))