scratchpad

5484.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/teeler/5f7da82a68025f19484e28dd911f26d8/scratchpad.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "lIYdn1woOS1n"
      },
      "outputs": [],
      "source": [
        "from dotenv import load_dotenv\n",
        "from agents import Agent, Runner, trace\n",
        "from agents.mcp import MCPServerStdio\n",
        "import os\n",
        "from mcp import ClientSession, StdioServerParameters, Tool as MCPTool, stdio_client\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from contextlib import AbstractAsyncContextManager, AsyncExitStack\n",
        "from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream\n",
        "from mcp.shared.message import SessionMessage\n",
        "from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client\n",
        "\n",
        "class ColabMCPServerStdio(MCPServerStdio):\n",
        "    def __init__(self, *args, **kwargs):\n",
        "        super().__init__(*args, **kwargs)\n",
        "        self.errlog = open(\"/tmp/err.log\", \"w+\")\n",
        "\n",
        "\n",
        "    def create_streams(\n",
        "        self,\n",
        "    ) -> AbstractAsyncContextManager[\n",
        "        tuple[\n",
        "            MemoryObjectReceiveStream[SessionMessage | Exception],\n",
        "            MemoryObjectSendStream[SessionMessage],\n",
        "            GetSessionIdCallback | None,\n",
        "        ]\n",
        "    ]:\n",
        "        \"\"\"Create the streams for the server.\"\"\"\n",
        "        return stdio_client(self.params, errlog=self.errlog)\n"
      ],
      "metadata": {
        "id": "K6l8tJcRLwqL"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "params = {\"command\": \"npx\", \"args\": [\"-y\", \"@modelcontextprotocol/server-memory\"]}\n",
        "\n",
        "async with ColabMCPServerStdio(params=puppeteer_params) as server:\n",
        "    tools = await server.list_tools()\n",
        "\n",
        "print(tools)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gSoYgWHB5yDg",
        "outputId": "16a99a80-d2f3-47b9-cff0-085ebab0045c"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[Tool(name='create_entities', title=None, description='Create multiple new entities in the knowledge graph', inputSchema={'type': 'object', 'properties': {'entities': {'type': 'array', 'items': {'type': 'object', 'properties': {'name': {'type': 'string', 'description': 'The name of the entity'}, 'entityType': {'type': 'string', 'description': 'The type of the entity'}, 'observations': {'type': 'array', 'items': {'type': 'string'}, 'description': 'An array of observation contents associated with the entity'}}, 'required': ['name', 'entityType', 'observations'], 'additionalProperties': False}}}, 'required': ['entities'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='create_relations', title=None, description='Create multiple new relations between entities in the knowledge graph. Relations should be in active voice', inputSchema={'type': 'object', 'properties': {'relations': {'type': 'array', 'items': {'type': 'object', 'properties': {'from': {'type': 'string', 'description': 'The name of the entity where the relation starts'}, 'to': {'type': 'string', 'description': 'The name of the entity where the relation ends'}, 'relationType': {'type': 'string', 'description': 'The type of the relation'}}, 'required': ['from', 'to', 'relationType'], 'additionalProperties': False}}}, 'required': ['relations'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='add_observations', title=None, description='Add new observations to existing entities in the knowledge graph', inputSchema={'type': 'object', 'properties': {'observations': {'type': 'array', 'items': {'type': 'object', 'properties': {'entityName': {'type': 'string', 'description': 'The name of the entity to add the observations to'}, 'contents': {'type': 'array', 'items': {'type': 'string'}, 'description': 'An array of observation contents to add'}}, 'required': ['entityName', 'contents'], 'additionalProperties': False}}}, 'required': ['observations'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='delete_entities', title=None, description='Delete multiple entities and their associated relations from the knowledge graph', inputSchema={'type': 'object', 'properties': {'entityNames': {'type': 'array', 'items': {'type': 'string'}, 'description': 'An array of entity names to delete'}}, 'required': ['entityNames'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='delete_observations', title=None, description='Delete specific observations from entities in the knowledge graph', inputSchema={'type': 'object', 'properties': {'deletions': {'type': 'array', 'items': {'type': 'object', 'properties': {'entityName': {'type': 'string', 'description': 'The name of the entity containing the observations'}, 'observations': {'type': 'array', 'items': {'type': 'string'}, 'description': 'An array of observations to delete'}}, 'required': ['entityName', 'observations'], 'additionalProperties': False}}}, 'required': ['deletions'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='delete_relations', title=None, description='Delete multiple relations from the knowledge graph', inputSchema={'type': 'object', 'properties': {'relations': {'type': 'array', 'items': {'type': 'object', 'properties': {'from': {'type': 'string', 'description': 'The name of the entity where the relation starts'}, 'to': {'type': 'string', 'description': 'The name of the entity where the relation ends'}, 'relationType': {'type': 'string', 'description': 'The type of the relation'}}, 'required': ['from', 'to', 'relationType'], 'additionalProperties': False}, 'description': 'An array of relations to delete'}}, 'required': ['relations'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='read_graph', title=None, description='Read the entire knowledge graph', inputSchema={'type': 'object', 'properties': {}, 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='search_nodes', title=None, description='Search for nodes in the knowledge graph based on a query', inputSchema={'type': 'object', 'properties': {'query': {'type': 'string', 'description': 'The search query to match against entity names, types, and observation content'}}, 'required': ['query'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None), Tool(name='open_nodes', title=None, description='Open specific nodes in the knowledge graph by their names', inputSchema={'type': 'object', 'properties': {'names': {'type': 'array', 'items': {'type': 'string'}, 'description': 'An array of entity names to retrieve'}}, 'required': ['names'], 'additionalProperties': False}, outputSchema=None, icons=None, annotations=None, meta=None)]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "oChsy1XINRT3"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "name": "scratchpad",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}