|
| 1 | +"""Integration tests for the Cohere wire mappings against the live Cohere API. |
| 2 | +
|
| 3 | +Gated by the ``COHERE_API_KEY`` env var: the tests run only when it is set (a |
| 4 | +live Cohere API key). Skipped in CI and local runs that don't have it in scope; |
| 5 | +runs end-to-end against the hosted Cohere endpoint when invoked with the key set. |
| 6 | +
|
| 7 | +``base_url`` is read from ``COHERE_BASE_URL`` (default ``https://api.cohere.com``), |
| 8 | +and the bound rerank model from ``COHERE_RERANK_MODEL`` (a sensible Cohere |
| 9 | +default, overridable via env). Nothing is hardcoded to a specific deployment. |
| 10 | +""" |
| 11 | + |
| 12 | +from __future__ import annotations |
| 13 | + |
| 14 | +import os |
| 15 | + |
| 16 | +import pytest |
| 17 | + |
| 18 | +from openarmature.retrieval import CohereRerankProvider, RerankRuntimeConfig |
| 19 | + |
| 20 | +_API_KEY = os.environ.get("COHERE_API_KEY") |
| 21 | +_BASE_URL = os.environ.get("COHERE_BASE_URL", "https://api.cohere.com") |
| 22 | +_RERANK_MODEL = os.environ.get("COHERE_RERANK_MODEL", "rerank-v3.5") |
| 23 | + |
| 24 | +requires_cohere = pytest.mark.skipif(not _API_KEY, reason="Requires COHERE_API_KEY (live Cohere API key)") |
| 25 | + |
| 26 | + |
| 27 | +@pytest.mark.integration |
| 28 | +@requires_cohere |
| 29 | +async def test_cohere_rerank_returns_sorted_results_with_search_units() -> None: |
| 30 | + """rerank() on a small pool returns results sorted by relevance descending, |
| 31 | + each index valid into the input documents, with a search_units usage record |
| 32 | + (Cohere meters rerank by search units, not tokens).""" |
| 33 | + documents = [ |
| 34 | + "The Sea of Tranquility was the Apollo 11 landing site.", |
| 35 | + "Cheese is made from milk.", |
| 36 | + "The lunar south pole holds water ice in permanently shadowed craters.", |
| 37 | + ] |
| 38 | + provider = CohereRerankProvider(model=_RERANK_MODEL, api_key=str(_API_KEY), base_url=_BASE_URL) |
| 39 | + try: |
| 40 | + response = await provider.rerank("Where on the moon is there water ice?", documents) |
| 41 | + finally: |
| 42 | + await provider.aclose() |
| 43 | + |
| 44 | + scores = [r.relevance_score for r in response.results] |
| 45 | + assert scores == sorted(scores, reverse=True) |
| 46 | + assert all(0 <= r.index < len(documents) for r in response.results) |
| 47 | + assert len({r.index for r in response.results}) == len(response.results) |
| 48 | + # Cohere meters rerank by search_units -> search_units; input_tokens stays null. |
| 49 | + assert response.usage is not None |
| 50 | + assert response.usage.search_units is not None |
| 51 | + assert response.usage.input_tokens is None |
| 52 | + |
| 53 | + |
| 54 | +@pytest.mark.integration |
| 55 | +@requires_cohere |
| 56 | +async def test_cohere_rerank_return_documents_is_wire_noop() -> None: |
| 57 | + """return_documents=True is a silent no-op on the Cohere wire: the results |
| 58 | + still come back (no error), and ScoredDocument.document is null on every |
| 59 | + result -- Cohere never echoes document text on /v2/rerank.""" |
| 60 | + documents = [ |
| 61 | + "The lunar maria are dark basaltic plains.", |
| 62 | + "Photosynthesis occurs in chloroplasts.", |
| 63 | + ] |
| 64 | + provider = CohereRerankProvider(model=_RERANK_MODEL, api_key=str(_API_KEY), base_url=_BASE_URL) |
| 65 | + try: |
| 66 | + response = await provider.rerank( |
| 67 | + "What are the dark plains on the moon?", |
| 68 | + documents, |
| 69 | + config=RerankRuntimeConfig(return_documents=True), |
| 70 | + ) |
| 71 | + finally: |
| 72 | + await provider.aclose() |
| 73 | + |
| 74 | + assert len(response.results) == len(documents) |
| 75 | + assert all(r.document is None for r in response.results) |
0 commit comments