graphs/middleware/__init__.py

"""Factories for LangChain agent middlewares used across the V2 graph."""
from __future__ import annotations

from typing import Any

try:
    from langchain.agents.middleware import ClearToolUsesEdit, ContextEditingMiddleware
    _MIDDLEWARES_AVAILABLE = True
except Exception:  # pragma: no cover - compatibility fallback for older envs.
    ClearToolUsesEdit = None  # type: ignore[assignment]
    ContextEditingMiddleware = None  # type: ignore[assignment]
    _MIDDLEWARES_AVAILABLE = False

try:
    from langchain.agents.middleware import SummarizationMiddleware
    _SUMMARIZATION_AVAILABLE = False
except Exception:  # pragma: no cover - compatibility fallback for older envs.
    SummarizationMiddleware = None  # type: ignore[assignment]
    _SUMMARIZATION_AVAILABLE = False

# Defaults for the main conversational agent.
# Tuned for "fast TTFT" on Mistral large with paginated project tool.
_MAIN_TRIGGER_TOKENS = 4_000
_MAIN_CLEAR_AT_LEAST = 1_500
_MAIN_KEEP_TOOL_RESULTS = 2
_MAIN_EXCLUDE_TOOLS: tuple[str, ...] = (
    "check_project_id",
    "generate_summary_pdf",
)

# Defaults for the project summarizer sub-agent: paginates more aggressively,
# so we trigger later and keep more recent pages.
_SUBAGENT_TRIGGER_TOKENS = 8_000
_SUBAGENT_CLEAR_AT_LEAST = 2_500
_SUBAGENT_KEEP_TOOL_RESULTS = 3
_SUBAGENT_EXCLUDE_TOOLS: tuple[str, ...] = ("check_project_id",)

_MAIN_PLACEHOLDER = "[chunks projet precedents - deja exploites dans la reponse]"
_SUBAGENT_PLACEHOLDER = "[page precedente du projet - deja resumee]"
_SUMMARY_MODEL = "mistral-small-latest"
_MAIN_SUMMARY_TRIGGER_TOKENS = 7_000
_MAIN_SUMMARY_KEEP_MESSAGES = 12
_SUBAGENT_SUMMARY_TRIGGER_TOKENS = 5_000
_SUBAGENT_SUMMARY_KEEP_MESSAGES = 8


def build_main_agent_middlewares() -> list[Any]:
    """Middlewares for the main conversational agent."""
    if not _MIDDLEWARES_AVAILABLE:
        return []
    middlewares: list[Any] = [
        ContextEditingMiddleware(
            edits=[
                ClearToolUsesEdit(
                    trigger=_MAIN_TRIGGER_TOKENS,
                    clear_at_least=_MAIN_CLEAR_AT_LEAST,
                    keep=_MAIN_KEEP_TOOL_RESULTS,
                    clear_tool_inputs=False,
                    exclude_tools=_MAIN_EXCLUDE_TOOLS,
                    placeholder=_MAIN_PLACEHOLDER,
                ),
            ],
            token_count_method="approximate",
        ),
    ]
    if _SUMMARIZATION_AVAILABLE:
        middlewares.append(
            SummarizationMiddleware(
                model=_SUMMARY_MODEL,
                trigger=("tokens", _MAIN_SUMMARY_TRIGGER_TOKENS),
                keep=("messages", _MAIN_SUMMARY_KEEP_MESSAGES),
            )
        )
    return middlewares


def build_project_subagent_middlewares() -> list[Any]:
    """Middlewares for the project summarizer sub-agent (paginates more)."""
    if not _MIDDLEWARES_AVAILABLE:
        return []
    middlewares: list[Any] = [
        ContextEditingMiddleware(
            edits=[
                ClearToolUsesEdit(
                    trigger=_SUBAGENT_TRIGGER_TOKENS,
                    clear_at_least=_SUBAGENT_CLEAR_AT_LEAST,
                    keep=_SUBAGENT_KEEP_TOOL_RESULTS,
                    clear_tool_inputs=False,
                    exclude_tools=_SUBAGENT_EXCLUDE_TOOLS,
                    placeholder=_SUBAGENT_PLACEHOLDER,
                ),
            ],
            token_count_method="approximate",
        ),
    ]
    if _SUMMARIZATION_AVAILABLE:
        middlewares.append(
            SummarizationMiddleware(
                model=_SUMMARY_MODEL,
                trigger=("tokens", _SUBAGENT_SUMMARY_TRIGGER_TOKENS),
                keep=("messages", _SUBAGENT_SUMMARY_KEEP_MESSAGES),
            )
        )
    return middlewares


__all__ = [
    "_MIDDLEWARES_AVAILABLE",
    "_SUMMARIZATION_AVAILABLE",
    "ClearToolUsesEdit",
    "ContextEditingMiddleware",
    "SummarizationMiddleware",
    "build_main_agent_middlewares",
    "build_project_subagent_middlewares",
]