AI工作流神器：Airflow中调用大模型

banq

这个代码库包含了一个用于在Apache Airflow中操作大语言模型（LLMs）的软件开发工具包（SDK），它基于Pydantic AI构建。

这个工具让用户可以直接在Airflow工作流中调用大语言模型，并通过装饰器任务来协调智能代理的调用。

该SDK沿用了Airflow常用的@task语法，并扩展了像@task.llm（大语言模型任务）、@task.llm_branch（大语言模型分支任务）和@task.agent（智能代理任务）这样的新功能。

特征：

  1. 带@task.llm的任务：这是让AI（比如GPT-3.5）帮你处理文字的任务，就像让学霸帮你写作文。
  2. 带@task.agent的任务：这是让AI当小队长，指挥其他AI工具一步步完成复杂任务，就像小组长分配工作给组员。
  3. 自动解析结果：只要提前说清楚要什么格式（包括用Pydantic模型），AI给出的答案就能自动整理得整整齐齐，就像自动批改选择题的阅卷机。
  4. 带@task.llm_branch的分支任务：根据AI的回答来决定下一步怎么做，就像玩游戏时根据NPC的对话选择不同剧情路线。
  5. 支持的AI模型：能用市面上所有主流AI（包括OpenAI、Anthropic、谷歌Gemini、Ollama、Groq、Mistral、Cohere、亚马逊Bedrock这些），就像游戏机兼容各种品牌的手柄。

设计原则
我们按照Airflow的任务流程模式，主要用三个"装饰器"（就是给函数加特效的小工具）：

  1. @task.llm：创建一个让AI大模型干活的任务。背后其实是用了一个没有额外工具的Pydantic AI小助手。
  2. @task.agent：创建一个让智能代理干活的任务。你可以直接把自己的Pydantic AI小助手传给它用。
  3. @task.llm_branch：创建一个会根据AI回答自动选择下一步路线的任务。强制要求AI必须从后续任务ID中选一个答案。
每个装饰器下面的函数就是个"翻译官"，负责把Airflow任务收到的原始输入转换成AI能听懂的话。如果你懒得翻译，直接原样返回输入内容也行。

为什么这么做？
现在越来越多的公司想用AI大模型解决实际问题，这就需要有个灵活又好用的工作流程管家。

Airflow就是目前最受欢迎的"数据流水线管家"，它有这些超能力：

  • 灵活排班：可以定时、随叫随到、或者根据突发事件启动任务
  • 自动分身术：能同时处理大批量任务，出错会自动处理还能全程监控
  • 智能岔路口：根据任务结果自动选择不同执行路线
  • 防崩溃机制：自带重试、智能等待、超时打断等功能
  • 资源分配：用Airflow资源池控制同时运行的任务数量
  • 全天候监控：详细记录运行日志，随时查看状态
  • 工业级强度：专为真实工作场景设计
我们这个工具包就是帮你把AI工作流轻松塞进Airflow管家的工具箱，不管是简单的AI问答还是复杂的智能决策流程都能搞定。