Skip to content

feat: Add support for serialization and tracing metadata params#452

Open
jacoblee93 wants to merge 2 commits into
databricks:mainfrom
jacoblee93:jacob/tracing
Open

feat: Add support for serialization and tracing metadata params#452
jacoblee93 wants to merge 2 commits into
databricks:mainfrom
jacoblee93:jacob/tracing

Conversation

@jacoblee93

@jacoblee93 jacoblee93 commented Jul 12, 2026

Copy link
Copy Markdown

Adds support for serialization of runs and populates useful metadata parameters for tracing

@prasadkona

Copy link
Copy Markdown

Just documenting one behavioral effect of this PR for anyone reading later — no action needed:

This change causes the Databricks workspace host to be emitted into trace/run metadata on every invocation. The path:

  • _get_invocation_params now adds a databricks_host key (via _get_databricks_host()).
  • BaseChatModel.generate() passes that dict to callback_manager.on_chat_model_start(..., invocation_params=params, ...).
  • on_chat_model_start folds invocation_params into the run's extra metadata, which tracers persist — e.g. LangSmith records it under the run's extra.invocation_params, and MLflow's LangChain tracer captures it onto the span.

Two clarifications:

  • It flows through the general invocation_params metadata channel, not _get_ls_params — the new _get_ls_params only sets ls_provider="databricks" and does not include the host.
  • It applies to existing workspace_client users too, not just callers of the new databricks_host/databricks_token fields, because _get_databricks_host() falls back to self.workspace_client.config.host. So integrations that previously did not surface the workspace URL in traces will start doing so after this change.

The host is workspace-identifying but not a secret (the token is correctly kept out of this path), so this is just a heads-up on the new trace content rather than a concern.

@aravind-segu aravind-segu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a quick question. Did we do an E2E test here to see if everything works here?

"""Name of Databricks Model Serving endpoint to query."""
target_uri: Optional[str] = None
"""The target MLflow deployment URI to use. Deprecated: use workspace_client instead."""
databricks_host: Optional[str] = None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We ask users to pass in the workspace_client, is there a specific reason to explicitly pass in databricks_host and token ?

@jacoblee93 jacoblee93 Jul 15, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Our service must recreate chat model instances from serialized chat model runs. We cannot easily represent complex objects like client instances, so it's convenient here to add an initialization path that uses only simple primitives like strings.

@jacoblee93

Copy link
Copy Markdown
Author

Yes, have tested end to end

@aravind-segu aravind-segu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the change, lgtm!

@aravind-segu

Copy link
Copy Markdown
Contributor

Thanks for the clarifications offline!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants