Tool CallingでUnity Catalogに登録されているツールを呼んでみる

以下記事の続きです。

https://qiita.com/isanakamishiro2/items/f889a2be243eb972327f

# 導入

Databricks Unity Catalogに登録されたツールを、前に作ったTool Calling Agentから呼んで実行してみます。

とはいえ、以下のdbdemosで解説されていることなので、ほとんどその内容のウォークスルーみたいなものです。
以下の内容は自作のTool Callingエージェントでしか動作しないわけではなく、OpenAI APIのTool Callingでも問題なく動作します。

https://www.databricks.com/resources/demos/tutorials/data-science-and-ai/llm-tools-functions

# Step1. ツールの登録

上記のdbdemosに含まれているツールをいくつかUnity Catalogに登録します。
なお、今回は全て```training.llm```カタログ・スキーマにツールを登録しています。

まずはインチをセンチ(cm)に変換するSQL関数を登録。

```sql
%sql
USE CATALOG training;
USE SCHEMA llm;

CREATE OR REPLACE FUNCTION convert_inch_to_cm(size_in_inch FLOAT)
RETURNS FLOAT
LANGUAGE SQL
COMMENT 'convert size from inch to cm'
RETURN size_in_inch * 2.54;

-- let's test our function:
SELECT convert_inch_to_cm(10) as 10_inches_in_cm;
```

次に緯度経度を指定して天気情報を取得するPython関数を登録。

:::note info
ダミー処理になっています。
実際に天気APIを呼び出した場合はNoteの内容を確認の上、コード中のコメントを外してください。
:::


```python
%sql
CREATE OR REPLACE FUNCTION get_weather(latitude DOUBLE, longitude DOUBLE)
RETURNS STRUCT<temperature_in_celsius DOUBLE, rain_in_mm DOUBLE>
LANGUAGE PYTHON
COMMENT 'This function retrieves the current temperature and rain information for a given latitude and longitude using the Open-Meteo API.'
AS
$$
  try:
    import requests as r
    # Note: this is provided for education only, non commercial - please get a license for real usage: https://api.open-meteo.com. Let s comment it to avoid issues for now
    # weather = r.get(f'https://api.open-meteo.com/v1/forecast?latitude={latitude}&longitude={longitude}&current=temperature_2m,rain&forecast_days=1').json()
    return {
      "temperature_in_celsius": weather["current"]["temperature_2m"],
      "rain_in_mm": weather["current"]["rain"]
    }
  except:
    return {"temperature_in_celsius": 25.0, "rain_in_mm": 0.0}
$$;

-- let's test our function:
SELECT get_weather(52.52, 13.41) as weather;
```

最後に数値計算ツール。

```python
%sql
USE CATALOG training;
USE SCHEMA llm;

CREATE OR REPLACE FUNCTION compute_math(expr STRING)
RETURNS STRING
LANGUAGE PYTHON
COMMENT 'Run any mathematical function and returns the result as output. Supports python syntax like math.sqrt(13)'
AS
$$
  import ast
  import operator
  import math
  operators = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, ast.Pow: operator.pow, ast.Mod: operator.mod, ast.FloorDiv: operator.floordiv, ast.UAdd: operator.pos, ast.USub: operator.neg}
    
  # Supported functions from the math module
  functions = {name: getattr(math, name) for name in dir(math) if callable(getattr(math, name))}

  def eval_node(node):
    if isinstance(node, ast.Num):  # <number>
      return node.n
    elif isinstance(node, ast.BinOp):  # <left> <operator> <right>
      return operators[type(node.op)](eval_node(node.left), eval_node(node.right))
    elif isinstance(node, ast.UnaryOp):  # <operator> <operand> e.g., -1
      return operators[type(node.op)](eval_node(node.operand))
    elif isinstance(node, ast.Call):  # <func>(<args>)
      func = node.func.id
      if func in functions:
        args = [eval_node(arg) for arg in node.args]
        return functions[func](*args)
      else:
        raise TypeError(f"Unsupported function: {func}")
    else:
      raise TypeError(f"Unsupported type: {type(node)}")  
  try:
    if expr.startswith('```') and expr.endswith('```'):
      expr = expr[3:-3].strip()      
    node = ast.parse(expr, mode='eval').body
    return eval_node(node)
  except Exception as ex:
    return str(ex)
$$;

-- let's test our function:
SELECT compute_math("(2+2)/3") as result;
```

これでUnity Catalog上に３つの関数が登録されました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/b7c0f2a3-5dd2-b800-758f-c62160392a78.png)

