agentWrapper.py

from sqlAgent import build_sql_agent
from chatAgent import build_graph,ChatState
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

import os
import logging
logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
logger = logging.getLogger(__name__)


_FALLBACK = (
    "I'm sorry, I wasn't able to retrieve that information right now. "
    "Please try again or contact our support team for immediate assistance."
)


class OrderChatbot:
    """
    High-level wrapper around the LangGraph pipeline.
    Each customer session keeps a unique thread_id so MemorySaver
    maintains per-customer conversation history automatically.
    """

    def __init__(self):
        # init_database(db_path)
        agent       = build_sql_agent()
        self._graph = build_graph(agent)
        self._sessions: dict[str, str] = {}   # cust_id → thread_id
        logger.info("OrderChatbot ready ✓")


    def chat(self, cust_id: str, user_message: str,session_id :str) -> dict:
        """
        Send a message and receive a structured response.

        Returns:
            {
                "response":   str,           # bot reply
                "guard":      str,           # SAFE | BLOCKED | ESCALATE
                "sql":        str | None,    # SQL that was executed
                "escalated":  bool,
                "history":    list[dict],    # full conversation so far
            }
        """
        # thread_id = self.start_session(cust_id)
        thread_id = session_id
        self._sessions[cust_id] = thread_id
        config    = {"configurable": {"thread_id": thread_id}}

        # Retrieve current state to build initial messages list
        current   = self._graph.get_state(config)
        prior_msgs = current.values.get("messages", []) if current.values else []

        # Append the new human message
        new_messages = list(prior_msgs) + [HumanMessage(content=user_message)]

        input_state: ChatState = {
            "messages":          new_messages,
            "session_id":        thread_id,
            "cust_id":           cust_id,
            "last_guard_status": "SAFE",
            "last_sql":          None,
            "escalated":         False,
        }

        # Run the graph
        output = self._graph.invoke(input_state, config=config)

        # Extract the latest AI message
        ai_msgs = [m for m in output["messages"] if isinstance(m, AIMessage)]
        response_text = ai_msgs[-1].content if ai_msgs else _FALLBACK

        # Build a human-readable history
        history = []
        for m in output["messages"]:
            if isinstance(m, HumanMessage):
                history.append({"role": "user",      "content": m.content})
            elif isinstance(m, AIMessage):
                history.append({"role": "assistant",  "content": m.content})

        return {
            "response":  response_text,
            "guard":     output.get("last_guard_status", "SAFE"),
            "sql":       output.get("last_sql"),
            "escalated": output.get("escalated", False),
            "history":   history,
        }

    def get_history(self, cust_id: str) -> list[dict]:
        """Return the full conversation history for a customer."""
        thread_id = self._sessions.get(cust_id)
        if not thread_id:
            return []
        config  = {"configurable": {"thread_id": thread_id}}
        state   = self._graph.get_state(config)
        if not state or not state.values:
            return []
        history = []
        for m in state.values.get("messages", []):
            if isinstance(m, HumanMessage):
                history.append({"role": "user",      "content": m.content})
            elif isinstance(m, AIMessage):
                history.append({"role": "assistant",  "content": m.content})
        return history

    def clear_session(self, cust_id: str) -> None:
      """Remove the customer's session (forces a fresh conversation)."""
      self._sessions.pop(cust_id, None)
      logger.info("Session cleared for cust_id=%s", cust_id)