Add OpenAI Provider

.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml

@@ -82,6 +82,7 @@ body:
         - odbc
         - openfaas
         - openlineage
+        - openai
         - opensearch
         - opsgenie
         - oracle

CONTRIBUTING.rst

@@ -676,11 +676,11 @@ datadog, dbt.cloud, deprecated_api, devel, devel_all, devel_ci, devel_hadoop, di
 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,
-oracle, otel, pagerduty, pandas, papermill, password, pgvector, pinecone, pinot, plexus, postgres,
-presto, rabbitmq, redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp, singularity,
-slack, smtp, snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino, vertica,
-virtualenv, weaviate, webhdfs, winrm, yandex, zendesk
+microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openai, openfaas, openlineage, opensearch,
+opsgenie, oracle, otel, pagerduty, pandas, papermill, password, pgvector, pinecone, pinot, plexus,
+postgres, presto, rabbitmq, redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp,
+singularity, slack, smtp, snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino,
+vertica, virtualenv, weaviate, webhdfs, winrm, yandex, zendesk
   .. END EXTRAS HERE
 
 Provider packages

INSTALL

@@ -103,11 +103,11 @@ datadog, dbt.cloud, deprecated_api, devel, devel_all, devel_ci, devel_hadoop, di
 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,
-oracle, otel, pagerduty, pandas, papermill, password, pgvector, pinecone, pinot, plexus, postgres,
-presto, rabbitmq, redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp, singularity,
-slack, smtp, snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino, vertica,
-virtualenv, weaviate, webhdfs, winrm, yandex, zendesk
+microsoft.winrm, mongo, mssql, mysql, neo4j, odbc, openai, openfaas, openlineage, opensearch,
+opsgenie, oracle, otel, pagerduty, pandas, papermill, password, pgvector, pinecone, pinot, plexus,
+postgres, presto, rabbitmq, redis, s3, s3fs, salesforce, samba, segment, sendgrid, sentry, sftp,
+singularity, slack, smtp, snowflake, spark, sqlite, ssh, statsd, tableau, tabular, telegram, trino,
+vertica, virtualenv, weaviate, webhdfs, winrm, yandex, zendesk
 # END EXTRAS HERE
 
 # For installing Airflow in development environments - see CONTRIBUTING.rst

airflow/providers/openai/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-openai``
+
+Changelog
+---------
+
+1.0.0
+.....
+
+Initial version of the provider.

airflow/providers/openai/__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.

airflow/providers/openai/hooks/__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.

airflow/providers/openai/hooks/openai.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 typing import Any
+
+import openai
+
+from airflow.hooks.base import BaseHook
+
+
+class OpenAIHook(BaseHook):
+    """
+    Use OpenAI SDK to interact with OpenAI APIs.
+
+    .. seealso:: https://platform.openai.com/docs/introduction/overview
+
+    :param conn_id: :ref:`OpenAI connection id <howto/connection:openai>`
+    """
+
+    conn_name_attr = "conn_id"
+    default_conn_name = "openai_default"
+    conn_type = "openai"
+    hook_name = "OpenAI"
+
+    def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.conn_id = conn_id
+        openai.api_key = self._get_api_key()
+        api_base = self._get_api_base()
+        if api_base:
+            openai.api_base = api_base
+
+    @staticmethod
+    def get_ui_field_behaviour() -> dict[str, Any]:
+        """Return custom field behaviour."""
+        return {
+            "hidden_fields": ["schema", "port", "login", "extra"],
+            "relabeling": {"password": "API Key"},
+            "placeholders": {},
+        }
+
+    def test_connection(self) -> tuple[bool, str]:
+        try:
+            openai.Model.list()
+            return True, "Connection established!"
+        except Exception as e:
+            return False, str(e)
+
+    def _get_api_key(self) -> str:
+        """Get the OpenAI API key from the connection."""
+        conn = self.get_connection(self.conn_id)
+        if not conn.password:
+            raise ValueError("OpenAI API key not found in connection")
+        return str(conn.password)
+
+    def _get_api_base(self) -> None | str:
+        conn = self.get_connection(self.conn_id)
+        return conn.host
+
+    def create_embeddings(
+        self, text: str | list[Any], model: str = "text-embedding-ada-002", **kwargs: Any
+    ) -> list[float]:
+        """Generate embeddings for the given text using the given model.
+
+        :param text: The text to generate embeddings for.
+        :param model: The model to use for generating embeddings.
+        """
+        response = openai.Embedding.create(model=model, input=text, **kwargs)
+        embeddings: list[float] = response["data"][0]["embedding"]
+        return embeddings

airflow/providers/openai/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.

airflow/providers/openai/operators/openai.py

@@ -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.
+
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.openai.hooks.openai import OpenAIHook
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class OpenAIEmbeddingOperator(BaseOperator):
+    """
+    Operator that accepts input text to generate OpenAI embeddings using the specified model.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:OpenAIEmbeddingOperator`
+
+    :param conn_id: The OpenAI connection.
+    :param input_text: The text to generate OpenAI embeddings on. Either input_text or input_callable
+        should be provided.
+    :param model: The OpenAI model to be used for generating the embeddings.
+    :param embedding_kwargs: For possible option check
+        .. seealso:: https://platform.openai.com/docs/api-reference/embeddings/create
+    """
+
+    template_fields: Sequence[str] = ("input_text",)
+
+    def __init__(
+        self,
+        conn_id: str,
+        input_text: str | list[Any],
+        model: str = "text-embedding-ada-002",
+        embedding_kwargs: dict | None = None,
+        **kwargs: Any,
+    ):
+        self.embedding_kwargs = embedding_kwargs or {}
+        super().__init__(**kwargs)
+        self.conn_id = conn_id
+        self.input_text = input_text
+        self.model = model
+
+    @cached_property
+    def hook(self) -> OpenAIHook:
+        """Return an instance of the OpenAIHook."""
+        return OpenAIHook(conn_id=self.conn_id)
+
+    def execute(self, context: Context) -> list[float]:
+        self.log.info("Input text: %s", self.input_text)
+        embeddings = self.hook.create_embeddings(self.input_text, model=self.model, **self.embedding_kwargs)
+        self.log.info("Embeddings: %s", embeddings)
+        return embeddings

airflow/providers/openai/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-openai
+
+name: OpenAI
+
+description: |
+    `OpenAI <https://platform.openai.com/docs/introduction>`__
+
+suspended: false
+
+versions:
+  - 1.0.0
+
+integrations:
+  - integration-name: OpenAI
+    external-doc-url: https://platform.openai.com/docs/introduction
+    how-to-guide:
+      - /docs/apache-airflow-providers-openai/operators/openai.rst
+    tags: [software]
+
+dependencies:
+  - apache-airflow>=2.5.0
+  - openai[datalib]>=0.28.1,<1.0
+
+hooks:
+  - integration-name: OpenAI
+    python-modules:
+      - airflow.providers.openai.hooks.openai
+
+operators:
+  - integration-name: OpenAI
+    python-modules:
+      - airflow.providers.openai.operators.openai
+
+connection-types:
+  - hook-class-name: airflow.providers.openai.hooks.openai.OpenAIHook
+    connection-type: openai

docs/apache-airflow-providers-openai/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/odbc/CHANGELOG.rst

docs/apache-airflow-providers-openai/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-openai
+----------------------------------------