agent.py

import os
from typing import TypedDict, Annotated
from dotenv import load_dotenv
import asyncio
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import InMemorySaver
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.prompt import Prompt
from rich.markdown import Markdown

# Load env vars
load_dotenv()
# Note: 
# The credentials included in this example are for demonstration only.
# You must use your own API keys, database credentials, and LLM configuration when running the project.
# os.environ["OPENAI_API_KEY"] = os.getenv(
#     "GITHUB_TOKEN", "YOUR_GITHUB_TOKEN"
# )
# os.environ["OPENAI_API_BASE"] = "https://models.github.ai/inference"

# LLM
llm = ChatOpenAI(
    model="openai/gpt-4.1",
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ["OPENAI_API_BASE"],
    temperature=0.2
)

# Rich console
console = Console()

# MCP client setup (absolute path to server)
mcp_server_path = os.path.join(
    os.path.dirname(__file__), "mcp-database-server", "dist", "src", "index.js"
)
mcp_client = MultiServerMCPClient({
    "database": {
        "command": "node",
        "args": [
            mcp_server_path,
            "--mysql",
            "--host", os.getenv("DB_HOST", "localhost"),
            "--database", os.getenv("DB_NAME", "YOUR_DATABASE_NAME"),
            "--port", os.getenv("DB_PORT", "3306"),
            "--user", os.getenv("DB_USER", "root"),
            "--password", os.getenv("DB_PASSWORD", "YOUR_PASSWORD")
        ],
        "transport": "stdio"
    }
})

# Conversation state type
class ChatState(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]

async def main():
    # Load system role prompt
    prompt_path = os.path.join(os.path.dirname(__file__), "prompts", "hospital_receptionist.txt")
    try:
        with open(prompt_path, "r", encoding="utf-8") as f:
            system_prompt = f.read().strip()
    except FileNotFoundError:
        system_prompt = (
            "You are CareDesk, the hospital's AI Receptionist. Enforce privacy; require patient_id for any "
            "patient-specific info. Help with appointments and general info."
        )
    system_messages = [SystemMessage(content=system_prompt)]
    # Acquire MCP tools
    try:
        tools = await mcp_client.get_tools()
        console.print(Panel("[bold green]MCP Database Server connected[/bold green]", border_style="green"))

        table = Table(title="Available Database Tools", title_style="bold cyan")
        table.add_column("Tool", style="bold", no_wrap=True)
        table.add_column("Description", style="dim")
        for t in tools:
            table.add_row(t.name, t.description or "")
        console.print(table)
    except Exception as e:
        console.print(Panel(f"[red]Failed to get MCP tools:[/red] {e}", border_style="red"))
        return

    # Bind LLM to tools
    llm_with_tools = llm.bind_tools(tools)

    # Nodes
    async def chat_node(state: ChatState):
        try:
            # Prepend system prompt to every turn
            resp = await llm_with_tools.ainvoke(system_messages + state["messages"])
            return {"messages": [resp]}
        except Exception as e:
            return {"messages": [AIMessage(content=f"Error: {str(e)}")]}

    tool_node = ToolNode(tools)

    def should_continue(state: ChatState):
        last = state["messages"][-1]
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        return END

    # Graph
    graph = StateGraph(ChatState)
    graph.add_node("chat", chat_node)
    graph.add_node("tools", tool_node)
    graph.add_edge(START, "chat")
    graph.add_conditional_edges("chat", should_continue, {"tools": "tools", END: END})
    graph.add_edge("tools", "chat")

    app = graph.compile(checkpointer=InMemorySaver())

    # Interactive loop
    console.print(Panel(
        "🤖 [bold magenta]Customer Support Agent (DB-enabled)[/bold magenta]\n[dim]Ask: 'List tables', 'Show doctors', 'Describe table patients'[/dim]\n[green]Type 'exit' to quit[/green]",
        border_style="magenta"
    ))
    console.print(Panel("Loaded role prompt: [italic]hospital_receptionist.txt[/italic]", border_style="blue"))

    # Friendly welcome message for first-time users
    console.print(Panel(
        "[bold]Welcome to CareDesk![/bold]\n\n"
        "I am the hospital's AI Receptionist. I can:\n"
        "• Book, reschedule, or cancel appointments\n"
        "• Share your own appointment, billing, or report info (with your patient_id)\n"
        "• Answer general questions about departments, doctors, and services\n\n"
        "[italic]Privacy first:[/italic] I will never share any patient's details without a valid patient_id.",
        border_style="cyan"
    ))

    config = {"configurable": {"thread_id": "config-thread"}}
    try:
        while True:
            user = Prompt.ask("[bold cyan]You[/bold cyan]").strip()
            if user.lower() in {"exit", "quit", "bye"}:
                console.print("👋 Goodbye!", style="bold yellow")
                break
            if not user:
                continue
            console.print(Panel(Markdown(user), title="You", title_align="left", border_style="cyan"))
            with console.status("Thinking...", spinner="dots"):
                result = await app.ainvoke({"messages": [HumanMessage(content=user)]}, config=config)
            msg = result["messages"][-1]
            content = getattr(msg, "content", "(no content)")
            console.print(Panel(Markdown(content), title="Agent", title_align="left", border_style="green"))
    finally:
        # Safely close MCP client if supported
        try:
            closer = getattr(mcp_client, "close", None)
            if callable(closer):
                maybe_coro = closer()
                if asyncio.iscoroutine(maybe_coro):
                    await maybe_coro
        except Exception:
            pass
        console.print("🔌 MCP connection closed.", style="dim")


if __name__ == "__main__":
    asyncio.run(main())