feat(integrations): add discovery engine v1 instrumentation - #776
Merged
Merged
Conversation
Trace Discovery Engine answer generation, conversations, ranking, and grounding
checks through Braintrust's integrations API.
## Automatic instrumentation
Install google-cloud-discoveryengine alongside braintrust, configure Google Cloud
credentials, then enable automatic instrumentation before making calls:
```python
import braintrust
from google.cloud import discoveryengine_v1 as discoveryengine
braintrust.init_logger(project="my-braintrust-project")
braintrust.auto_instrument()
client = discoveryengine.ConversationalSearchServiceClient()
response = client.answer_query(request={
"serving_config": serving_config,
"query": {"text": "What was last year's revenue?"},
})
```
Discovery Engine is enabled by default. Importing the provider before or after
setup works. Repeated setup is idempotent. To opt out while enabling other
integrations, use braintrust.auto_instrument(discoveryengine=False).
For explicit package-level setup or one client instance, use these alternatives:
```python
from braintrust.integrations.discoveryengine import (
setup_discoveryengine,
wrap_discoveryengine,
)
setup_discoveryengine() # Instrument supported Discovery Engine v1 clients.
# Alternatively, wrap only this instance:
client = wrap_discoveryengine(discoveryengine.RankServiceClient())
```
## Supported versions and APIs
The integration requires google-cloud-discoveryengine >= 0.20.3. The provider
matrix currently tests 0.20.3; later releases pass the minimum-version gate but
are not covered by this commit's recordings. Targets are the v1 clients only:
| Client (and corresponding AsyncClient) | Methods |
| --- | --- |
| ConversationalSearchServiceClient | answer_query, stream_answer_query, converse_conversation |
| RankServiceClient | rank |
| GroundedGenerationServiceClient | check_grounding |
Patch exactly these ten public sync/async methods. Search, Assistant, CRUD,
v1alpha/v1beta, and answer_query(asynchronous_mode=True) are excluded.
Defer generate_grounded_content and stream_generate_grounded_content because
live v1 endpoints returned method-not-found errors.
## Span shape and Google integration coexistence
Each supported call emits one llm span with provider="google" and origin
"discoveryengine-auto". Log a model only when requested and omit unavailable
token metrics. Streaming aggregates into one output and records time to first
text; Google's final complete answer replaces accumulated deltas. Explicit
closure/cancellation finalizes partial output without leaving the span current
between reads. Retain structured citations and cap logged ranking output at 100
records without changing the provider's returned result.
With a user-created task span around separate calls:
```text
answer workflow [task; created by the application]
|-- discoveryengine.rank [llm]
|-- discoveryengine.stream_answer_query [llm; one span for the entire stream]
`-- generate_content [llm; separate google.genai call, if made]
```
Discovery Engine targets google.cloud.discoveryengine_v1; the existing
Google GenAI integration targets google.genai. Discovery Engine calls its own
service transport rather than invoking the local GenAI client, so enabling both
does not create a duplicate GenAI span for the same Discovery Engine request.
Server-side retrieval/model work is not exposed as local child spans.
Other Google integrations retain their own patch scopes. This integration does
not patch shared Google authentication, HTTP, or gRPC code. Calls made inside an
instrumented agent/tool follow its active parent span; simply enabling all
integrations does not create an agent parent. The independent discoveryengine
and google_genai flags allow either integration to be disabled at setup.
## Validation and recording
Use real REST VCR and async gRPC recordings. The test-only gRPC helper replaces
one callable on one transport instance and restores it after each test. Replay
matches requests and reconstructs responses/errors; it does not validate network
behavior, retry timing, or real cancellation.
The final suite has 26 passing offline tests, including automatic import order,
manual/setup idempotence, span shape, stream lifecycle, and error propagation.
The broader core suite passed with 853 tests; pylint and pre-commit checks pass.
Earlier red/green checks exposed duplicate final streaming snapshots and repeated
conversation summaries before their fixes.
Document recording prerequisites and commands in the SDK VCR workflow skill.
Live recording requires BRAINTRUST_DISCOVERYENGINE_PROJECT and, for generation
resources, BRAINTRUST_DISCOVERYENGINE_APP and BRAINTRUST_DISCOVERYENGINE_DATASTORE,
plus gcloud ADC. Playback derives resource paths from cassettes and needs no
access to the recorded project. No private-resource defaults remain in test code.
Refs #773
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 528ee7c611
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| span_input, metadata = _safe_extract(_prepare, method, request, default=(None, {"provider": "google"})) | ||
| return start_span( | ||
| name=f"google_discoveryengine.{method}", | ||
| type="llm", |
There was a problem hiding this comment.
For default answer_query, converse_conversation, and check_grounding calls, _prepare() does not populate metadata.model, and none of these paths records token metrics, yet this line classifies every operation as an llm span. Such spans violate the integration span contract and cannot be attributed or priced correctly; operations for which the provider exposes neither model nor usage should use a non-LLM type, or the required LLM metadata and metrics must be populated.
Useful? React with 👍 / 👎.
Classify google_discoveryengine.rank and check_grounding as task spans: these APIs rank or assess existing content rather than generate answers. Keep answer_query, stream_answer_query, and converse_conversation as llm. Automatic setup and manual wrapping use the same classification for sync and async clients. ```text answer workflow [task] |-- google_discoveryengine.rank [task] |-- google_discoveryengine.stream_answer_query [llm] `-- google_discoveryengine.check_grounding [task] ``` Update existing REST/gRPC cassette tests: six assertions fail before the change, and all 26 integration tests pass afterward. Pylint and pre-commit pass. Reuse recordings because provider requests and responses are unchanged.
Andrew Kent (realark)
approved these changes
Sep 17, 2026
| if method == "answer_query" and _get(request, "asynchronous_mode", False): | ||
| return wrapped(*args, **kwargs) | ||
| span = _start(method, request) | ||
| if method == "stream_answer_query": |
Contributor
There was a problem hiding this comment.
my large-language friend flagged this as problematic if users abandon their stream. it seems like it's valid to just discard the stream and upon GC it auto-closes, in which case our span never ends (i don't think the "leak" classification is right though. it just seems like we don't report ending data)
AI summary:
- Abandoned streams leak an unended span (tracing.py:289) — the span only ends via _AnswerStreamState.finish, reached on exhaustion/close/cancel/error. _AnswerStream is a wrapt.ObjectProxy, not a generator, with no __del__ and an __exit__ that delegates to the wrapped object. So for chunk in stream: ... break (very common) logs a row with metrics.start and never an output or metrics.end. Tests only cover full consumption and explicit close.
gRPC's own response object cleans itself up on GC. grpc/_channel.py:564:
def __del__(self) -> None:
with self._state.condition:
if self._state.code is None:
self._state.code = grpc.StatusCode.CANCELLED
self._state.details = "Cancelled upon garbage collection!"
self._call.cancel(...)
So for chunk in stream: break is a supported pattern — dropping the reference cancels the RPC. The wrapt proxy doesn't break that (the proxy holding the last ref means GC'ing the proxy still GC's the rendezvous), but there's no equivalent hook for the span. Verified against the real 0.20.3 install with a stand-in rendezvous:
Member
Author
There was a problem hiding this comment.
yeah I should really add a note to the skill about this! thanks for flagging
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resolves https://linear.app/braintrustdata/issue/SDK-381/bot-add-google-cloud-discoveryengine-vertex-ai-search-instrumentation
Resolves #773
https://www.braintrust.dev/app/Braintrust%20SDKs/p/Google%20Discovery%20Engine%20Demo%202026-09-16/logs?r=7e04e23d2be71b8acb90051a2e656c08&s=86ab77fd0b492004
Setup
For explicit package-level setup:
Alternatively, instrument one client instance:
Supported versions and APIs
The integration accepts
google-cloud-discoveryengine >= 0.20.3; the provider matrix and recordings currently cover 0.20.3. Onlygoogle.cloud.discoveryengine_v1clients are instrumented.AsyncClient)ConversationalSearchServiceClientanswer_query,stream_answer_query,converse_conversationRankServiceClientrankGroundedGenerationServiceClientcheck_grounding