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