From 2a14475efc9cfc39472e8c7be181c7f7725c1771 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 09:27:46 +0000 Subject: [PATCH 01/22] Plugin skeleton --- plugins/airflow-wingman/__init__.py | 0 plugins/airflow-wingman/plugin.py | 43 +++++++++++++++++++ plugins/airflow-wingman/pyproject.toml | 20 +++++++++ .../templates/wingman_chat.html | 5 +++ 4 files changed, 68 insertions(+) create mode 100644 plugins/airflow-wingman/__init__.py create mode 100644 plugins/airflow-wingman/plugin.py create mode 100644 plugins/airflow-wingman/pyproject.toml create mode 100644 plugins/airflow-wingman/templates/wingman_chat.html diff --git a/plugins/airflow-wingman/__init__.py b/plugins/airflow-wingman/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/plugins/airflow-wingman/plugin.py b/plugins/airflow-wingman/plugin.py new file mode 100644 index 0000000..03f0405 --- /dev/null +++ b/plugins/airflow-wingman/plugin.py @@ -0,0 +1,43 @@ +from airflow.plugins_manager import AirflowPlugin +from flask_appbuilder import BaseView as AppBuilderBaseView, expose +from flask import Blueprint + + +bp = Blueprint( + "wingman", + __name__, + template_folder="templates", + static_folder="static", + static_url_path="/static/wingman", +) + + +class WingmanView(AppBuilderBaseView): + route_base = "/wingman" + default_view = "chat" + + @expose("/") + def chat(self): + """ + Chat interface for Airflow Wingman. + """ + return self.render_template( + "wingman_chat.html", + title="Airflow Wingman" + ) + + +# Create AppBuilder View +v_appbuilder_view = WingmanView() +v_appbuilder_package = { + "name": "Airflow Wingman", + "category": "AI", + "view": v_appbuilder_view, +} + + +# Create Plugin +class WingmanPlugin(AirflowPlugin): + name = "wingman" + flask_blueprints = [bp] + appbuilder_views = [v_appbuilder_package] diff --git a/plugins/airflow-wingman/pyproject.toml b/plugins/airflow-wingman/pyproject.toml new file mode 100644 index 0000000..b98e4e3 --- /dev/null +++ b/plugins/airflow-wingman/pyproject.toml @@ -0,0 +1,20 @@ +[build-system] +requires = ["setuptools", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "airflow-wingman" +version = "0.1.0" +description = "Airflow plugin to enable LLMs chat" +authors = [ + {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"} +] +dependencies = [ + "apache-airflow>=2.10.0", +] + +[project.urls] +repository = "https://github.com/abhishekbhakat/airflow-mcp-server" + +[tool.setuptools] +packages = ["airflow-wingman"] diff --git a/plugins/airflow-wingman/templates/wingman_chat.html b/plugins/airflow-wingman/templates/wingman_chat.html new file mode 100644 index 0000000..072a34f --- /dev/null +++ b/plugins/airflow-wingman/templates/wingman_chat.html @@ -0,0 +1,5 @@ +{% extends "appbuilder/base.html" %} + +{% block content %} + +{% endblock %} From 0b50ed0b85613518ef6ceff87af8fcec199d672a Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 09:47:28 +0000 Subject: [PATCH 02/22] Models list --- plugins/airflow-wingman/plugin.py | 48 ++++- .../templates/wingman_chat.html | 188 ++++++++++++++++++ 2 files changed, 233 insertions(+), 3 deletions(-) diff --git a/plugins/airflow-wingman/plugin.py b/plugins/airflow-wingman/plugin.py index 03f0405..9b5f6cb 100644 --- a/plugins/airflow-wingman/plugin.py +++ b/plugins/airflow-wingman/plugin.py @@ -16,21 +16,63 @@ class WingmanView(AppBuilderBaseView): route_base = "/wingman" default_view = "chat" + AVAILABLE_MODELS = { + "anthropic": { + "name": "Anthropic", + "endpoint": "https://api.anthropic.com/v1/messages", + "models": [ + { + "id": "claude-3.5-sonnet", + "name": "Claude 3.5 Sonnet", + "default": True, + "context_window": 200000, + "description": "Input $3/M tokens, Output $15/M tokens", + }, + { + "id": "claude-3.5-haiku", + "name": "Claude 3.5 Haiku", + "default": False, + "context_window": 200000, + "description": "Input $0.80/M tokens, Output $4/M tokens", + }, + ], + }, + "openrouter": { + "name": "OpenRouter", + "endpoint": "https://openrouter.ai/api/v1/chat/completions", + "models": [ + { + "id": "anthropic/claude-3.5-sonnet", + "name": "Claude 3.5 Sonnet", + "default": False, + "context_window": 200000, + "description": "Input $3/M tokens, Output $15/M tokens", + }, + { + "id": "anthropic/claude-3.5-haiku", + "name": "Claude 3.5 Haiku", + "default": False, + "context_window": 200000, + "description": "Input $0.80/M tokens, Output $4/M tokens", + }, + ], + }, + } + @expose("/") def chat(self): """ Chat interface for Airflow Wingman. """ return self.render_template( - "wingman_chat.html", - title="Airflow Wingman" + "wingman_chat.html", title="Airflow Wingman", models=self.AVAILABLE_MODELS ) # Create AppBuilder View v_appbuilder_view = WingmanView() v_appbuilder_package = { - "name": "Airflow Wingman", + "name": "Wingman", "category": "AI", "view": v_appbuilder_view, } diff --git a/plugins/airflow-wingman/templates/wingman_chat.html b/plugins/airflow-wingman/templates/wingman_chat.html index 072a34f..dafa921 100644 --- a/plugins/airflow-wingman/templates/wingman_chat.html +++ b/plugins/airflow-wingman/templates/wingman_chat.html @@ -1,5 +1,193 @@ {% extends "appbuilder/base.html" %} {% block content %} +
+ +
+
+
+
+

Airflow Wingman

+
+
+
+
+
+ +
+
+
+

Model Selection

+
+
+ {% for provider_id, provider in models.items() %} +
+

{{ provider.name }}

+ {% for model in provider.models %} +
+ +
+ {% endfor %} +
+ {% endfor %} +
+ + +
+ + +
+
+

API Key

+
+
+
+ + + Your API key will be used for the selected provider + +
+
+
+
+ + +
+
+
+ +
+
+
+ + + + +
+
+
+
+
+
+ + + + {% endblock %} From f1fb9ac925cef60aa23a798ee7f681128e33c696 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:15:48 +0000 Subject: [PATCH 03/22] Spliting out from local install --- plugins/airflow-wingman/__init__.py | 0 plugins/airflow-wingman/plugin.py | 85 -------- plugins/airflow-wingman/pyproject.toml | 20 -- .../templates/wingman_chat.html | 193 ------------------ 4 files changed, 298 deletions(-) delete mode 100644 plugins/airflow-wingman/__init__.py delete mode 100644 plugins/airflow-wingman/plugin.py delete mode 100644 plugins/airflow-wingman/pyproject.toml delete mode 100644 plugins/airflow-wingman/templates/wingman_chat.html diff --git a/plugins/airflow-wingman/__init__.py b/plugins/airflow-wingman/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/plugins/airflow-wingman/plugin.py b/plugins/airflow-wingman/plugin.py deleted file mode 100644 index 9b5f6cb..0000000 --- a/plugins/airflow-wingman/plugin.py +++ /dev/null @@ -1,85 +0,0 @@ -from airflow.plugins_manager import AirflowPlugin -from flask_appbuilder import BaseView as AppBuilderBaseView, expose -from flask import Blueprint - - -bp = Blueprint( - "wingman", - __name__, - template_folder="templates", - static_folder="static", - static_url_path="/static/wingman", -) - - -class WingmanView(AppBuilderBaseView): - route_base = "/wingman" - default_view = "chat" - - AVAILABLE_MODELS = { - "anthropic": { - "name": "Anthropic", - "endpoint": "https://api.anthropic.com/v1/messages", - "models": [ - { - "id": "claude-3.5-sonnet", - "name": "Claude 3.5 Sonnet", - "default": True, - "context_window": 200000, - "description": "Input $3/M tokens, Output $15/M tokens", - }, - { - "id": "claude-3.5-haiku", - "name": "Claude 3.5 Haiku", - "default": False, - "context_window": 200000, - "description": "Input $0.80/M tokens, Output $4/M tokens", - }, - ], - }, - "openrouter": { - "name": "OpenRouter", - "endpoint": "https://openrouter.ai/api/v1/chat/completions", - "models": [ - { - "id": "anthropic/claude-3.5-sonnet", - "name": "Claude 3.5 Sonnet", - "default": False, - "context_window": 200000, - "description": "Input $3/M tokens, Output $15/M tokens", - }, - { - "id": "anthropic/claude-3.5-haiku", - "name": "Claude 3.5 Haiku", - "default": False, - "context_window": 200000, - "description": "Input $0.80/M tokens, Output $4/M tokens", - }, - ], - }, - } - - @expose("/") - def chat(self): - """ - Chat interface for Airflow Wingman. - """ - return self.render_template( - "wingman_chat.html", title="Airflow Wingman", models=self.AVAILABLE_MODELS - ) - - -# Create AppBuilder View -v_appbuilder_view = WingmanView() -v_appbuilder_package = { - "name": "Wingman", - "category": "AI", - "view": v_appbuilder_view, -} - - -# Create Plugin -class WingmanPlugin(AirflowPlugin): - name = "wingman" - flask_blueprints = [bp] - appbuilder_views = [v_appbuilder_package] diff --git a/plugins/airflow-wingman/pyproject.toml b/plugins/airflow-wingman/pyproject.toml deleted file mode 100644 index b98e4e3..0000000 --- a/plugins/airflow-wingman/pyproject.toml +++ /dev/null @@ -1,20 +0,0 @@ -[build-system] -requires = ["setuptools", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "airflow-wingman" -version = "0.1.0" -description = "Airflow plugin to enable LLMs chat" -authors = [ - {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"} -] -dependencies = [ - "apache-airflow>=2.10.0", -] - -[project.urls] -repository = "https://github.com/abhishekbhakat/airflow-mcp-server" - -[tool.setuptools] -packages = ["airflow-wingman"] diff --git a/plugins/airflow-wingman/templates/wingman_chat.html b/plugins/airflow-wingman/templates/wingman_chat.html deleted file mode 100644 index dafa921..0000000 --- a/plugins/airflow-wingman/templates/wingman_chat.html +++ /dev/null @@ -1,193 +0,0 @@ -{% extends "appbuilder/base.html" %} - -{% block content %} -
- -
-
-
-
-

Airflow Wingman

-
-
-
-
- -
- -
-
-
-

Model Selection

-
-
- {% for provider_id, provider in models.items() %} -
-

{{ provider.name }}

- {% for model in provider.models %} -
- -
- {% endfor %} -
- {% endfor %} -
- - -
- - -
-
-

API Key

