81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
from langchain_mistralai import ChatMistralAI
|
|
from langgraph.graph import MessagesState
|
|
from langgraph.prebuilt import ToolNode
|
|
from langchain.chat_models import init_chat_model
|
|
from langgraph.graph import START, END
|
|
|
|
from .tools import getTools
|
|
|
|
# LLM principal
|
|
llm = ChatMistralAI( # LLM sans outils
|
|
model="mistral-large-latest",
|
|
temperature=0,
|
|
max_retries=2,
|
|
)
|
|
|
|
# NODES
|
|
def call_to_LLM(state: MessagesState):
|
|
"""Noeud qui s'occupe de gérer les appels au LLM"""
|
|
# Initialisation du LLM
|
|
model = llm.bind_tools(getTools())
|
|
|
|
# Appel du LLM
|
|
return {"messages": [model.invoke(state["messages"])]}
|
|
|
|
# fonction de routage : Après reponse_question, si le LLM veut appeler un outil, on va au tool_node
|
|
def should_continue(state: MessagesState):
|
|
"""
|
|
Vérifier s'il y a un appel aux outils dans le dernier message
|
|
"""
|
|
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 "no_tools"
|
|
|
|
def task_ended(state: MessagesState):
|
|
"""
|
|
Vérifier si l'agent a terminé son cycle, ou s'il faut le relancer
|
|
"""
|
|
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 "terminé" in ai_message.content.lower():
|
|
return END
|
|
return "continue"
|
|
|
|
|
|
class BasicToolNode: # De mon ancien projet, https://github.com/LJ5O/Assistant/blob/main/modules/Brain/src/LLM/graph/nodes/BasicToolNode.py
|
|
"""A node that runs the tools requested in the last AIMessage."""
|
|
|
|
def __init__(self, tools: list) -> None:
|
|
self.tools_by_name = {tool.name: tool for tool in tools}
|
|
|
|
def __call__(self, inputs: dict):
|
|
if messages := inputs.get("messages", []):
|
|
message = messages[-1]
|
|
else:
|
|
raise ValueError("No message found in input")
|
|
outputs = []
|
|
for tool_call in message.tool_calls:
|
|
#print(tool_call["args"])
|
|
tool_result = self.tools_by_name[tool_call["name"]].invoke(
|
|
tool_call["args"]
|
|
)
|
|
outputs.append(
|
|
ToolMessage(
|
|
content=json.dumps(tool_result),
|
|
name=tool_call["name"],
|
|
tool_call_id=tool_call["id"],
|
|
)
|
|
)
|
|
return {"messages": outputs} |