diff --git a/examples/agent/openapi.py b/examples/agent/openapi.py new file mode 100644 index 000000000..6cace4b44 --- /dev/null +++ b/examples/agent/openapi.py @@ -0,0 +1,119 @@ +import aiohttp +import json +import sys +import urllib.parse + +class OpenAPIMethod: + def __init__(self, url, name, descriptor, catalog): + ''' + Wraps a remote OpenAPI method as an async Python function. + ''' + self.url = url + self.__name__ = name + + assert 'post' in descriptor, 'Only POST methods are supported' + post_descriptor = descriptor['post'] + + self.__doc__ = post_descriptor.get('description', '') + parameters = post_descriptor.get('parameters', []) + request_body = post_descriptor.get('requestBody') + + self.parameters = {p['name']: p for p in parameters} + assert all(param['in'] == 'query' for param in self.parameters.values()), f'Only query path parameters are supported (path: {url}, descriptor: {json.dumps(descriptor)})' + + self.body = None + if request_body: + assert 'application/json' in request_body['content'], f'Only application/json is supported for request body (path: {url}, descriptor: {json.dumps(descriptor)})' + + body_name = 'body' + i = 2 + while body_name in self.parameters: + body_name = f'body{i}' + i += 1 + + self.body = dict( + name=body_name, + required=request_body['required'], + schema=request_body['content']['application/json']['schema'], + ) + + self.parameters_schema = dict( + type='object', + properties={ + **({ + self.body['name']: self.body['schema'] + } if self.body else {}), + **{ + name: param['schema'] + for name, param in self.parameters.items() + } + }, + required=[name for name, param in self.parameters.items() if param['required']] + ([self.body['name']] if self.body and self.body['required'] else []) + ) + + if (components := catalog.get('components', {})) is not None: + if (schemas := components.get('schemas')) is not None: + del schemas['HTTPValidationError'] + del schemas['ValidationError'] + if not schemas: + del components['schemas'] + if components: + self.parameters_schema['components'] = components + + async def __call__(self, **kwargs): + if self.body: + body = kwargs.pop(self.body['name'], None) + if self.body['required']: + assert body is not None, f'Missing required body parameter: {self.body["name"]}' + else: + body = None + + query_params = {} + for name, param in self.parameters.items(): + value = kwargs.pop(name, None) + if param['required']: + assert value is not None, f'Missing required parameter: {name}' + + assert param['in'] == 'query', 'Only query parameters are supported' + query_params[name] = value + + params = '&'.join(f'{name}={urllib.parse.quote(str(value))}' for name, value in query_params.items() if value is not None) + url = f'{self.url}?{params}' + async with aiohttp.ClientSession() as session: + async with session.post(url, json=body) as response: + if response.status == 500: + raise Exception(await response.text()) + response.raise_for_status() + response_json = await response.json() + + return response_json + +async def discover_tools(tool_endpoints: list[str], verbose) -> tuple[dict, list]: + tool_map = {} + tools = [] + + async with aiohttp.ClientSession() as session: + for url in tool_endpoints: + assert url.startswith('http://') or url.startswith('https://'), f'Tools must be URLs, not local files: {url}' + + catalog_url = f'{url}/openapi.json' + async with session.get(catalog_url) as response: + response.raise_for_status() + catalog = await response.json() + + for path, descriptor in catalog['paths'].items(): + fn = OpenAPIMethod(url=f'{url}{path}', name=path.replace('/', ' ').strip().replace(' ', '_'), descriptor=descriptor, catalog=catalog) + tool_map[fn.__name__] = fn + if verbose: + print(f'Function {fn.__name__}: params schema: {fn.parameters_schema}', file=sys.stderr) + tools.append(dict( + type='function', + function=dict( + name=fn.__name__, + description=fn.__doc__ or '', + parameters=fn.parameters_schema, + ) + ) + ) + + return tool_map, tools diff --git a/examples/agent/run.py b/examples/agent/run.py index bc2322bc4..5a47ebe68 100644 --- a/examples/agent/run.py +++ b/examples/agent/run.py @@ -12,6 +12,7 @@ import aiohttp import asyncio from functools import wraps import json +from openapi import discover_tools import os from pydantic import BaseModel import sys @@ -19,120 +20,7 @@ import typer from typing import Annotated, Literal, Optional import urllib.parse -class OpenAPIMethod: - def __init__(self, url, name, descriptor, catalog): - ''' - Wraps a remote OpenAPI method as an async Python function. - ''' - self.url = url - self.__name__ = name - assert 'post' in descriptor, 'Only POST methods are supported' - post_descriptor = descriptor['post'] - - self.__doc__ = post_descriptor.get('description', '') - parameters = post_descriptor.get('parameters', []) - request_body = post_descriptor.get('requestBody') - - self.parameters = {p['name']: p for p in parameters} - assert all(param['in'] == 'query' for param in self.parameters.values()), f'Only query path parameters are supported (path: {url}, descriptor: {json.dumps(descriptor)})' - - self.body = None - if request_body: - assert 'application/json' in request_body['content'], f'Only application/json is supported for request body (path: {url}, descriptor: {json.dumps(descriptor)})' - - body_name = 'body' - i = 2 - while body_name in self.parameters: - body_name = f'body{i}' - i += 1 - - self.body = dict( - name=body_name, - required=request_body['required'], - schema=request_body['content']['application/json']['schema'], - ) - - self.parameters_schema = dict( - type='object', - properties={ - **({ - self.body['name']: self.body['schema'] - } if self.body else {}), - **{ - name: param['schema'] - for name, param in self.parameters.items() - } - }, - required=[name for name, param in self.parameters.items() if param['required']] + ([self.body['name']] if self.body and self.body['required'] else []) - ) - - if (components := catalog.get('components', {})) is not None: - if (schemas := components.get('schemas')) is not None: - del schemas['HTTPValidationError'] - del schemas['ValidationError'] - if not schemas: - del components['schemas'] - if components: - self.parameters_schema['components'] = components - - async def __call__(self, **kwargs): - if self.body: - body = kwargs.pop(self.body['name'], None) - if self.body['required']: - assert body is not None, f'Missing required body parameter: {self.body["name"]}' - else: - body = None - - query_params = {} - for name, param in self.parameters.items(): - value = kwargs.pop(name, None) - if param['required']: - assert value is not None, f'Missing required parameter: {name}' - - assert param['in'] == 'query', 'Only query parameters are supported' - query_params[name] = value - - params = '&'.join(f'{name}={urllib.parse.quote(str(value))}' for name, value in query_params.items() if value is not None) - url = f'{self.url}?{params}' - async with aiohttp.ClientSession() as session: - async with session.post(url, json=body) as response: - if response.status == 500: - raise Exception(await response.text()) - response.raise_for_status() - response_json = await response.json() - - return response_json - -async def discover_tools(tool_endpoints: list[str], verbose) -> tuple[dict, list]: - tool_map = {} - tools = [] - - async with aiohttp.ClientSession() as session: - for url in tool_endpoints: - assert url.startswith('http://') or url.startswith('https://'), f'Tools must be URLs, not local files: {url}' - - catalog_url = f'{url}/openapi.json' - async with session.get(catalog_url) as response: - response.raise_for_status() - catalog = await response.json() - - for path, descriptor in catalog['paths'].items(): - fn = OpenAPIMethod(url=f'{url}{path}', name=path.replace('/', ' ').strip().replace(' ', '_'), descriptor=descriptor, catalog=catalog) - tool_map[fn.__name__] = fn - if verbose: - print(f'Function {fn.__name__}: params schema: {fn.parameters_schema}', file=sys.stderr) - tools.append(dict( - type='function', - function=dict( - name=fn.__name__, - description=fn.__doc__ or '', - parameters=fn.parameters_schema, - ) - ) - ) - - return tool_map, tools def typer_async_workaround(): diff --git a/examples/agent/serve_tools_inside_docker.sh b/examples/agent/serve_tools_inside_docker.sh index 898241c79..5fca28edc 100755 --- a/examples/agent/serve_tools_inside_docker.sh +++ b/examples/agent/serve_tools_inside_docker.sh @@ -9,25 +9,29 @@ set -euo pipefail PORT=${PORT:-8088} BRAVE_SEARCH_API_KEY=${BRAVE_SEARCH_API_KEY:-} +DATA_DIR=${DATA_DIR:-$HOME/.llama.cpp/agent/tools/data} +UV_CACHE_DIR=${UV_CACHE_DIR:-$HOME/.llama.cpp/agent/tools/uv_cache} -excludes=() -if [[ -z "${BRAVE_SEARCH_API_KEY:-}" ]]; then - echo "#" >&2 - echo "# Please set BRAVE_SEARCH_API_KEY environment variable in order to enable the brave_search tool" >&2 - echo "#" >&2 - excludes+=( "brave_search" ) -fi +mkdir -p "$DATA_DIR" +mkdir -p "$UV_CACHE_DIR" args=( --port $PORT "$@" ) -if [[ "${#excludes[@]}" -gt 0 ]]; then - args+=( --exclude="$(IFS=\|; echo "${excludes[*]}")" ) -fi +echo "# Warming up the uv cache" +docker run \ + -w /src \ + -v $PWD/examples/agent:/src \ + -v "$UV_CACHE_DIR":/root/.cache/uv:rw \ + --rm -it ghcr.io/astral-sh/uv:python3.12-alpine \ + uv run serve_tools.py --help echo "# Running inside docker: serve_tools.py ${args[*]}" docker run \ -p $PORT:$PORT \ -w /src \ -v $PWD/examples/agent:/src \ - --env BRAVE_SEARCH_API_KEY=$BRAVE_SEARCH_API_KEY \ + -v "$UV_CACHE_DIR":/root/.cache/uv \ + -v "$DATA_DIR":/data:rw \ + --env "MEMORY_SQLITE_DB=/data/memory.db" \ + --env "BRAVE_SEARCH_API_KEY=$BRAVE_SEARCH_API_KEY" \ --rm -it ghcr.io/astral-sh/uv:python3.12-alpine \ uv run serve_tools.py "${args[@]}"