Add Cohere Provider (#34921)

* Add Cohere Provider * Add Cohere Provider * Move link to seealso sphinx directive * Updated check for parameters * Update dependency of the cohere * Move the dag out of rst and into system tests * Add dependency to of cohere python sdk * Add cache_property for the cohere client * Remove unwanted get_conn method * Add correct label to password field * Expose timeout, max_retries and api_url to user * Fix documentation * Update interface of CohereEmbeddingOperator operator * Updated testcases * Updated testcases * Fix static check and docs build * Update CONTRIBUTING.rst Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> * Update docs/apache-airflow-providers-cohere/operators/embedding.rst Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> * Update airflow/providers/cohere/operators/embedding.py Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> * Update airflow/providers/cohere/hooks/cohere.py Co-authored-by: Hussein Awala <hussein@awala.fr> * Update airflow/providers/cohere/CHANGELOG.rst Co-authored-by: Hussein Awala <hussein@awala.fr> * Address the PR comments * Resolve conflicts * Fix breaking tests * Fix static checks * Update airflow/providers/cohere/operators/embedding.py Co-authored-by: Pankaj Singh <98807258+pankajastro@users.noreply.github.com> * Fix docstring * Add note for initial release * Add security.rst file * Update airflow/providers/cohere/hooks/cohere.py * Update airflow/providers/cohere/operators/embedding.py Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> * Update docs/apache-airflow-providers-cohere/operators/embedding.rst Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> * Add ref to sequrity.rst * Update docs/apache-airflow-providers-cohere/security.rst * Add /changelog.rst * Resolve conflicts --------- Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> Co-authored-by: Hussein Awala <hussein@awala.fr> Co-authored-by: Pankaj Singh <98807258+pankajastro@users.noreply.github.com> Co-authored-by: Pankaj <pankaj.singh@astronomer.io>
apache · Nov 6, 2023 · 7fc19d4 · 7fc19d4
1 parent 9a6ce3a
commit 7fc19d4
Show file tree

Hide file tree

Showing 35 changed files with 810 additions and 47 deletions.
diff --git a/.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml b/.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml
@@ -48,6 +48,7 @@ body:
         - celery
         - cloudant
         - cncf-kubernetes
+        - cohere
         - common-io
         - common-sql
         - daskexecutor

diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
@@ -671,9 +671,9 @@ aiobotocore, airbyte, alibaba, all, all_dbs, amazon, apache.atlas, apache.beam,
 apache.drill, apache.druid, apache.flink, apache.hdfs, apache.hive, apache.impala, apache.kafka,
 apache.kylin, apache.livy, apache.pig, apache.pinot, apache.spark, apache.sqoop, apache.webhdfs,
 apprise, arangodb, asana, async, atlas, atlassian.jira, aws, azure, cassandra, celery, cgroups,
-cloudant, cncf.kubernetes, common.io, common.sql, crypto, dask, daskexecutor, databricks, datadog,
-dbt.cloud, deprecated_api, devel, devel_all, devel_ci, devel_hadoop, dingding, discord, doc,
-doc_gen, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api, github,
+cloudant, cncf.kubernetes, cohere, common.io, common.sql, crypto, dask, daskexecutor, databricks,
+datadog, dbt.cloud, deprecated_api, devel, devel_all, devel_ci, devel_hadoop, dingding, discord,
+doc, doc_gen, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api, github,
 github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive, http, imap, influxdb, jdbc,
 jenkins, kerberos, kubernetes, ldap, leveldb, microsoft.azure, microsoft.mssql, microsoft.psrp,
 microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, openlineage, opensearch, opsgenie,

diff --git a/INSTALL b/INSTALL
@@ -98,9 +98,9 @@ aiobotocore, airbyte, alibaba, all, all_dbs, amazon, apache.atlas, apache.beam,
 apache.drill, apache.druid, apache.flink, apache.hdfs, apache.hive, apache.impala, apache.kafka,
 apache.kylin, apache.livy, apache.pig, apache.pinot, apache.spark, apache.sqoop, apache.webhdfs,
 apprise, arangodb, asana, async, atlas, atlassian.jira, aws, azure, cassandra, celery, cgroups,
-cloudant, cncf.kubernetes, common.io, common.sql, crypto, dask, daskexecutor, databricks, datadog,
-dbt.cloud, deprecated_api, devel, devel_all, devel_ci, devel_hadoop, dingding, discord, doc,
-doc_gen, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api, github,
+cloudant, cncf.kubernetes, cohere, common.io, common.sql, crypto, dask, daskexecutor, databricks,
+datadog, dbt.cloud, deprecated_api, devel, devel_all, devel_ci, devel_hadoop, dingding, discord,
+doc, doc_gen, docker, druid, elasticsearch, exasol, facebook, ftp, gcp, gcp_api, github,
 github_enterprise, google, google_auth, grpc, hashicorp, hdfs, hive, http, imap, influxdb, jdbc,
 jenkins, kerberos, kubernetes, ldap, leveldb, microsoft.azure, microsoft.mssql, microsoft.psrp,
 microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openfaas, openlineage, opensearch, opsgenie,

diff --git a/airflow/providers/cohere/.latest-doc-only-change.txt b/airflow/providers/cohere/.latest-doc-only-change.txt
@@ -0,0 +1 @@
+c645d8e40c167ea1f6c332cdc3ea0ca5a9363205
diff --git a/airflow/providers/cohere/CHANGELOG.rst b/airflow/providers/cohere/CHANGELOG.rst
@@ -0,0 +1,26 @@
+ .. 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.
+
+``apache-airflow-providers-cohere``
+
+Changelog
+---------
+
+1.0.0
+.....
+
+Initial version of the provider.
diff --git a/airflow/providers/cohere/__init__.py b/airflow/providers/cohere/__init__.py
@@ -0,0 +1,24 @@
+#
+# 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.
+#
+# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
+# OVERWRITTEN WHEN PREPARING DOCUMENTATION FOR THE PACKAGES.
+#
+# IF YOU WANT TO MODIFY IT, YOU SHOULD MODIFY THE TEMPLATE
+# `PROVIDER__INIT__PY_TEMPLATE.py.jinja2` IN the `dev/provider_packages` DIRECTORY
+#
diff --git a/airflow/providers/cohere/hooks/__init__.py b/airflow/providers/cohere/hooks/__init__.py
@@ -0,0 +1,24 @@
+#
+# 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.
+#
+# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
+# OVERWRITTEN WHEN PREPARING DOCUMENTATION FOR THE PACKAGES.
+#
+# IF YOU WANT TO MODIFY IT, YOU SHOULD MODIFY THE TEMPLATE
+# `PROVIDER__INIT__PY_TEMPLATE.py.jinja2` IN the `dev/provider_packages` DIRECTORY
+#
diff --git a/airflow/providers/cohere/hooks/cohere.py b/airflow/providers/cohere/hooks/cohere.py
@@ -0,0 +1,83 @@
+#
+# 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 functools import cached_property
+from typing import Any
+
+import cohere
+
+from airflow.hooks.base import BaseHook
+
+
+class CohereHook(BaseHook):
+    """
+    Use Cohere Python SDK to interact with Cohere platform.
+
+    .. seealso:: https://docs.cohere.com/docs
+
+    :param conn_id: :ref:`Cohere connection id <howto/connection:cohere>`
+    :param timeout: Request timeout in seconds.
+    :param max_retries: Maximal number of retries for requests.
+    """
+
+    conn_name_attr = "conn_id"
+    default_conn_name = "cohere_default"
+    conn_type = "cohere"
+    hook_name = "Cohere"
+
+    def __init__(
+        self,
+        conn_id: str = default_conn_name,
+        timeout: int | None = None,
+        max_retries: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.conn_id = conn_id
+        self.timeout = timeout
+        self.max_retries = max_retries
+
+    @cached_property
+    def get_conn(self) -> cohere.Client:
+        conn = self.get_connection(self.conn_id)
+        return cohere.Client(
+            api_key=conn.password, timeout=self.timeout, max_retries=self.max_retries, api_url=conn.host
+        )
+
+    def create_embeddings(
+        self, texts: list[str], model: str = "embed-multilingual-v2.0"
+    ) -> list[list[float]]:
+        response = self.get_conn.embed(texts=texts, model=model)
+        embeddings = response.embeddings
+        return embeddings
+
+    @staticmethod
+    def get_ui_field_behaviour() -> dict[str, Any]:
+        return {
+            "hidden_fields": ["schema", "login", "port", "extra"],
+            "relabeling": {
+                "password": "API Key",
+            },
+        }
+
+    def test_connection(self) -> tuple[bool, str]:
+        try:
+            self.get_conn.generate("Test", max_tokens=10)
+            return True, "Connection established"
+        except Exception as e:
+            return False, str(e)
diff --git a/airflow/providers/cohere/operators/__init__.py b/airflow/providers/cohere/operators/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/airflow/providers/cohere/operators/embedding.py b/airflow/providers/cohere/operators/embedding.py
@@ -0,0 +1,69 @@
+# 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 functools import cached_property
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.cohere.hooks.cohere import CohereHook
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class CohereEmbeddingOperator(BaseOperator):
+    """Creates the embedding base by interacting with cohere hosted services.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:CohereEmbeddingOperator`
+
+    :param input_text: single string text or list of text items that need to be embedded.
+    :param conn_id: Optional. The name of the Airflow connection to get connection
+        information for Cohere. Defaults to "cohere_default".
+    :param timeout: Timeout in seconds for Cohere API.
+    :param max_retries: Number of times to retry before failing.
+    """
+
+    template_fields: Sequence[str] = ("input_text",)
+
+    def __init__(
+        self,
+        input_text: list[str] | str,
+        conn_id: str = CohereHook.default_conn_name,
+        timeout: int | None = None,
+        max_retries: int | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        if isinstance(input_text, str):
+            input_text = [input_text]
+        self.conn_id = conn_id
+        self.input_text = input_text
+        self.timeout = timeout
+        self.max_retries = max_retries
+
+    @cached_property
+    def hook(self) -> CohereHook:
+        """Return an instance of the CohereHook."""
+        return CohereHook(conn_id=self.conn_id, timeout=self.timeout, max_retries=self.max_retries)
+
+    def execute(self, context: Context) -> list[list[float]]:
+        """Embed texts using Cohere embed services."""
+        return self.hook.create_embeddings(self.input_text)
diff --git a/airflow/providers/cohere/provider.yaml b/airflow/providers/cohere/provider.yaml
@@ -0,0 +1,54 @@
+# 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.
+
+---
+package-name: apache-airflow-providers-cohere
+
+name: Cohere
+
+description: |
+    `Cohere <https://docs.cohere.com/docs>`__
+
+suspended: false
+
+versions:
+  - 1.0.0
+
+integrations:
+  - integration-name: Cohere
+    external-doc-url: https://docs.cohere.com/docs
+    how-to-guide:
+      - /docs/apache-airflow-providers-cohere/operators/embedding.rst
+    tags: [software]
+
+dependencies:
+  - apache-airflow>=2.5.0
+  - cohere>=4.27
+
+hooks:
+  - integration-name: Cohere
+    python-modules:
+      - airflow.providers.cohere.hooks.cohere
+
+operators:
+  - integration-name: Cohere
+    python-modules:
+      - airflow.providers.cohere.operators.embedding
+
+connection-types:
+  - hook-class-name: airflow.providers.cohere.hooks.cohere.CohereHook
+    connection-type: cohere
diff --git a/docs/apache-airflow-providers-cohere/changelog.rst b/docs/apache-airflow-providers-cohere/changelog.rst
@@ -0,0 +1,18 @@
+ .. 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.
+
+.. include:: ../../airflow/providers/cohere/CHANGELOG.rst
diff --git a/docs/apache-airflow-providers-cohere/commits.rst b/docs/apache-airflow-providers-cohere/commits.rst
@@ -0,0 +1,19 @@
+ .. 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.
+
+Package apache-airflow-providers-cohere
+-------------------------------------------