MCP function calling

mcp-run-python/deno.json → mcp-run-python/deno.jsonc

@@ -32,7 +32,7 @@
       "src/*.ts",
       "src/prepareEnvCode.ts", // required to override gitignore
       "README.md",
-      "deno.json"
+      "deno.jsonc"
     ]
   }
 }

mcp-run-python/src/main.ts

@@ -15,26 +15,29 @@ const VERSION = '0.0.13'
 
 export async function main() {
   const { args } = Deno
-  if (args.length === 1 && args[0] === 'stdio') {
-    await runStdio()
-  } else if (args.length >= 1 && args[0] === 'sse') {
-    const flags = parseArgs(Deno.args, {
-      string: ['port'],
-      default: { port: '3001' },
-    })
-    const port = parseInt(flags.port)
-    runSse(port)
-  } else if (args.length === 1 && args[0] === 'warmup') {
-    await warmup()
+  const flags = parseArgs(args, {
+    string: ['port', 'callbacks'],
+    default: { port: '3001' },
+  })
+  const { _: [task], callbacks, port } = flags
+  if (task === 'stdio') {
+    await runStdio(callbacks)
+  } else if (task === 'sse' || task === 'http') {
+    runSse(parseInt(port), callbacks)
+  } else if (task === 'warmup') {
+    await warmup(callbacks)
   } else {
     console.error(
       `\
 Invalid arguments.
 
-Usage: deno run -N -R=node_modules -W=node_modules --node-modules-dir=auto jsr:@pydantic/mcp-run-python [stdio|sse|warmup]
+Usage:
+  deno run -N -R=node_modules -W=node_modules --node-modules-dir=auto \\
+    jsr:@pydantic/mcp-run-python [stdio|sse|warmup]
 
 options:
-  --port <port>  Port to run the SSE server on (default: 3001)`,
+  --port <port>                    Port to run the SSE server on (default: 3001).
+  --callbacks <python-signatures>  Python code representing the signatures of client functions the server can call.`,
     )
     Deno.exit(1)
   }
@@ -43,7 +46,8 @@ options:
 /*
  * Create an MCP server with the `run_python_code` tool registered.
  */
-function createServer(): McpServer {
+function createServer(callbacks?: string): McpServer {
+  const functions = _extractFunctions(callbacks)
   const server = new McpServer(
     {
       name: 'MCP Run Python',
@@ -57,20 +61,39 @@ function createServer(): McpServer {
     },
   )
 
-  const toolDescription = `Tool to execute Python code and return stdout, stderr, and return value.
+  let toolDescription = `Tool to execute Python code and return stdout, stderr, and return value.
 
-The code may be async, and the value on the last line will be returned as the return value.
+The code may be async, and the value on the last line will be returned as the return.
 
 The code will be executed with Python 3.12.
 
-Dependencies may be defined via PEP 723 script metadata, e.g. to install "pydantic", the script should start
-with a comment of the form:
+Dependencies may be defined via PEP 723 script metadata.
+
+To make HTTP requests, you must use the "httpx" library in async mode.
 
+For example:
+
+\`\`\`python
 # /// script
-# dependencies = ['pydantic']
+# dependencies = ['httpx']
 # ///
-print('python code here')
+import httpx
+
+async with httpx.AsyncClient() as client:
+    response = await client.get('https://example.com')
+# return the text of the page
+response.text
+\`\`\`
 `
+  if (callbacks) {
+    toolDescription += `
+The following functions are already defined globally and available to call from within your code:
+
+\`\`\`python
+${callbacks}
+\`\`\`
+    `
+  }
 
   let setLogLevel: LoggingLevel = 'emergency'
 
@@ -85,29 +108,46 @@ print('python code here')
     { python_code: z.string().describe('Python code to run') },
     async ({ python_code }: { python_code: string }) => {
       const logPromises: Promise<void>[] = []
-      const result = await runCode([{
+      const mainPy = {
         name: 'main.py',
         content: python_code,
         active: true,
-      }], (level, data) => {
+      }
+      const codeLog = (level: LoggingLevel, data: string) => {
         if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
           logPromises.push(server.server.sendLoggingMessage({ level, data }))
         }
-      })
+      }
+      async function clientCallback(func: string, args?: string, kwargs?: string) {
+        const { content } = await server.server.createMessage({
+          messages: [],
+          maxTokens: 0,
+          systemPrompt: '',
+          metadata: { pydantic_custom_use: '__python_function_call__', func, args, kwargs },
+        })
+        if (content.type !== 'text') {
+          throw new Error('Expected return content type to be "text"')
+        } else {
+          return content.text
+        }
+      }
+
+      const result = await runCode([mainPy], codeLog, functions, clientCallback)
       await Promise.all(logPromises)
       return {
         content: [{ type: 'text', text: asXml(result) }],
       }
     },
   )
+
   return server
 }
 
 /*
  * Run the MCP server using the SSE transport, e.g. over HTTP.
  */
-function runSse(port: number) {
-  const mcpServer = createServer()
+function runSse(port: number, callbacks?: string) {
+  const mcpServer = createServer(callbacks)
   const transports: { [sessionId: string]: SSEServerTransport } = {}
 
   const server = http.createServer(async (req, res) => {
@@ -162,16 +202,20 @@ function runSse(port: number) {
 /*
  * Run the MCP server using the Stdio transport.
  */
-async function runStdio() {
-  const mcpServer = createServer()
+async function runStdio(callbacks?: string) {
+  const mcpServer = createServer(callbacks)
   const transport = new StdioServerTransport()
   await mcpServer.connect(transport)
 }
 
 /*
  * Run pyodide to download packages which can otherwise interrupt the server
  */
-async function warmup() {
+async function warmup(callbacks?: string) {
+  if (callbacks) {
+    const functions = _extractFunctions(callbacks)
+    console.error(`Functions extracted from callbacks: ${JSON.stringify(functions)}`)
+  }
   console.error(
     `Running warmup script for MCP Run Python version ${VERSION}...`,
   )
@@ -193,6 +237,10 @@ a
   console.log('\nwarmup successful 🎉')
 }
 
+function _extractFunctions(callbacks?: string): string[] {
+  return callbacks ? [...callbacks.matchAll(/^async def (\w+)/g).map(([, f]) => f)] : []
+}
+
 // list of log levels to use for level comparison
 const LogLevels: LoggingLevel[] = [
   'debug',

mcp-run-python/src/prepare_env.py

@@ -10,15 +10,17 @@
 import re
 import sys
 import traceback
-from collections.abc import Iterable, Iterator
+from collections.abc import Awaitable, Iterable, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Literal, TypedDict
+from typing import Any, Callable, Literal, TypedDict
 
 import micropip
 import pyodide_js
 import tomllib
+from pydantic import ConfigDict, TypeAdapter
+from pydantic_core import to_json
 from pyodide.code import find_imports
 
 __all__ = 'prepare_env', 'dump_json'
@@ -31,18 +33,18 @@ class File(TypedDict):
 
 
 @dataclass
-class Success:
+class PrepSuccess:
     dependencies: list[str] | None
     kind: Literal['success'] = 'success'
 
 
 @dataclass
-class Error:
+class PrepError:
     message: str
     kind: Literal['error'] = 'error'
 
 
-async def prepare_env(files: list[File]) -> Success | Error:
+async def prepare_env(files: list[File]) -> PrepSuccess | PrepError:
     sys.setrecursionlimit(400)
 
     cwd = Path.cwd()
@@ -68,14 +70,12 @@ async def prepare_env(files: list[File]) -> Success | Error:
             except Exception:
                 with open(logs_filename) as f:
                     logs = f.read()
-                return Error(message=f'{logs} {traceback.format_exc()}')
+                return PrepError(message=f'{logs} {traceback.format_exc()}')
 
-    return Success(dependencies=dependencies)
+    return PrepSuccess(dependencies=dependencies)
 
 
 def dump_json(value: Any) -> str | None:
-    from pydantic_core import to_json
-
     if value is None:
         return None
     if isinstance(value, str):
@@ -84,6 +84,50 @@ def dump_json(value: Any) -> str | None:
         return to_json(value, indent=2, fallback=_json_fallback).decode()
 
 
+class CallSuccess(TypedDict):
+    return_value: Any
+    kind: Literal['success']
+
+
+class CallError(TypedDict):
+    exc_type: str
+    message: str
+    kind: Literal['error']
+
+
+call_result_ta: TypeAdapter[CallSuccess | CallError] = TypeAdapter(
+    CallSuccess | CallError, config=ConfigDict(defer_build=True)
+)
+
+
+@dataclass(slots=True)
+class RegisterFunction:
+    _func_name: str
+    _callback: Callable[[str, str | None, str | None], Awaitable[str]]
+
+    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        result_json = await self._callback(self._func_name, _dump_args(args), _dump_args(kwargs))
+        result = call_result_ta.validate_json(result_json)
+        if result['kind'] == 'success':
+            return result['return_value']
+
+        exc_type, message = result['exc_type'], result['message']
+        try:
+            exc_type_ = __builtins__[exc_type]
+        except KeyError:
+            raise Exception(f'{message}\n(Raised exception type: {exc_type})')
+        else:
+            raise exc_type_(message)
+
+    def __repr__(self) -> str:
+        return f'<client callback {self._func_name}>'
+
+
+def _dump_args(value: Any) -> str | None:
+    if value:
+        return to_json(value, fallback=_json_fallback).decode()
+
+
 def _json_fallback(value: Any) -> Any:
     tp: Any = type(value)
     module = tp.__module__
@@ -95,7 +139,7 @@ def _json_fallback(value: Any) -> Any:
     elif module == 'pyodide.ffi':
         return value.to_py()
     else:
-        return repr(value)
+        return str(value)
 
 
 def _add_extra_dependencies(dependencies: list[str]) -> list[str]:

mcp-run-python/src/runCode.ts

@@ -1,4 +1,3 @@
-/* eslint @typescript-eslint/no-explicit-any: off */
 import { loadPyodide } from 'pyodide'
 import { preparePythonCode } from './prepareEnvCode.ts'
 import type { LoggingLevel } from '@modelcontextprotocol/sdk/types.js'
@@ -12,6 +11,8 @@ export interface CodeFile {
 export async function runCode(
   files: CodeFile[],
   log: (level: LoggingLevel, data: string) => void,
+  functionNames?: string[],
+  clientCallback?: (func: string, args?: string, kwargs?: string) => Promise<string>,
 ): Promise<RunSuccess | RunError> {
   // remove once https://github.com/pyodide/pyodide/pull/5514 is released
   const realConsoleLog = console.log
@@ -58,6 +59,14 @@ export async function runCode(
 
   const prepareStatus = await preparePyEnv.prepare_env(pyodide.toPy(files))
 
+  const globals: Record<string, unknown> = { __name__: '__main__' }
+
+  if (functionNames && clientCallback) {
+    for (const functionName of functionNames) {
+      globals[functionName] = preparePyEnv.RegisterFunction(functionName, clientCallback)
+    }
+  }
+
   let runResult: RunSuccess | RunError
   if (prepareStatus.kind == 'error') {
     runResult = {
@@ -70,7 +79,7 @@ export async function runCode(
     const activeFile = files.find((f) => f.active)! || files[0]
     try {
       const rawValue = await pyodide.runPythonAsync(activeFile.content, {
-        globals: pyodide.toPy({ __name__: '__main__' }),
+        globals: pyodide.toPy(globals),
         filename: activeFile.name,
       })
       runResult = {
@@ -99,7 +108,7 @@ interface RunSuccess {
   // we could record stdout and stderr separately, but I suspect simplicity is more important
   output: string[]
   dependencies: string[]
-  returnValueJson: string | null
+  returnValueJson: string | undefined
 }
 
 interface RunError {
@@ -153,6 +162,13 @@ function formatError(err: any): string {
     / {2}File "\/lib\/python\d+\.zip\/_pyodide\/.*\n {4}.*\n(?: {4,}\^+\n)?/g,
     '',
   )
+  // remove frames from _prepare_env.py
+  errStr = errStr.replace(
+    / {2}File "\/tmp\/mcp_run_python\/_prepare_env.py".*\n {4,}.+\n/g,
+    '',
+  )
+  // remove trailing newlines
+  errStr = errStr.replace(/\n+$/, '')
   return errStr
 }
 
@@ -164,8 +180,12 @@ interface PrepareError {
   kind: 'error'
   message: string
 }
+
 interface PreparePyEnv {
   prepare_env: (files: CodeFile[]) => Promise<PrepareSuccess | PrepareError>
-  // deno-lint-ignore no-explicit-any
-  dump_json: (value: any) => string | null
+  RegisterFunction: (
+    func_name: string,
+    callback: (func_name: string, args?: string, kwargs?: string) => Promise<string>,
+  ) => unknown
+  dump_json: (value: unknown) => string | undefined
 }