unitycatalog-ai

import asyncio from unitycatalog.ai.core.client import UnitycatalogFunctionClient from unitycatalog.client import ApiClient, Configuration import os # --- Configuration --- # Replace with your Unity Catalog server host or Databricks host # For local OSS server, typically 'http://localhost:8080/api/2.1/unity-catalog' UC_HOST = os.environ.get('UC_HOST', 'http://localhost:8080/api/2.1/unity-catalog') # For Databricks, typically 'https://<your-databricks-instance>/api/2.1/unity-catalog' # Ensure you have DBR 16.4+ and serverless compute enabled for Databricks. # Replace with your Databricks token if connecting to a Databricks-managed UC DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN', 'dapi...') CATALOG = os.environ.get('UC_CATALOG', 'my_catalog') SCHEMA = os.environ.get('UC_SCHEMA', 'my_schema') # Configure the Unity Catalog API client config = Configuration( host=UC_HOST, access_token=DATABRICKS_TOKEN if 'databricks.net' in UC_HOST else None ) api_client = ApiClient(configuration=config) # Use the UnityCatalog client to create an instance of the AI function client client = UnitycatalogFunctionClient(api_client=api_client) # --- Define a Python function to be registered --- # Requirements: type hints for all arguments and return, Google-style docstring. # External imports must be within the function body for execution in sandbox mode. def add_numbers(number_1: float, number_2: float) -> float: """ A function that accepts two floating point numbers, adds them, and returns the resulting sum as a float. Args: number_1 (float): The first number. number_2 (float): The second number. Returns: float: The sum of the two numbers. """ return number_1 + number_2 async def main(): try: # Register the Python function in Unity Catalog print(f"Registering function {CATALOG}.{SCHEMA}.add_numbers...") created_function = await client.create_python_function( func=add_numbers, catalog=CATALOG, schema=SCHEMA, name="add_numbers", replace=True # Overwrite if already exists ) print(f"Function registered: {created_function.name}") # Execute the registered function print(f"Executing function {created_function.name}...") result = await client.execute_function( name=created_function.name, catalog=CATALOG, schema=SCHEMA, parameters={'number_1': 10.5, 'number_2': 20.3} ) print(f"Execution result: {result.result}") # Clean up (optional: delete the function) # print(f"Deleting function {created_function.name}...") # await client.delete_function( # name=created_function.name, # catalog=CATALOG, # schema=SCHEMA # ) # print("Function deleted.") except Exception as e: print(f"An error occurred: {e}") if __name__ == "__main__": # The client is async-based, so run in an event loop asyncio.run(main())