# Step2. ツールの呼び出し

Step1.で作成したツールをTool Callingで実行判断をさせて、さらに実際に実行してみます。

ノートブックを作成し、langchainなどのパッケージをインストール。

```python
%pip install -q -U langchain-core==0.3.8 langgraph==0.2.34 langchain-databricks==0.1.0 langchain_community==0.3.1

dbutils.library.restartPython()
```

Unity Catalogからツール情報を取得する関数```get_tools```を定義。
```warehouse_id```には、Unity Catalogから関数のメタ情報を参照するために使用するSQLウェアハウスのIDを設定する必要があります。
(SQLウェアハウスの画面からIDを確認できます）

```python
from langchain_community.tools.databricks.tool import UCFunctionToolkit

catalog = "training"
db = "llm"
warehouse_id = "xxxxxxxxxxx"

def get_tools():
    return (
        UCFunctionToolkit(warehouse_id=warehouse_id)
        .include(f"{catalog}.{db}.*")
        .get_tools())

get_tools()
```

```:出力
[StructuredTool(name='training__llm__compute_math', description='Run any mathematical function and returns the result as output. Supports python syntax like math.sqrt(13)', args_schema=<class 'langchain_community.tools.databricks.tool.training__llm__compute_math__params'>, func=<function UCFunctionToolkit._make_tool.<locals>.func at 0x7faae7f9ac20>),
 StructuredTool(name='training__llm__convert_inch_to_cm', description='convert size from inch to cm', args_schema=<class 'langchain_community.tools.databricks.tool.training__llm__convert_inch_to_cm__params'>, func=<function UCFunctionToolkit._make_tool.<locals>.func at 0x7faae80b03a0>),
 StructuredTool(name='training__llm__get_weather', description='This function retrieves the current temperature and rain information for a given latitude and longitude using the Open-Meteo API.', args_schema=<class 'langchain_community.tools.databricks.tool.training__llm__get_weather__params'>, func=<function UCFunctionToolkit._make_tool.<locals>.func at 0x7faae74bd900>)]
```

3つのツールが取得できているのがわかります。

あとはLangChain/LangGraphの通常のツール呼び出しの仕方と同じです。

まずは計算ツールが使われることを想定したクエリをエージェントに渡してみます。

```python
import json
from langchain_databricks import ChatDatabricks
from langgraph.prebuilt import ToolNode

endpoint = "test_tool_calling_agent"

chat = ChatDatabricks(endpoint=endpoint).bind_tools(get_tools())
message_with_single_tool_call = chat.invoke(
    [{"role": "user", "content": "100+25*2はいくつ？"}],
)

tool_node = ToolNode(get_tools())
print(tool_node.invoke({"messages": [message_with_single_tool_call]}))
```

```:出力
{'messages': [ToolMessage(content='{"format": "SCALAR", "value": "150", "truncated": false}', name='training__llm__compute_math', tool_call_id='1')]}
```

想定通りに適切なツールが選択されて、計算も実行されました。

他にも何個か実行してみます。

```python
message_with_single_tool_call = chat.invoke(
    [{"role": "user", "content": "サンフランシスコの天気は？"}],
)

tool_node = ToolNode(get_tools())
print(tool_node.invoke({"messages": [message_with_single_tool_call]}))
```

```:出力
{'messages': [ToolMessage(content='{"format": "SCALAR", "value": "{\\"temperature_in_celsius\\":\\"25.0\\",\\"rain_in_mm\\":\\"0.0\\"}", "truncated": false}', name='training__llm__get_weather', tool_call_id='1')]}
```


```python
message_with_single_tool_call = chat.invoke(
    [{"role": "user", "content": "10インチは何センチ？"}],
)

tool_node = ToolNode(get_tools())
print(tool_node.invoke({"messages": [message_with_single_tool_call]}))
```

```:出力
{'messages': [ToolMessage(content='{"format": "SCALAR", "value": "25.4", "truncated": false}', name='training__llm__convert_inch_to_cm', tool_call_id='1')]}
```

想定通りの関数が呼べていますね。

# まとめ

LLMのTool Callingを行うにあたってツールをどのように管理するべきか？というのは結構難しい課題だと思います。Unity Catalogの関数管理はそのための十分な機能を提供しており、理にかなってるなと感じています。

こういった機能が充足してくると、エージェント構築が楽しくなってきますね。
だから東京リージョンにも早くMosaic AI Agent Frameworkを(ry