databricks-tellr-app 0.4.3.dev32__tar.gz → 0.4.3.dev33__tar.gz

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-tellr-app
-Version: 0.4.3.dev32
+Version: 0.4.3.dev33
 Summary: Tellr application package for Databricks Apps
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/databricks_tellr_app.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-tellr-app
-Version: 0.4.3.dev32
+Version: 0.4.3.dev33
 Summary: Tellr application package for Databricks Apps
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/databricks_tellr_app.egg-info/SOURCES.txt

@@ -147,6 +147,7 @@ setup.py
 ./src/services/design_system_templates.py
 ./src/services/drive_uploader.py
 ./src/services/foreman_service.py
+./src/services/gateway_openai.py
 ./src/services/genie_service.py
 ./src/services/google_slides_auth.py
 ./src/services/google_slides_prompts_defaults.py
@@ -371,6 +372,7 @@ src/services/design_system_service.py
 src/services/design_system_templates.py
 src/services/drive_uploader.py
 src/services/foreman_service.py
+src/services/gateway_openai.py
 src/services/genie_service.py
 src/services/google_slides_auth.py
 src/services/google_slides_prompts_defaults.py

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "databricks-tellr-app"
-version = "0.4.3.dev32"
+version = "0.4.3.dev33"
 description = "Tellr application package for Databricks Apps"
 readme = "README.md"
 requires-python = ">=3.10"

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/api/services/chat_service.py

@@ -1445,16 +1445,9 @@ class ChatService:
         def run_title_gen():
             """Generate a session title in parallel with the main agent."""
             try:
-                from databricks_langchain import ChatDatabricks
-                from src.core.databricks_client import get_user_client
-
-                from src.core.defaults import DEFAULT_CONFIG
-                naming_model = ChatDatabricks(
-                    endpoint=DEFAULT_CONFIG["llm"]["endpoint"],
-                    max_tokens=50,
-                    temperature=0.3,
-                    workspace_client=get_user_client(),
-                )
+                from src.api.services.session_naming import build_session_title_model
+
+                naming_model = build_session_title_model()
                 generated_title = generate_session_title(message, naming_model)
                 if generated_title:
                     session_manager.rename_session(session_id, generated_title)
