Support AI Gateway URLs in DatabricksOpenAI#373
Conversation
jennsun
changed the title
| from typing_extensions import override | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
There was a problem hiding this comment.
logging was to get more observability during notebook debugging (see pr description) but I can remove logging instances before merging
| logger.debug("AI Gateway V2 returned no endpoints") | ||
| return None | ||
| # Extract gateway base URL from the first endpoint's ai_gateway_url | ||
| gateway_url = endpoints[0].get("ai_gateway_url") |
There was a problem hiding this comment.
using direct http request to list ai gateway endpoints to check ai gateway status + get url so we don't assume hardcoded gateway url value here
There was a problem hiding this comment.
for context example response to calling https://dogfood.staging.databricks.com/ajax-api/ai-gateway/v2/endpoints looks like:
"endpoints": [
{
"name": "databricks-bge-large-en",
"id": "41eaa162-d05a-3401-bf27-2a040480e808",
"created_by": "Databricks",
"config": {
"destinations": [
{
"name": "system.ai.bge_large_en_v1_5",
"type": "PAY_PER_TOKEN_FOUNDATION_MODEL",
"traffic_percentage": 100,
"id": "01a22d18007038459db772d12ef43b3d"
}
],
"usage_tracking": {
"enabled": true
}
},
"created_timestamp": "1699610000000",
"last_updated_timestamp": "1699610000000",
"ai_gateway_url": "https://6051921418418893.ai-gateway.staging.cloud.databricks.com"
},
| # Default: Serving endpoints | ||
| target_base_url = f"{workspace_client.config.host}/serving-endpoints" | ||
| # Use a sync http client for the gateway probe (init is synchronous) | ||
| sync_http_client = _get_authorized_http_client(workspace_client) |
There was a problem hiding this comment.
using a sync http client here to make endpoint calls to gateway since it does a blocking http_client.get call in init
| self, | ||
| workspace_client: WorkspaceClient | None = None, | ||
| base_url: str | None = None, | ||
| use_ai_gateway: bool = True, |
There was a problem hiding this comment.
should set this to be False so existing queries to model serving endpoint agents don't break, see testing in notebook: https://dogfood.staging.databricks.com/editor/notebooks/713384057499604?o=6051921418418893#command/5345838133151630
…endpoints cuj doesnt break
aravind-segu
left a comment
aravind-segu
left a comment
There was a problem hiding this comment.
quick comment, and will approve after the logging messages are removed
| if use_ai_gateway: | ||
| gateway_url = _get_ai_gateway_base_url(http_client, workspace_client.config.host) | ||
| if gateway_url: | ||
| return gateway_url |
There was a problem hiding this comment.
nit: Instead of falling through, lets throw an error here. If the user is specifically asing for use_ai_gateway and its not enabled or something, we should not silently swallow it
There was a problem hiding this comment.
good catch - i'll raise a value error here
There was a problem hiding this comment.
Can we say:
"Please ensure AI Gateway V2 is enabled for the workspace when use_ai_gateway is set to True
…le for this workspace, remove logs
2 participants
test notebooks:
current databricks-openai package not detecting ai gateway and routing through serving endpoints on an ai gateway enabled workspace: https://dogfood.staging.databricks.com/editor/notebooks/713384057499601?o=6051921418418893
updated databricks-openai package (This branch) to auto detect ai gateway and route request through that endpoint on an ai gateway enabled workspace: https://dogfood.staging.databricks.com/editor/notebooks/713384057499604?o=6051921418418893
testing on non-gateway enabled workspace, changes enabled but queries model serving endpoint instead of ai gateway: https://eng-ml-agent-platform.staging.cloud.databricks.com/editor/notebooks/1312429564150317?o=2850744067564480
if ai gateway is not enabled and we set use_ai_gateway=true
