
Apache Airflow Cohere Provider 详解安装、连接配置与 Embedding / Rerank 算子实战【免费下载链接】airflowApache Airflow - A platform to programmatically author, schedule, and monitor workflows项目地址: https://gitcode.com/GitHub_Trending/ai/airflowApache Airflow 的apache-airflow-providers-cohere是一个生产级lifecycle: production的 Provider 包用于在 DAG 中调用 Cohere 平台的 Embedding 与 Rerank API。本篇基于该包的 README、provider.yaml声明及核心源码完整讲解其安装要求、Cohere 连接配置方式以及CohereHook、CohereEmbeddingOperator、CohereRerankOperator三个核心组件的参数、默认值与底层实现机制读完后可直接在生产 DAG 中集成 Cohere 多语言模型能力。包概览apache-airflow-providers-cohere根据 providers/cohere/README.rst该 Provider 包当前发布版本为1.7.0其全部类都位于airflow.providers.coherePython 包下。provider.yaml 中声明了包的生命周期状态state: ready lifecycle: production从源码结构看该 Provider 提供三类可复用组件均在 provider.yaml 的hooks/operators/connection-types段注册组件类型模块路径CohereHookHookairflow.providers.cohere.hooks.cohereCohereEmbeddingOperatorOperatorairflow.providers.cohere.operators.embeddingCohereRerankOperatorOperatorairflow.providers.cohere.operators.rerankcohere连接类型Hook 类airflow.providers.cohere.hooks.cohere.CohereHookREADME 头部还注明该文件由模板PROVIDER_README_TEMPLATE.rst.jinja2位于dev/breeze/src/airflow_breeze/templates目录自动生成因此其中的版本、依赖信息与实际发布包严格一致。安装方式与版本要求安装命令在已有 Airflow 安装基础上可通过以下命令叠加安装该 Provider 包pip install apache-airflow-providers-coherePython 版本支持根据 README 与 pyproject.tomlrequires-python 3.10及 classifiers该包支持 Python3.10、3.11、3.12、3.13、3.14。依赖要求Requirements完整继承 README 中的依赖版本要求表并与 pyproject.toml 中dependencies段一一对应PIP 包版本要求apache-airflow2.11.0apache-airflow-providers-common-compat1.8.0cohere5.13.4fastavro1.10.0条件Python 3.13.x即3.13 and 3.14fastavro1.12.1条件Python 3.14几个值得注意的点最低 Airflow 版本为 2.11.0这是使用本 Provider 的前提版本过低的 Airflow 环境需要先升级fastavro采用按 Python 版本分叉的版本约束Python 3.13 要求1.10.0而 Python 3.14 提升到1.12.1这属于典型的按解释器版本分叉的条件依赖写法包的 changelog 可在官方文档站对应 Provider 文档的 changelog 页查阅README 中给出的入口为apache-airflow-providers-cohere1.7.0 文档。配置 Cohere 连接在编写 DAG 前需要先在 Airflow 中配置一个连接类型为cohere的连接。连接字段语义参考 providers/cohere/docs/connections.rst默认连接 ID所有 Cohere Hook 默认指向cohere_default源码中由CohereHook.default_conn_name cohere_default定义见 hooks/cohere.pyPassword必填存放 Cohere API Key。注意连接表单中该字段会被重命名显示为 “API Key”Host可选指定 API 主机base_url用于指向自托管或代理的 Cohere 兼容端点。从源码结构看hooks/cohere.py 的get_ui_field_behaviour连接 UI 表单会隐藏schema、login、port、extra四个无关字段并将password重标记为 “API Key”与provider.yaml中的ui-field-behaviour声明一致。这实际上明确了 Cohere 连接取用连接对象的哪些字段API Key 取自passwordbase_url取自host。CohereHook客户端封装与核心方法CohereHookproviders/cohere/src/airflow/providers/cohere/hooks/cohere.py基于 Cohere 官方 Python SDK 的ClientV2API v2构建是全部能力的底层入口。构造参数CohereHook( conn_idcohere_default, # 默认连接 ID timeoutNone, # 请求超时秒 max_retriesNone, # 已弃用见下文 request_optionsNone, # 函数级请求配置 )timeout透传给cohere.ClientV2的请求超时秒request_optionsSDK 的RequestOptions字典支持timeout_in_seconds单次 API 调用超时、max_retries失败重试次数、additional_headers、additional_query_parameters、additional_body_parameters等字段字段说明见 operators/embedding.py 的 docstringmax_retries已弃用参数。构造时会发出AirflowProviderDeprecationWarning并将其归并进request_options[max_retries]源码 L69-L78新代码应直接使用request_options。get_conn延迟创建并缓存客户端def get_conn(self) - cohere.ClientV2: if self._client is None: conn self.get_connection(self.conn_id) self._client cohere.ClientV2( api_keyconn.password, timeoutself.timeout, base_urlconn.host or None, ) return self._clientget_conn采用“首次调用才创建、之后复用缓存实例”的策略hooks/cohere.py。单元测试TestCohereHook::test__get_api_keytests/unit/cohere/hooks/test_cohere.py验证了映射关系Connection(password...)→api_keyConnection(host...)→base_urltimeout原样透传。create_embeddings文本向量化def create_embeddings( self, texts: list[str], model: str embed-multilingual-v3.0 ) - list[list[float]]: response self.get_conn().embed( textstexts, modelmodel, input_typesearch_document, embedding_types[float], request_optionsself.request_options, ) if response.embeddings.float_ is None: raise ValueError(Embeddings response is missing float_ field) return response.embeddings.float_要点hooks/cohere.py默认模型为embed-multilingual-v3.0input_type固定为search_documentembedding_types固定为[float]返回list[list[float]]。源码注释特别说明原本想返回 SDK 的EmbedByTypeResponseEmbeddings类型但因 Airflow XCom 对复杂类型Cohere 嵌入对象、Pydantic 模型序列化/反序列化的限制而暂时改为返回纯 float 列表——这是一个重要的实战细节意味着下游任务拿到的是可被 XCom 安全传递的原始向量而非 SDK 对象若响应缺少float_字段会直接抛ValueError而非静默失败。rerank按相关性重排文档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]:全部为关键字参数keyword-only默认模型rerank-v3.5hooks/cohere.pytop_n与max_tokens_per_doc仅在非None时才会加入请求 kwargs避免向 SDK 传入显式的None返回值通过response.model_dump(modejson)转换为XCom 可序列化的纯字典包含id和results每项为indexrelevance_score。test_connection连接可用性验证def test_connection( self, model: str command-r-plus-08-2024, messages: ChatMessages | None None, ) - tuple[bool, str]:该classmethodhooks/cohere.py会真实发起一次chat调用默认模型command-r-plus-08-2024默认消息hello world!成功返回(True, Connection successfully established.)异常则返回(False, Unexpected error: ...)——这就是在 Airflow UI 连接编辑页点击 “Test” 按钮时的底层行为。CohereEmbeddingOperator 实战CohereEmbeddingOperatoroperators/embedding.py用于调用 Cohere Embedding API 生成文本向量官方 How-To 指南见 docs/operators/embedding.rst。参数参数说明input_text必填。单个字符串或字符串列表即需要向量化的一段或多段文本conn_idCohere 连接 ID默认cohere_defaulttimeoutCohere API 请求超时秒max_retries失败重试次数已弃用见 Hook 说明request_options请求级配置字段含timeout_in_seconds、max_retries、additional_headers、additional_query_parameters、additional_body_parametersinput_text是模板字段template_fields (input_text,)支持 Jinja 渲染。一个容易忽略的实现细节execute中只有在渲染之后才把单个字符串规范化为单元素列表operators/embedding.py 注释明确说明不能在__init__中做这件事否则会作用于未渲染的值。官方示例取自系统测试 DAG example_cohere_embedding_operator.pyfrom airflow import DAG from airflow.providers.cohere.operators.embedding import CohereEmbeddingOperator with DAG(example_cohere_embedding, scheduleNone, start_datedatetime(2023, 1, 1), catchupFalse) as dag: texts [ On Kernel-Target Alignment. We describe a family of global optimization procedures, that automatically decompose optimization problems into smaller loosely coupled, problems, then combine the solutions of these with message passing algorithms., ] CohereEmbeddingOperator(input_texttexts, task_idembedding_via_text) CohereEmbeddingOperator(input_texttexts[0], task_idembedding_via_task)示例同时演示了两种用法传入文本列表批量向量化以及传入单个字符串。execute的返回值为list[list[float]]即与输入文本一一对应的浮点向量列表可直接存入 XCom 供下游任务消费。CohereRerankOperator 实战CohereRerankOperatoroperators/rerank.py基于 Cohere Rerank API 对文档按与查询词的相关度重新排序How-To 指南见 docs/operators/rerank.rst。参数参数必填说明query是用于评估相关性的搜索查询词documents是待排序的文本文档列表model否重排模型缺省时使用 Hook 默认的rerank-v3.5也可传端点专属部署名top_n否最多返回多少条排序结果默认返回全部max_tokens_per_doc否每个文档最多保留处理的 token 数conn_id否Cohere 连接 ID默认cohere_defaulttimeout否请求超时秒request_options否请求级配置同 Hook模板字段为(query, documents, top_n, max_tokens_per_doc)且documents在 UI 中按json渲染template_fields_renderers。从源码看operators/rerank.pyexecute会把渲染后的top_n/max_tokens_per_doc强制转为int再传给 Hook——因此这两个字段支持传入模板渲染出的字符串数字如 Airflow 参数注入而model仅在非None时才参与请求构造。官方示例取自 example_cohere_rerank_operator.pyfrom airflow import DAG from airflow.providers.cohere.operators.rerank import CohereRerankOperator with DAG(example_cohere_rerank, scheduleNone, start_datedatetime(2023, 1, 1), catchupFalse) as dag: CohereRerankOperator( task_idrerank_documents, queryWhat 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_n2, )输出格式Operator 把 Cohere 响应转换为 XCom 可序列化的字典其中results列表按相关性从高到低排列每项包含原文档的零基index与relevance_score——用index即可把结果映射回输入documents列表中的对应条目。简化的响应示例{ id: rerank-request-id, results: [ {index: 1, relevance_score: 0.99}, {index: 0, relevance_score: 0.12} ] }验证与源码导航单元测试tests/unit/cohere/hooks/test_cohere.py 覆盖客户端构造映射password→api_key、host→base_url以及rerank在默认模型、省略未设置限制项时的 SDK 调用形态Operator 测试位于 tests/unit/cohere/operators/ 下的test_embedding.py、test_rerank.py系统测试 DAGtests/system/cohere/ 中的两个 example DAG 既作为真实端到端用例运行也通过[START ...]/[END ...]标记片段被官方文档embedding.rst、rerank.rst以exampleinclude方式直接引用保证文档示例与实际代码同源包元数据provider.yaml 是 Provider 在 Airflow 元数据连接类型注册、集成信息、生命周期中的唯一声明源pyproject.toml 定义打包flit、依赖与入口点apache_airflow_provider注册get_provider_info。小结apache-airflow-providers-cohere1.7.0以CohereHook为统一入口围绕 Cohere API v2 提供了文本向量化默认embed-multilingual-v3.0与文档重排默认rerank-v3.5两类算子。使用时的关键约定连接中password存 API Key、host存可选的 API 主机新代码优先用request_options替代已弃用的max_retries两个 Operator 的返回值都经过显式的 XCom 友好化转换float 向量列表、model_dump(modejson)字典可以直接在 DAG 的下游任务中消费。【免费下载链接】airflowApache Airflow - A platform to programmatically author, schedule, and monitor workflows项目地址: https://gitcode.com/GitHub_Trending/ai/airflow创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考