Skip to content

Commit cc400a1

Browse files
authored
fix(cohere): capture rerank relevant_spans (#1010)
> 🤖 **This text was AI-generated.** Cohere v2 rerank dropped the relevant_spans request option from trace metadata and discarded the matching passages returned by v5 models: ```text before: output[] = {index, relevance_score} ``` Capture the option in the metadata allowlist and include each returned passage's character offsets and relevance score: ```text after: metadata.relevant_spans = {mode, threshold?, max_spans?} output[] = {index, relevance_score, relevant_spans?: [{start, end, score}]} ``` Omit relevant_spans when absent from a result. Add a recorded v5 example covering both cases and skip it on SDK versions without the request option. Fixes #1004 Linear: [SDK-630](https://linear.app/braintrustdata/issue/SDK-630)
1 parent d747b96 commit cc400a1

4 files changed

Lines changed: 130 additions & 20 deletions

File tree

‎.agents/skills/commit-message/SKILL.md‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -20,23 +20,21 @@ Use Conventional Commit form:
2020
```text
2121
<type>(<scope>): <imperative summary>
2222
23-
<optional body>
23+
<body>
2424
25-
<optional footer>
25+
<footer(optional)>
2626
```
2727

2828
- Choose the most specific suitable type: `feat`, `fix`, `ref` (repository convention for behavior-preserving restructuring), `perf`, `test`, `docs`, `chore`, `ci`, `style`, or `revert`.
2929
- Add a scope only when it meaningfully narrows the change.
3030
- Keep the subject at 72 characters or fewer, lowercase, imperative, and without a trailing period.
31-
- Add a body when the subject alone does not explain the purpose or a meaningful behavior change. Keep it focused on what changed and why; wrap prose at about 100 characters.
32-
- Add `BREAKING CHANGE: ...` or an issue-closing footer only when the diff or conversation supports it.
33-
34-
Use direct, specific language. Name relevant APIs, versions, errors, or files when they clarify the change. Avoid restating the diff, unsupported claims, decorative examples, and filler.
35-
36-
For substantive changes, prefer a compact code snippet or ASCII diagram in the body when it explains behavior or data flow more clearly than prose. Keep it focused and omit it when it adds no clarity.
37-
38-
## Output and commit execution
39-
40-
For a message suggestion, return one copyable message in a fenced code block. Add at most one brief sentence explaining a non-obvious type or scope choice. Ask a focused question only when the intent cannot reasonably be inferred from the diff and conversation.
31+
- always add a body. The body must:
32+
- start with a clear AI-generated notice: `> 🤖 **This text was AI-generated.**`
33+
- explain the problem and the solution like you are writing docs for a user. Use simple english and link to issues/prs/docs as appropriate
34+
- for both the problem and solution, try to include code snippets or ASCII diagrams that illustrates the change. ASCII diagrams should be in code snippet blocks
35+
- mention limitations or next steps if applicable
36+
- do not add validation or testing sections unless there is an extremely good reason to do so
37+
- Add `BREAKING CHANGE: ...` or an issue-closing footer only when the diff or conversation supports it. If the change is breaking, make sure to append `!` to the type, like `feat!`
38+
- add relevant github or linear issues to the footer if applicable
4139

4240
Do not run `git commit` or `git commit --amend` unless the user explicitly asks you to commit. When committing a multiline message, preserve real newline characters: prefer writing the message to a temporary file and using `git commit -F <file>` (or `git commit --amend -F <file>`). Multiple `-m` flags or ANSI-C quoting are acceptable for simple cases. Never assume literal `\n` inside ordinary quotes becomes a newline. After committing, verify the stored message with `git log -1 --format=medium`.

‎py/src/braintrust/integrations/cohere/cassettes/latest/test_wrap_cohere_rerank_relevant_spans.yaml‎

Lines changed: 76 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎py/src/braintrust/integrations/cohere/test_cohere.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -452,6 +452,37 @@ def test_wrap_cohere_rerank_sync(memory_logger):
452452
assert span["metrics"].get("search_units") == 1
453453

454454

455+
@pytest.mark.vcr
456+
def test_wrap_cohere_rerank_relevant_spans(memory_logger):
457+
client = _v2_client(require_methods=("rerank",))
458+
if "relevant_spans" not in inspect.signature(client.rerank).parameters:
459+
pytest.skip("Cohere rerank relevant_spans requires a newer SDK")
460+
client = wrap_cohere(client)
461+
462+
model = "rerank-v5.0-fast"
463+
documents = [
464+
"Paris is the capital of France. The Eiffel Tower is a famous landmark in Paris.",
465+
"Vienna is the capital of Austria. The city is known for its classical music.",
466+
]
467+
relevant_spans = {"mode": "enabled", "threshold": 0.3, "max_spans": 10}
468+
response = client.rerank(
469+
model=model,
470+
query="What is the capital of France?",
471+
documents=documents,
472+
top_n=2,
473+
relevant_spans=relevant_spans,
474+
)
475+
assert any(result.relevant_spans for result in response.results)
476+
assert any(result.relevant_spans is None for result in response.results)
477+
478+
spans = memory_logger.pop()
479+
assert len(spans) == 1
480+
span = spans[0]
481+
assert span["context"]["span_origin"]["instrumentation"]["name"] == "cohere-auto"
482+
assert span["metadata"]["relevant_spans"] == relevant_spans
483+
assert span["output"] == [result.model_dump(exclude_none=True) for result in response.results]
484+
485+
455486
@pytest.mark.vcr
456487
def test_wrap_cohere_chat_stream_v1_sync(memory_logger):
457488
assert not memory_logger.pop()

‎py/src/braintrust/integrations/cohere/tracing.py‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ def start_span(*args, **kwargs):
8585
"max_chunks_per_doc",
8686
"max_tokens_per_doc",
8787
"rank_fields",
88+
"relevant_spans",
8889
"priority",
8990
)
9091
_AUDIO_TRANSCRIPTION_METADATA_KEYS = (
@@ -337,8 +338,8 @@ def _embed_output(result: Any) -> dict[str, Any] | None:
337338
def _rerank_output(result: Any) -> list[dict[str, Any]] | None:
338339
"""Return the top-N rerank results (capped at 100) from a rerank response.
339340
340-
Each result is summarized to ``{index, relevance_score}`` — the document
341-
payload (if present via ``return_documents=True``) is dropped on purpose.
341+
Each result includes its index, relevance score, and relevant spans when
342+
present. The document payload (via ``return_documents=True``) is dropped.
342343
"""
343344
results = _get_field(result, "results")
344345
if not isinstance(results, list):
@@ -348,12 +349,16 @@ def _rerank_output(result: Any) -> list[dict[str, Any]] | None:
348349
for item in results[:100]:
349350
if item is None:
350351
continue
351-
out.append(
352-
{
353-
"index": _get_field(item, "index"),
354-
"relevance_score": _get_field(item, "relevance_score"),
355-
}
356-
)
352+
entry = {
353+
"index": _get_field(item, "index"),
354+
"relevance_score": _get_field(item, "relevance_score"),
355+
}
356+
relevant_spans = _get_field(item, "relevant_spans")
357+
if relevant_spans is not None:
358+
entry["relevant_spans"] = [
359+
{key: _get_field(span, key) for key in ("start", "end", "score")} for span in relevant_spans
360+
]
361+
out.append(entry)
357362
return out
358363

359364

0 commit comments

Comments
 (0)