{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example of using Agent node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from langchain_core.tools import tool\n",
    "from langchain_mcp_adapters.client import MultiServerMCPClient\n",
    "from langgraph.graph import END, START, StateGraph\n",
    "from orcheo.graph.state import State\n",
    "from orcheo.nodes.ai import Agent\n",
    "from orcheo.nodes.code import PythonCode\n",
    "from orcheo.nodes.telegram import MessageTelegram\n",
    "\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "\n",
    "model_settings = {\n",
    "    \"model\": \"gpt-4o-mini\",\n",
    "    \"api_key\": os.getenv(\"OPENAI_API_KEY\"),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured output with vanilla JSON dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'John Doe'}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "so_schema = \"\"\"\n",
    "json_dict = {\n",
    "    \"type\": \"object\",\n",
    "    \"title\": \"Person\",\n",
    "    \"description\": \"A person\",\n",
    "    \"properties\": {\n",
    "        \"name\": {\"type\": \"string\"},\n",
    "    },\n",
    "    \"required\": [\"name\"],\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "so_config = {\n",
    "    \"schema_type\": \"json_dict\",\n",
    "    \"schema_str\": so_schema,\n",
    "}\n",
    "\n",
    "agent_node = Agent(\n",
    "    name=\"agent\",\n",
    "    model_settings=model_settings,\n",
    "    structured_output=so_config,\n",
    "    system_prompt=\"Your name is John Doe.\",\n",
    "    checkpointer=\"memory\",\n",
    ")\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = await agent_node(  # noqa: F704, PLE1142\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": \"What's your name?\"}]}, config\n",
    ")\n",
    "\n",
    "result[\"structured_response\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured output with OpenAI JSON dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'John Doe'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "so_schema = \"\"\"\n",
    "oai_json_schema = {\n",
    "    \"name\": \"get_person\",\n",
    "    \"strict\": True,\n",
    "    \"schema\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\"name\": {\"type\": \"string\"}},\n",
    "        \"additionalProperties\": False,\n",
    "        \"required\": [\"name\"],\n",
    "    },\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "so_config = {\n",
    "    \"schema_type\": \"json_dict\",\n",
    "    \"schema_str\": so_schema,\n",
    "}\n",
    "\n",
    "agent_node = Agent(\n",
    "    name=\"agent\",\n",
    "    model_settings=model_settings,\n",
    "    structured_output=so_config,\n",
    "    system_prompt=\"Your name is John Doe.\",\n",
    "    checkpointer=\"memory\",\n",
    ")\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = await agent_node(  # noqa: F704, PLE1142\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": \"What's your name?\"}]}, config\n",
    ")\n",
    "\n",
    "result[\"structured_response\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured output with Pydantic Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Person(name='John Doe')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "so_schema = \"\"\"\n",
    "class Person(BaseModel):\n",
    "    \\\"\\\"\\\"A Person.\\\"\\\"\\\"\n",
    "\n",
    "    name: str\n",
    "\"\"\"\n",
    "\n",
    "so_config = {\n",
    "    \"schema_type\": \"json_dict\",\n",
    "    \"schema_str\": so_schema,\n",
    "}\n",
    "\n",
    "agent_node = Agent(\n",
    "    name=\"agent\",\n",
    "    model_settings=model_settings,\n",
    "    structured_output=so_config,\n",
    "    system_prompt=\"Your name is John Doe.\",\n",
    "    checkpointer=\"memory\",\n",
    ")\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = await agent_node(  # noqa: F704, PLE1142\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": \"What's your name?\"}]}, config\n",
    ")\n",
    "\n",
    "result[\"structured_response\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured output with Typed dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'John Doe'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "so_schema = \"\"\"\n",
    "class Person(TypedDict):\n",
    "    \\\"\\\"\\\"A Person.\\\"\\\"\\\"\n",
    "\n",
    "    name: str\n",
    "\"\"\"\n",
    "\n",
    "so_config = {\n",
    "    \"schema_type\": \"typed_dict\",\n",
    "    \"schema_str\": so_schema,\n",
    "}\n",
    "\n",
    "agent_node = Agent(\n",
    "    name=\"agent\",\n",
    "    model_settings=model_settings,\n",
    "    structured_output=so_config,\n",
    "    system_prompt=\"Your name is John Doe.\",\n",
    "    checkpointer=\"memory\",\n",
    ")\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = await agent_node(  # noqa: F704, PLE1142\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": \"What's your name?\"}]}, config\n",
    ")\n",
    "\n",
    "result[\"structured_response\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use a node as a tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Main differences of a graph node and a function tool: A graph node only\n",
    "receives `state: dict` as input, and returns a dict, while a function tool\n",
    "can have arbitrary inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ToolMessage(content='{\"message_id\": 533, \"status\": \"sent\"}', name='MessageTelegram', id='1d985ce7-0209-4709-88f0-5447b5e70257', tool_call_id='call_kZHXVposdO3eO9NI79tps6fR')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "telegram_node = MessageTelegram(\n",
    "    name=\"MessageTelegram\",\n",
    "    token=os.getenv(\"TELEGRAM_TOKEN\"),\n",
    ")\n",
    "\n",
    "agent_node = Agent(\n",
    "    name=\"agent\",\n",
    "    model_settings=model_settings,\n",
    "    tools=[telegram_node],\n",
    "    system_prompt=\"Your name is John Doe.\",\n",
    "    checkpointer=\"memory\",\n",
    ")\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = await agent_node(  # noqa: F704, PLE1142\n",
    "    {\n",
    "        \"messages\": [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": (\n",
    "                    f\"Say hello to {os.getenv('TELEGRAM_CHAT_ID')} using \"\n",
    "                    \"message_telegram tool\"\n",
    "                ),\n",
    "            }\n",
    "        ]\n",
    "    },\n",
    "    config,\n",
    ")\n",
    "\n",
    "result[\"messages\"][-2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use tools from MCP servers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize the MCP client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "mcp_servers = {\n",
    "    \"filesystem\": {\n",
    "        \"command\": \"npx\",\n",
    "        \"args\": [\n",
    "            \"-y\",\n",
    "            \"@modelcontextprotocol/server-filesystem\",\n",
    "            \"~/Desktop\",\n",
    "        ],\n",
    "        \"transport\": \"stdio\",\n",
    "    },\n",
    "    \"git\": {\n",
    "        \"command\": \"uvx\",\n",
    "        \"args\": [\"mcp-server-git\"],\n",
    "        \"transport\": \"stdio\",\n",
    "    },\n",
    "}\n",
    "\n",
    "client = MultiServerMCPClient(mcp_servers)\n",
    "\n",
    "tools = await client.get_tools()  # noqa: F704, PLE1142"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pass the tools to the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Here's what's on your desktop:\n",
      "\n",
      "- **.DS_Store** (file)\n",
      "- **.localized** (file)\n",
      "\n",
      "It seems that there is no git repository on your desktop, so I was unable to retrieve the git status for it. If you have a specific directory with a git repository, please let me know!\n"
     ]
    }
   ],
   "source": [
    "agent_node = Agent(\n",
    "    name=\"agent\",\n",
    "    model_settings=model_settings,\n",
    "    tools=tools,\n",
    "    system_prompt=\"You have access to the filesystem and git.\",\n",
    "    checkpointer=\"memory\",\n",
    ")\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = await agent_node(  # noqa: F704, PLE1142\n",
    "    {\n",
    "        \"messages\": [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": (\n",
    "                    \"Hello, tell me what's on my desktop and what's the git \"\n",
    "                    \"status of it.\"\n",
    "                ),\n",
    "            }\n",
    "        ]\n",
    "    },\n",
    "    config,\n",
    ")\n",
    "\n",
    "print(result[\"messages\"][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use sub-graph as a tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a sub-graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "python_code_node = PythonCode(\n",
    "    name=\"PythonCode\",\n",
    "    code=\"return {'messages': [{'role': 'ai', 'content': 'Hello, ' + state['outputs']['initial'] + '.'}]}\",  # noqa: E501\n",
    ")\n",
    "\n",
    "tool_graph = StateGraph(State)\n",
    "tool_graph.add_node(\"python_code\", python_code_node)\n",
    "tool_graph.add_edge(START, \"python_code\")\n",
    "tool_graph.add_edge(\"python_code\", END)\n",
    "\n",
    "python_code_graph = tool_graph.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "### Wrap the sub-graph as a tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool(parse_docstring=True)\n",
    "def greet(name: str) -> dict:\n",
    "    \"\"\"Greet the user.\n",
    "\n",
    "    Args:\n",
    "        name: The name of the user to greet.\n",
    "    \"\"\"\n",
    "    result = asyncio.run(\n",
    "        python_code_graph.ainvoke(\n",
    "            {\"messages\": [], \"outputs\": {\"initial\": name}}, config={}\n",
    "        )\n",
    "    )\n",
    "    return result[\"outputs\"][\"PythonCode\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use the tool in an agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='Hello, my name is John Doe', additional_kwargs={}, response_metadata={}, id='33de491c-5b7c-4a69-bf32-aa04eb53fb7e'),\n",
       " AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_ljMrBXo50wphcXt5v4YORZIK', 'function': {'arguments': '{\"name\":\"John Doe\"}', 'name': 'greet'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 15, 'prompt_tokens': 63, 'total_tokens': 78, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_62a23a81ef', 'id': 'chatcmpl-BaMDDCuQ24e8EhJe7XqKFpeIhT0TA', 'service_tier': 'default', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--780cb5b7-fd3c-4455-9985-b81a91824ffa-0', tool_calls=[{'name': 'greet', 'args': {'name': 'John Doe'}, 'id': 'call_ljMrBXo50wphcXt5v4YORZIK', 'type': 'tool_call'}], usage_metadata={'input_tokens': 63, 'output_tokens': 15, 'total_tokens': 78, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}}),\n",
       " ToolMessage(content='{\"messages\": [{\"role\": \"ai\", \"content\": \"Hello, John Doe.\"}]}', name='greet', id='546fd998-7357-4b51-a040-3ec759abe839', tool_call_id='call_ljMrBXo50wphcXt5v4YORZIK'),\n",
       " AIMessage(content='Hello, John Doe. How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 13, 'prompt_tokens': 106, 'total_tokens': 119, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_62a23a81ef', 'id': 'chatcmpl-BaMDEafZVkb6SK8TDA042Ts2rfE1k', 'service_tier': 'default', 'finish_reason': 'stop', 'logprobs': None}, id='run--8bb511fd-efe7-43e0-9038-c405a3f3ea61-0', usage_metadata={'input_tokens': 106, 'output_tokens': 13, 'total_tokens': 119, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_node = Agent(\n",
    "    name=\"agent\",\n",
    "    model_settings=model_settings,\n",
    "    tools=[greet],\n",
    "    system_prompt=\"Your name is John Doe.\",\n",
    "    checkpointer=\"memory\",\n",
    ")\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"123\"}}\n",
    "result = await agent_node(  # noqa: F704, PLE1142\n",
    "    {\n",
    "        \"messages\": [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": (\"Hello, my name is John Doe\"),\n",
    "            }\n",
    "        ]\n",
    "    },\n",
    "    config,\n",
    ")\n",
    "\n",
    "result[\"messages\"]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
