Use tool calling and MCP servers with LLMs

%pip install -qU pixeltable openai mcp

import getpass
import os

if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key: ')

import pixeltable as pxt
from pixeltable.functions import openai

# Create a fresh directory
pxt.drop_dir('tools_demo', force=True)
pxt.create_dir('tools_demo')

# Define tool functions as Pixeltable UDFs
@pxt.udf
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    # In production, call a real weather API
    weather_data = {
        'new york': 'Sunny, 72°F',
        'london': 'Cloudy, 58°F',
        'tokyo': 'Rainy, 65°F',
        'paris': 'Partly cloudy, 68°F',
    }
    return weather_data.get(
        city.lower(), f'Weather data not available for {city}'
    )


@pxt.udf
def get_stock_price(symbol: str) -> str:
    """Get the current stock price for a symbol."""
    # In production, call a real stock API
    prices = {
        'AAPL': '$178.50',
        'GOOGL': '$141.25',
        'MSFT': '$378.90',
        'AMZN': '$185.30',
    }
    return prices.get(symbol.upper(), f'Price not available for {symbol}')

# Create a Tools object with our functions
tools = pxt.tools(get_weather, get_stock_price)

# Create table for queries
queries = pxt.create_table('tools_demo/queries', {'query': pxt.String})

# Add LLM call with tools
queries.add_computed_column(
    response=openai.chat_completions(
        messages=[{'role': 'user', 'content': queries.query}],
        model='gpt-4o-mini',
        tools=tools,  # Pass tools to the LLM
    )
)

# Automatically execute tool calls and get results
queries.add_computed_column(
    tool_results=openai.invoke_tools(tools, queries.response)
)

# Insert queries that require tool calls
sample_queries = [
    {'query': "What's the weather in Tokyo?"},
    {'query': "What's the stock price of Apple?"},
    {
        'query': "What's the weather in Paris and the price of Microsoft stock?"
    },
]

queries.insert(sample_queries)

# View results
queries.select(queries.query, queries.tool_results).collect()

from mcp.server.fastmcp import FastMCP

mcp = FastMCP('PixeltableDemo', stateless_http=True)

@mcp.tool()
def calculate_discount(price: float, discount_percent: float) -> float:
    """Calculate the discounted price."""
    return price * (1 - discount_percent / 100)

@mcp.tool()
def check_inventory(product_id: str) -> str:
    """Check inventory status for a product."""
    # In production, query your inventory database
    inventory = {
        'SKU001': 'In stock (42 units)',
        'SKU002': 'Low stock (3 units)',
        'SKU003': 'Out of stock',
    }
    return inventory.get(product_id, f'Unknown product: {product_id}')

if __name__ == '__main__':
    mcp.run(transport='streamable-http')

# Connect to the MCP server using pxt.mcp_udfs()
# This creates a Pixeltable UDF for each tool exposed by the server
# See: https://docs.pixeltable.com/platform/custom-functions#5-mcp-udfs
mcp_tools = pxt.mcp_udfs('https://docs.pixeltable.com/mcp')

# View available tools - each is now a callable Pixeltable function
for tool in mcp_tools:
    print(f'- {tool.name}: {tool.comment()}')

# Bundle MCP tools for LLM use
mcp_toolset = pxt.tools(*mcp_tools)

# Create a table with MCP tool-calling pipeline
mcp_queries = pxt.create_table(
    'tools_demo/mcp_queries', {'query': pxt.String}
)

# Add LLM call with MCP tools
mcp_queries.add_computed_column(
    response=openai.chat_completions(
        messages=[{'role': 'user', 'content': mcp_queries.query}],
        model='gpt-4o-mini',
        tools=mcp_toolset,
    )
)

# Execute MCP tool calls
mcp_queries.add_computed_column(
    tool_results=openai.invoke_tools(mcp_toolset, mcp_queries.response)
)

# View the schema - note that mcp_toolset is stored as persistent metadata
# Every subsequent insert will use these same tools automatically
mcp_queries.describe()

# Test with e-commerce queries
mcp_queries.insert(
    [
        {'query': 'What is Pixeltable?'},
        {'query': 'How to use OpenAI in Pixeltable?'},
    ]
)

mcp_queries.select(mcp_queries.query, mcp_queries.tool_results).collect()

# Extract the search result with a named column
mcp_queries.select(
    search_result=mcp_queries.tool_results[
        'SearchPixeltableDocumentation'
    ][0]
).collect()