From e465be5941f7820a72085cd55fb7ad94ddd2b93f Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Mon, 27 Mar 2023 21:46:44 -0300 Subject: [PATCH] feat: agents as tools working --- src/backend/langflow/api/endpoints.py | 4 +- src/backend/langflow/interface/run.py | 49 ++++++++++ src/backend/langflow/utils/graph.py | 130 +++++++++++++++++++++++++- 3 files changed, 179 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/api/endpoints.py b/src/backend/langflow/api/endpoints.py index 7214ec6a6..7211eb8d6 100644 --- a/src/backend/langflow/api/endpoints.py +++ b/src/backend/langflow/api/endpoints.py @@ -2,7 +2,7 @@ from typing import Any, Dict from fastapi import APIRouter, HTTPException -from langflow.interface.run import process_data_graph +from langflow.interface.run import process_graph from langflow.interface.types import build_langchain_types_dict # build router @@ -17,6 +17,6 @@ def get_all(): @router.post("/predict") def get_load(data: Dict[str, Any]): try: - return process_data_graph(data) + return process_graph(data) except Exception as e: return HTTPException(status_code=500, detail=str(e)) diff --git a/src/backend/langflow/interface/run.py b/src/backend/langflow/interface/run.py index 0844e2073..469c64807 100644 --- a/src/backend/langflow/interface/run.py +++ b/src/backend/langflow/interface/run.py @@ -4,6 +4,8 @@ import re from typing import Any, Dict from langflow.interface import loading +from langflow.utils import payload +from langflow.utils.graph import Graph def process_data_graph(data_graph: Dict[str, Any]): @@ -27,6 +29,52 @@ def process_data_graph(data_graph: Dict[str, Any]): } +def process_graph(data_graph: Dict[str, Any]): + """ + Process graph by extracting input variables and replacing ZeroShotPrompt + with PromptTemplate,then run the graph and return the result and thought. + """ + nodes = data_graph["nodes"] + # Add input variables + # ? Is this necessary? + nodes = payload.extract_input_variables(nodes) + # Nodes, edges and root node + edges = data_graph["edges"] + graph = Graph(nodes, edges) + + langchain_object = graph.build() + message = data_graph["message"] + # Process json + result, thought = get_result_and_thought_using_graph(langchain_object, message) + + return { + "result": result, + "thought": re.sub( + r"\x1b\[([0-9,A-Z]{1,2}(;[0-9,A-Z]{1,2})?)?[m|K]", "", thought + ).strip(), + } + + +def get_result_and_thought_using_graph(loaded_langchain, message: str): + """Get result and thought from extracted json""" + loaded_langchain.verbose = True + try: + with io.StringIO() as output_buffer, contextlib.redirect_stdout(output_buffer): + result = loaded_langchain(message) + + result = ( + result.get(loaded_langchain.output_keys[0]) + if isinstance(result, dict) + else result + ) + thought = output_buffer.getvalue() + + except Exception as e: + result = f"Error: {str(e)}" + thought = "" + return result, thought + + def get_result_and_thought(extracted_json: Dict[str, Any], message: str): """Get result and thought from extracted json""" try: @@ -41,6 +89,7 @@ def get_result_and_thought(extracted_json: Dict[str, Any], message: str): else result ) thought = output_buffer.getvalue() + except Exception as e: result = f"Error: {str(e)}" thought = "" diff --git a/src/backend/langflow/utils/graph.py b/src/backend/langflow/utils/graph.py index 59fe6cffd..2f8729538 100644 --- a/src/backend/langflow/utils/graph.py +++ b/src/backend/langflow/utils/graph.py @@ -1,6 +1,12 @@ +# Description: Graph class for building a graph of nodes and edges +# Insights: +# - Defer prompts building to the last moment or when they have all the tools +# - Build each inner agent first, then build the outer agent +from langchain.agents.load_tools import get_all_tool_names + from copy import deepcopy import types -from typing import Any, Dict, List, Union +from typing import Any, Dict, List, Optional, Union from langflow.interface import loading from langflow.utils import payload, util from langflow.interface.listing import ALL_TYPES_DICT @@ -166,6 +172,35 @@ class Node: return id(self) +class AgentNode(Node): + def __init__(self, data: Dict): + super().__init__(data) + self.tools: List[ToolNode] = [] + self.chains: List[ChainNode] = [] + + def _set_tools_and_chains(self) -> None: + for edge in self.edges: + source_node = edge.source + if isinstance(source_node, ToolNode): + self.tools.append(source_node) + elif isinstance(source_node, ChainNode): + self.chains.append(source_node) + + def build(self, force: bool = False) -> Any: + if not self._built or force: + self._set_tools_and_chains() + # First, build the tools + for tool_node in self.tools: + tool_node.build() + + # Next, build the chains and the rest + for chain_node in self.chains: + chain_node.build(tools=self.tools) + + self._build() + return deepcopy(self._built_object) + + class Edge: def __init__(self, source: "Node", target: "Node"): self.source: "Node" = source @@ -204,6 +239,51 @@ class Edge: ) +class PromptNode(Node): + def __init__(self, data: Dict): + super().__init__(data) + + def build(self, tools: Optional[List[Node]] = None, force: bool = False) -> Any: + if not self._built or force: + # Check if it is a ZeroShotPrompt and needs a tool + if self.node_type == "ZeroShotPrompt": + tools = ( + [tool_node.build() for tool_node in tools] + if tools is not None + else [] + ) + self.params["tools"] = tools + + self._build() + return deepcopy(self._built_object) + + +class ToolNode(Node): + def __init__(self, data: Dict): + super().__init__(data) + + def build(self, force: bool = False) -> Any: + if not self._built or force: + self._build() + return deepcopy(self._built_object) + + +class ChainNode(Node): + def __init__(self, data: Dict): + super().__init__(data) + + def build(self, tools: Optional[List[Node]] = None, force: bool = False) -> Any: + if not self._built or force: + # Check if the chain requires a PromptNode + for key, value in self.params.items(): + if isinstance(value, PromptNode): + # Build the PromptNode, passing the tools if available + self.params[key] = value.build(tools=tools or [], force=force) + + self._build() + return deepcopy(self._built_object) + + class Graph: def __init__( self, @@ -270,7 +350,24 @@ class Graph: return edges def _build_nodes(self) -> List[Node]: - return [Node(node) for node in self._nodes] + nodes = [] + all_tool_names = set(get_all_tool_names()) + for node in self._nodes: + node_data = node["data"] + node_type = node_data["type"] + node_lc_type = node_data["node"]["template"]["_type"] + + if node_type in ["ZeroShotPrompt", "PromptTemplate"]: + nodes.append(PromptNode(node)) + elif "agent" in node_type.lower(): + nodes.append(AgentNode(node)) + elif "chain" in node_type.lower(): + nodes.append(ChainNode(node)) + elif "tool" in node_type.lower() or node_lc_type in all_tool_names: + nodes.append(ToolNode(node)) + else: + nodes.append(Node(node)) + return nodes def get_children_by_node_type(self, node: Node, node_type: str) -> List[Node]: children = [] @@ -280,3 +377,32 @@ class Graph: if node_type in node_types: children.append(node) return children + + def _build_agent(self, agent_node: Node) -> None: + # Identify the ZeroShotPrompt node and any inner ZeroShotAgent nodes + zero_shot_prompt_node = None + inner_agent_nodes = [] + for edge in agent_node.edges: + if edge.source == agent_node: + source_node = edge.target + if ( + isinstance(source_node, DeferredNode) + and source_node.node_type == "ZeroShotPrompt" + ): + zero_shot_prompt_node = source_node + elif source_node.node_type == "ZeroShotAgent": + inner_agent_nodes.append(source_node) + + # First, build any inner ZeroShotAgent nodes + for inner_agent_node in inner_agent_nodes: + self._build_agent(inner_agent_node) + + # Build the ZeroShotAgent node itself + agent_built = agent_node.build() + + if zero_shot_prompt_node: + # Set the tools parameter in the ZeroShotPrompt node + zero_shot_prompt_node.params["tools"] = agent_built.tools + + # Build the ZeroShotPrompt node + zero_shot_prompt_node.build()