-
-
-
- - - Your API key will be used for the selected provider - -
-
-
-
- - -
-
-
- -
-
-
- - - - -
-
-
-
-
-
- - - - -{% endblock %} From 724579a337b5edeaf0c33b45055bd3476e29e691 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:16:19 +0000 Subject: [PATCH 04/22] fix uv build error --- airflow-mcp-server/pyproject.toml | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/airflow-mcp-server/pyproject.toml b/airflow-mcp-server/pyproject.toml index 671a747..76422c4 100644 --- a/airflow-mcp-server/pyproject.toml +++ b/airflow-mcp-server/pyproject.toml @@ -3,6 +3,9 @@ name = "airflow-mcp-server" version = "0.2.0" description = "MCP Server for Airflow" requires-python = ">=3.11" +authors = [ + {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"} +] dependencies = [ "aiofiles>=24.1.0", "aiohttp>=3.11.11", @@ -14,6 +17,9 @@ dependencies = [ "pyyaml>=6.0.0", ] +[project.urls] +repository = "https://github.com/abhishekbhakat/airflow-mcp-server" + [project.scripts] airflow-mcp-server = "airflow_mcp_server.__main__:main" @@ -42,9 +48,7 @@ exclude = [ packages = ["src/airflow_mcp_server"] [tool.hatch.build.targets.wheel.sources] -"src/airflow_mcp_server" = [ - "*.yaml", -] +"src/airflow_mcp_server" = "airflow_mcp_server" [tool.pytest.ini_options] pythonpath = ["src"] From 8f8dfeb3b1fe6014df8efba5ee8a2deaa2842aee Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:18:15 +0000 Subject: [PATCH 05/22] Airflow Wingman separate project --- airflow-wingman/__init__.py | 0 airflow-wingman/llms_models.py | 42 +++++ airflow-wingman/plugin.py | 44 +++++ airflow-wingman/pyproject.toml | 62 +++++++ airflow-wingman/templates/wingman_chat.html | 193 ++++++++++++++++++++ 5 files changed, 341 insertions(+) create mode 100644 airflow-wingman/__init__.py create mode 100644 airflow-wingman/llms_models.py create mode 100644 airflow-wingman/plugin.py create mode 100644 airflow-wingman/pyproject.toml create mode 100644 airflow-wingman/templates/wingman_chat.html diff --git a/airflow-wingman/__init__.py b/airflow-wingman/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/airflow-wingman/llms_models.py b/airflow-wingman/llms_models.py new file mode 100644 index 0000000..fd328ab --- /dev/null +++ b/airflow-wingman/llms_models.py @@ -0,0 +1,42 @@ +MODELS = { + "anthropic": { + "name": "Anthropic", + "endpoint": "https://api.anthropic.com/v1/messages", + "models": [ + { + "id": "claude-3.5-sonnet", + "name": "Claude 3.5 Sonnet", + "default": True, + "context_window": 200000, + "description": "Input $3/M tokens, Output $15/M tokens", + }, + { + "id": "claude-3.5-haiku", + "name": "Claude 3.5 Haiku", + "default": False, + "context_window": 200000, + "description": "Input $0.80/M tokens, Output $4/M tokens", + }, + ], + }, + "openrouter": { + "name": "OpenRouter", + "endpoint": "https://openrouter.ai/api/v1/chat/completions", + "models": [ + { + "id": "anthropic/claude-3.5-sonnet", + "name": "Claude 3.5 Sonnet", + "default": False, + "context_window": 200000, + "description": "Input $3/M tokens, Output $15/M tokens", + }, + { + "id": "anthropic/claude-3.5-haiku", + "name": "Claude 3.5 Haiku", + "default": False, + "context_window": 200000, + "description": "Input $0.80/M tokens, Output $4/M tokens", + }, + ], + }, + } \ No newline at end of file diff --git a/airflow-wingman/plugin.py b/airflow-wingman/plugin.py new file mode 100644 index 0000000..78b01fd --- /dev/null +++ b/airflow-wingman/plugin.py @@ -0,0 +1,44 @@ +from airflow.plugins_manager import AirflowPlugin +from flask_appbuilder import BaseView as AppBuilderBaseView, expose +from flask import Blueprint + +from airflow_wingman.llms_models import MODELS + + +bp = Blueprint( + "wingman", + __name__, + template_folder="templates", + static_folder="static", + static_url_path="/static/wingman", +) + + +class WingmanView(AppBuilderBaseView): + route_base = "/wingman" + default_view = "chat" + + @expose("/") + def chat(self): + """ + Chat interface for Airflow Wingman. + """ + return self.render_template( + "wingman_chat.html", title="Airflow Wingman", models=MODELS + ) + + +# Create AppBuilder View +v_appbuilder_view = WingmanView() +v_appbuilder_package = { + "name": "Wingman", + "category": "AI", + "view": v_appbuilder_view, +} + + +# Create Plugin +class WingmanPlugin(AirflowPlugin): + name = "wingman" + flask_blueprints = [bp] + appbuilder_views = [v_appbuilder_package] diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml new file mode 100644 index 0000000..b94fbaa --- /dev/null +++ b/airflow-wingman/pyproject.toml @@ -0,0 +1,62 @@ + +[project] +name = "airflow-wingman" +version = "0.1.0" +description = "Airflow plugin to enable LLMs chat" +requires-python = ">=3.11" +authors = [ + {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"} +] +dependencies = [ + "apache-airflow>=2.10.0", + "airflow-mcp-server>=0.2.0" +] + +[project.urls] +repository = "https://github.com/abhishekbhakat/airflow-mcp-server" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["airflow-wingman"] + +[tool.ruff] +line-length = 200 +indent-width = 4 +fix = true +preview = true + +lint.select = [ + "E", # pycodestyle errors + "F", # pyflakes + "I", # isort + "W", # pycodestyle warnings + "C90", # Complexity + "C", # flake8-comprehensions + "ISC", # flake8-implicit-str-concat + "T10", # flake8-debugger + "A", # flake8-builtins + "UP", # pyupgrade +] + +lint.ignore = [ + "C416", # Unnecessary list comprehension - rewrite as a generator expression + "C408", # Unnecessary `dict` call - rewrite as a literal + "ISC001" # Single line implicit string concatenation +] + +lint.fixable = ["ALL"] +lint.unfixable = [] + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" +skip-magic-trailing-comma = false + +[tool.ruff.lint.isort] +combine-as-imports = true + +[tool.ruff.lint.mccabe] +max-complexity = 12 diff --git a/airflow-wingman/templates/wingman_chat.html b/airflow-wingman/templates/wingman_chat.html new file mode 100644 index 0000000..dafa921 --- /dev/null +++ b/airflow-wingman/templates/wingman_chat.html @@ -0,0 +1,193 @@ +{% extends "appbuilder/base.html" %} + +{% block content %} +
+ +
+
+
+
+

Airflow Wingman

+
+
+
+
+ +
+ +
+
+
+

Model Selection

+
+
+ {% for provider_id, provider in models.items() %} +
+

{{ provider.name }}

+ {% for model in provider.models %} +
+ +
+ {% endfor %} +
+ {% endfor %} +
+ + +
+ + +
+
+

API Key

+
+
+
+ + + Your API key will be used for the selected provider + +
+
+
+
+ + +
+
+
+ +
+
+
+ + + + +
+
+
+
+
+
+ + + + +{% endblock %} From fac5a8f843561dbdb0aa875ce80d4388fbc3897c Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:18:26 +0000 Subject: [PATCH 06/22] Update Airflow base image --- Dockerfile | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index 7eca009..bc79e1c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1 +1,2 @@ -FROM quay.io/astronomer/astro-runtime:12.6.0 +FROM quay.io/astronomer/astro-runtime:12.7.1 +RUN cd airflow-wingman && pip install -e . \ No newline at end of file From beb5cfc092dd5950d9f739f3eb375bcb46f53b00 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:25:27 +0000 Subject: [PATCH 07/22] inlcude package-data --- airflow-mcp-server/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow-mcp-server/pyproject.toml b/airflow-mcp-server/pyproject.toml index 76422c4..0ffd0c9 100644 --- a/airflow-mcp-server/pyproject.toml +++ b/airflow-mcp-server/pyproject.toml @@ -46,6 +46,7 @@ exclude = [ [tool.hatch.build.targets.wheel] packages = ["src/airflow_mcp_server"] +package-data = {"airflow_mcp_server"= ["*.yaml"]} [tool.hatch.build.targets.wheel.sources] "src/airflow_mcp_server" = "airflow_mcp_server" From 9f3290a59fc5609de60b2df89d13e0f958e5b7a2 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:39:16 +0000 Subject: [PATCH 08/22] Improvised packaging --- LICENSE | 2 +- airflow-mcp-server/LICENSE | 21 +++++++++++++++++++++ airflow-mcp-server/pyproject.toml | 7 +++++++ airflow-wingman/LICENSE | 21 +++++++++++++++++++++ airflow-wingman/README.md | 2 ++ airflow-wingman/pyproject.toml | 7 +++++++ 6 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 airflow-mcp-server/LICENSE create mode 100644 airflow-wingman/LICENSE create mode 100644 airflow-wingman/README.md diff --git a/LICENSE b/LICENSE index 8036231..b17cff9 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2025 Abhishek +Copyright (c) 2025 Abhishek Bhakat Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/airflow-mcp-server/LICENSE b/airflow-mcp-server/LICENSE new file mode 100644 index 0000000..b17cff9 --- /dev/null +++ b/airflow-mcp-server/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Abhishek Bhakat + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/airflow-mcp-server/pyproject.toml b/airflow-mcp-server/pyproject.toml index 0ffd0c9..82dace7 100644 --- a/airflow-mcp-server/pyproject.toml +++ b/airflow-mcp-server/pyproject.toml @@ -2,6 +2,7 @@ name = "airflow-mcp-server" version = "0.2.0" description = "MCP Server for Airflow" +readme = "README.md" requires-python = ">=3.11" authors = [ {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"} @@ -16,6 +17,12 @@ dependencies = [ "pydantic>=2.10.5", "pyyaml>=6.0.0", ] +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] +license = "MIT" +license-files = ["LICEN[CS]E*"] [project.urls] repository = "https://github.com/abhishekbhakat/airflow-mcp-server" diff --git a/airflow-wingman/LICENSE b/airflow-wingman/LICENSE new file mode 100644 index 0000000..b17cff9 --- /dev/null +++ b/airflow-wingman/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Abhishek Bhakat + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/airflow-wingman/README.md b/airflow-wingman/README.md new file mode 100644 index 0000000..a3eb0bf --- /dev/null +++ b/airflow-wingman/README.md @@ -0,0 +1,2 @@ +# Airflow Wingman +Airflow plugin to enable LLMs chat in Airflow Webserver. \ No newline at end of file diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml index b94fbaa..fa0f504 100644 --- a/airflow-wingman/pyproject.toml +++ b/airflow-wingman/pyproject.toml @@ -3,6 +3,7 @@ name = "airflow-wingman" version = "0.1.0" description = "Airflow plugin to enable LLMs chat" +readme = "README.md" requires-python = ">=3.11" authors = [ {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"} @@ -11,6 +12,12 @@ dependencies = [ "apache-airflow>=2.10.0", "airflow-mcp-server>=0.2.0" ] +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] +license = "MIT" +license-files = ["LICEN[CS]E*"] [project.urls] repository = "https://github.com/abhishekbhakat/airflow-mcp-server" From 4ec58e3bb0737a45f8224350e7be918b707ef874 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:39:38 +0000 Subject: [PATCH 09/22] Local install --- .dockerignore | 1 - Dockerfile | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/.dockerignore b/.dockerignore index 9ec1580..46f5212 100644 --- a/.dockerignore +++ b/.dockerignore @@ -6,7 +6,6 @@ logs/ .venv airflow.db airflow.cfg -airflow-mcp-server/ resources/ assets/ README.md diff --git a/Dockerfile b/Dockerfile index bc79e1c..327a17d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,2 +1,3 @@ FROM quay.io/astronomer/astro-runtime:12.7.1 +RUN cd airflow-mcp-server && pip install -e . RUN cd airflow-wingman && pip install -e . \ No newline at end of file From 7b59a454c6bd5c69287cfa48365a5d5d2a18bf5a Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:45:11 +0000 Subject: [PATCH 10/22] Fine grained classifiers and repo urls --- airflow-mcp-server/pyproject.toml | 8 ++++++-- airflow-wingman/pyproject.toml | 9 +++++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/airflow-mcp-server/pyproject.toml b/airflow-mcp-server/pyproject.toml index 82dace7..49e76c5 100644 --- a/airflow-mcp-server/pyproject.toml +++ b/airflow-mcp-server/pyproject.toml @@ -18,14 +18,18 @@ dependencies = [ "pyyaml>=6.0.0", ] classifiers = [ - "Programming Language :: Python :: 3", + "Development Status :: 3 - Alpha", "Operating System :: OS Independent", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.10", ] license = "MIT" license-files = ["LICEN[CS]E*"] [project.urls] -repository = "https://github.com/abhishekbhakat/airflow-mcp-server" +GitHub = "https://github.com/abhishekbhakat/airflow-mcp-server" +Issues = "https://github.com/abhishekbhakat/airflow-mcp-server/issues" [project.scripts] airflow-mcp-server = "airflow_mcp_server.__main__:main" diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml index fa0f504..7cbae6d 100644 --- a/airflow-wingman/pyproject.toml +++ b/airflow-wingman/pyproject.toml @@ -13,14 +13,19 @@ dependencies = [ "airflow-mcp-server>=0.2.0" ] classifiers = [ - "Programming Language :: Python :: 3", + "Development Status :: 3 - Alpha", + "Environment :: Plugins", "Operating System :: OS Independent", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.10", ] license = "MIT" license-files = ["LICEN[CS]E*"] [project.urls] -repository = "https://github.com/abhishekbhakat/airflow-mcp-server" +GitHub = "https://github.com/abhishekbhakat/airflow-mcp-server" +Issues = "https://github.com/abhishekbhakat/airflow-mcp-server/issues" [build-system] requires = ["hatchling"] From ea0af4e1fc2812c2f5b4b6af6c828aaa0fe04f81 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 10:52:20 +0000 Subject: [PATCH 11/22] fix project structure --- airflow-wingman/__init__.py | 0 airflow-wingman/pyproject.toml | 2 +- airflow-wingman/src/airflow_wingman/__init__.py | 6 ++++++ airflow-wingman/{ => src/airflow_wingman}/llms_models.py | 0 airflow-wingman/{ => src/airflow_wingman}/plugin.py | 0 .../{ => src/airflow_wingman}/templates/wingman_chat.html | 0 6 files changed, 7 insertions(+), 1 deletion(-) delete mode 100644 airflow-wingman/__init__.py create mode 100644 airflow-wingman/src/airflow_wingman/__init__.py rename airflow-wingman/{ => src/airflow_wingman}/llms_models.py (100%) rename airflow-wingman/{ => src/airflow_wingman}/plugin.py (100%) rename airflow-wingman/{ => src/airflow_wingman}/templates/wingman_chat.html (100%) diff --git a/airflow-wingman/__init__.py b/airflow-wingman/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml index 7cbae6d..1455a0a 100644 --- a/airflow-wingman/pyproject.toml +++ b/airflow-wingman/pyproject.toml @@ -32,7 +32,7 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] -packages = ["airflow-wingman"] +packages = ["src/airflow_wingman"] [tool.ruff] line-length = 200 diff --git a/airflow-wingman/src/airflow_wingman/__init__.py b/airflow-wingman/src/airflow_wingman/__init__.py new file mode 100644 index 0000000..a1c4dad --- /dev/null +++ b/airflow-wingman/src/airflow_wingman/__init__.py @@ -0,0 +1,6 @@ +from importlib.metadata import version + +from airflow_wingman.plugin import WingmanPlugin + +__version__ = version("airflow-wingman") +__all__ = ["WingmanPlugin"] diff --git a/airflow-wingman/llms_models.py b/airflow-wingman/src/airflow_wingman/llms_models.py similarity index 100% rename from airflow-wingman/llms_models.py rename to airflow-wingman/src/airflow_wingman/llms_models.py diff --git a/airflow-wingman/plugin.py b/airflow-wingman/src/airflow_wingman/plugin.py similarity index 100% rename from airflow-wingman/plugin.py rename to airflow-wingman/src/airflow_wingman/plugin.py diff --git a/airflow-wingman/templates/wingman_chat.html b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html similarity index 100% rename from airflow-wingman/templates/wingman_chat.html rename to airflow-wingman/src/airflow_wingman/templates/wingman_chat.html From a93e16c151d6d6cfd68ce8edff9eae30e57272f3 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 11:02:46 +0000 Subject: [PATCH 12/22] fix register airflow plugin --- airflow-wingman/pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml index 1455a0a..b04c01b 100644 --- a/airflow-wingman/pyproject.toml +++ b/airflow-wingman/pyproject.toml @@ -27,6 +27,9 @@ license-files = ["LICEN[CS]E*"] GitHub = "https://github.com/abhishekbhakat/airflow-mcp-server" Issues = "https://github.com/abhishekbhakat/airflow-mcp-server/issues" +[project.entry-points."airflow.plugins"] +wingman = "airflow_wingman:WingmanPlugin" + [build-system] requires = ["hatchling"] build-backend = "hatchling.build" From a4dc402b50f61162860307047643ca999cb3e466 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 11:03:52 +0000 Subject: [PATCH 13/22] version bump --- airflow-wingman/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml index b04c01b..37d22a1 100644 --- a/airflow-wingman/pyproject.toml +++ b/airflow-wingman/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "airflow-wingman" -version = "0.1.0" +version = "0.2.0" description = "Airflow plugin to enable LLMs chat" readme = "README.md" requires-python = ">=3.11" From 820d84df8811246852bcd60670a77e1134fdf3ac Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 11:11:54 +0000 Subject: [PATCH 14/22] version bump for better project details --- airflow-mcp-server/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-mcp-server/pyproject.toml b/airflow-mcp-server/pyproject.toml index 49e76c5..e3f292a 100644 --- a/airflow-mcp-server/pyproject.toml +++ b/airflow-mcp-server/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "airflow-mcp-server" -version = "0.2.0" +version = "0.3.0" description = "MCP Server for Airflow" readme = "README.md" requires-python = ">=3.11" From 48d556b1e1acfee963d36e1d318dd28361082334 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 12:00:15 +0000 Subject: [PATCH 15/22] Helper Note --- .../src/airflow_wingman/templates/wingman_chat.html | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html index dafa921..12d69be 100644 --- a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html +++ b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html @@ -9,6 +9,9 @@

Airflow Wingman

+
+

Note: For best results with function/tool calling capabilities, we recommend using models like Claude-3.5 Sonnet or GPT-4o. These models excel at understanding and using complex tools effectively.

+
From f3cc2381302eb3fdb88981e793d324bc1fd65861 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 12:19:58 +0000 Subject: [PATCH 16/22] Custom Models from openrouter --- .../src/airflow_wingman/llms_models.py | 28 +++++--- .../templates/wingman_chat.html | 68 ++++++++++++++++--- 2 files changed, 76 insertions(+), 20 deletions(-) diff --git a/airflow-wingman/src/airflow_wingman/llms_models.py b/airflow-wingman/src/airflow_wingman/llms_models.py index fd328ab..2e154b4 100644 --- a/airflow-wingman/src/airflow_wingman/llms_models.py +++ b/airflow-wingman/src/airflow_wingman/llms_models.py @@ -1,4 +1,17 @@ MODELS = { + "openai": { + "name": "OpenAI", + "endpoint": "https://api.openai.com/v1/chat/completions", + "models": [ + { + "id": "gpt-4o", + "name": "GPT-4o", + "default": True, + "context_window": 128000, + "description": "Input $5/M tokens, Output $15/M tokens", + } + ], + }, "anthropic": { "name": "Anthropic", "endpoint": "https://api.anthropic.com/v1/messages", @@ -24,18 +37,11 @@ MODELS = { "endpoint": "https://openrouter.ai/api/v1/chat/completions", "models": [ { - "id": "anthropic/claude-3.5-sonnet", - "name": "Claude 3.5 Sonnet", + "id": "custom", + "name": "Custom Model", "default": False, - "context_window": 200000, - "description": "Input $3/M tokens, Output $15/M tokens", - }, - { - "id": "anthropic/claude-3.5-haiku", - "name": "Claude 3.5 Haiku", - "default": False, - "context_window": 200000, - "description": "Input $0.80/M tokens, Output $4/M tokens", + "context_window": 128000, # Default context window, will be updated based on model + "description": "Enter any model name supported by OpenRouter (e.g., 'anthropic/claude-3-opus', 'meta-llama/llama-2-70b')", }, ], }, diff --git a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html index 12d69be..6c6360b 100644 --- a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html +++ b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html @@ -21,7 +21,7 @@
-

Model Selection

+

Provider Selection

{% for provider_id, provider in models.items() %} @@ -29,13 +29,15 @@

{{ provider.name }}

{% for model in provider.models %}
-
@@ -44,6 +46,14 @@ {% endfor %}
+ +
+
+ + +
+
+
- - -
-
-

API Key

-
-
-
- - - Your API key will be used for the selected provider - -
-
-
-
+
@@ -180,7 +177,7 @@ document.addEventListener('DOMContentLoaded', function() { const modelName = this.getAttribute('data-model-name'); console.log('Selected provider:', provider); console.log('Model name:', modelName); - + if (provider === 'openrouter') { console.log('Enabling model name input'); modelNameInput.disabled = false; @@ -201,7 +198,7 @@ document.addEventListener('DOMContentLoaded', function() { const modelName = defaultSelected.getAttribute('data-model-name'); console.log('Initial provider:', provider); console.log('Initial model name:', modelName); - + if (provider === 'openrouter') { console.log('Initially enabling model name input'); modelNameInput.disabled = false; @@ -218,20 +215,105 @@ document.addEventListener('DOMContentLoaded', function() { const sendButton = document.getElementById('send-button'); const chatMessages = document.getElementById('chat-messages'); + let currentMessageDiv = null; + function addMessage(content, isUser) { const messageDiv = document.createElement('div'); messageDiv.className = `message ${isUser ? 'message-user' : 'message-assistant'}`; messageDiv.textContent = content; chatMessages.appendChild(messageDiv); chatMessages.scrollTop = chatMessages.scrollHeight; + return messageDiv; } - function sendMessage() { + async function sendMessage() { const message = messageInput.value.trim(); - if (message) { - addMessage(message, true); - messageInput.value = ''; - // TODO: Add API call to send message and get response + if (!message) return; + + // Get selected model + const selectedModel = document.querySelector('input[name="model"]:checked'); + if (!selectedModel) { + alert('Please select a model'); + return; + } + + const provider = selectedModel.getAttribute('data-provider'); + const modelId = selectedModel.value.split(':')[1]; + const modelName = provider === 'openrouter' ? modelNameInput.value : modelId; + + // Clear input and add user message + messageInput.value = ''; + addMessage(message, true); + + try { + // Create messages array with system message + const messages = [ + { + role: 'system', + content: 'You are a helpful AI assistant integrated into Apache Airflow.' + }, + { + role: 'user', + content: message + } + ]; + + // Create assistant message div + currentMessageDiv = addMessage('', false); + + // Get API key + const apiKey = document.getElementById('api-key').value.trim(); + if (!apiKey) { + alert('Please enter an API key'); + return; + } + + // Send request + const response = await fetch('/wingman/chat', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + provider: provider, + model: modelName, + messages: messages, + api_key: apiKey, + stream: true, + temperature: 0.7 + }) + }); + + if (!response.ok) { + const error = await response.json(); + throw new Error(error.error || 'Failed to get response'); + } + + // Handle streaming response + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + + while (true) { + const { value, done } = await reader.read(); + if (done) break; + + const chunk = decoder.decode(value); + const lines = chunk.split('\n'); + + for (const line of lines) { + if (line.startsWith('data: ')) { + const content = line.slice(6); + if (content) { + currentMessageDiv.textContent += content; + chatMessages.scrollTop = chatMessages.scrollHeight; + } + } + } + } + } catch (error) { + console.error('Error:', error); + currentMessageDiv.textContent = `Error: ${error.message}`; + currentMessageDiv.style.color = 'red'; } } diff --git a/airflow-wingman/src/airflow_wingman/views.py b/airflow-wingman/src/airflow_wingman/views.py new file mode 100644 index 0000000..0dc6934 --- /dev/null +++ b/airflow-wingman/src/airflow_wingman/views.py @@ -0,0 +1,76 @@ +"""Views for Airflow Wingman plugin.""" + +from flask import Response, request, stream_with_context +from flask.json import jsonify +from flask_appbuilder import BaseView as AppBuilderBaseView, expose + +from airflow_wingman.llm_client import LLMClient +from airflow_wingman.llms_models import MODELS + + +class WingmanView(AppBuilderBaseView): + """View for Airflow Wingman plugin.""" + + route_base = "/wingman" + default_view = "chat" + + @expose("/") + def chat(self): + """Render chat interface.""" + providers = {provider: info["name"] for provider, info in MODELS.items()} + return self.render_template("wingman_chat.html", title="Airflow Wingman", models=MODELS, providers=providers) + + @expose("/chat", methods=["POST"]) + async def chat_completion(self): + """Handle chat completion requests.""" + try: + data = self._validate_chat_request(request.get_json()) + + # Create a new client for this request + client = LLMClient(data["api_key"]) + + if data["stream"]: + return self._handle_streaming_response(client, data) + else: + return await self._handle_regular_response(client, data) + + except ValueError as e: + return jsonify({"error": str(e)}), 400 + except Exception as e: + return jsonify({"error": str(e)}), 500 + + def _validate_chat_request(self, data: dict) -> dict: + """Validate chat request data.""" + if not data: + raise ValueError("No data provided") + + required_fields = ["provider", "model", "messages", "api_key"] + missing = [f for f in required_fields if not data.get(f)] + if missing: + raise ValueError(f"Missing required fields: {', '.join(missing)}") + + return { + "provider": data["provider"], + "model": data["model"], + "messages": data["messages"], + "api_key": data["api_key"], + "stream": data.get("stream", False), + "temperature": data.get("temperature", 0.7), + "max_tokens": data.get("max_tokens"), + } + + def _handle_streaming_response(self, client: LLMClient, data: dict) -> Response: + """Handle streaming response.""" + + async def generate(): + async for chunk in await client.chat_completion( + messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=True + ): + yield f"data: {chunk}\n\n" + + return Response(stream_with_context(generate()), mimetype="text/event-stream") + + async def _handle_regular_response(self, client: LLMClient, data: dict) -> Response: + """Handle regular response.""" + response = await client.chat_completion(messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=False) + return jsonify(response) From 093577dd96d29d783cad75e2fe6cde282394a283 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 13:25:13 +0000 Subject: [PATCH 18/22] resolve csrf --- .../templates/wingman_chat.html | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html index 3187fde..f1ba825 100644 --- a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html +++ b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html @@ -1,5 +1,10 @@ {% extends "appbuilder/base.html" %} +{% block head_meta %} + {{ super() }} + +{% endblock %} + {% block content %}
@@ -237,8 +242,7 @@ document.addEventListener('DOMContentLoaded', function() { return; } - const provider = selectedModel.getAttribute('data-provider'); - const modelId = selectedModel.value.split(':')[1]; + const [provider, modelId] = selectedModel.value.split(':'); const modelName = provider === 'openrouter' ? modelNameInput.value : modelId; // Clear input and add user message @@ -268,11 +272,29 @@ document.addEventListener('DOMContentLoaded', function() { return; } + // Debug log the request + const requestData = { + provider: provider, + model: modelName, + messages: messages, + api_key: apiKey, + stream: true, + temperature: 0.7 + }; + console.log('Sending request:', {...requestData, api_key: '***'}); + + // Get CSRF token + const csrfToken = document.querySelector('meta[name="csrf-token"]')?.getAttribute('content'); + if (!csrfToken) { + throw new Error('CSRF token not found. Please refresh the page.'); + } + // Send request const response = await fetch('/wingman/chat', { method: 'POST', headers: { 'Content-Type': 'application/json', + 'X-CSRFToken': csrfToken }, body: JSON.stringify({ provider: provider, From f89f6e7d5e821bf879bf61fadcf11cebef8be369 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 13:35:36 +0000 Subject: [PATCH 19/22] Use synchronous calls --- .../src/airflow_wingman/llm_client.py | 48 +++++++++---------- airflow-wingman/src/airflow_wingman/views.py | 20 ++++---- 2 files changed, 35 insertions(+), 33 deletions(-) diff --git a/airflow-wingman/src/airflow_wingman/llm_client.py b/airflow-wingman/src/airflow_wingman/llm_client.py index dc2543c..0c51316 100644 --- a/airflow-wingman/src/airflow_wingman/llm_client.py +++ b/airflow-wingman/src/airflow_wingman/llm_client.py @@ -2,10 +2,10 @@ Client for making API calls to various LLM providers using their official SDKs. """ -from collections.abc import AsyncGenerator +from collections.abc import Generator -from anthropic import AsyncAnthropic -from openai import AsyncOpenAI +from anthropic import Anthropic +from openai import OpenAI class LLMClient: @@ -16,9 +16,9 @@ class LLMClient: api_key: API key for the provider """ self.api_key = api_key - self.openai_client = AsyncOpenAI(api_key=api_key) - self.anthropic_client = AsyncAnthropic(api_key=api_key) - self.openrouter_client = AsyncOpenAI( + self.openai_client = OpenAI(api_key=api_key) + self.anthropic_client = Anthropic(api_key=api_key) + self.openrouter_client = OpenAI( base_url="https://openrouter.ai/api/v1", api_key=api_key, default_headers={ @@ -27,9 +27,9 @@ class LLMClient: }, ) - async def chat_completion( + def chat_completion( self, messages: list[dict[str, str]], model: str, provider: str, temperature: float = 0.7, max_tokens: int | None = None, stream: bool = False - ) -> AsyncGenerator[str, None] | dict: + ) -> Generator[str, None, None] | dict: """Send a chat completion request to the specified provider. Args: @@ -41,29 +41,29 @@ class LLMClient: stream: Whether to stream the response Returns: - If stream=True, returns an async generator yielding response chunks + If stream=True, returns a generator yielding response chunks If stream=False, returns the complete response """ try: if provider == "openai": - return await self._openai_chat_completion(messages, model, temperature, max_tokens, stream) + return self._openai_chat_completion(messages, model, temperature, max_tokens, stream) elif provider == "anthropic": - return await self._anthropic_chat_completion(messages, model, temperature, max_tokens, stream) + return self._anthropic_chat_completion(messages, model, temperature, max_tokens, stream) elif provider == "openrouter": - return await self._openrouter_chat_completion(messages, model, temperature, max_tokens, stream) + return self._openrouter_chat_completion(messages, model, temperature, max_tokens, stream) else: return {"error": f"Unknown provider: {provider}"} except Exception as e: return {"error": f"API request failed: {str(e)}"} - async def _openai_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): + def _openai_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): """Handle OpenAI chat completion requests.""" - response = await self.openai_client.chat.completions.create(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=stream) + response = self.openai_client.chat.completions.create(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=stream) if stream: - async def response_generator(): - async for chunk in response: + def response_generator(): + for chunk in response: if chunk.choices[0].delta.content: yield chunk.choices[0].delta.content @@ -71,7 +71,7 @@ class LLMClient: else: return {"content": response.choices[0].message.content} - async def _anthropic_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): + def _anthropic_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): """Handle Anthropic chat completion requests.""" # Convert messages to Anthropic format system_message = next((m["content"] for m in messages if m["role"] == "system"), None) @@ -80,12 +80,12 @@ class LLMClient: if m["role"] != "system": conversation.append({"role": "assistant" if m["role"] == "assistant" else "user", "content": m["content"]}) - response = await self.anthropic_client.messages.create(model=model, messages=conversation, system=system_message, temperature=temperature, max_tokens=max_tokens, stream=stream) + response = self.anthropic_client.messages.create(model=model, messages=conversation, system=system_message, temperature=temperature, max_tokens=max_tokens, stream=stream) if stream: - async def response_generator(): - async for chunk in response: + def response_generator(): + for chunk in response: if chunk.delta.text: yield chunk.delta.text @@ -93,14 +93,14 @@ class LLMClient: else: return {"content": response.content[0].text} - async def _openrouter_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): + def _openrouter_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): """Handle OpenRouter chat completion requests.""" - response = await self.openrouter_client.chat.completions.create(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=stream) + response = self.openrouter_client.chat.completions.create(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=stream) if stream: - async def response_generator(): - async for chunk in response: + def response_generator(): + for chunk in response: if chunk.choices[0].delta.content: yield chunk.choices[0].delta.content diff --git a/airflow-wingman/src/airflow_wingman/views.py b/airflow-wingman/src/airflow_wingman/views.py index 0dc6934..a56b7dc 100644 --- a/airflow-wingman/src/airflow_wingman/views.py +++ b/airflow-wingman/src/airflow_wingman/views.py @@ -21,7 +21,7 @@ class WingmanView(AppBuilderBaseView): return self.render_template("wingman_chat.html", title="Airflow Wingman", models=MODELS, providers=providers) @expose("/chat", methods=["POST"]) - async def chat_completion(self): + def chat_completion(self): """Handle chat completion requests.""" try: data = self._validate_chat_request(request.get_json()) @@ -32,7 +32,7 @@ class WingmanView(AppBuilderBaseView): if data["stream"]: return self._handle_streaming_response(client, data) else: - return await self._handle_regular_response(client, data) + return self._handle_regular_response(client, data) except ValueError as e: return jsonify({"error": str(e)}), 400 @@ -62,15 +62,17 @@ class WingmanView(AppBuilderBaseView): def _handle_streaming_response(self, client: LLMClient, data: dict) -> Response: """Handle streaming response.""" - async def generate(): - async for chunk in await client.chat_completion( - messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=True - ): + def generate(): + for chunk in client.chat_completion(messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=True): yield f"data: {chunk}\n\n" - return Response(stream_with_context(generate()), mimetype="text/event-stream") + response = Response(stream_with_context(generate()), mimetype="text/event-stream") + response.headers["Content-Type"] = "text/event-stream" + response.headers["Cache-Control"] = "no-cache" + response.headers["Connection"] = "keep-alive" + return response - async def _handle_regular_response(self, client: LLMClient, data: dict) -> Response: + def _handle_regular_response(self, client: LLMClient, data: dict) -> Response: """Handle regular response.""" - response = await client.chat_completion(messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=False) + response = client.chat_completion(messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=False) return jsonify(response) From 3a03f08a4a251626bbf4a97aeef25a3bc0592824 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 13:47:51 +0000 Subject: [PATCH 20/22] Openrouter app name --- airflow-wingman/src/airflow_wingman/llm_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-wingman/src/airflow_wingman/llm_client.py b/airflow-wingman/src/airflow_wingman/llm_client.py index 0c51316..dd58370 100644 --- a/airflow-wingman/src/airflow_wingman/llm_client.py +++ b/airflow-wingman/src/airflow_wingman/llm_client.py @@ -22,7 +22,7 @@ class LLMClient: base_url="https://openrouter.ai/api/v1", api_key=api_key, default_headers={ - "HTTP-Referer": "http://localhost:8080", # Required by OpenRouter + "HTTP-Referer": "Airflow Wingman", # Required by OpenRouter "X-Title": "Airflow Wingman", # Required by OpenRouter }, ) From 5d199ba154a845a4d745f142e1284199e3c0a883 Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 14:23:21 +0000 Subject: [PATCH 21/22] Chat history and growing context --- .../src/airflow_wingman/mcp_tools.py | 7 ++ airflow-wingman/src/airflow_wingman/notes.py | 13 ++++ .../src/airflow_wingman/prompt_engineering.py | 40 +++++++++++ .../templates/wingman_chat.html | 72 ++++++++++++++----- airflow-wingman/src/airflow_wingman/views.py | 10 ++- 5 files changed, 121 insertions(+), 21 deletions(-) create mode 100644 airflow-wingman/src/airflow_wingman/mcp_tools.py create mode 100644 airflow-wingman/src/airflow_wingman/notes.py create mode 100644 airflow-wingman/src/airflow_wingman/prompt_engineering.py diff --git a/airflow-wingman/src/airflow_wingman/mcp_tools.py b/airflow-wingman/src/airflow_wingman/mcp_tools.py new file mode 100644 index 0000000..7d5d866 --- /dev/null +++ b/airflow-wingman/src/airflow_wingman/mcp_tools.py @@ -0,0 +1,7 @@ +import asyncio + +from airflow_mcp_server.tools.tool_manager import get_airflow_tools + +# Get tools with their parameters +tools = asyncio.run(get_airflow_tools(mode="safe")) +TOOLS = {tool.name: {"description": tool.description, "parameters": tool.inputSchema} for tool in tools} diff --git a/airflow-wingman/src/airflow_wingman/notes.py b/airflow-wingman/src/airflow_wingman/notes.py new file mode 100644 index 0000000..5616650 --- /dev/null +++ b/airflow-wingman/src/airflow_wingman/notes.py @@ -0,0 +1,13 @@ +INTERFACE_MESSAGES = { + "model_recommendation": {"title": "Note", "content": "For best results with function/tool calling capabilities, we recommend using models like Claude-3.5 Sonnet or GPT-4."}, + "security_note": { + "title": "Security", + "content": "For your security, API keys are required for each session and are never stored. If you refresh the page or close the browser, you'll need to enter your API key again.", + }, + "context_window": { + "title": "Context Window", + "content": "Each model has a maximum context window size that determines how much text it can process. " + "For long conversations or large code snippets, consider using models with larger context windows like Claude-3 Opus (200K tokens) or GPT-4 Turbo (128K tokens). " + "For better results try to keep the context size as low as possible. Try using new chats instead of reusing the same chat.", + }, +} diff --git a/airflow-wingman/src/airflow_wingman/prompt_engineering.py b/airflow-wingman/src/airflow_wingman/prompt_engineering.py new file mode 100644 index 0000000..6317971 --- /dev/null +++ b/airflow-wingman/src/airflow_wingman/prompt_engineering.py @@ -0,0 +1,40 @@ +""" +Prompt engineering for the Airflow Wingman plugin. +Contains prompts and instructions for the AI assistant. +""" + +import json + +from airflow_wingman.mcp_tools import TOOLS + +INSTRUCTIONS = { + "default": f"""You are Airflow Wingman, a helpful AI assistant integrated into Apache Airflow. +You have deep knowledge of Apache Airflow's architecture, DAGs, operators, and best practices. +The Airflow version being used is >=2.10. + +You have access to the following Airflow API tools: + +{json.dumps(TOOLS, indent=2)} + +You can use these tools to fetch information and help users understand and manage their Airflow environment. +""" +} + + +def prepare_messages(messages: list[dict[str, str]], instruction_key: str = "default") -> list[dict[str, str]]: + """Prepare messages for the chat completion request. + + Args: + messages: List of messages in the conversation + instruction_key: Key for the instruction template to use + + Returns: + List of message dictionaries ready for the chat completion API + """ + instruction = INSTRUCTIONS.get(instruction_key, INSTRUCTIONS["default"]) + + # Add instruction as first system message if not present + if not messages or messages[0].get("role") != "system": + messages.insert(0, {"role": "system", "content": instruction}) + + return messages diff --git a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html index f1ba825..72657f0 100644 --- a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html +++ b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html @@ -15,9 +15,11 @@

Airflow Wingman

-

Note: For best results with function/tool calling capabilities, we recommend using models like Claude-3.5 Sonnet or GPT-4o. These models excel at understanding and using complex tools effectively.

+

{{ interface_messages.model_recommendation.title }}: {{ interface_messages.model_recommendation.content }}

-

Security: For your security, API keys are required for each session and are never stored. If you refresh the page or close the browser, you'll need to enter your API key again. This ensures your API keys remain secure in shared environments.

+

{{ interface_messages.security_note.title }}: {{ interface_messages.security_note.content }}

+
+

{{ interface_messages.context_window.title }}: {{ interface_messages.context_window.content }}

@@ -104,13 +106,22 @@
-
- - - - +
+
+
+ + + + +
+
@@ -218,9 +229,23 @@ document.addEventListener('DOMContentLoaded', function() { const messageInput = document.getElementById('message-input'); const sendButton = document.getElementById('send-button'); + const refreshButton = document.getElementById('refresh-button'); const chatMessages = document.getElementById('chat-messages'); let currentMessageDiv = null; + let messageHistory = []; + + function clearChat() { + // Clear the chat messages + chatMessages.innerHTML = ''; + // Reset message history + messageHistory = []; + // Clear the input field + messageInput.value = ''; + // Enable input if it was disabled + messageInput.disabled = false; + sendButton.disabled = false; + } function addMessage(content, isUser) { const messageDiv = document.createElement('div'); @@ -250,17 +275,14 @@ document.addEventListener('DOMContentLoaded', function() { addMessage(message, true); try { - // Create messages array with system message - const messages = [ - { - role: 'system', - content: 'You are a helpful AI assistant integrated into Apache Airflow.' - }, - { - role: 'user', - content: message - } - ]; + // Add user message to history + messageHistory.push({ + role: 'user', + content: message + }); + + // Use full message history for the request + const messages = [...messageHistory]; // Create assistant message div currentMessageDiv = addMessage('', false); @@ -314,6 +336,7 @@ document.addEventListener('DOMContentLoaded', function() { // Handle streaming response const reader = response.body.getReader(); const decoder = new TextDecoder(); + let fullResponse = ''; while (true) { const { value, done } = await reader.read(); @@ -327,11 +350,20 @@ document.addEventListener('DOMContentLoaded', function() { const content = line.slice(6); if (content) { currentMessageDiv.textContent += content; + fullResponse += content; chatMessages.scrollTop = chatMessages.scrollHeight; } } } } + + // Add assistant's response to history + if (fullResponse) { + messageHistory.push({ + role: 'assistant', + content: fullResponse + }); + } } catch (error) { console.error('Error:', error); currentMessageDiv.textContent = `Error: ${error.message}`; @@ -345,6 +377,8 @@ document.addEventListener('DOMContentLoaded', function() { sendMessage(); } }); + + refreshButton.addEventListener('click', clearChat); }); {% endblock %} diff --git a/airflow-wingman/src/airflow_wingman/views.py b/airflow-wingman/src/airflow_wingman/views.py index a56b7dc..34158d7 100644 --- a/airflow-wingman/src/airflow_wingman/views.py +++ b/airflow-wingman/src/airflow_wingman/views.py @@ -6,6 +6,8 @@ from flask_appbuilder import BaseView as AppBuilderBaseView, expose from airflow_wingman.llm_client import LLMClient from airflow_wingman.llms_models import MODELS +from airflow_wingman.notes import INTERFACE_MESSAGES +from airflow_wingman.prompt_engineering import prepare_messages class WingmanView(AppBuilderBaseView): @@ -18,7 +20,7 @@ class WingmanView(AppBuilderBaseView): def chat(self): """Render chat interface.""" providers = {provider: info["name"] for provider, info in MODELS.items()} - return self.render_template("wingman_chat.html", title="Airflow Wingman", models=MODELS, providers=providers) + return self.render_template("wingman_chat.html", title="Airflow Wingman", models=MODELS, providers=providers, interface_messages=INTERFACE_MESSAGES) @expose("/chat", methods=["POST"]) def chat_completion(self): @@ -49,10 +51,14 @@ class WingmanView(AppBuilderBaseView): if missing: raise ValueError(f"Missing required fields: {', '.join(missing)}") + # Prepare messages with system instruction while maintaining history + messages = data["messages"] + messages = prepare_messages(messages) + return { "provider": data["provider"], "model": data["model"], - "messages": data["messages"], + "messages": messages, "api_key": data["api_key"], "stream": data.get("stream", False), "temperature": data.get("temperature", 0.7), From 16cd3f48fec243ce44d4020228c25629e70a416e Mon Sep 17 00:00:00 2001 From: abhishekbhakat Date: Mon, 24 Feb 2025 16:50:08 +0000 Subject: [PATCH 22/22] Clean up for only MCP Server --- .astro/config.yaml | 2 - .astro/dag_integrity_exceptions.txt | 1 - .astro/test_dag_integrity_default.py | 141 ------- .dockerignore | 11 - .github/workflows/python-publish.yml | 3 +- Dockerfile | 30 +- README.md | 47 ++- airflow-mcp-server/Dockerfile | 27 -- airflow-mcp-server/LICENSE | 21 - airflow-mcp-server/README.md | 66 --- airflow-wingman/LICENSE | 21 - airflow-wingman/README.md | 2 - airflow-wingman/pyproject.toml | 79 ---- .../src/airflow_wingman/__init__.py | 6 - .../src/airflow_wingman/llm_client.py | 109 ----- .../src/airflow_wingman/llms_models.py | 48 --- .../src/airflow_wingman/mcp_tools.py | 7 - airflow-wingman/src/airflow_wingman/notes.py | 13 - airflow-wingman/src/airflow_wingman/plugin.py | 32 -- .../src/airflow_wingman/prompt_engineering.py | 40 -- .../templates/wingman_chat.html | 384 ------------------ airflow-wingman/src/airflow_wingman/views.py | 84 ---- airflow_settings.yaml | 25 -- dags/.airflowignore | 0 dags/exampledag.py | 100 ----- packages.txt | 0 .../pyproject.toml => pyproject.toml | 0 requirements.txt | 1 - .../airflow_mcp_server/__init__.py | 0 .../airflow_mcp_server/__main__.py | 0 .../airflow_mcp_server/client/__init__.py | 0 .../client/airflow_client.py | 0 .../airflow_mcp_server/parser/__init__.py | 0 .../parser/operation_parser.py | 0 .../airflow_mcp_server/resources/v1.yaml | 0 .../src => src}/airflow_mcp_server/server.py | 0 .../airflow_mcp_server/server_safe.py | 0 .../airflow_mcp_server/server_unsafe.py | 0 .../airflow_mcp_server/tools/__init__.py | 0 .../airflow_mcp_server/tools/airflow_tool.py | 0 .../airflow_mcp_server/tools/base_tools.py | 0 .../airflow_mcp_server/tools/tool_manager.py | 0 .../tests => tests}/__init__.py | 0 .../client/test_airflow_client.py | 0 .../tests => tests}/conftest.py | 0 tests/dags/test_dag_example.py | 83 ---- .../parser/test_operation_parser.py | 0 .../tests => tests}/tools/__init__.py | 0 .../tools/test_airflow_tool.py | 0 .../tests => tests}/tools/test_models.py | 0 .../tools/test_tool_manager.py | 0 airflow-mcp-server/uv.lock => uv.lock | 0 52 files changed, 66 insertions(+), 1317 deletions(-) delete mode 100644 .astro/config.yaml delete mode 100644 .astro/dag_integrity_exceptions.txt delete mode 100644 .astro/test_dag_integrity_default.py delete mode 100644 .dockerignore delete mode 100644 airflow-mcp-server/Dockerfile delete mode 100644 airflow-mcp-server/LICENSE delete mode 100644 airflow-mcp-server/README.md delete mode 100644 airflow-wingman/LICENSE delete mode 100644 airflow-wingman/README.md delete mode 100644 airflow-wingman/pyproject.toml delete mode 100644 airflow-wingman/src/airflow_wingman/__init__.py delete mode 100644 airflow-wingman/src/airflow_wingman/llm_client.py delete mode 100644 airflow-wingman/src/airflow_wingman/llms_models.py delete mode 100644 airflow-wingman/src/airflow_wingman/mcp_tools.py delete mode 100644 airflow-wingman/src/airflow_wingman/notes.py delete mode 100644 airflow-wingman/src/airflow_wingman/plugin.py delete mode 100644 airflow-wingman/src/airflow_wingman/prompt_engineering.py delete mode 100644 airflow-wingman/src/airflow_wingman/templates/wingman_chat.html delete mode 100644 airflow-wingman/src/airflow_wingman/views.py delete mode 100644 airflow_settings.yaml delete mode 100644 dags/.airflowignore delete mode 100644 dags/exampledag.py delete mode 100644 packages.txt rename airflow-mcp-server/pyproject.toml => pyproject.toml (100%) delete mode 100644 requirements.txt rename {airflow-mcp-server/src => src}/airflow_mcp_server/__init__.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/__main__.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/client/__init__.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/client/airflow_client.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/parser/__init__.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/parser/operation_parser.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/resources/v1.yaml (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/server.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/server_safe.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/server_unsafe.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/tools/__init__.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/tools/airflow_tool.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/tools/base_tools.py (100%) rename {airflow-mcp-server/src => src}/airflow_mcp_server/tools/tool_manager.py (100%) rename {airflow-mcp-server/tests => tests}/__init__.py (100%) rename {airflow-mcp-server/tests => tests}/client/test_airflow_client.py (100%) rename {airflow-mcp-server/tests => tests}/conftest.py (100%) delete mode 100644 tests/dags/test_dag_example.py rename {airflow-mcp-server/tests => tests}/parser/test_operation_parser.py (100%) rename {airflow-mcp-server/tests => tests}/tools/__init__.py (100%) rename {airflow-mcp-server/tests => tests}/tools/test_airflow_tool.py (100%) rename {airflow-mcp-server/tests => tests}/tools/test_models.py (100%) rename {airflow-mcp-server/tests => tests}/tools/test_tool_manager.py (100%) rename airflow-mcp-server/uv.lock => uv.lock (100%) diff --git a/.astro/config.yaml b/.astro/config.yaml deleted file mode 100644 index 72839af..0000000 --- a/.astro/config.yaml +++ /dev/null @@ -1,2 +0,0 @@ -project: - name: airflow-mcp-server diff --git a/.astro/dag_integrity_exceptions.txt b/.astro/dag_integrity_exceptions.txt deleted file mode 100644 index c9a2a63..0000000 --- a/.astro/dag_integrity_exceptions.txt +++ /dev/null @@ -1 +0,0 @@ -# Add dag files to exempt from parse test below. ex: dags/ \ No newline at end of file diff --git a/.astro/test_dag_integrity_default.py b/.astro/test_dag_integrity_default.py deleted file mode 100644 index e433703..0000000 --- a/.astro/test_dag_integrity_default.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Test the validity of all DAGs. **USED BY DEV PARSE COMMAND DO NOT EDIT**""" - -from contextlib import contextmanager -import logging -import os - -import pytest - -from airflow.models import DagBag, Variable, Connection -from airflow.hooks.base import BaseHook -from airflow.utils.db import initdb - -# init airflow database -initdb() - -# The following code patches errors caused by missing OS Variables, Airflow Connections, and Airflow Variables - - -# =========== MONKEYPATCH BaseHook.get_connection() =========== -def basehook_get_connection_monkeypatch(key: str, *args, **kwargs): - print( - f"Attempted to fetch connection during parse returning an empty Connection object for {key}" - ) - return Connection(key) - - -BaseHook.get_connection = basehook_get_connection_monkeypatch -# # =========== /MONKEYPATCH BASEHOOK.GET_CONNECTION() =========== - - -# =========== MONKEYPATCH OS.GETENV() =========== -def os_getenv_monkeypatch(key: str, *args, **kwargs): - default = None - if args: - default = args[0] # os.getenv should get at most 1 arg after the key - if kwargs: - default = kwargs.get( - "default", None - ) # and sometimes kwarg if people are using the sig - - env_value = os.environ.get(key, None) - - if env_value: - return env_value # if the env_value is set, return it - if ( - key == "JENKINS_HOME" and default is None - ): # fix https://github.com/astronomer/astro-cli/issues/601 - return None - if default: - return default # otherwise return whatever default has been passed - return f"MOCKED_{key.upper()}_VALUE" # if absolutely nothing has been passed - return the mocked value - - -os.getenv = os_getenv_monkeypatch -# # =========== /MONKEYPATCH OS.GETENV() =========== - -# =========== MONKEYPATCH VARIABLE.GET() =========== - - -class magic_dict(dict): - def __init__(self, *args, **kwargs): - self.update(*args, **kwargs) - - def __getitem__(self, key): - return {}.get(key, "MOCKED_KEY_VALUE") - - -_no_default = object() # allow falsey defaults - - -def variable_get_monkeypatch(key: str, default_var=_no_default, deserialize_json=False): - print( - f"Attempted to get Variable value during parse, returning a mocked value for {key}" - ) - - if default_var is not _no_default: - return default_var - if deserialize_json: - return magic_dict() - return "NON_DEFAULT_MOCKED_VARIABLE_VALUE" - - -Variable.get = variable_get_monkeypatch -# # =========== /MONKEYPATCH VARIABLE.GET() =========== - - -@contextmanager -def suppress_logging(namespace): - """ - Suppress logging within a specific namespace to keep tests "clean" during build - """ - logger = logging.getLogger(namespace) - old_value = logger.disabled - logger.disabled = True - try: - yield - finally: - logger.disabled = old_value - - -def get_import_errors(): - """ - Generate a tuple for import errors in the dag bag, and include DAGs without errors. - """ - with suppress_logging("airflow"): - dag_bag = DagBag(include_examples=False) - - def strip_path_prefix(path): - return os.path.relpath(path, os.environ.get("AIRFLOW_HOME")) - - # Initialize an empty list to store the tuples - result = [] - - # Iterate over the items in import_errors - for k, v in dag_bag.import_errors.items(): - result.append((strip_path_prefix(k), v.strip())) - - # Check if there are DAGs without errors - for file_path in dag_bag.dags: - # Check if the file_path is not in import_errors, meaning no errors - if file_path not in dag_bag.import_errors: - result.append((strip_path_prefix(file_path), "No import errors")) - - return result - - -@pytest.mark.parametrize( - "rel_path, rv", get_import_errors(), ids=[x[0] for x in get_import_errors()] -) -def test_file_imports(rel_path, rv): - """Test for import errors on a file""" - if os.path.exists(".astro/dag_integrity_exceptions.txt"): - with open(".astro/dag_integrity_exceptions.txt", "r") as f: - exceptions = f.readlines() - print(f"Exceptions: {exceptions}") - if (rv != "No import errors") and rel_path not in exceptions: - # If rv is not "No import errors," consider it a failed test - raise Exception(f"{rel_path} failed to import with message \n {rv}") - else: - # If rv is "No import errors," consider it a passed test - print(f"{rel_path} passed the import test") diff --git a/.dockerignore b/.dockerignore deleted file mode 100644 index 46f5212..0000000 --- a/.dockerignore +++ /dev/null @@ -1,11 +0,0 @@ -astro -.git -.env -airflow_settings.yaml -logs/ -.venv -airflow.db -airflow.cfg -resources/ -assets/ -README.md diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 355ebcf..d5eb3e2 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -31,7 +31,6 @@ jobs: pip install uv - name: Build release distributions - working-directory: airflow-mcp-server run: | uv pip install --system build python -m build @@ -40,7 +39,7 @@ jobs: uses: actions/upload-artifact@v4 with: name: release-dists - path: airflow-mcp-server/dist/ + path: dist/ pypi-publish: runs-on: ubuntu-latest diff --git a/Dockerfile b/Dockerfile index 327a17d..84b935e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,3 +1,27 @@ -FROM quay.io/astronomer/astro-runtime:12.7.1 -RUN cd airflow-mcp-server && pip install -e . -RUN cd airflow-wingman && pip install -e . \ No newline at end of file +# Use a Python image with uv pre-installed +FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim AS uv + +WORKDIR /app + +ENV UV_COMPILE_BYTECODE=1 +ENV UV_LINK_MODE=copy + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,source=uv.lock,target=uv.lock \ + --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ + uv sync --frozen --no-install-project --no-dev --no-editable + +ADD . /app +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --frozen --no-dev --no-editable + +FROM python:3.12-slim-bookworm + +WORKDIR /app + +COPY --from=uv /root/.local /root/.local +COPY --from=uv --chown=app:app /app/.venv /app/.venv + +ENV PATH="/app/.venv/bin:$PATH" + +ENTRYPOINT ["airflow-mcp-server"] diff --git a/README.md b/README.md index 5987c89..23bc815 100644 --- a/README.md +++ b/README.md @@ -2,12 +2,13 @@ ## Overview -A Model Context Protocol server for controlling Airflow via Airflow APIs. +A [Model Context Protocol](https://modelcontextprotocol.io/) server for controlling Airflow via Airflow APIs. ## Demo Video https://github.com/user-attachments/assets/f3e60fff-8680-4dd9-b08e-fa7db655a705 + ## Setup ### Usage with Claude Desktop @@ -29,13 +30,41 @@ https://github.com/user-attachments/assets/f3e60fff-8680-4dd9-b08e-fa7db655a705 } ``` +### Operation Modes -# Scope +The server supports two operation modes: -2 different streams in which Airflow MCP Server can be used: -- Adding Airflow to AI (_complete access to an Airflow deployment_) - - This will enable AI to be able to write DAGs and just do things in a schedule on its own. - - Use command `airflow-mcp-server` or `airflow-mcp-server --unsafe`. -- Adding AI to Airflow (_read-only access using Airflow Plugin_) - - This stream can enable Users to be able to get a better understanding about their deployment. Specially in cases where teams have hundreds, if not thousands of dags. - - Use command `airflow-mcp-server --safe`. +- **Safe Mode** (`--safe`): Only allows read-only operations (GET requests). This is useful when you want to prevent any modifications to your Airflow instance. +- **Unsafe Mode** (`--unsafe`): Allows all operations including modifications. This is the default mode. + +To start in safe mode: +```bash +airflow-mcp-server --safe +``` + +To explicitly start in unsafe mode (though this is default): +```bash +airflow-mcp-server --unsafe +``` + +### Considerations + +The MCP Server expects environment variables to be set: +- `AIRFLOW_BASE_URL`: The base URL of the Airflow API +- `AUTH_TOKEN`: The token to use for authorization (_This should be base64 encoded username:password_) +- `OPENAPI_SPEC`: The path to the OpenAPI spec file (_Optional_) (_defaults to latest stable release_) + +*Currently, only Basic Auth is supported.* + +**Page Limit** + +The default is 100 items, but you can change it using `maximum_page_limit` option in [api] section in the `airflow.cfg` file. + +## Tasks + +- [x] First API +- [x] Parse OpenAPI Spec +- [x] Safe/Unsafe mode implementation +- [ ] Parse proper description with list_tools. +- [ ] Airflow config fetch (_specifically for page limit_) +- [ ] Env variables optional (_env variables might not be ideal for airflow plugins_) diff --git a/airflow-mcp-server/Dockerfile b/airflow-mcp-server/Dockerfile deleted file mode 100644 index 84b935e..0000000 --- a/airflow-mcp-server/Dockerfile +++ /dev/null @@ -1,27 +0,0 @@ -# Use a Python image with uv pre-installed -FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim AS uv - -WORKDIR /app - -ENV UV_COMPILE_BYTECODE=1 -ENV UV_LINK_MODE=copy - -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,source=uv.lock,target=uv.lock \ - --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ - uv sync --frozen --no-install-project --no-dev --no-editable - -ADD . /app -RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --frozen --no-dev --no-editable - -FROM python:3.12-slim-bookworm - -WORKDIR /app - -COPY --from=uv /root/.local /root/.local -COPY --from=uv --chown=app:app /app/.venv /app/.venv - -ENV PATH="/app/.venv/bin:$PATH" - -ENTRYPOINT ["airflow-mcp-server"] diff --git a/airflow-mcp-server/LICENSE b/airflow-mcp-server/LICENSE deleted file mode 100644 index b17cff9..0000000 --- a/airflow-mcp-server/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2025 Abhishek Bhakat - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/airflow-mcp-server/README.md b/airflow-mcp-server/README.md deleted file mode 100644 index 722735c..0000000 --- a/airflow-mcp-server/README.md +++ /dev/null @@ -1,66 +0,0 @@ -# airflow-mcp-server: An MCP Server for controlling Airflow - - -## Overview -A [Model Context Protocol](https://modelcontextprotocol.io/) server for controlling Airflow via Airflow APIs. - - -## Setup - -### Usage with Claude Desktop - -```json -{ - "mcpServers": { - "airflow-mcp-server": { - "command": "uvx", - "args": [ - "airflow-mcp-server" - ], - "env": { - "AIRFLOW_BASE_URL": "http:///api/v1", - "AUTH_TOKEN": "" - } - } - } -} -``` - -### Operation Modes - -The server supports two operation modes: - -- **Safe Mode** (`--safe`): Only allows read-only operations (GET requests). This is useful when you want to prevent any modifications to your Airflow instance. -- **Unsafe Mode** (`--unsafe`): Allows all operations including modifications. This is the default mode. - -To start in safe mode: -```bash -airflow-mcp-server --safe -``` - -To explicitly start in unsafe mode (though this is default): -```bash -airflow-mcp-server --unsafe -``` - -### Considerations - -The MCP Server expects environment variables to be set: -- `AIRFLOW_BASE_URL`: The base URL of the Airflow API -- `AUTH_TOKEN`: The token to use for authorization (_This should be base64 encoded username:password_) -- `OPENAPI_SPEC`: The path to the OpenAPI spec file (_Optional_) (_defaults to latest stable release_) - -*Currently, only Basic Auth is supported.* - -**Page Limit** - -The default is 100 items, but you can change it using `maximum_page_limit` option in [api] section in the `airflow.cfg` file. - -## Tasks - -- [x] First API -- [x] Parse OpenAPI Spec -- [x] Safe/Unsafe mode implementation -- [ ] Parse proper description with list_tools. -- [ ] Airflow config fetch (_specifically for page limit_) -- [ ] Env variables optional (_env variables might not be ideal for airflow plugins_) diff --git a/airflow-wingman/LICENSE b/airflow-wingman/LICENSE deleted file mode 100644 index b17cff9..0000000 --- a/airflow-wingman/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2025 Abhishek Bhakat - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/airflow-wingman/README.md b/airflow-wingman/README.md deleted file mode 100644 index a3eb0bf..0000000 --- a/airflow-wingman/README.md +++ /dev/null @@ -1,2 +0,0 @@ -# Airflow Wingman -Airflow plugin to enable LLMs chat in Airflow Webserver. \ No newline at end of file diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml deleted file mode 100644 index 5dfb861..0000000 --- a/airflow-wingman/pyproject.toml +++ /dev/null @@ -1,79 +0,0 @@ - -[project] -name = "airflow-wingman" -version = "0.2.0" -description = "Airflow plugin to enable LLMs chat" -readme = "README.md" -requires-python = ">=3.11" -authors = [ - {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"} -] -dependencies = [ - "apache-airflow>=2.10.0", - "airflow-mcp-server>=0.2.0", - "openai>=1.64.0", - "anthropic>=0.46.0" -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Plugins", - "Operating System :: OS Independent", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.10", -] -license = "MIT" -license-files = ["LICEN[CS]E*"] - -[project.urls] -GitHub = "https://github.com/abhishekbhakat/airflow-mcp-server" -Issues = "https://github.com/abhishekbhakat/airflow-mcp-server/issues" - -[project.entry-points."airflow.plugins"] -wingman = "airflow_wingman:WingmanPlugin" - -[build-system] -requires = ["hatchling"] -build-backend = "hatchling.build" - -[tool.hatch.build.targets.wheel] -packages = ["src/airflow_wingman"] - -[tool.ruff] -line-length = 200 -indent-width = 4 -fix = true -preview = true - -lint.select = [ - "E", # pycodestyle errors - "F", # pyflakes - "I", # isort - "W", # pycodestyle warnings - "C90", # Complexity - "C", # flake8-comprehensions - "ISC", # flake8-implicit-str-concat - "T10", # flake8-debugger - "A", # flake8-builtins - "UP", # pyupgrade -] - -lint.ignore = [ - "C416", # Unnecessary list comprehension - rewrite as a generator expression - "C408", # Unnecessary `dict` call - rewrite as a literal - "ISC001" # Single line implicit string concatenation -] - -lint.fixable = ["ALL"] -lint.unfixable = [] - -[tool.ruff.format] -quote-style = "double" -indent-style = "space" -skip-magic-trailing-comma = false - -[tool.ruff.lint.isort] -combine-as-imports = true - -[tool.ruff.lint.mccabe] -max-complexity = 12 diff --git a/airflow-wingman/src/airflow_wingman/__init__.py b/airflow-wingman/src/airflow_wingman/__init__.py deleted file mode 100644 index a1c4dad..0000000 --- a/airflow-wingman/src/airflow_wingman/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from importlib.metadata import version - -from airflow_wingman.plugin import WingmanPlugin - -__version__ = version("airflow-wingman") -__all__ = ["WingmanPlugin"] diff --git a/airflow-wingman/src/airflow_wingman/llm_client.py b/airflow-wingman/src/airflow_wingman/llm_client.py deleted file mode 100644 index dd58370..0000000 --- a/airflow-wingman/src/airflow_wingman/llm_client.py +++ /dev/null @@ -1,109 +0,0 @@ -""" -Client for making API calls to various LLM providers using their official SDKs. -""" - -from collections.abc import Generator - -from anthropic import Anthropic -from openai import OpenAI - - -class LLMClient: - def __init__(self, api_key: str): - """Initialize the LLM client. - - Args: - api_key: API key for the provider - """ - self.api_key = api_key - self.openai_client = OpenAI(api_key=api_key) - self.anthropic_client = Anthropic(api_key=api_key) - self.openrouter_client = OpenAI( - base_url="https://openrouter.ai/api/v1", - api_key=api_key, - default_headers={ - "HTTP-Referer": "Airflow Wingman", # Required by OpenRouter - "X-Title": "Airflow Wingman", # Required by OpenRouter - }, - ) - - def chat_completion( - self, messages: list[dict[str, str]], model: str, provider: str, temperature: float = 0.7, max_tokens: int | None = None, stream: bool = False - ) -> Generator[str, None, None] | dict: - """Send a chat completion request to the specified provider. - - Args: - messages: List of message dictionaries with 'role' and 'content' - model: Model identifier - provider: Provider identifier (openai, anthropic, openrouter) - temperature: Sampling temperature (0-1) - max_tokens: Maximum tokens to generate - stream: Whether to stream the response - - Returns: - If stream=True, returns a generator yielding response chunks - If stream=False, returns the complete response - """ - try: - if provider == "openai": - return self._openai_chat_completion(messages, model, temperature, max_tokens, stream) - elif provider == "anthropic": - return self._anthropic_chat_completion(messages, model, temperature, max_tokens, stream) - elif provider == "openrouter": - return self._openrouter_chat_completion(messages, model, temperature, max_tokens, stream) - else: - return {"error": f"Unknown provider: {provider}"} - except Exception as e: - return {"error": f"API request failed: {str(e)}"} - - def _openai_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): - """Handle OpenAI chat completion requests.""" - response = self.openai_client.chat.completions.create(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=stream) - - if stream: - - def response_generator(): - for chunk in response: - if chunk.choices[0].delta.content: - yield chunk.choices[0].delta.content - - return response_generator() - else: - return {"content": response.choices[0].message.content} - - def _anthropic_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): - """Handle Anthropic chat completion requests.""" - # Convert messages to Anthropic format - system_message = next((m["content"] for m in messages if m["role"] == "system"), None) - conversation = [] - for m in messages: - if m["role"] != "system": - conversation.append({"role": "assistant" if m["role"] == "assistant" else "user", "content": m["content"]}) - - response = self.anthropic_client.messages.create(model=model, messages=conversation, system=system_message, temperature=temperature, max_tokens=max_tokens, stream=stream) - - if stream: - - def response_generator(): - for chunk in response: - if chunk.delta.text: - yield chunk.delta.text - - return response_generator() - else: - return {"content": response.content[0].text} - - def _openrouter_chat_completion(self, messages: list[dict[str, str]], model: str, temperature: float, max_tokens: int | None, stream: bool): - """Handle OpenRouter chat completion requests.""" - response = self.openrouter_client.chat.completions.create(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, stream=stream) - - if stream: - - def response_generator(): - for chunk in response: - if chunk.choices[0].delta.content: - yield chunk.choices[0].delta.content - - return response_generator() - else: - return {"content": response.choices[0].message.content} diff --git a/airflow-wingman/src/airflow_wingman/llms_models.py b/airflow-wingman/src/airflow_wingman/llms_models.py deleted file mode 100644 index 2e154b4..0000000 --- a/airflow-wingman/src/airflow_wingman/llms_models.py +++ /dev/null @@ -1,48 +0,0 @@ -MODELS = { - "openai": { - "name": "OpenAI", - "endpoint": "https://api.openai.com/v1/chat/completions", - "models": [ - { - "id": "gpt-4o", - "name": "GPT-4o", - "default": True, - "context_window": 128000, - "description": "Input $5/M tokens, Output $15/M tokens", - } - ], - }, - "anthropic": { - "name": "Anthropic", - "endpoint": "https://api.anthropic.com/v1/messages", - "models": [ - { - "id": "claude-3.5-sonnet", - "name": "Claude 3.5 Sonnet", - "default": True, - "context_window": 200000, - "description": "Input $3/M tokens, Output $15/M tokens", - }, - { - "id": "claude-3.5-haiku", - "name": "Claude 3.5 Haiku", - "default": False, - "context_window": 200000, - "description": "Input $0.80/M tokens, Output $4/M tokens", - }, - ], - }, - "openrouter": { - "name": "OpenRouter", - "endpoint": "https://openrouter.ai/api/v1/chat/completions", - "models": [ - { - "id": "custom", - "name": "Custom Model", - "default": False, - "context_window": 128000, # Default context window, will be updated based on model - "description": "Enter any model name supported by OpenRouter (e.g., 'anthropic/claude-3-opus', 'meta-llama/llama-2-70b')", - }, - ], - }, - } \ No newline at end of file diff --git a/airflow-wingman/src/airflow_wingman/mcp_tools.py b/airflow-wingman/src/airflow_wingman/mcp_tools.py deleted file mode 100644 index 7d5d866..0000000 --- a/airflow-wingman/src/airflow_wingman/mcp_tools.py +++ /dev/null @@ -1,7 +0,0 @@ -import asyncio - -from airflow_mcp_server.tools.tool_manager import get_airflow_tools - -# Get tools with their parameters -tools = asyncio.run(get_airflow_tools(mode="safe")) -TOOLS = {tool.name: {"description": tool.description, "parameters": tool.inputSchema} for tool in tools} diff --git a/airflow-wingman/src/airflow_wingman/notes.py b/airflow-wingman/src/airflow_wingman/notes.py deleted file mode 100644 index 5616650..0000000 --- a/airflow-wingman/src/airflow_wingman/notes.py +++ /dev/null @@ -1,13 +0,0 @@ -INTERFACE_MESSAGES = { - "model_recommendation": {"title": "Note", "content": "For best results with function/tool calling capabilities, we recommend using models like Claude-3.5 Sonnet or GPT-4."}, - "security_note": { - "title": "Security", - "content": "For your security, API keys are required for each session and are never stored. If you refresh the page or close the browser, you'll need to enter your API key again.", - }, - "context_window": { - "title": "Context Window", - "content": "Each model has a maximum context window size that determines how much text it can process. " - "For long conversations or large code snippets, consider using models with larger context windows like Claude-3 Opus (200K tokens) or GPT-4 Turbo (128K tokens). " - "For better results try to keep the context size as low as possible. Try using new chats instead of reusing the same chat.", - }, -} diff --git a/airflow-wingman/src/airflow_wingman/plugin.py b/airflow-wingman/src/airflow_wingman/plugin.py deleted file mode 100644 index b1e37d6..0000000 --- a/airflow-wingman/src/airflow_wingman/plugin.py +++ /dev/null @@ -1,32 +0,0 @@ -"""Plugin definition for Airflow Wingman.""" - -from airflow.plugins_manager import AirflowPlugin -from flask import Blueprint - -from airflow_wingman.views import WingmanView - -# Create Blueprint -bp = Blueprint( - "wingman", - __name__, - template_folder="templates", - static_folder="static", - static_url_path="/static/wingman", -) - -# Create AppBuilder View -v_appbuilder_view = WingmanView() -v_appbuilder_package = { - "name": "Wingman", - "category": "AI", - "view": v_appbuilder_view, -} - - -# Create Plugin -class WingmanPlugin(AirflowPlugin): - """Airflow plugin for Wingman chat interface.""" - - name = "wingman" - flask_blueprints = [bp] - appbuilder_views = [v_appbuilder_package] diff --git a/airflow-wingman/src/airflow_wingman/prompt_engineering.py b/airflow-wingman/src/airflow_wingman/prompt_engineering.py deleted file mode 100644 index 6317971..0000000 --- a/airflow-wingman/src/airflow_wingman/prompt_engineering.py +++ /dev/null @@ -1,40 +0,0 @@ -""" -Prompt engineering for the Airflow Wingman plugin. -Contains prompts and instructions for the AI assistant. -""" - -import json - -from airflow_wingman.mcp_tools import TOOLS - -INSTRUCTIONS = { - "default": f"""You are Airflow Wingman, a helpful AI assistant integrated into Apache Airflow. -You have deep knowledge of Apache Airflow's architecture, DAGs, operators, and best practices. -The Airflow version being used is >=2.10. - -You have access to the following Airflow API tools: - -{json.dumps(TOOLS, indent=2)} - -You can use these tools to fetch information and help users understand and manage their Airflow environment. -""" -} - - -def prepare_messages(messages: list[dict[str, str]], instruction_key: str = "default") -> list[dict[str, str]]: - """Prepare messages for the chat completion request. - - Args: - messages: List of messages in the conversation - instruction_key: Key for the instruction template to use - - Returns: - List of message dictionaries ready for the chat completion API - """ - instruction = INSTRUCTIONS.get(instruction_key, INSTRUCTIONS["default"]) - - # Add instruction as first system message if not present - if not messages or messages[0].get("role") != "system": - messages.insert(0, {"role": "system", "content": instruction}) - - return messages diff --git a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html b/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html deleted file mode 100644 index 72657f0..0000000 --- a/airflow-wingman/src/airflow_wingman/templates/wingman_chat.html +++ /dev/null @@ -1,384 +0,0 @@ -{% extends "appbuilder/base.html" %} - -{% block head_meta %} - {{ super() }} - -{% endblock %} - -{% block content %} -
- -
-
-
-
-

Airflow Wingman

-
-
-

{{ interface_messages.model_recommendation.title }}: {{ interface_messages.model_recommendation.content }}

-
-

{{ interface_messages.security_note.title }}: {{ interface_messages.security_note.content }}

-
-

{{ interface_messages.context_window.title }}: {{ interface_messages.context_window.content }}

-
-
-
-
- -
- -
-
-
-

Provider Selection

-
-
- {% for provider_id, provider in models.items() %} -
-

{{ provider.name }}

- {% for model in provider.models %} -
- -
- {% endfor %} -
- {% endfor %} -
- - -
-
- - - Only required for OpenRouter provider -
-
- - -
-
- - - Your API key will be used for the selected provider -
-
- - -
-
- - -
-
-
- -
-
-
-
- -
-
-
- - - - -
-
-
-
-
-
-
-
- - - - -{% endblock %} diff --git a/airflow-wingman/src/airflow_wingman/views.py b/airflow-wingman/src/airflow_wingman/views.py deleted file mode 100644 index 34158d7..0000000 --- a/airflow-wingman/src/airflow_wingman/views.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Views for Airflow Wingman plugin.""" - -from flask import Response, request, stream_with_context -from flask.json import jsonify -from flask_appbuilder import BaseView as AppBuilderBaseView, expose - -from airflow_wingman.llm_client import LLMClient -from airflow_wingman.llms_models import MODELS -from airflow_wingman.notes import INTERFACE_MESSAGES -from airflow_wingman.prompt_engineering import prepare_messages - - -class WingmanView(AppBuilderBaseView): - """View for Airflow Wingman plugin.""" - - route_base = "/wingman" - default_view = "chat" - - @expose("/") - def chat(self): - """Render chat interface.""" - providers = {provider: info["name"] for provider, info in MODELS.items()} - return self.render_template("wingman_chat.html", title="Airflow Wingman", models=MODELS, providers=providers, interface_messages=INTERFACE_MESSAGES) - - @expose("/chat", methods=["POST"]) - def chat_completion(self): - """Handle chat completion requests.""" - try: - data = self._validate_chat_request(request.get_json()) - - # Create a new client for this request - client = LLMClient(data["api_key"]) - - if data["stream"]: - return self._handle_streaming_response(client, data) - else: - return self._handle_regular_response(client, data) - - except ValueError as e: - return jsonify({"error": str(e)}), 400 - except Exception as e: - return jsonify({"error": str(e)}), 500 - - def _validate_chat_request(self, data: dict) -> dict: - """Validate chat request data.""" - if not data: - raise ValueError("No data provided") - - required_fields = ["provider", "model", "messages", "api_key"] - missing = [f for f in required_fields if not data.get(f)] - if missing: - raise ValueError(f"Missing required fields: {', '.join(missing)}") - - # Prepare messages with system instruction while maintaining history - messages = data["messages"] - messages = prepare_messages(messages) - - return { - "provider": data["provider"], - "model": data["model"], - "messages": messages, - "api_key": data["api_key"], - "stream": data.get("stream", False), - "temperature": data.get("temperature", 0.7), - "max_tokens": data.get("max_tokens"), - } - - def _handle_streaming_response(self, client: LLMClient, data: dict) -> Response: - """Handle streaming response.""" - - def generate(): - for chunk in client.chat_completion(messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=True): - yield f"data: {chunk}\n\n" - - response = Response(stream_with_context(generate()), mimetype="text/event-stream") - response.headers["Content-Type"] = "text/event-stream" - response.headers["Cache-Control"] = "no-cache" - response.headers["Connection"] = "keep-alive" - return response - - def _handle_regular_response(self, client: LLMClient, data: dict) -> Response: - """Handle regular response.""" - response = client.chat_completion(messages=data["messages"], model=data["model"], provider=data["provider"], temperature=data["temperature"], max_tokens=data["max_tokens"], stream=False) - return jsonify(response) diff --git a/airflow_settings.yaml b/airflow_settings.yaml deleted file mode 100644 index 1c16dc0..0000000 --- a/airflow_settings.yaml +++ /dev/null @@ -1,25 +0,0 @@ -# This file allows you to configure Airflow Connections, Pools, and Variables in a single place for local development only. -# NOTE: json dicts can be added to the conn_extra field as yaml key value pairs. See the example below. - -# For more information, refer to our docs: https://www.astronomer.io/docs/astro/cli/develop-project#configure-airflow_settingsyaml-local-development-only -# For questions, reach out to: https://support.astronomer.io -# For issues create an issue ticket here: https://github.com/astronomer/astro-cli/issues - -airflow: - connections: - - conn_id: - conn_type: - conn_host: - conn_schema: - conn_login: - conn_password: - conn_port: - conn_extra: - example_extra_field: example-value - pools: - - pool_name: - pool_slot: - pool_description: - variables: - - variable_name: - variable_value: diff --git a/dags/.airflowignore b/dags/.airflowignore deleted file mode 100644 index e69de29..0000000 diff --git a/dags/exampledag.py b/dags/exampledag.py deleted file mode 100644 index 8b08b7b..0000000 --- a/dags/exampledag.py +++ /dev/null @@ -1,100 +0,0 @@ -""" -## Astronaut ETL example DAG - -This DAG queries the list of astronauts currently in space from the -Open Notify API and prints each astronaut's name and flying craft. - -There are two tasks, one to get the data from the API and save the results, -and another to print the results. Both tasks are written in Python using -Airflow's TaskFlow API, which allows you to easily turn Python functions into -Airflow tasks, and automatically infer dependencies and pass data. - -The second task uses dynamic task mapping to create a copy of the task for -each Astronaut in the list retrieved from the API. This list will change -depending on how many Astronauts are in space, and the DAG will adjust -accordingly each time it runs. - -For more explanation and getting started instructions, see our Write your -first DAG tutorial: https://www.astronomer.io/docs/learn/get-started-with-airflow - -![Picture of the ISS](https://www.esa.int/var/esa/storage/images/esa_multimedia/images/2010/02/space_station_over_earth/10293696-3-eng-GB/Space_Station_over_Earth_card_full.jpg) -""" - -from airflow import Dataset -from airflow.decorators import dag, task -from pendulum import datetime -import requests - - -# Define the basic parameters of the DAG, like schedule and start_date -@dag( - start_date=datetime(2024, 1, 1), - schedule="@daily", - catchup=False, - doc_md=__doc__, - default_args={"owner": "Astro", "retries": 3}, - tags=["example"], -) -def example_astronauts(): - # Define tasks - @task( - # Define a dataset outlet for the task. This can be used to schedule downstream DAGs when this task has run. - outlets=[Dataset("current_astronauts")] - ) # Define that this task updates the `current_astronauts` Dataset - def get_astronauts(**context) -> list[dict]: - """ - This task uses the requests library to retrieve a list of Astronauts - currently in space. The results are pushed to XCom with a specific key - so they can be used in a downstream pipeline. The task returns a list - of Astronauts to be used in the next task. - """ - try: - r = requests.get("http://api.open-notify.org/astros.json") - r.raise_for_status() - number_of_people_in_space = r.json()["number"] - list_of_people_in_space = r.json()["people"] - except: - print("API currently not available, using hardcoded data instead.") - number_of_people_in_space = 12 - list_of_people_in_space = [ - {"craft": "ISS", "name": "Oleg Kononenko"}, - {"craft": "ISS", "name": "Nikolai Chub"}, - {"craft": "ISS", "name": "Tracy Caldwell Dyson"}, - {"craft": "ISS", "name": "Matthew Dominick"}, - {"craft": "ISS", "name": "Michael Barratt"}, - {"craft": "ISS", "name": "Jeanette Epps"}, - {"craft": "ISS", "name": "Alexander Grebenkin"}, - {"craft": "ISS", "name": "Butch Wilmore"}, - {"craft": "ISS", "name": "Sunita Williams"}, - {"craft": "Tiangong", "name": "Li Guangsu"}, - {"craft": "Tiangong", "name": "Li Cong"}, - {"craft": "Tiangong", "name": "Ye Guangfu"}, - ] - - context["ti"].xcom_push( - key="number_of_people_in_space", value=number_of_people_in_space - ) - return list_of_people_in_space - - @task - def print_astronaut_craft(greeting: str, person_in_space: dict) -> None: - """ - This task creates a print statement with the name of an - Astronaut in space and the craft they are flying on from - the API request results of the previous task, along with a - greeting which is hard-coded in this example. - """ - craft = person_in_space["craft"] - name = person_in_space["name"] - - print(f"{name} is currently in space flying on the {craft}! {greeting}") - - # Use dynamic task mapping to run the print_astronaut_craft task for each - # Astronaut in space - print_astronaut_craft.partial(greeting="Hello! :)").expand( - person_in_space=get_astronauts() # Define dependencies using TaskFlow API syntax - ) - - -# Instantiate the DAG -example_astronauts() diff --git a/packages.txt b/packages.txt deleted file mode 100644 index e69de29..0000000 diff --git a/airflow-mcp-server/pyproject.toml b/pyproject.toml similarity index 100% rename from airflow-mcp-server/pyproject.toml rename to pyproject.toml diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 1bb359b..0000000 --- a/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -# Astro Runtime includes the following pre-installed providers packages: https://www.astronomer.io/docs/astro/runtime-image-architecture#provider-packages diff --git a/airflow-mcp-server/src/airflow_mcp_server/__init__.py b/src/airflow_mcp_server/__init__.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/__init__.py rename to src/airflow_mcp_server/__init__.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/__main__.py b/src/airflow_mcp_server/__main__.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/__main__.py rename to src/airflow_mcp_server/__main__.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/client/__init__.py b/src/airflow_mcp_server/client/__init__.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/client/__init__.py rename to src/airflow_mcp_server/client/__init__.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/client/airflow_client.py b/src/airflow_mcp_server/client/airflow_client.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/client/airflow_client.py rename to src/airflow_mcp_server/client/airflow_client.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/parser/__init__.py b/src/airflow_mcp_server/parser/__init__.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/parser/__init__.py rename to src/airflow_mcp_server/parser/__init__.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/parser/operation_parser.py b/src/airflow_mcp_server/parser/operation_parser.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/parser/operation_parser.py rename to src/airflow_mcp_server/parser/operation_parser.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/resources/v1.yaml b/src/airflow_mcp_server/resources/v1.yaml similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/resources/v1.yaml rename to src/airflow_mcp_server/resources/v1.yaml diff --git a/airflow-mcp-server/src/airflow_mcp_server/server.py b/src/airflow_mcp_server/server.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/server.py rename to src/airflow_mcp_server/server.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/server_safe.py b/src/airflow_mcp_server/server_safe.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/server_safe.py rename to src/airflow_mcp_server/server_safe.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/server_unsafe.py b/src/airflow_mcp_server/server_unsafe.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/server_unsafe.py rename to src/airflow_mcp_server/server_unsafe.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/tools/__init__.py b/src/airflow_mcp_server/tools/__init__.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/tools/__init__.py rename to src/airflow_mcp_server/tools/__init__.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/tools/airflow_tool.py b/src/airflow_mcp_server/tools/airflow_tool.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/tools/airflow_tool.py rename to src/airflow_mcp_server/tools/airflow_tool.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/tools/base_tools.py b/src/airflow_mcp_server/tools/base_tools.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/tools/base_tools.py rename to src/airflow_mcp_server/tools/base_tools.py diff --git a/airflow-mcp-server/src/airflow_mcp_server/tools/tool_manager.py b/src/airflow_mcp_server/tools/tool_manager.py similarity index 100% rename from airflow-mcp-server/src/airflow_mcp_server/tools/tool_manager.py rename to src/airflow_mcp_server/tools/tool_manager.py diff --git a/airflow-mcp-server/tests/__init__.py b/tests/__init__.py similarity index 100% rename from airflow-mcp-server/tests/__init__.py rename to tests/__init__.py diff --git a/airflow-mcp-server/tests/client/test_airflow_client.py b/tests/client/test_airflow_client.py similarity index 100% rename from airflow-mcp-server/tests/client/test_airflow_client.py rename to tests/client/test_airflow_client.py diff --git a/airflow-mcp-server/tests/conftest.py b/tests/conftest.py similarity index 100% rename from airflow-mcp-server/tests/conftest.py rename to tests/conftest.py diff --git a/tests/dags/test_dag_example.py b/tests/dags/test_dag_example.py deleted file mode 100644 index 6ff3552..0000000 --- a/tests/dags/test_dag_example.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Example DAGs test. This test ensures that all Dags have tags, retries set to two, and no import errors. This is an example pytest and may not be fit the context of your DAGs. Feel free to add and remove tests.""" - -import os -import logging -from contextlib import contextmanager -import pytest -from airflow.models import DagBag - - -@contextmanager -def suppress_logging(namespace): - logger = logging.getLogger(namespace) - old_value = logger.disabled - logger.disabled = True - try: - yield - finally: - logger.disabled = old_value - - -def get_import_errors(): - """ - Generate a tuple for import errors in the dag bag - """ - with suppress_logging("airflow"): - dag_bag = DagBag(include_examples=False) - - def strip_path_prefix(path): - return os.path.relpath(path, os.environ.get("AIRFLOW_HOME")) - - # prepend "(None,None)" to ensure that a test object is always created even if it's a no op. - return [(None, None)] + [ - (strip_path_prefix(k), v.strip()) for k, v in dag_bag.import_errors.items() - ] - - -def get_dags(): - """ - Generate a tuple of dag_id, in the DagBag - """ - with suppress_logging("airflow"): - dag_bag = DagBag(include_examples=False) - - def strip_path_prefix(path): - return os.path.relpath(path, os.environ.get("AIRFLOW_HOME")) - - return [(k, v, strip_path_prefix(v.fileloc)) for k, v in dag_bag.dags.items()] - - -@pytest.mark.parametrize( - "rel_path,rv", get_import_errors(), ids=[x[0] for x in get_import_errors()] -) -def test_file_imports(rel_path, rv): - """Test for import errors on a file""" - if rel_path and rv: - raise Exception(f"{rel_path} failed to import with message \n {rv}") - - -APPROVED_TAGS = {} - - -@pytest.mark.parametrize( - "dag_id,dag,fileloc", get_dags(), ids=[x[2] for x in get_dags()] -) -def test_dag_tags(dag_id, dag, fileloc): - """ - test if a DAG is tagged and if those TAGs are in the approved list - """ - assert dag.tags, f"{dag_id} in {fileloc} has no tags" - if APPROVED_TAGS: - assert not set(dag.tags) - APPROVED_TAGS - - -@pytest.mark.parametrize( - "dag_id,dag, fileloc", get_dags(), ids=[x[2] for x in get_dags()] -) -def test_dag_retries(dag_id, dag, fileloc): - """ - test if a DAG has retries set - """ - assert ( - dag.default_args.get("retries", None) >= 2 - ), f"{dag_id} in {fileloc} must have task retries >= 2." diff --git a/airflow-mcp-server/tests/parser/test_operation_parser.py b/tests/parser/test_operation_parser.py similarity index 100% rename from airflow-mcp-server/tests/parser/test_operation_parser.py rename to tests/parser/test_operation_parser.py diff --git a/airflow-mcp-server/tests/tools/__init__.py b/tests/tools/__init__.py similarity index 100% rename from airflow-mcp-server/tests/tools/__init__.py rename to tests/tools/__init__.py diff --git a/airflow-mcp-server/tests/tools/test_airflow_tool.py b/tests/tools/test_airflow_tool.py similarity index 100% rename from airflow-mcp-server/tests/tools/test_airflow_tool.py rename to tests/tools/test_airflow_tool.py diff --git a/airflow-mcp-server/tests/tools/test_models.py b/tests/tools/test_models.py similarity index 100% rename from airflow-mcp-server/tests/tools/test_models.py rename to tests/tools/test_models.py diff --git a/airflow-mcp-server/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py similarity index 100% rename from airflow-mcp-server/tests/tools/test_tool_manager.py rename to tests/tools/test_tool_manager.py diff --git a/airflow-mcp-server/uv.lock b/uv.lock similarity index 100% rename from airflow-mcp-server/uv.lock rename to uv.lock