Na początku był chaos…
A zaraz potem powstała funkcja python. Zdefiniowana całkiem normalnie, przy pomocy pythonowego def. Gdybyśmy nic więcej nie zrobili to życie tej funkcji skończyłoby się wraz z zakończeniem sesji,
def total_revenue_by_customer(customer_id: str) -> float:
query = f"SELECT SUM(amount) FROM sales.transactions WHERE customer_id = '{customer_id}'"
result = spark.sql(query).collect()[0][0]
return result
ale… my wpadliśmy na pomysł aby funkcję opublikować w UC.
Można to zrobić wykorzystując moduł DatabricksFunctionClient(), a w nim funkcje:
- client.create_python_function(func, catalog, schema, replace) – ta metoda jako argument przyjmuje wcześniej utworzoną lokalną w skali sesji funkcję – więc to jest nasz przypadek (!!!)
from databricks.sdk import WorkspaceClient
workspace_client = WorkspaceClient()
workspace_client.functions.create_python_function(
func=total_revenue_by_customer, # function name from prvious step
catalog="system",
schema="ai",
replace=True
)
W ten sposób mamy już funkcję opublikowaną w UC, ale ta funkcja jeszcze nie jest toolem agenta. Żeby funkcja stała się toolem, należy ją zarejestrować w postaci toola. Służy do tego funkcja register()
from databricks.sdk import DatabricksFunctionClient
client = DatabricksFunctionClient()
client.register(
func=total_revenue_by_customer, # this is the function name
name="total_revenue_by_customer",
description="Calculate total revenue for the client",
parameters=[{"name": "customer_id", "type": "string"}],
returns={"type": "double"}
)
Zamiast tych trzech osobnych kroków: definicja, publikacja, rejestracja, można po prostu podczas definiowania funkcji oznaczyć ją dekoratorem register()
from databricks.sdk import DatabricksFunctionClient
client = DatabricksFunctionClient()
@client.register(
name="total_revenue_by_customer",
description="Calculate total revenue for the client",
parameters=[{"name": "customer_id", "type": "string"}],
returns={"type": "double"}
)
def total_revenue_by_customer(customer_id: str) -> float:
query = f"SELECT SUM(amount) FROM sales.transactions WHERE customer_id = '{customer_id}'"
result = spark.sql(query).collect()[0][0]
return result
