Clean up for only MCP Server
This commit is contained in:
@@ -1,2 +0,0 @@
|
|||||||
project:
|
|
||||||
name: airflow-mcp-server
|
|
||||||
@@ -1 +0,0 @@
|
|||||||
# Add dag files to exempt from parse test below. ex: dags/<test-file>
|
|
||||||
@@ -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")
|
|
||||||
@@ -1,11 +0,0 @@
|
|||||||
astro
|
|
||||||
.git
|
|
||||||
.env
|
|
||||||
airflow_settings.yaml
|
|
||||||
logs/
|
|
||||||
.venv
|
|
||||||
airflow.db
|
|
||||||
airflow.cfg
|
|
||||||
resources/
|
|
||||||
assets/
|
|
||||||
README.md
|
|
||||||
3
.github/workflows/python-publish.yml
vendored
3
.github/workflows/python-publish.yml
vendored
@@ -31,7 +31,6 @@ jobs:
|
|||||||
pip install uv
|
pip install uv
|
||||||
|
|
||||||
- name: Build release distributions
|
- name: Build release distributions
|
||||||
working-directory: airflow-mcp-server
|
|
||||||
run: |
|
run: |
|
||||||
uv pip install --system build
|
uv pip install --system build
|
||||||
python -m build
|
python -m build
|
||||||
@@ -40,7 +39,7 @@ jobs:
|
|||||||
uses: actions/upload-artifact@v4
|
uses: actions/upload-artifact@v4
|
||||||
with:
|
with:
|
||||||
name: release-dists
|
name: release-dists
|
||||||
path: airflow-mcp-server/dist/
|
path: dist/
|
||||||
|
|
||||||
pypi-publish:
|
pypi-publish:
|
||||||
runs-on: ubuntu-latest
|
runs-on: ubuntu-latest
|
||||||
|
|||||||
30
Dockerfile
30
Dockerfile
@@ -1,3 +1,27 @@
|
|||||||
FROM quay.io/astronomer/astro-runtime:12.7.1
|
# Use a Python image with uv pre-installed
|
||||||
RUN cd airflow-mcp-server && pip install -e .
|
FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim AS uv
|
||||||
RUN cd airflow-wingman && pip install -e .
|
|
||||||
|
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"]
|
||||||
|
|||||||
47
README.md
47
README.md
@@ -2,12 +2,13 @@
|
|||||||
|
|
||||||
|
|
||||||
## Overview
|
## 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
|
## Demo Video
|
||||||
|
|
||||||
https://github.com/user-attachments/assets/f3e60fff-8680-4dd9-b08e-fa7db655a705
|
https://github.com/user-attachments/assets/f3e60fff-8680-4dd9-b08e-fa7db655a705
|
||||||
|
|
||||||
|
|
||||||
## Setup
|
## Setup
|
||||||
|
|
||||||
### Usage with Claude Desktop
|
### 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:
|
- **Safe Mode** (`--safe`): Only allows read-only operations (GET requests). This is useful when you want to prevent any modifications to your Airflow instance.
|
||||||
- Adding Airflow to AI (_complete access to an Airflow deployment_)
|
- **Unsafe Mode** (`--unsafe`): Allows all operations including modifications. This is the default mode.
|
||||||
- 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`.
|
To start in safe mode:
|
||||||
- Adding AI to Airflow (_read-only access using Airflow Plugin_)
|
```bash
|
||||||
- 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.
|
airflow-mcp-server --safe
|
||||||
- Use command `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_)
|
||||||
|
|||||||
@@ -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"]
|
|
||||||
@@ -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.
|
|
||||||
@@ -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://<host:port>/api/v1",
|
|
||||||
"AUTH_TOKEN": "<base64_encoded_username_password>"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
### 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_)
|
|
||||||
@@ -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.
|
|
||||||
@@ -1,2 +0,0 @@
|
|||||||
# Airflow Wingman
|
|
||||||
Airflow plugin to enable LLMs chat in Airflow Webserver.
|
|
||||||
@@ -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
|
|
||||||
@@ -1,6 +0,0 @@
|
|||||||
from importlib.metadata import version
|
|
||||||
|
|
||||||
from airflow_wingman.plugin import WingmanPlugin
|
|
||||||
|
|
||||||
__version__ = version("airflow-wingman")
|
|
||||||
__all__ = ["WingmanPlugin"]
|
|
||||||
@@ -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}
|
|
||||||
@@ -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')",
|
|
||||||
},
|
|
||||||
],
|
|
||||||
},
|
|
||||||
}
|
|
||||||
@@ -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}
|
|
||||||
@@ -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.",
|
|
||||||
},
|
|
||||||
}
|
|
||||||
@@ -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]
|
|
||||||
@@ -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
|
|
||||||
@@ -1,384 +0,0 @@
|
|||||||
{% extends "appbuilder/base.html" %}
|
|
||||||
|
|
||||||
{% block head_meta %}
|
|
||||||
{{ super() }}
|
|
||||||
<meta name="csrf-token" content="{{ csrf_token() }}">
|
|
||||||
{% endblock %}
|
|
||||||
|
|
||||||
{% block content %}
|
|
||||||
<div class="container-fluid">
|
|
||||||
<!-- Banner -->
|
|
||||||
<div class="row">
|
|
||||||
<div class="col-md-12">
|
|
||||||
<div class="panel panel-primary">
|
|
||||||
<div class="panel-heading">
|
|
||||||
<h3 class="panel-title">Airflow Wingman</h3>
|
|
||||||
</div>
|
|
||||||
<div class="alert alert-info" style="margin: 15px;">
|
|
||||||
<p><strong>{{ interface_messages.model_recommendation.title }}:</strong> {{ interface_messages.model_recommendation.content }}</p>
|
|
||||||
<hr style="margin: 10px 0;">
|
|
||||||
<p><strong>{{ interface_messages.security_note.title }}:</strong> {{ interface_messages.security_note.content }}</p>
|
|
||||||
<hr style="margin: 10px 0;">
|
|
||||||
<p><strong>{{ interface_messages.context_window.title }}:</strong> {{ interface_messages.context_window.content }}</p>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
|
|
||||||
<div class="row">
|
|
||||||
<!-- Sidebar -->
|
|
||||||
<div class="col-md-3">
|
|
||||||
<div class="panel panel-default">
|
|
||||||
<div class="panel-heading">
|
|
||||||
<h3 class="panel-title">Provider Selection</h3>
|
|
||||||
</div>
|
|
||||||
<div class="panel-body">
|
|
||||||
{% for provider_id, provider in models.items() %}
|
|
||||||
<div class="provider-section mb-3">
|
|
||||||
<h4 class="provider-name">{{ provider.name }}</h4>
|
|
||||||
{% for model in provider.models %}
|
|
||||||
<div class="radio model-option">
|
|
||||||
<label class="model-label" title="{{ model.description }}">
|
|
||||||
<input type="radio"
|
|
||||||
name="model"
|
|
||||||
value="{{ provider_id }}:{{ model.id }}"
|
|
||||||
{% if model.default %}checked{% endif %}
|
|
||||||
data-context-window="{{ model.context_window }}"
|
|
||||||
data-provider="{{ provider_id }}"
|
|
||||||
data-model-name="{{ model.name }}">
|
|
||||||
{{ model.name }}
|
|
||||||
</label>
|
|
||||||
</div>
|
|
||||||
{% endfor %}
|
|
||||||
</div>
|
|
||||||
{% endfor %}
|
|
||||||
</div>
|
|
||||||
|
|
||||||
<!-- Model Name Input -->
|
|
||||||
<div class="panel-body" style="border-top: 1px solid #ddd; padding-top: 15px;">
|
|
||||||
<div class="form-group">
|
|
||||||
<label for="modelName">Model Name</label>
|
|
||||||
<input type="text" class="form-control" id="modelName" placeholder="Enter model name for OpenRouter" disabled>
|
|
||||||
<small class="form-text text-muted">Only required for OpenRouter provider</small>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
|
|
||||||
<!-- API Key Input -->
|
|
||||||
<div class="panel-body" style="border-top: 1px solid #ddd; padding-top: 15px;">
|
|
||||||
<div class="form-group">
|
|
||||||
<label for="api-key">API Key</label>
|
|
||||||
<input type="password"
|
|
||||||
class="form-control"
|
|
||||||
id="api-key"
|
|
||||||
placeholder="Enter API key for selected provider"
|
|
||||||
required
|
|
||||||
autocomplete="off">
|
|
||||||
<small class="text-muted">Your API key will be used for the selected provider</small>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
|
|
||||||
<style>
|
|
||||||
.provider-section {
|
|
||||||
margin-bottom: 20px;
|
|
||||||
}
|
|
||||||
.provider-name {
|
|
||||||
font-size: 16px;
|
|
||||||
font-weight: bold;
|
|
||||||
margin-bottom: 10px;
|
|
||||||
color: #666;
|
|
||||||
}
|
|
||||||
.model-option {
|
|
||||||
margin-left: 15px;
|
|
||||||
margin-bottom: 8px;
|
|
||||||
}
|
|
||||||
.model-option label {
|
|
||||||
display: block;
|
|
||||||
cursor: pointer;
|
|
||||||
}
|
|
||||||
</style>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
|
|
||||||
<!-- Main Chat Window -->
|
|
||||||
<div class="col-md-9">
|
|
||||||
<div class="panel panel-default" style="height: calc(80vh - 250px); display: flex; flex-direction: column;">
|
|
||||||
<div class="panel-body" style="flex-grow: 1; overflow-y: auto; padding: 15px;" id="chat-messages">
|
|
||||||
<!-- Messages will be dynamically added here -->
|
|
||||||
</div>
|
|
||||||
<div class="panel-footer" style="padding: 15px; background-color: white;">
|
|
||||||
<div class="row">
|
|
||||||
<div class="col-md-2">
|
|
||||||
<button class="btn btn-default btn-block" type="button" id="refresh-button" title="Start a new chat">
|
|
||||||
<i class="fa fa-refresh"></i> New Chat
|
|
||||||
</button>
|
|
||||||
</div>
|
|
||||||
<div class="col-md-10">
|
|
||||||
<div class="input-group">
|
|
||||||
<input type="text" class="form-control" id="message-input" placeholder="Type your message...">
|
|
||||||
<span class="input-group-btn">
|
|
||||||
<button class="btn btn-primary" type="button" id="send-button">
|
|
||||||
<i class="fa fa-paper-plane"></i> Send
|
|
||||||
</button>
|
|
||||||
</span>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
|
|
||||||
<style>
|
|
||||||
.message {
|
|
||||||
margin-bottom: 15px;
|
|
||||||
max-width: 80%;
|
|
||||||
clear: both;
|
|
||||||
}
|
|
||||||
|
|
||||||
.message-user {
|
|
||||||
float: right;
|
|
||||||
background-color: #f0f7ff;
|
|
||||||
border: 1px solid #d1e6ff;
|
|
||||||
border-radius: 15px 15px 0 15px;
|
|
||||||
padding: 10px 15px;
|
|
||||||
}
|
|
||||||
|
|
||||||
.message-assistant {
|
|
||||||
float: left;
|
|
||||||
background-color: #f8f9fa;
|
|
||||||
border: 1px solid #e9ecef;
|
|
||||||
border-radius: 15px 15px 15px 0;
|
|
||||||
padding: 10px 15px;
|
|
||||||
}
|
|
||||||
|
|
||||||
#chat-messages::after {
|
|
||||||
content: "";
|
|
||||||
clear: both;
|
|
||||||
display: table;
|
|
||||||
}
|
|
||||||
|
|
||||||
.panel-body::-webkit-scrollbar {
|
|
||||||
width: 8px;
|
|
||||||
}
|
|
||||||
|
|
||||||
.panel-body::-webkit-scrollbar-track {
|
|
||||||
background: #f1f1f1;
|
|
||||||
}
|
|
||||||
|
|
||||||
.panel-body::-webkit-scrollbar-thumb {
|
|
||||||
background: #888;
|
|
||||||
border-radius: 4px;
|
|
||||||
}
|
|
||||||
|
|
||||||
.panel-body::-webkit-scrollbar-thumb:hover {
|
|
||||||
background: #555;
|
|
||||||
}
|
|
||||||
</style>
|
|
||||||
|
|
||||||
<script>
|
|
||||||
document.addEventListener('DOMContentLoaded', function() {
|
|
||||||
// Add title attributes for tooltips
|
|
||||||
document.querySelectorAll('[data-bs-toggle="tooltip"]').forEach(function(el) {
|
|
||||||
el.title = el.getAttribute('title') || el.getAttribute('data-bs-original-title');
|
|
||||||
});
|
|
||||||
|
|
||||||
// Handle model selection and model name input
|
|
||||||
const modelNameInput = document.getElementById('modelName');
|
|
||||||
const modelRadios = document.querySelectorAll('input[name="model"]');
|
|
||||||
|
|
||||||
modelRadios.forEach(function(radio) {
|
|
||||||
radio.addEventListener('change', function() {
|
|
||||||
const provider = this.value.split(':')[0]; // Get provider from value instead of data attribute
|
|
||||||
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;
|
|
||||||
modelNameInput.value = '';
|
|
||||||
modelNameInput.placeholder = 'Enter model name for OpenRouter';
|
|
||||||
} else {
|
|
||||||
console.log('Disabling model name input');
|
|
||||||
modelNameInput.disabled = true;
|
|
||||||
modelNameInput.value = modelName;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
});
|
|
||||||
|
|
||||||
// Set initial state based on default selection
|
|
||||||
const defaultSelected = document.querySelector('input[name="model"]:checked');
|
|
||||||
if (defaultSelected) {
|
|
||||||
const provider = defaultSelected.value.split(':')[0]; // Get provider from value instead of data attribute
|
|
||||||
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;
|
|
||||||
modelNameInput.value = '';
|
|
||||||
modelNameInput.placeholder = 'Enter model name for OpenRouter';
|
|
||||||
} else {
|
|
||||||
console.log('Initially disabling model name input');
|
|
||||||
modelNameInput.disabled = true;
|
|
||||||
modelNameInput.value = modelName;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
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');
|
|
||||||
messageDiv.className = `message ${isUser ? 'message-user' : 'message-assistant'}`;
|
|
||||||
messageDiv.textContent = content;
|
|
||||||
chatMessages.appendChild(messageDiv);
|
|
||||||
chatMessages.scrollTop = chatMessages.scrollHeight;
|
|
||||||
return messageDiv;
|
|
||||||
}
|
|
||||||
|
|
||||||
async function sendMessage() {
|
|
||||||
const message = messageInput.value.trim();
|
|
||||||
if (!message) return;
|
|
||||||
|
|
||||||
// Get selected model
|
|
||||||
const selectedModel = document.querySelector('input[name="model"]:checked');
|
|
||||||
if (!selectedModel) {
|
|
||||||
alert('Please select a model');
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
const [provider, modelId] = selectedModel.value.split(':');
|
|
||||||
const modelName = provider === 'openrouter' ? modelNameInput.value : modelId;
|
|
||||||
|
|
||||||
// Clear input and add user message
|
|
||||||
messageInput.value = '';
|
|
||||||
addMessage(message, true);
|
|
||||||
|
|
||||||
try {
|
|
||||||
// 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);
|
|
||||||
|
|
||||||
// Get API key
|
|
||||||
const apiKey = document.getElementById('api-key').value.trim();
|
|
||||||
if (!apiKey) {
|
|
||||||
alert('Please enter an API key');
|
|
||||||
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,
|
|
||||||
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();
|
|
||||||
let fullResponse = '';
|
|
||||||
|
|
||||||
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;
|
|
||||||
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}`;
|
|
||||||
currentMessageDiv.style.color = 'red';
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
sendButton.addEventListener('click', sendMessage);
|
|
||||||
messageInput.addEventListener('keypress', function(e) {
|
|
||||||
if (e.key === 'Enter') {
|
|
||||||
sendMessage();
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
refreshButton.addEventListener('click', clearChat);
|
|
||||||
});
|
|
||||||
</script>
|
|
||||||
{% endblock %}
|
|
||||||
@@ -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)
|
|
||||||
@@ -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:
|
|
||||||
@@ -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
|
|
||||||
|
|
||||||
|
|
||||||
"""
|
|
||||||
|
|
||||||
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()
|
|
||||||
@@ -1 +0,0 @@
|
|||||||
# Astro Runtime includes the following pre-installed providers packages: https://www.astronomer.io/docs/astro/runtime-image-architecture#provider-packages
|
|
||||||
@@ -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, <DAG objects> 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."
|
|
||||||
0
airflow-mcp-server/uv.lock → uv.lock
generated
0
airflow-mcp-server/uv.lock → uv.lock
generated
Reference in New Issue
Block a user