graphs/workflows/agent.py

from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.language_models.chat_models import BaseChatModel
from graphs.state import CustomState
from graphs.prompts_v2 import load_v2_prompt
from graphs.middleware import build_main_agent_middlewares
from datetime import datetime

from graphs.tools.retrieval_tools import (
    search_formations,
    search_prestations,
    search_project_docs,
)
from graphs.tools.check_project import check_project_id
from graphs.tools.summarize import build_generate_summary_pdf

# Prompts
chat_prompt_v2_template = load_v2_prompt("chat_system.md")
today_date = datetime.now().strftime("%d/%m/%Y")
system_prompt = chat_prompt_v2_template.replace("{TODAY_DATE}", today_date)


def create_agent_graph(llm: BaseChatModel, checkpointer=None):
    tools = [
        search_formations,
        search_prestations,
        search_project_docs,
        check_project_id,
        build_generate_summary_pdf(llm),
    ]

    # Agent
    agent = create_agent(
        model="mistral-small-2603",
        tools=tools,
        system_prompt=system_prompt,
        state_schema=CustomState,
        checkpointer=checkpointer or InMemorySaver(),
        middleware=build_main_agent_middlewares(),
    )

    return agent