from pydantic import BaseModel, Field from typing_extensions import TypedDict from typing import Annotated, List, Dict, Literal, Optional, Tuple, Union from app.config import settings from langchain_core.tools import Tool from langchain_openai.chat_models.base import ChatOpenAI from langgraph.prebuilt import ToolNode from langgraph.graph import StateGraph, START, END, add_messages from langgraph.graph.state import CompiledStateGraph from langchain_core.messages import AIMessageChunk from langchain_core.output_parsers.json import parse_partial_json from langchain_core.messages import SystemMessage class State(TypedDict): messages: Annotated[list, add_messages] class RetrieverTool(Tool): """Custom Tool Just For Override _to_args_and_kwargs function.""" def _to_args_and_kwargs(self, tool_input: Union[str, Dict], tool_call_id: Optional[str] = None) -> Tuple[Tuple, Dict]: if isinstance(tool_input, str): return (tool_input,), {} else: return (), tool_input class SearchRoomsSchema(BaseModel): from_date: str = Field(description="The start date of the stay") to_date: str = Field(description="The end date of the stay") adults: int = Field(description="The number of adults") children: int = Field(description="The number of children") def __init__(self, from_date: str, to_date: str, adults: int, children: int) -> None: self.from_date = from_date self.to_date = to_date self.adults = adults self.children = children def __setattr__(self, name, value) -> None: self.__dict__[name] = value class ClassifyResponse(BaseModel): answer: str = Field(description="Final answer to the user's question in markdown format") related_questions: list[str] = Field(description="Give 2-3 related questions from user's perspective to the user's question") class HotelAgent: def __init__(self): self.llm = ChatOpenAI(model="gpt-5.2", api_key=settings.OPENAI_API_KEY) async def data_retrieval(self, query:str, callbacks = None, **kwargs): return "Room search completed" async def get_tools(self) -> List[RetrieverTool]: return [ RetrieverTool( name='data_retrieval', description=( "Retrieves detailed and up-to-date information about Plaza Beach Resorts, " "including room availability for specified dates, occupancy limits " "(adults and children), room types, amenities, and property-specific details. " "Use this tool only when the requested information is not already available " "in the system prompt. Do not use it for greetings or unrelated questions." ), func=self.data_retrieval, coroutine=self.data_retrieval, args_schema=SearchRoomsSchema ) ] async def should_continue(self, state: State) -> Literal["continue", "end"]: messages = state["messages"] last_message = messages[-1] tool_calls = last_message.tool_calls if not tool_calls: return 'end' tool_name = tool_calls[0].get("name") return "continue" async def build_graph(self) -> "CompiledStateGraph": """ Builds and compiles the state graph defining agent behavior and tool usage. Returns: CompiledStateGraph: Compiled graph ready for execution. """ tools = await self.get_tools() builder = StateGraph(state_schema=State) builder.add_node("agent",self.agent) builder.add_node("tools", ToolNode(tools=tools)) builder.add_edge(START, "agent") builder.add_conditional_edges( "agent", self.should_continue, { "continue":"tools", "end":END }, ) builder.add_edge("tools", "agent") builder.add_edge("agent", END) return builder.compile() async def agent(self, state: State): tools = await self.get_tools() self.llm.use_responses_api = True llm_with_tools = self.llm.bind_tools(tools=tools) messages = await self.get_agent_prompt() messages.extend(state["messages"]) response = await llm_with_tools.ainvoke( messages, config={"run_name":"Agent"}, response_format = ClassifyResponse ) return {"messages": [response]} async def get_agent_prompt(self, inputs={}) -> list: system_template = """ You are the official AI assistant for Plaza Beach Resorts (St. Pete Beach, Florida), featuring three hotels: Plaza Beach Hotel (beachfront), Bayview Plaza Waterfront Resort & Marina (waterfront, pet-friendly, couples-friendly), and Bay Palms Waterfront Resort & Marina (family-friendly suites with kitchens). Your role: - Provide helpful, friendly, and professional information only about these three hotels. - You are only here for guide users about the hotels, you can't perform any actions like booking, cancellation, payment, availability check or reservation changes etc. - Politely refuse unrelated questions. - Don't tell user about the tools you have or what you can't do(until user ask about it), you are only here to answer questions about the hotels. - Use the data retrieval tool only if needed and never mention it. - Never guess. If information is unavailable, say so politely. - Don't give big examples in your response, just answer the question. We have related questions for user to ask you. - Keep responses clear, concise, and welcoming. Important: Your response should be in markdown formatting. """ messages = [ SystemMessage(content=system_template) ] messages.extend(inputs.get('chat_history', [])) return messages async def astream(self, inputs: dict): graph = await self.build_graph() final_answer = "" answer = "" escaped_text = None async for event in graph.astream(inputs ,config={"run_name": "graph"}, stream_mode="messages"): chunk,meta = event if chunk and isinstance(chunk, AIMessageChunk) and hasattr(chunk, 'content') and chunk.content and len(chunk.content): token = chunk.content[0].get('text') if not token: continue final_answer += token escaped_text = parse_partial_json(final_answer) if escaped_text is not None and 'answer' in escaped_text: new_token = escaped_text['answer'][len(answer):] answer += new_token if new_token == "": continue yield {"event": "stream", "data": new_token} self.answer = answer yield {"event": "final", "data": final_answer}