From 5f8f17cac36b3634e325198f5ba4990c8b375642 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Sat, 6 Jan 2024 14:42:28 -0800
Subject: [PATCH] Improve run names

---
 examples/langgraph.ipynb        | 20 ++------------------
 permchain/langgraph/__init__.py | 19 +++++++++++--------
 permchain/pregel/read.py        |  1 +
 permchain/pregel/write.py       |  1 +
 4 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/examples/langgraph.ipynb b/examples/langgraph.ipynb
index b502b1e33..b179a5fda 100644
--- a/examples/langgraph.ipynb
+++ b/examples/langgraph.ipynb
@@ -13,30 +13,14 @@
    "execution_count": 1,
    "id": "d642e6af-217a-4414-a78c-509b44155eca",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "input_variables=['agent_scratchpad', 'input'] input_types={'chat_history': typing.List[typing.Union[langchain_core.messages.ai.AIMessage, langchain_core.messages.human.HumanMessage, langchain_core.messages.chat.ChatMessage, langchain_core.messages.system.SystemMessage, langchain_core.messages.function.FunctionMessage, langchain_core.messages.tool.ToolMessage]], 'agent_scratchpad': typing.List[typing.Union[langchain_core.messages.ai.AIMessage, langchain_core.messages.human.HumanMessage, langchain_core.messages.chat.ChatMessage, langchain_core.messages.system.SystemMessage, langchain_core.messages.function.FunctionMessage, langchain_core.messages.tool.ToolMessage]]} messages=[SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=[], template='You are a helpful assistant')), MessagesPlaceholder(variable_name='chat_history', optional=True), HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], template='{input}')), MessagesPlaceholder(variable_name='agent_scratchpad')]\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/nuno/dev/langchain/libs/core/langchain_core/_api/deprecation.py:191: LangChainDeprecationWarning: The class `ChatOpenAI` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use langchain_openai.ChatOpenAI instead.\n",
-      "  warn_deprecated(\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from langchain.chat_models import ChatOpenAI\n",
     "from langchain import hub\n",
     "from langchain.agents import create_openai_functions_agent\n",
     "from langchain_community.chat_models import ChatOpenAI\n",
     "from langchain_community.tools.tavily_search import TavilySearchResults\n",
-    "from langchain_core.runnables import RunnablePassthrough, RunnableLambda\n",
+    "from langchain_core.runnables import RunnablePassthrough\n",
     "from permchain.langgraph import Graph, END\n",
     "\n",
     "tools = [TavilySearchResults(max_results=1)]\n",
diff --git a/permchain/langgraph/__init__.py b/permchain/langgraph/__init__.py
index 6f97bf491..e082ee8ca 100644
--- a/permchain/langgraph/__init__.py
+++ b/permchain/langgraph/__init__.py
@@ -3,7 +3,11 @@ from collections import defaultdict
 from typing import Any, Callable, Dict, NamedTuple
 
 from langchain_core.runnables import Runnable
-from langchain_core.runnables.base import RunnableLike, coerce_to_runnable
+from langchain_core.runnables.base import (
+    RunnableLambda,
+    RunnableLike,
+    coerce_to_runnable,
+)
 
 from permchain.pregel import Channel, Pregel
 
@@ -111,17 +115,16 @@ class Graph:
             outgoing_edges[self.finish_point].append(END)
 
         nodes = {
-            key: (
-                Channel.subscribe_to(key)
-                | node
-                | Channel.write_to(*outgoing_edges[key])
-            )
-            for key, node in self.nodes.items()
+            key: Channel.subscribe_to(key) | node for key, node in self.nodes.items()
         }
 
+        for key, edges in outgoing_edges.items():
+            if edges:
+                nodes[key] |= Channel.write_to(*edges)
+
         for key, branches in self.branches.items():
             for branch in branches:
-                nodes[key] |= branch.runnable
+                nodes[key] |= RunnableLambda(branch.runnable, name=f"{key}_condition")
 
         return Pregel(
             nodes=nodes,
diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py
index 1c9b1dc13..f5184f821 100644
--- a/permchain/pregel/read.py
+++ b/permchain/pregel/read.py
@@ -40,6 +40,7 @@ class ChannelRead(RunnableLambda):
     def __init__(self, channel: str) -> None:
         super().__init__(func=self._read, afunc=self._aread)
         self.channel = channel
+        self.name = f"ChannelRead<{channel}>"
 
     def _read(self, _: Any, config: RunnableConfig) -> Any:
         try:
diff --git a/permchain/pregel/write.py b/permchain/pregel/write.py
index 3c21d39c8..7135e5a76 100644
--- a/permchain/pregel/write.py
+++ b/permchain/pregel/write.py
@@ -30,6 +30,7 @@ class ChannelWrite(RunnablePassthrough):
         channels: Sequence[tuple[str, Runnable | None]],
     ):
         super().__init__(func=self._write, afunc=self._awrite, channels=channels)
+        self.name = f"ChannelWrite<{','.join(chan for chan, _ in self.channels)}>"
 
     @property
     def config_specs(self) -> list[ConfigurableFieldSpec]: