メインコンテンツまでスキップ

LlamaIndex と Databricks Unity Catalog ツールの統合

Databricks Unity Catalog を使用して、SQL 関数と Python 関数を LlamaIndex ワークフローのツールとして統合します。この統合により、Unity Catalog のガバナンスと LlamaIndex の機能が組み合わされ、LLM の大規模なデータセットのインデックス作成とクエリが可能になります。

必要条件

  • Python 3.10以降をインストールします。

Unity Catalog ツールを LlamaIndex と統合する

ノートブックまたは Python スクリプトで次のコードを実行して、Unity Catalog ツールを作成し、LlamaIndex エージェントで使用します。

  1. LlamaIndex の Databricks Unity Catalog 統合パッケージをインストールします。

    Python
    %pip install unitycatalog-llamaindex[databricks]
    dbutils.library.restartPython()
  2. Unity Catalog 関数クライアントのインスタンスを作成します。

    Python
    from unitycatalog.ai.core.base import get_uc_function_client

    client = get_uc_function_client()
  3. Python で記述された Unity Catalog 関数を作成します。

    Python
    CATALOG = "your_catalog"
    SCHEMA = "your_schema"

    func_name = f"{CATALOG}.{SCHEMA}.code_function"

    def code_function(code: str) -> str:
    """
    Runs Python code.

    Args:
    code (str): The Python code to run.
    Returns:
    str: The result of running the Python code.
    """
    import sys
    from io import StringIO
    stdout = StringIO()
    sys.stdout = stdout
    exec(code)
    return stdout.getvalue()

    client.create_python_function(
    func=code_function,
    catalog=CATALOG,
    schema=SCHEMA,
    replace=True
    )
  4. Unity Catalog 関数のインスタンスをツールキットとして作成し、それを実行してツールが正しく動作することを確認します。

    Python
    from unitycatalog.ai.llama_index.toolkit import UCFunctionToolkit
    import mlflow

    # Enable traces
    mlflow.llama_index.autolog()

    # Create a UCFunctionToolkit that includes the UC function
    toolkit = UCFunctionToolkit(function_names=[func_name])

    # Fetch the tools stored in the toolkit
    tools = toolkit.tools
    python_exec_tool = tools[0]

    # Run the tool directly
    result = python_exec_tool.call(code="print(1 + 1)")
    print(result) # Outputs: {"format": "SCALAR", "value": "2\n"}
  5. LlamaIndex ReActAgent でツールを使用するには、Unity Catalog 関数を LlamaIndex ツール コレクションの一部として定義します。次に、エージェントが LlamaIndex ツール コレクションを呼び出して、エージェントが正しく動作することを確認します。

    Python
    from llama_index.llms.openai import OpenAI
    from llama_index.core.agent import ReActAgent

    llm = OpenAI()

    agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)

    agent.chat("Please run the following python code: `print(1 + 1)`")
この記事は役に立ちましたか?