New graph representation

docker · Nov 13, 2024 · def0e7f · def0e7f
1 parent 5a5002a
commit def0e7f
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 24 deletions.
diff --git a/deps-lock.json b/deps-lock.json
@@ -930,6 +930,16 @@
       "mvn-repo": "https://repo.maven.apache.org/maven2/",
       "hash": "sha256-OeNB9nv+85PkeDkNSYjxGad5ykSQZssNM/gLQv8E9D0="
     },
+    {
+      "mvn-path": "org/clojure/core.match/1.1.0/core.match-1.1.0.jar",
+      "mvn-repo": "https://repo.maven.apache.org/maven2/",
+      "hash": "sha256-10V6tjEIWae9cTmEM4IEX6PN7A0T97qSEpfy8/uZj1M="
+    },
+    {
+      "mvn-path": "org/clojure/core.match/1.1.0/core.match-1.1.0.pom",
+      "mvn-repo": "https://repo.maven.apache.org/maven2/",
+      "hash": "sha256-NnHYN2UlIwq6Ah8fYmx54g86ELYrXfgXIiWJDsSv4EU="
+    },
     {
       "mvn-path": "org/clojure/core.memoize/1.0.253/core.memoize-1.0.253.jar",
       "mvn-repo": "https://repo.maven.apache.org/maven2/",

diff --git a/deps.edn b/deps.edn
@@ -1,5 +1,6 @@
 {:paths ["src" "dev"]
  :deps {org.clojure/clojure {:mvn/version "1.11.4"}
+        org.clojure/core.match {:mvn/version "1.1.0"}
         markdown-clj/markdown-clj {:mvn/version "1.12.1"}
         pogonos/pogonos {:mvn/version "0.2.1"}
         dev.weavejester/medley {:mvn/version "1.8.0"}

diff --git a/src/graph.clj b/src/graph.clj
@@ -3,6 +3,7 @@
    [babashka.fs :as fs]
    [clojure.core.async :as async]
    [clojure.pprint :refer [pprint]]
+   [clojure.core.match :refer [match]]
    git
    jsonrpc
    openai
@@ -141,14 +142,14 @@
 
 (declare stream chat-with-tools)
 
-(defn add-last-message-as-tool-call 
+(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 
+(defn append-new-messages
   [state sub-graph-state]
   {:messages (->> (:messages sub-graph-state)
                   (filter (complement (fn [m] (some #(= m %) (:messages state))))))})
@@ -220,6 +221,36 @@
          ;; TODO handling missing edges
          (recur new-state ((get-in graph [:edges node]) new-state)))))))
 
+(defn update-graph [m l x]
+  (match [x]
+
+         ; create an edge from the last node to this one 
+    [[(n1 :guard string?)]] (cond-> m
+                              true (add-edge l n1))
+
+         ; add a conditional edge from the last node to this function 
+    [[:edge (n2 :guard (comp not string?))]] (cond-> m
+                                               l (add-conditional-edges
+                                                  l
+                                                  n2))
+
+         ; 
+    [[n1 (n2 :guard (comp not string?))]] (cond-> m
+                                            true (add-node n1 n2)
+                                            l (add-edge l n1))
+
+    :else m))
+
+(defn path-item [{:keys [m l]} [n1 :as item]]
+  {:m (update-graph m l item)
+   :l (when (not (= :edge n1)) n1)})
+
+(defn paths [agg p]
+  (:m (reduce path-item {:m agg} p)))
+
+(defn construct-graph [x]
+  (reduce paths {} x))
+
 ; ============================================================
 ; this is the graph we tend to use in our experiments thus far
 ; ============================================================
@@ -238,6 +269,17 @@
       (add-edge "tools-query" "completion")
       (add-conditional-edges "completion" tool-or-end)))
 
+(def chat-with-tools-representation
+  [[["start" start]
+    ["tools-query" tools-query]
+    ["completion" completion]
+    [:edge tool-or-end]]
+   [["sub-graph" (sub-graph-node nil)]
+    ["tools-query"]]
+   [["tool" (tool-node nil)]
+    ["tools-query"]]
+   [["end" end]]])
+
 (defn one-tool-call [_]
   (-> {}
       (add-node "start" start)
@@ -250,6 +292,16 @@
       (add-edge "tool" "end")
       (add-conditional-edges "completion" tool-or-end)))
 
+(def one-tool-call-representation
+  [[["start" start]
+    ["completion" completion]
+    [:edge tool-or-end]]
+   [["sub-graph" (sub-graph-node nil)]
+    ["completion"]]
+   [["tool" (tool-node nil)]
+    ["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")

diff --git a/src/graphs/sql.clj b/src/graphs/sql.clj
@@ -106,26 +106,19 @@
       (update-in [:messages] concat [(last (:messages state))])))
 
 (defn graph [_]
-  (-> {}
-      (graph/add-node "start" graph/start)
-      (graph/add-node "list-tables-inject-tool" list-tables-inject-tool)
-      (graph/add-edge "start" "list-tables-inject-tool")
+  (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})]
+     ["query-gen"]]
+    [["end" graph/end]]]))
 
-      (graph/add-node "list-tables-tool" (graph/tool-node nil))
-      (graph/add-edge "list-tables-inject-tool" "list-tables-tool")
-
-      ; 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 "query-gen" query-gen)
-      (graph/add-edge "model-get-schema" "query-gen")
-
-      (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" "query-gen")))
diff --git a/test/graph_t.clj b/test/graph_t.clj
@@ -0,0 +1,8 @@
+(ns graph_t
+  (:require [clojure.test :as t]
+            [graph]
+            [graphs.sql]))
+
+(t/deftest
+  (graph/construct-graph graph/chat-with-tools-representation)
+  (graph/construct-graph graphs.sql/graph-data))