from typing import TypedDict, Annotated, List, Union
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langchain.agents import create_react_agent
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langgraph.prebuilt.tool_node import ToolNode
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
# --------------------------
# 1. Définition du Tool
# --------------------------
@tool
def smart_get_info(app: str) -> str:
"""Fournit des informations sur une application."""
return f"L'application {app} est gérée par l'équipe IT Sécurité."
tools = [smart_get_info]
# --------------------------
# 2. Prompt système spécial Mistral (non agent-aware)
# --------------------------
prompt = ChatPromptTemplate.from_messages([
("system", """Tu es un agent intelligent pour la gestion des applications.
Tu dois raisonner étape par étape, et utiliser des outils quand c'est nécessaire pour répondre à des questions précises.
Utilise le format suivant :
Thought: décris ce que tu penses faire.
Action: le nom d’un outil à appeler (par exemple, smart_get_info)
Action Input: l’entrée à donner à cet outil.
Quand tu as obtenu l’information et que tu peux répondre, termine avec :
Thought: Je connais maintenant la réponse.
Final Answer: ta réponse finale à afficher à l’utilisateur.
Toujours terminer par Final Answer: quand tu veux mettre fin à l'échange."""),
MessagesPlaceholder(variable_name="messages")
])
# --------------------------
# 3. Initialisation du LLM (Mistral via vLLM ou Ollama)
# --------------------------
llm = ChatOpenAI(
base_url="http://localhost:11434/v1", # Ollama ou vLLM local
model_name="mistral",
api_key="ollama", # requis par LangChain mais ignoré par Ollama
temperature=0,
streaming=True
)
# --------------------------
# 4. Création de l'agent et du ToolNode
# --------------------------
agent_node = create_react_agent(llm, tools, prompt)
tool_node = ToolNode(tools=tools)
# --------------------------
# 5. Définition du state
# --------------------------
class AgentState(TypedDict):
messages: Annotated[List[Union[HumanMessage]], add_messages]
# --------------------------
# 6. Condition d'arrêt (quand Final Answer est généré)
# --------------------------
def should_continue(state: AgentState) -> str:
last_msg = state["messages"][-1].content
return "false" if "Final Answer:" in last_msg else "true"
# --------------------------
# 7. Construction du graphe LangGraph
# --------------------------
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
"agent",
condition=should_continue,
path_map={"true": "tools", "false": END}
)
workflow.add_edge("tools", "agent")
# Mémoire de session
memory = MemorySaver()
app = workflow.compile(checkpointer=memory)
# --------------------------
# 8. Exécution d'un cas de test avec thread_id obligatoire
# --------------------------
inputs = {"messages": [HumanMessage(content="Quelle équipe gère l'application AAA ?")]}
config = {"configurable": {"thread_id": "session-aaa-001"}}
print("---- Début du stream ----")
for step in app.stream(inputs, config=config, stream_mode="values"):
for msg in step["messages"]:
print(f"[{msg.type.upper()}] {msg.content}")
print("---- Fin du stream ----")
