import json from typing import Annotated, TypedDict, Sequence, Literal, Any from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage, ToolMessage from langgraph.graph.message import add_messages from langgraph.graph import StateGraph, START, END from langgraph.prebuilt import ToolNode from langchain_core.tools import tool from langchain_openai import ChatOpenAI from app.core.config import settings from app.core.prompts import BASE_SYSTEM_TEMPLATE from app.core.constants import ( MAX_TOKEN_LIMIT, HOTEL_RETRIEVE_DOCUMENTATION_DESCRIPTION, HUMAN_ESCALATION_DESCRIPTION, GET_RESERVATION_DETAILS_DESCRIPTION, GET_AVAILABLE_ADDONS_DESCRIPTION, ) from app.services.hotel_rag_service import HotelRAGService from app.models.chat import ( AgentRetrieverResponse, GetReservationDetailsArgs, GetAvailableAddonsArgs, HumanEscalationArgs, HotelAgentResponse, ) from app.utils.extra import get_reservation, get_available_addons as fetch_available_addons class AgentState(TypedDict): messages: Annotated[Sequence[BaseMessage], add_messages] class HotelLangGraphAgentService: MAX_TOKEN_LIMIT = MAX_TOKEN_LIMIT def __init__(self): self.reservation_details: dict | None = None self.llm = ChatOpenAI( model=settings.LLM_MODEL, temperature=0.5, openai_api_key=settings.OPENAI_API_KEY, streaming=False, ) self.tools = self._build_tools() self.llm_with_tools = self.llm.bind_tools( tools=self.tools, parallel_tool_calls=True, strict=True, response_format=HotelAgentResponse, ) self.graph = self._build_graph() def get_agent_prompt(self) -> list[BaseMessage]: return [SystemMessage(content=BASE_SYSTEM_TEMPLATE)] def _build_tools(self): @tool( description=GET_RESERVATION_DETAILS_DESCRIPTION, args_schema=GetReservationDetailsArgs, return_direct=True, ) async def get_reservation_details() -> str: details = self.reservation_details or {} print("details", details) email = details.get("email") confirmation_number = details.get("confirmation_number") hotel_id = details.get("hotel_id") sub_hotel_id = details.get("sub_hotel_id") if not email or not confirmation_number: return "No reservation details are available for this guest session." lookup_kwargs = { "confirmation_number": str(confirmation_number), "email": str(email), } if hotel_id is not None: lookup_kwargs["hotel_id"] = str(hotel_id) if sub_hotel_id is not None: lookup_kwargs["sub_hotel_id"] = str(sub_hotel_id) result = await get_reservation(**lookup_kwargs) print("result", result) return result @tool( description=GET_AVAILABLE_ADDONS_DESCRIPTION, args_schema=GetAvailableAddonsArgs, ) async def get_available_addons() -> str: details = self.reservation_details or {} hotel_id = details.get("hotel_id") sub_hotel_id = details.get("sub_hotel_id") if not hotel_id: return json.dumps( { "error": ( "Hotel ID is missing from this guest session, " "so available products and add-ons cannot be loaded." ) } ) result = await fetch_available_addons( hotel_id=str(hotel_id), sub_hotel_id=str(sub_hotel_id) if sub_hotel_id is not None else None, ) return json.dumps(result, default=str) @tool( description=HOTEL_RETRIEVE_DOCUMENTATION_DESCRIPTION, args_schema=AgentRetrieverResponse, ) async def retrieve_documentation(query: str, question: str): all_queries = [query, question] rag_service = HotelRAGService() documents = await rag_service._safe_retrieve_batch(all_queries, k=4) if not documents: return "No relevant hotel information was found for this query." return rag_service.format_documents(documents) @tool( description=HUMAN_ESCALATION_DESCRIPTION, args_schema=HumanEscalationArgs, return_direct=True, ) async def human_escalation(query: str): return json.dumps( { "escalated": True, "staff_query": query, "message": "Guest request has been escalated to a hotel representative.", } ) return [ retrieve_documentation, human_escalation, get_reservation_details, get_available_addons, ] def _num_tokens_from_messages(self, messages: list[BaseMessage]) -> int: try: return self.llm.get_num_tokens_from_messages(messages) except Exception: total_chars = 0 for msg in messages: content = msg.content if isinstance(content, str): total_chars += len(content) elif isinstance(content, list): for item in content: if isinstance(item, dict): text = item.get("text") if isinstance(text, str): total_chars += len(text) return max(1, total_chars // 4) def _history_to_base_messages(self, history: list[dict], question_token: int = 0) -> list[BaseMessage]: if not history: return [] normalized: list[BaseMessage] = [] for item in history: if not isinstance(item, dict): continue role = str(item.get("type") or item.get("role") or "").lower() text = item.get("text") or item.get("content") or "" if not isinstance(text, str) or not text.strip(): continue if role in {"human", "user"}: normalized.append(HumanMessage(content=text)) elif role in {"bot", "ai", "assistant"}: normalized.append(AIMessage(content=text)) elif role in {"hotel_admin", "admin", "system"}: # Hotel admin replies are authoritative staff messages. admin_text = text if text.lower().startswith("admin:") else f"Admin: {text}" normalized.append(SystemMessage(content=admin_text)) if not normalized: return [] token_count = self._num_tokens_from_messages(normalized) max_allowed = max(0, self.MAX_TOKEN_LIMIT - question_token) while normalized and token_count > max_allowed: normalized.pop(0) token_count = self._num_tokens_from_messages(normalized) return normalized async def call_agent(self, state: AgentState): messages = self.get_agent_prompt() messages.extend(state["messages"]) response = await self.llm_with_tools.ainvoke( messages, config={"run_name": "HotelAgent"}, ) return {"messages": [response]} def should_continue(self, state: AgentState) -> Literal["tools", END]: messages = state["messages"] last_message = messages[-1] if getattr(last_message, "tool_calls", None): return "tools" return END def _build_graph(self): workflow = StateGraph(AgentState) workflow.add_node("agent", self.call_agent) workflow.add_node("tools", ToolNode(self.tools)) workflow.add_edge(START, "agent") workflow.add_conditional_edges("agent", self.should_continue) workflow.add_edge("tools", "agent") return workflow.compile() def _normalize_content(self, content: Any) -> str: if content is None: return "" if isinstance(content, str): return content if isinstance(content, list): parts = [] for item in content: if isinstance(item, str): parts.append(item) elif isinstance(item, dict): text = item.get("text") if isinstance(text, str): parts.append(text) return "".join(parts) if isinstance(content, dict): text = content.get("text") return text if isinstance(text, str) else "" return str(content) def _parse_final_answer(self, content: Any) -> dict: text = self._normalize_content(content).strip() if not text: return {"answer": "", "escalated": False} if text.startswith("{"): try: parsed = json.loads(text) if isinstance(parsed, dict): return { "answer": parsed.get("answer", "") if isinstance(parsed.get("answer"), str) else "", "escalated": bool(parsed.get("escalated", False)), } except Exception: pass return {"answer": text, "escalated": False} def _extract_tool_usage(self, messages: Sequence[BaseMessage]) -> tuple[bool, bool, str | None]: """Return (used_retrieve, used_escalation, escalation_query).""" used_retrieve = False used_escalation = False escalation_query = None for message in messages: tool_calls = getattr(message, "tool_calls", None) or [] for tool_call in tool_calls: name = tool_call.get("name") if isinstance(tool_call, dict) else getattr(tool_call, "name", None) if name in { "retrieve_documentation", "get_reservation_details", "get_available_addons", }: used_retrieve = True elif name == "human_escalation": used_escalation = True args = tool_call.get("args") if isinstance(tool_call, dict) else getattr(tool_call, "args", {}) or {} query = args.get("query") if isinstance(args, dict) else None if isinstance(query, str) and query.strip(): escalation_query = query.strip() if isinstance(message, ToolMessage): name = getattr(message, "name", None) or "" if name in { "retrieve_documentation", "get_reservation_details", "get_available_addons", }: used_retrieve = True elif name == "human_escalation": used_escalation = True raw = self._normalize_content(message.content) try: payload = json.loads(raw) staff_query = payload.get("staff_query") if isinstance(staff_query, str) and staff_query.strip(): escalation_query = staff_query.strip() except Exception: if raw.strip() and not escalation_query: escalation_query = raw.strip() return used_retrieve, used_escalation, escalation_query def _resolve_event(self, used_retrieve: bool, used_escalation: bool) -> str: if used_escalation: return "support_escalation" if used_retrieve: return "lookup_answer" return "answer" def _build_history(self, history: list[dict] | None, query: str, answer: str) -> list[dict]: updated = [] for item in history or []: if isinstance(item, dict): updated.append(item) updated.append({"type": "user", "text": query}) if answer: updated.append({"type": "ai", "text": answer}) return updated async def answer( self, query: str, history: list[dict] | None = None, reservation_details: dict | None = None, ) -> dict: self.reservation_details = reservation_details question_token = self._num_tokens_from_messages([HumanMessage(content=query)]) history_messages = self._history_to_base_messages(history or [], question_token=question_token) history_messages.append(HumanMessage(content=query)) result = await self.graph.ainvoke({"messages": history_messages}) messages = result.get("messages") or [] used_retrieve, used_escalation, escalation_query = self._extract_tool_usage(messages) event = self._resolve_event(used_retrieve, used_escalation) last_ai = None for message in reversed(messages): if isinstance(message, AIMessage) and not getattr(message, "tool_calls", None): last_ai = message break parsed = self._parse_final_answer(getattr(last_ai, "content", "") if last_ai else "") answer_text = parsed.get("answer", "") data = { "answer": answer_text, "history": self._build_history(history, query, answer_text), } if event == "support_escalation" and escalation_query: data["escalation_query"] = escalation_query return { "event": event, "data": data, } def get_hotel_langgraph_service(): return HotelLangGraphAgentService()