@@ -1999,17 +1992,9 @@ class ChatService:
             session must never fail the turn, so this catches and logs.
             """
             try:
-                from databricks_langchain import ChatDatabricks
+                from src.api.services.session_naming import build_session_title_model
 
-                from src.core.databricks_client import get_user_client
-                from src.core.defaults import DEFAULT_CONFIG
-
-                naming_model = ChatDatabricks(
-                    endpoint=DEFAULT_CONFIG["llm"]["endpoint"],
-                    max_tokens=50,
-                    temperature=0.3,
-                    workspace_client=get_user_client(),
-                )
+                naming_model = build_session_title_model()
                 generated_title = generate_session_title(message, naming_model)
                 if generated_title:
                     session_manager.rename_session(session_id, generated_title)

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/api/services/session_naming.py

@@ -88,6 +88,22 @@ def _was_truncated(response: object) -> bool:
     return False
 
 
+def build_session_title_model():
+    """The title model: Gateway-routed, on the requesting user's OBO identity."""
+    from databricks_langchain import ChatDatabricks
+
+    from src.core.databricks_client import get_user_client
+    from src.core.defaults import SESSION_TITLE_MODEL
+
+    return ChatDatabricks(
+        model=SESSION_TITLE_MODEL,
+        use_ai_gateway=True,
+        max_tokens=50,
+        temperature=0.3,
+        workspace_client=get_user_client(),
+    )
+
+
 def generate_session_title(
     user_message: str,
     model,

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/core/defaults.py

@@ -232,3 +232,7 @@ This means the user wants to modify these specific slides. Your response should:
    For these operations, respond with a helpful message guiding users to the UI or workaround - do NOT attempt to return HTML.""",
     },
 }
+
+#: Session-title generation (both engine paths) goes through Unity AI Gateway (ws2a).
+#: Kept separate from DEFAULT_CONFIG["llm"], which the out-of-scope monolith still reads.
+SESSION_TITLE_MODEL = "system.ai.claude-opus-4-6"

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/services/agent_runtime.py

@@ -231,7 +231,7 @@ def bind_structured_output_model(
     the plain structured-output binding of ``schema``.
     """
     model = model_factory(
-        endpoint=configuration.endpoint_name,
+        model=configuration.endpoint_name,
         temperature=configuration.temperature,
         max_tokens=configuration.max_tokens,
         top_p=configuration.top_p,
@@ -322,7 +322,9 @@ class DatabricksModelAdapter:
     def _default_model_factory(**kwargs: Any) -> Any:
         from databricks_langchain import ChatDatabricks  # type: ignore[import-untyped]
 
-        return ChatDatabricks(**kwargs)
+        # Workstream 2a: every graph model call goes through Unity AI Gateway
+        # ({host}/ai-gateway/mlflow/v1).  There is no /serving-endpoints fallback.
+        return ChatDatabricks(use_ai_gateway=True, **kwargs)
 
     @staticmethod
     def _default_client_factory() -> Any:

databricks_tellr_app-0.4.3.dev33/src/services/gateway_openai.py

@@ -0,0 +1,12 @@
+"""The one OpenAI-compatible client for app LLM calls through Unity AI Gateway (ws2a)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def gateway_openai_client(workspace_client: Any) -> Any:
+    """``DatabricksOpenAI`` bound to ``{host}/ai-gateway/mlflow/v1`` with the caller's identity."""
+    from databricks_openai import DatabricksOpenAI
+
+    return DatabricksOpenAI(workspace_client=workspace_client, use_ai_gateway=True)

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/services/graph_configuration_draft.py

@@ -44,6 +44,7 @@ from src.services.model_endpoint_catalog import (
     EndpointValidationFailure,
     ModelEndpointCatalog,
     validate_endpoint_name_policy,
+    validate_gateway_model_name,
 )
 from src.services.prompt_assembler import PromptAssembler, PromptAssemblyRejected
 
@@ -285,6 +286,24 @@ def _endpoint_name_policy_validator(
     return ()
 
 
+def _gateway_model_name_validator(
+    current: DefinitionContent, candidate: DefinitionContent
+) -> tuple[DraftValidationIssue, ...]:
+    """ws2a: a save that CHANGES the model must name a system.ai.* Gateway model.
+
+    Save-only by design.  It is never composed into ``_save_local_validators``,
+    which publication and rollback reuse, because rollback re-validates historical
+    releases whose roles name ``databricks-*`` endpoints (spec §5.2).
+    """
+    if candidate.model.endpoint_name == current.model.endpoint_name:
+        return ()
+    try:
+        validate_gateway_model_name(candidate.model.endpoint_name)
+    except EndpointValidationFailure as failure:
+        return (_endpoint_issue(failure),)
+    return ()
+
+
 class RemoteEndpointDraftValidator(Protocol):
     """Remote exact-endpoint check over one complete reconstructed candidate.
 
@@ -449,6 +468,14 @@ class _GraphConfigurationDraft(_GraphConfigurationWorkbench):
                     client_candidate=candidate,
                     server=self._draft_aggregate_snapshot(locked.snapshot),
                 )
+            # Post-lock (ruling I-1): the locked draft is now the client's own
+            # base, so "changed" means changed by this request, never by a
+            # concurrent save that a stale client has not seen.
+            gateway_issues = _gateway_model_name_validator(
+                locked.selected.draft.content, content
+            )
+            if gateway_issues:
+                raise DraftContentRejected(*gateway_issues)
             self._run_candidate_validators(self.post_stale_validators, content)
             self._validate_remote_endpoint(content)
             return self._write_locked_content(
@@ -508,6 +535,14 @@ class _GraphConfigurationDraft(_GraphConfigurationWorkbench):
                     client_candidate=validated,
                     server=self._draft_aggregate_snapshot(locked.snapshot),
                 )
+            # Post-lock (ruling I-1): the locked draft is now the client's own
+            # base, so "changed" means changed by this request, never by a
+            # concurrent save that a stale client has not seen.
+            gateway_issues = _gateway_model_name_validator(
+                locked.selected.draft.content, validated
+            )
+            if gateway_issues:
+                raise DraftContentRejected(*gateway_issues)
             self._run_candidate_validators(self.post_stale_validators, validated)
             self._validate_remote_endpoint(validated)
             return self._write_locked_content(

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/services/html_to_google_slides.py

@@ -16,6 +16,7 @@ from typing import Callable, Dict, List, Optional, Tuple
 from databricks.sdk import WorkspaceClient
 
 from src.services.converter_jail.codeprep import _fix_apostrophe_strings
+from src.services.gateway_openai import gateway_openai_client
 from src.services.google_slides_auth import GoogleSlidesAuth, GoogleSlidesAuthError
 from src.services.google_slides_prompts_defaults import (
     DEFAULT_GSLIDES_SYSTEM_PROMPT,
@@ -339,7 +340,7 @@ class GoogleSlidesConversionError(Exception):
 class HtmlToGoogleSlidesConverter:
     """LLM-powered converter: HTML → Google Slides API code → execution."""
 
-    DEFAULT_MODEL = "databricks-claude-sonnet-4-5"
+    DEFAULT_MODEL = "system.ai.claude-sonnet-4-5"
 
     def __init__(
         self,
@@ -355,7 +356,7 @@ class HtmlToGoogleSlidesConverter:
             from src.core.databricks_client import get_databricks_client
             self.ws_client = get_databricks_client()
 
-        self.llm_client = self.ws_client.serving_endpoints.get_open_ai_client()
+        self.llm_client = gateway_openai_client(self.ws_client)
         self.auth = google_auth or GoogleSlidesAuth()
         self.SYSTEM_PROMPT = DEFAULT_GSLIDES_SYSTEM_PROMPT
         self.USER_PROMPT = DEFAULT_GSLIDES_USER_PROMPT

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/services/html_to_pptx.py

@@ -13,6 +13,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple
 from bs4 import BeautifulSoup
 from databricks.sdk import WorkspaceClient
 
+from src.services.gateway_openai import gateway_openai_client
 from src.services.pptx_prompts_defaults import (
     DEFAULT_SYSTEM_PROMPT,
     DEFAULT_USER_PROMPT_TEMPLATE,
@@ -48,7 +49,7 @@ class HtmlToPptxConverterV3:
     """
     
     # Model configuration
-    DEFAULT_MODEL = "databricks-claude-sonnet-4-5"
+    DEFAULT_MODEL = "system.ai.claude-sonnet-4-5"
     
 
     def __init__(
@@ -60,19 +61,19 @@ class HtmlToPptxConverterV3:
         
         Args:
             workspace_client: Databricks client (optional, uses singleton if not provided)
-            model_endpoint: LLM model name (default: databricks-claude-sonnet-4-5)
+            model_endpoint: LLM model name (default: system.ai.claude-sonnet-4-5)
         """
         self.model_endpoint = model_endpoint or self.DEFAULT_MODEL
-        
+
         # Use provided client or get singleton from databricks_client
         if workspace_client:
             self.ws_client = workspace_client
         else:
             from src.core.databricks_client import get_databricks_client
             self.ws_client = get_databricks_client()
-        
-        # Initialize OpenAI-compatible client for LLM calls
-        self.llm_client = self.ws_client.serving_endpoints.get_open_ai_client()
+
+        # Initialize OpenAI-compatible client for LLM calls through Unity AI Gateway
+        self.llm_client = gateway_openai_client(self.ws_client)
         
         # Load default prompts
         self.SYSTEM_PROMPT = DEFAULT_SYSTEM_PROMPT

{databricks_tellr_app-0.4.3.dev32 → databricks_tellr_app-0.4.3.dev33}/src/services/model_endpoint_catalog.py

@@ -1,4 +1,4 @@
-"""Exact Databricks foundation-model endpoint discovery and validation."""
+"""Unity AI Gateway model discovery and validation for the workbench (ws2a)."""
 
 from __future__ import annotations
 
@@ -9,8 +9,7 @@ from typing import Any, Literal, Protocol
 
 import requests
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import DatabricksError, PermissionDenied, ResourceDoesNotExist
-from databricks.sdk.service.serving import EndpointStateConfigUpdate, EndpointStateReady
+from databricks.sdk.errors import DatabricksError, NotFound, PermissionDenied
 
 CatalogFailureCode = Literal["catalog_forbidden", "catalog_unavailable"]
 EndpointValidationCode = Literal[
@@ -23,6 +22,8 @@ EndpointValidationCode = Literal[
     "endpoint_update_in_progress",
     "endpoint_update_failed",
     "endpoint_update_canceled",
+    "endpoint_not_gateway_model",
+    "endpoint_not_chat_model",
 ]
 
 
@@ -72,6 +73,27 @@ _TRANSPORT_FAILURES: tuple[type[BaseException], ...] = (
 )
 
 
+GATEWAY_ENDPOINTS_PATH = "/api/ai-gateway/v2/endpoints"
+GATEWAY_CHAT_API_TYPE = "mlflow/v1/chat/completions"
+_GATEWAY_PREFIX = "databricks-"
+_SYSTEM_AI_PREFIX = "system.ai."
+
+
+def gateway_invocable_name(gateway_endpoint_name: str) -> str | None:
+    """``databricks-<m>`` -> ``system.ai.<m>`` (spec §2 naming rule); otherwise ``None``."""
+    if not gateway_endpoint_name.startswith(_GATEWAY_PREFIX):
+        return None
+    model = gateway_endpoint_name[len(_GATEWAY_PREFIX):]
+    return f"{_SYSTEM_AI_PREFIX}{model}" if model else None
+
+
+def gateway_endpoint_name(model_name: str) -> str:
+    """``system.ai.<m>`` -> ``databricks-<m>``; any other (legacy) name is looked up as-is."""
+    if model_name.startswith(_SYSTEM_AI_PREFIX):
+        return f"{_GATEWAY_PREFIX}{model_name[len(_SYSTEM_AI_PREFIX):]}"
+    return model_name
+
+
 class ModelEndpointCatalog(Protocol):
     def list_system_models(self) -> SystemModelDiscovery: ...
 
@@ -103,6 +125,20 @@ def validate_endpoint_name_policy(name: str) -> None:
         )
 
 
+_GATEWAY_MODEL_NAME = re.compile(r"^system\.ai\.[a-z0-9][a-z0-9._-]*[a-z0-9]$")
+
+
+def validate_gateway_model_name(name: str) -> None:
+    """A changed draft model must be a ``system.ai.*`` Gateway name (spec §5.2)."""
+    if not _GATEWAY_MODEL_NAME.fullmatch(name):
+        raise EndpointValidationFailure(
+            "endpoint_not_gateway_model",
+            "Model name must start with `system.ai.` and contain only lowercase letters, "
+            "digits, hyphens, underscores and periods.",
+            False,
+        )
+
+
 #: The remote check runs while the save holds the exclusive draft lock, so its
 #: client must not inherit the SDK's 300 s default retry window.  With a 5 s
 #: window and a 3 s per-request timeout, measured transport exhaustion ends
@@ -162,7 +198,7 @@ class DatabricksModelEndpointCatalog:
 
     def list_system_models(self) -> SystemModelDiscovery:
         try:
-            endpoints = tuple(self._workspace_client.serving_endpoints.list())
+            listed = self._workspace_client.api_client.do("GET", GATEWAY_ENDPOINTS_PATH)
         except PermissionDenied as error:
             raise ModelEndpointCatalogFailure(
                 "catalog_forbidden",
@@ -177,45 +213,31 @@ class DatabricksModelEndpointCatalog:
             ) from error
 
         discovered: list[SystemModelEndpoint] = []
-        for endpoint in endpoints:
-            served_entities = getattr(endpoint.config, "served_entities", None) or ()
-            foundation_model = next(
-                (
-                    getattr(entity, "foundation_model", None)
-                    for entity in served_entities
-                    if getattr(entity, "foundation_model", None) is not None
-                ),
-                None,
-            )
-            if foundation_model is None:
-                continue
-
-            name = getattr(endpoint, "name", None)
+        for entry in (listed or {}).get("endpoints") or ():
+            name = entry.get("name") if isinstance(entry, dict) else None
             if not isinstance(name, str) or not name.strip():
                 raise ModelEndpointCatalogFailure(
                     "catalog_unavailable",
                     "Model endpoint discovery returned an endpoint without a name.",
                     True,
                 )
-
+            invocable = gateway_invocable_name(name)
+            if invocable is None:
+                continue
             discovered.append(
-                SystemModelEndpoint(
-                    name=name,
-                    display_name=getattr(foundation_model, "display_name", None),
-                    description=getattr(foundation_model, "description", None),
-                    docs=getattr(foundation_model, "docs", None),
-                )
+                SystemModelEndpoint(name=invocable, display_name=None, description=None, docs=None)
             )
 
-        discovered.sort(key=lambda item: ((item.display_name or item.name).casefold(), item.name))
+        discovered.sort(key=lambda item: item.name)
         return SystemModelDiscovery(endpoints=tuple(discovered))
 
     def validate_custom_endpoint_remote(self, name: str) -> None:
         # Defensive: never let a path-shaped name reach the interpolated request.
         validate_endpoint_name_policy(name)
+        path = f"{GATEWAY_ENDPOINTS_PATH}/{gateway_endpoint_name(name)}"
         try:
-            detail = self._workspace_client.serving_endpoints.get(name)
-        except ResourceDoesNotExist as error:
+            detail = self._workspace_client.api_client.do("GET", path)
+        except NotFound as error:
             raise _validation_failure(
                 "endpoint_unknown", "Endpoint name was not found.", False
             ) from error
@@ -232,41 +254,10 @@ class DatabricksModelEndpointCatalog:
                 True,
             ) from error
 
-        if getattr(detail, "name", None) != name:
-            raise _validation_failure(
-                "endpoint_name_mismatch",
-                "Endpoint validation did not return the exact requested name.",
-                False,
-            )
-
-        state = getattr(detail, "state", None)
-        config_update = getattr(state, "config_update", None)
-        if config_update == EndpointStateConfigUpdate.IN_PROGRESS:
-            raise _validation_failure(
-                "endpoint_update_in_progress",
-                "Endpoint configuration update is in progress.",
-                True,
-            )
-        if config_update == EndpointStateConfigUpdate.UPDATE_FAILED:
-            raise _validation_failure(
-                "endpoint_update_failed",
-                "Endpoint configuration update failed.",
-                False,
-            )
-        if config_update == EndpointStateConfigUpdate.UPDATE_CANCELED:
+        api_types = (detail or {}).get("supported_api_types") or ()
+        if GATEWAY_CHAT_API_TYPE not in api_types:
             raise _validation_failure(
-                "endpoint_update_canceled",
-                "Endpoint configuration update was canceled.",
-                False,
-            )
-        if (
-            getattr(state, "ready", None) != EndpointStateReady.READY
-            or config_update != EndpointStateConfigUpdate.NOT_UPDATING
-        ):
-            raise _validation_failure(
-                "endpoint_not_ready",
-                "Endpoint is not ready for invocation.",
-                True,
+                "endpoint_not_chat_model", "Endpoint is not a chat model.", False
             )