This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fd896f247fe Add rerank support to Cohere provider (#72501)
fd896f247fe is described below
commit fd896f247fe9cdc2b15026b2385ff2c878b4e844
Author: Jeff(Wei-Hao) Lu <[email protected]>
AuthorDate: Wed Sep 9 19:04:07 2026 +0800
Add rerank support to Cohere provider (#72501)
* Add CohereRerankOperator for cohere provider
* add rerank, reranking to wordlist
* Fix Cohere rerank parameters and fix tests
---
docs/spelling_wordlist.txt | 2 +
providers/cohere/docs/index.rst | 7 +-
providers/cohere/docs/operators/rerank.rst | 71 ++++++++++++++++
providers/cohere/provider.yaml | 2 +
.../airflow/providers/cohere/get_provider_info.py | 13 ++-
.../src/airflow/providers/cohere/hooks/cohere.py | 26 ++++++
.../airflow/providers/cohere/operators/rerank.py | 95 ++++++++++++++++++++++
.../cohere/example_cohere_rerank_operator.py} | 44 +++++-----
.../cohere/tests/unit/cohere/hooks/test_cohere.py | 49 ++++++++++-
.../tests/unit/cohere/operators/test_rerank.py | 86 ++++++++++++++++++++
10 files changed, 366 insertions(+), 29 deletions(-)
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index f5d77467181..140ca2e93aa 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -1447,6 +1447,8 @@ requeue
requeued
requeues
requeuing
+rerank
+reranking
reserialize
Reserializing
resetdb
diff --git a/providers/cohere/docs/index.rst b/providers/cohere/docs/index.rst
index e3b1570d93c..262b709ff9a 100644
--- a/providers/cohere/docs/index.rst
+++ b/providers/cohere/docs/index.rst
@@ -19,16 +19,18 @@
``apache-airflow-providers-cohere``
======================================
-The ``cohere`` provider gives Dags direct access to Cohere's own Embed API —
this page
+The ``cohere`` provider gives Dags direct access to Cohere's Embed and Rerank
APIs — this page
compares that choice against ``common.ai``.
When to use this provider
--------------------------
-Use ``cohere`` when a Dag needs Cohere's native embedding models specifically:
+Use ``cohere`` when a Dag needs Cohere's native embedding or reranking models:
* ``CohereEmbeddingOperator`` — call Cohere's
`Embed API <https://docs.cohere.com/docs/embeddings>`__ directly via
``CohereHook``.
+* ``CohereRerankOperator`` — reorder documents with Cohere's
+ `Rerank API <https://docs.cohere.com/docs/rerank-overview>`__.
Use :doc:`apache-airflow-providers-common-ai:index` instead when the embedding
step should
stay vendor-neutral:
@@ -53,6 +55,7 @@ stay vendor-neutral:
Connection types <connections>
Operators <operators/embedding>
+ Rerank operator <operators/rerank>
.. toctree::
:hidden:
diff --git a/providers/cohere/docs/operators/rerank.rst
b/providers/cohere/docs/operators/rerank.rst
new file mode 100644
index 00000000000..44ae215098e
--- /dev/null
+++ b/providers/cohere/docs/operators/rerank.rst
@@ -0,0 +1,71 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/operator:CohereRerankOperator:
+
+CohereRerankOperator
+====================
+
+Use :class:`~airflow.providers.cohere.operators.rerank.CohereRerankOperator`
to reorder
+documents by their relevance to a query with Cohere's
+`Rerank API <https://docs.cohere.com/docs/rerank-overview>`__.
+
+Before you begin
+^^^^^^^^^^^^^^^^
+
+Configure a :ref:`Cohere connection <howto/connection:cohere>`. The operator
uses the
+``cohere_default`` connection unless another ``conn_id`` is provided.
+
+The operator requires:
+
+* ``query``: The search query used to evaluate relevance.
+* ``documents``: A list of text documents to rank.
+
+The operator uses ``rerank-v3.5`` by default. Set ``model`` to use another
Cohere model or
+an endpoint-specific deployment name. Use ``top_n`` to limit the number of
results and
+``max_tokens_per_doc`` to control how much of each document Cohere processes.
The query,
+documents, and both limits are templated fields; rendered limit values are
converted to integers
+before they are sent to Cohere.
+
+Using the operator
+^^^^^^^^^^^^^^^^^^
+
+.. exampleinclude::
/../../cohere/tests/system/cohere/example_cohere_rerank_operator.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_cohere_rerank]
+ :end-before: [END howto_operator_cohere_rerank]
+
+Output
+^^^^^^
+
+The operator converts the Cohere response to an XCom-serializable dictionary.
Its ``results``
+list is ordered from most to least relevant. Each result contains the original
document's
+zero-based ``index`` and its ``relevance_score``. Use the index to associate a
result with the
+corresponding item in the input ``documents`` list.
+
+An abbreviated response looks like this:
+
+.. code-block:: json
+
+ {
+ "id": "rerank-request-id",
+ "results": [
+ {"index": 1, "relevance_score": 0.99},
+ {"index": 0, "relevance_score": 0.12}
+ ]
+ }
diff --git a/providers/cohere/provider.yaml b/providers/cohere/provider.yaml
index 97e057f6f69..8e6d67ca638 100644
--- a/providers/cohere/provider.yaml
+++ b/providers/cohere/provider.yaml
@@ -62,6 +62,7 @@ integrations:
external-doc-url: https://docs.cohere.com/docs
how-to-guide:
- /docs/apache-airflow-providers-cohere/operators/embedding.rst
+ - /docs/apache-airflow-providers-cohere/operators/rerank.rst
tags: [software]
hooks:
@@ -73,6 +74,7 @@ operators:
- integration-name: Cohere
python-modules:
- airflow.providers.cohere.operators.embedding
+ - airflow.providers.cohere.operators.rerank
connection-types:
- hook-class-name: airflow.providers.cohere.hooks.cohere.CohereHook
diff --git a/providers/cohere/src/airflow/providers/cohere/get_provider_info.py
b/providers/cohere/src/airflow/providers/cohere/get_provider_info.py
index ad445d1afbd..a4f9eef9a91 100644
--- a/providers/cohere/src/airflow/providers/cohere/get_provider_info.py
+++ b/providers/cohere/src/airflow/providers/cohere/get_provider_info.py
@@ -30,7 +30,10 @@ def get_provider_info():
{
"integration-name": "Cohere",
"external-doc-url": "https://docs.cohere.com/docs";,
- "how-to-guide":
["/docs/apache-airflow-providers-cohere/operators/embedding.rst"],
+ "how-to-guide": [
+
"/docs/apache-airflow-providers-cohere/operators/embedding.rst",
+
"/docs/apache-airflow-providers-cohere/operators/rerank.rst",
+ ],
"tags": ["software"],
}
],
@@ -38,7 +41,13 @@ def get_provider_info():
{"integration-name": "Cohere", "python-modules":
["airflow.providers.cohere.hooks.cohere"]}
],
"operators": [
- {"integration-name": "Cohere", "python-modules":
["airflow.providers.cohere.operators.embedding"]}
+ {
+ "integration-name": "Cohere",
+ "python-modules": [
+ "airflow.providers.cohere.operators.embedding",
+ "airflow.providers.cohere.operators.rerank",
+ ],
+ }
],
"connection-types": [
{
diff --git a/providers/cohere/src/airflow/providers/cohere/hooks/cohere.py
b/providers/cohere/src/airflow/providers/cohere/hooks/cohere.py
index 848a70edd74..b8c583519b1 100644
--- a/providers/cohere/src/airflow/providers/cohere/hooks/cohere.py
+++ b/providers/cohere/src/airflow/providers/cohere/hooks/cohere.py
@@ -110,6 +110,32 @@ class CohereHook(BaseHook):
raise ValueError("Embeddings response is missing float_ field")
return response.embeddings.float_
+ def rerank(
+ self,
+ *,
+ query: str,
+ documents: list[str],
+ model: str = "rerank-v3.5",
+ top_n: int | None = None,
+ max_tokens_per_doc: int | None = None,
+ ) -> dict[str, Any]:
+ """Rerank documents by their relevance to a query."""
+ rerank_kwargs: dict[str, Any] = {
+ "query": query,
+ "documents": documents,
+ "model": model,
+ "request_options": self.request_options,
+ }
+ if top_n is not None:
+ rerank_kwargs["top_n"] = top_n
+ if max_tokens_per_doc is not None:
+ rerank_kwargs["max_tokens_per_doc"] = max_tokens_per_doc
+
+ response = self.get_conn().rerank(
+ **rerank_kwargs,
+ )
+ return response.model_dump(mode="json")
+
@classmethod
def get_ui_field_behaviour(cls) -> dict[str, Any]:
return {
diff --git a/providers/cohere/src/airflow/providers/cohere/operators/rerank.py
b/providers/cohere/src/airflow/providers/cohere/operators/rerank.py
new file mode 100644
index 00000000000..a77358e3f61
--- /dev/null
+++ b/providers/cohere/src/airflow/providers/cohere/operators/rerank.py
@@ -0,0 +1,95 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.cohere.hooks.cohere import CohereHook
+from airflow.providers.common.compat.sdk import BaseOperator
+
+if TYPE_CHECKING:
+ from cohere.core.request_options import RequestOptions
+
+ from airflow.providers.common.compat.sdk import Context
+
+
+class CohereRerankOperator(BaseOperator):
+ """
+ Rerank documents by their relevance to a query using Cohere's Rerank API.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:CohereRerankOperator`
+
+ :param query: The search query used to rank the documents.
+ :param documents: Text documents to compare with the query.
+ :param model: Model to use for reranking. Uses the hook's default when
omitted.
+ :param top_n: Maximum number of ranked documents to return. By default,
all are returned.
+ :param max_tokens_per_doc: Maximum number of tokens retained from each
document.
+ :param conn_id: Cohere connection id.
+ :param timeout: Request timeout in seconds.
+ :param request_options: Request-specific configuration passed to the
Cohere client.
+ """
+
+ template_fields: Sequence[str] = ("query", "documents", "top_n",
"max_tokens_per_doc")
+ template_fields_renderers = {"documents": "json"}
+
+ def __init__(
+ self,
+ *,
+ query: str,
+ documents: list[str],
+ model: str | None = None,
+ top_n: int | None = None,
+ max_tokens_per_doc: int | None = None,
+ conn_id: str = CohereHook.default_conn_name,
+ timeout: int | None = None,
+ request_options: RequestOptions | None = None,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.query = query
+ self.documents = documents
+ self.model = model
+ self.top_n = top_n
+ self.max_tokens_per_doc = max_tokens_per_doc
+ self.conn_id = conn_id
+ self.timeout = timeout
+ self.request_options = request_options
+
+ @cached_property
+ def hook(self) -> CohereHook:
+ """Return a Cohere hook."""
+ return CohereHook(
+ conn_id=self.conn_id,
+ timeout=self.timeout,
+ request_options=self.request_options,
+ )
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ """Rerank the documents and return an XCom-serializable response."""
+ rerank_kwargs: dict[str, Any] = {"query": self.query, "documents":
self.documents}
+ if self.model is not None:
+ rerank_kwargs["model"] = self.model
+ if self.top_n is not None:
+ rerank_kwargs["top_n"] = int(self.top_n)
+ if self.max_tokens_per_doc is not None:
+ rerank_kwargs["max_tokens_per_doc"] = int(self.max_tokens_per_doc)
+ return self.hook.rerank(**rerank_kwargs)
diff --git a/providers/cohere/tests/unit/cohere/hooks/test_cohere.py
b/providers/cohere/tests/system/cohere/example_cohere_rerank_operator.py
similarity index 51%
copy from providers/cohere/tests/unit/cohere/hooks/test_cohere.py
copy to providers/cohere/tests/system/cohere/example_cohere_rerank_operator.py
index f4e7a218b13..2be7067abe8 100644
--- a/providers/cohere/tests/unit/cohere/hooks/test_cohere.py
+++ b/providers/cohere/tests/system/cohere/example_cohere_rerank_operator.py
@@ -14,33 +14,29 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+
from __future__ import annotations
-from unittest.mock import patch
+from datetime import datetime
+
+from airflow import DAG
+from airflow.providers.cohere.operators.rerank import CohereRerankOperator
-from airflow.models import Connection
-from airflow.providers.cohere.hooks.cohere import (
- CohereHook,
-)
+with DAG("example_cohere_rerank", schedule=None, start_date=datetime(2023, 1,
1), catchup=False) as dag:
+ # [START howto_operator_cohere_rerank]
+ CohereRerankOperator(
+ task_id="rerank_documents",
+ query="What is the capital of the United States?",
+ documents=[
+ "Carson City is the capital city of Nevada.",
+ "Washington, D.C. is the capital of the United States.",
+ "The capital city of France is Paris.",
+ ],
+ top_n=2,
+ )
+ # [END howto_operator_cohere_rerank]
-class TestCohereHook:
- """
- Test for CohereHook
- """
+from tests_common.test_utils.system_tests import get_test_run
- def test__get_api_key(self):
- api_key = "test"
- base_url = "http://some_host.com";
- timeout = 150
- with (
- patch.object(
- CohereHook,
- "get_connection",
- return_value=Connection(conn_type="cohere", password=api_key,
host=base_url),
- ),
- patch("cohere.ClientV2") as client,
- ):
- hook = CohereHook(timeout=timeout)
- _ = hook.get_conn()
- client.assert_called_once_with(api_key=api_key, timeout=timeout,
base_url=base_url)
+test_run = get_test_run(dag)
diff --git a/providers/cohere/tests/unit/cohere/hooks/test_cohere.py
b/providers/cohere/tests/unit/cohere/hooks/test_cohere.py
index f4e7a218b13..4306113d448 100644
--- a/providers/cohere/tests/unit/cohere/hooks/test_cohere.py
+++ b/providers/cohere/tests/unit/cohere/hooks/test_cohere.py
@@ -16,7 +16,7 @@
# under the License.
from __future__ import annotations
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
from airflow.models import Connection
from airflow.providers.cohere.hooks.cohere import (
@@ -44,3 +44,50 @@ class TestCohereHook:
hook = CohereHook(timeout=timeout)
_ = hook.get_conn()
client.assert_called_once_with(api_key=api_key, timeout=timeout,
base_url=base_url)
+
+ @patch.object(CohereHook, "get_conn", autospec=True)
+ def test_rerank(self, mock_get_conn):
+ response = MagicMock(spec=["model_dump"])
+ response.model_dump.return_value = {
+ "id": "rerank-id",
+ "results": [{"index": 1, "relevance_score": 0.9}],
+ }
+ mock_get_conn.return_value.rerank.return_value = response
+ request_options = {"timeout_in_seconds": 10}
+ hook = CohereHook(request_options=request_options)
+
+ result = hook.rerank(
+ query="Where is the capital?",
+ documents=["first", "second"],
+ model="rerank-v3.5",
+ top_n=1,
+ max_tokens_per_doc=512,
+ )
+
+ mock_get_conn.return_value.rerank.assert_called_once_with(
+ query="Where is the capital?",
+ documents=["first", "second"],
+ model="rerank-v3.5",
+ top_n=1,
+ max_tokens_per_doc=512,
+ request_options=request_options,
+ )
+ response.model_dump.assert_called_once_with(mode="json")
+ assert result == {"id": "rerank-id", "results": [{"index": 1,
"relevance_score": 0.9}]}
+
+ @patch.object(CohereHook, "get_conn", autospec=True)
+ def test_rerank_uses_default_model_and_omits_unset_limits(self,
mock_get_conn):
+ response = MagicMock(spec=["model_dump"])
+ response.model_dump.return_value = {"results": []}
+ mock_get_conn.return_value.rerank.return_value = response
+ hook = CohereHook()
+
+ result = hook.rerank(query="Where is the capital?",
documents=["first", "second"])
+
+ mock_get_conn.return_value.rerank.assert_called_once_with(
+ query="Where is the capital?",
+ documents=["first", "second"],
+ model="rerank-v3.5",
+ request_options=None,
+ )
+ assert result == {"results": []}
diff --git a/providers/cohere/tests/unit/cohere/operators/test_rerank.py
b/providers/cohere/tests/unit/cohere/operators/test_rerank.py
new file mode 100644
index 00000000000..d57dba6c3de
--- /dev/null
+++ b/providers/cohere/tests/unit/cohere/operators/test_rerank.py
@@ -0,0 +1,86 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest import mock
+
+from airflow.providers.cohere.hooks.cohere import CohereHook
+from airflow.providers.cohere.operators.rerank import CohereRerankOperator
+
+
[email protected](CohereRerankOperator, "hook",
new_callable=mock.PropertyMock)
+def
test_execute_coerces_templated_limits_and_forwards_model(mock_hook_property):
+ expected = {"id": "rerank-id", "results": [{"index": 1, "relevance_score":
0.9}]}
+ mock_hook = mock.create_autospec(CohereHook, instance=True)
+ mock_hook_property.return_value = mock_hook
+ mock_hook.rerank.return_value = expected
+ documents = ["first", "second"]
+ operator = CohereRerankOperator(
+ task_id="rerank",
+ query="Where is the capital?",
+ documents=documents,
+ model="rerank-v3.5",
+ top_n="1",
+ max_tokens_per_doc="512",
+ )
+
+ result = operator.execute(context={})
+
+ mock_hook.rerank.assert_called_once_with(
+ query="Where is the capital?",
+ documents=documents,
+ model="rerank-v3.5",
+ top_n=1,
+ max_tokens_per_doc=512,
+ )
+ assert result == expected
+
+
[email protected](CohereRerankOperator, "hook",
new_callable=mock.PropertyMock)
+def test_execute_omits_unset_optional_arguments(mock_hook_property):
+ mock_hook = mock.create_autospec(CohereHook, instance=True)
+ mock_hook_property.return_value = mock_hook
+ operator = CohereRerankOperator(
+ task_id="rerank",
+ query="Where is the capital?",
+ documents=["first", "second"],
+ )
+
+ operator.execute(context={})
+
+ mock_hook.rerank.assert_called_once_with(query="Where is the capital?",
documents=["first", "second"])
+
+
[email protected]("airflow.providers.cohere.operators.rerank.CohereHook",
autospec=True)
+def test_hook_uses_operator_connection_options(mock_hook_class):
+ request_options = {"timeout_in_seconds": 10}
+ operator = CohereRerankOperator(
+ task_id="rerank",
+ query="Where is the capital?",
+ documents=["first", "second"],
+ conn_id="cohere_custom",
+ timeout=30,
+ request_options=request_options,
+ )
+
+ assert operator.hook == mock_hook_class.return_value
+ mock_hook_class.assert_called_once_with(
+ conn_id="cohere_custom",
+ timeout=30,
+ request_options=request_options,
+ )
