graphs/agents/summarizer_agent.py

"""Summarizer agents: split into two nodes (LLM -> Export PDF/Upload)."""
from io import BytesIO
from typing import Callable, Union
from datetime import datetime, timezone

from langchain_core.messages import SystemMessage
from langchain_core.language_models.chat_models import BaseChatModel

from graphs.state import AgentState
from graphs.prompts import SUMMARIZER_SYSTEM_PROMPT

from langchain_core.documents import Document
from langchain_core.messages import BaseMessage, SystemMessage, AIMessage


def summarizer_llm_node(
    llm: BaseChatModel,
):
    """Node: ask the LLM to generate a Markdown summary from the conversation/context.

    Produces state keys:
    - summary_markdown
    - messages (appends the AIMessage containing the summary)
    """
    def _run(state: AgentState) -> AgentState:
        messages = list(state.get("messages", []))

        sys = SystemMessage(content=SUMMARIZER_SYSTEM_PROMPT)
        response = llm.invoke([sys] + messages)
        summary_markdown = response.content or ""


        return {
            "summary_markdown": summary_markdown,
            "messages": messages + [AIMessage(content=summary_markdown)],
        }

    return _run


def summarizer_export_node(
    markdown_to_pdf: Callable[[str], Union[bytes, BytesIO]],
    upload_pdf: Callable[[bytes, str], str],
):
    """Node: takes `summary_markdown` from state, builds PDF and uploads it.

    Produces state keys:
    - summary_pdf_path (URL or storage path)
    - messages (appends an AIMessage with a link/attachment metadata)
    """
    def _run(state: AgentState) -> AgentState:
        import os

        messages = list(state.get("messages", []))
        summary_markdown = state.get("summary_markdown", "") or ""

        # Convert Markdown to PDF
        pdf_result = markdown_to_pdf(summary_markdown)

        # Ensure pdf_bytes is of type bytes (not BytesIO)
        if hasattr(pdf_result, "read"):
            pdf_bytes = pdf_result.read()
        else:
            pdf_bytes = pdf_result

        now = datetime.now(timezone.utc)
        # Build a project-relative path: summaries/YYYY/MM/summary-<timestamp>.pdf
        filename = f"summary-{now.strftime('%Y%m%d-%H%M%S')}.pdf"
        dir_path = os.path.join("tmp_summaries")
        if not os.path.exists(dir_path):
            os.makedirs(dir_path, exist_ok=True)
        object_path = os.path.join(dir_path, filename)

        # Always write locally first so the file exists even if upload fails
        with open(object_path, "wb") as f:
            f.write(pdf_bytes)

        # Try remote upload, but do not fail the node if upload fails
        uploaded_path = None
        try:
            uploaded_path = upload_pdf(object_path, filename)
        except Exception as e:
            # Swallow upload error but keep local file
            import logging
            logging.exception("Upload of summary PDF failed")
            uploaded_path = None

        

        metadata = {
            "file_name": filename,
            "type": "conversation_summary_pdf",
            "generated_at": now.isoformat(),
            "link": uploaded_path if uploaded_path else "",
        }

        return {
            "summary_pdf_path": uploaded_path if uploaded_path else object_path,
            "messages": messages + [AIMessage(content=f"Synthèse prête: {uploaded_path or object_path}", metadata={"document": metadata})],
        }

    return _run


def summarizer_node(
    llm: BaseChatModel,
    markdown_to_pdf: Callable[[str], Union[bytes, BytesIO]],
    upload_pdf: Callable[[bytes, str], str],
):
    """Compatibility wrapper: runs LLM summarization then export in one node."""
    llm_runner = summarizer_llm_node(llm)
    export_runner = summarizer_export_node(markdown_to_pdf, upload_pdf)

    def _run(state: AgentState) -> AgentState:
        state_after_llm = llm_runner(state)
        state_after_export = export_runner(state_after_llm)
        return state_after_export

    return _run