databricks-openai

import os from databricks_openai import DatabricksOpenAI, VectorSearchRetrieverTool from openai.types.chat import ChatCompletionMessageParam # Ensure your Databricks host and token are set as environment variables # DATABRICKS_HOST='https://<your-workspace-url>.cloud.databricks.com' # DATABRICKS_TOKEN='dapi...' # Initialize the Databricks OpenAI client # It automatically picks up DATABRICKS_HOST and DATABRICKS_TOKEN from environment variables # or Databricks CLI configuration. client = DatabricksOpenAI( host=os.environ.get('DATABRICKS_HOST', ''), token=os.environ.get('DATABRICKS_TOKEN', '') ) # Example 1: Simple chat completion using a Databricks-hosted model try: chat_completion = client.chat.completions.create( model="databricks-gpt-5-mini", # Or your specific Databricks-hosted endpoint name messages=[ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is Databricks Unity Catalog?"} ] ) print("\n--- Simple Chat Completion ---") print(f"Assistant: {chat_completion.choices[0].message.content}") except Exception as e: print(f"Error during simple chat completion: {e}") # Example 2: Chat completion with a Vector Search Retriever Tool # Replace 'catalog.schema.my_index_name' with your actual Vector Search index name index_name = os.environ.get('DATABRICKS_VECTOR_SEARCH_INDEX', 'catalog.schema.my_index_name') try: dbvs_tool = VectorSearchRetrieverTool(index_name=index_name) messages: list[ChatCompletionMessageParam] = [ {"role": "system", "content": "You are a helpful assistant that uses provided tools."}, {"role": "user", "content": "Using the Databricks documentation, what is Spark?"} ] first_response = client.chat.completions.create( model="databricks-gpt-5-mini", messages=messages, tools=[dbvs_tool.tool] ) print("\n--- Chat Completion with Tool ---") tool_call = first_response.choices[0].message.tool_calls[0] if tool_call.function.name == dbvs_tool.tool.function.name: args = json.loads(tool_call.function.arguments) # In a real scenario, this would execute on Databricks # For demonstration, we simulate a response # result = dbvs_tool.execute(query=args["query"]) result = {"docs": [{"text": "Apache Spark is an open-source, distributed processing system used for big data workloads."}]} messages.append(first_response.choices[0].message) messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": json.dumps(result)}) second_response = client.chat.completions.create( model="databricks-gpt-5-mini", messages=messages, tools=[dbvs_tool.tool] # Tools should still be passed for the second turn ) print(f"Assistant (using tool): {second_response.choices[0].message.content}") else: print(f"Assistant: {first_response.choices[0].message.content}") except Exception as e: print(f"Error during tool usage: {e}")