Spaces:
Sleeping
Sleeping
File size: 6,954 Bytes
c509967 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 | from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import ToolMessage
import json
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from chemgraph.tools.generic_tools import repl_tool
from chemgraph.tools.generic_tools import calculator
from chemgraph.prompt.single_agent_prompt import single_agent_prompt
from chemgraph.utils.logging_config import setup_logger
logger = setup_logger(__name__)
class State(TypedDict):
"""Type definition for the state dictionary used in the graph.
Attributes
----------
messages : list
List of messages in the conversation, annotated with add_messages
"""
messages: Annotated[list, add_messages]
class BasicToolNode:
"""A node that executes tools requested in the last AIMessage.
This class processes tool calls from AI messages and executes the corresponding
tools, handling their results and any potential errors.
Parameters
----------
tools : list
List of tool objects that can be called by the node
Attributes
----------
tools_by_name : dict
Dictionary mapping tool names to their corresponding tool objects
"""
def __init__(self, tools: list) -> None:
"""Initialize the tool node.
Parameters
----------
tools : list
Tool objects keyed by their ``name`` attribute.
"""
self.tools_by_name = {tool.name: tool for tool in tools}
def __call__(self, inputs: State) -> State:
"""Execute tools requested in the last message.
Parameters
----------
inputs : State
The current state containing messages
Returns
-------
State
Updated state containing tool execution results
Raises
------
ValueError
If no message is found in the input state
"""
if messages := inputs.get("messages", []):
message = messages[-1]
else:
raise ValueError("No message found in input")
outputs = []
for tool_call in message.tool_calls:
try:
tool_name = tool_call.get("name")
if not tool_name or tool_name not in self.tools_by_name:
raise ValueError(f"Invalid tool name: {tool_name}")
tool_result = self.tools_by_name[tool_name].invoke(tool_call.get("args", {}))
# Handle different types of tool results
result_content = (
tool_result.dict()
if hasattr(tool_result, "dict")
else (tool_result if isinstance(tool_result, dict) else str(tool_result))
)
outputs.append(
ToolMessage(
content=json.dumps(result_content),
name=tool_name,
tool_call_id=tool_call.get("id", ""),
)
)
except Exception as e:
outputs.append(
ToolMessage(
content=json.dumps({"error": str(e)}),
name=tool_name if tool_name else "unknown_tool",
tool_call_id=tool_call.get("id", ""),
)
)
return {"messages": outputs}
def route_tools(state: State):
"""Route to the 'tools' node if the last message has tool calls; otherwise, route to END.
Parameters
----------
state : State
The current state containing messages
Returns
-------
str
Either 'tools' or END based on the presence of tool calls
Raises
------
ValueError
If no messages are found in the input state
"""
if isinstance(state, list):
ai_message = state[-1]
elif messages := state.get("messages", []):
ai_message = messages[-1]
else:
raise ValueError(f"No messages found in input state to tool_edge: {state}")
if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
return "tools"
return END
def CompChemAgent(state: State, llm: ChatOpenAI, system_prompt=single_agent_prompt, tools=None):
"""LLM node that processes messages and decides next actions.
Parameters
----------
state : State
The current state containing messages
llm : ChatOpenAI
The language model to use for processing
system_prompt : str, optional
The system prompt to guide the LLM's behavior,
by default single_agent_prompt
tools : list, optional
List of tools available to the agent, by default None
Returns
-------
dict
Updated state containing the LLM's response
"""
if tools is None:
tools = []
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": f"{state['messages']}"},
]
llm_with_tools = llm.bind_tools(tools=tools)
return {"messages": [llm_with_tools.invoke(messages)]}
def construct_relp_graph(llm: ChatOpenAI, system_prompt=single_agent_prompt):
"""Construct a graph for REPL-based Python execution workflow.
This function creates a state graph that implements a workflow for executing
Python code through a REPL interface, using LLM agents and tools.
Parameters
----------
llm : ChatOpenAI
The language model to use in the workflow
system_prompt : str, optional
The system prompt to guide the LLM's behavior,
by default single_agent_prompt
Returns
-------
StateGraph
A compiled state graph implementing the REPL workflow
Raises
------
Exception
If there is an error during graph construction
"""
try:
logger.info("Constructing geometry optimization graph")
checkpointer = MemorySaver()
tools = [
repl_tool,
calculator,
]
tool_node = BasicToolNode(tools=tools)
graph_builder = StateGraph(State)
graph_builder.add_node(
"CompChemAgent",
lambda state: CompChemAgent(state, llm, system_prompt=system_prompt, tools=tools),
)
graph_builder.add_node("tools", tool_node)
graph_builder.add_conditional_edges(
"CompChemAgent",
route_tools,
{"tools": "tools", END: END},
)
graph_builder.add_edge("tools", "CompChemAgent")
graph_builder.add_edge(START, "CompChemAgent")
graph = graph_builder.compile(checkpointer=checkpointer)
logger.info("Graph construction completed")
return graph
except Exception as e:
logger.error(f"Error constructing graph: {str(e)}")
raise
|