diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..7a1acfe --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,15 @@ +repos: +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.1.11 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.5.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-added-large-files diff --git a/LICENSE b/LICENSE index 8036231..b17cff9 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2025 Abhishek +Copyright (c) 2025 Abhishek Bhakat Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index 0fa64bd..931d612 100644 --- a/README.md +++ b/README.md @@ -1 +1,2 @@ -# airflow-wingman \ No newline at end of file +# Airflow Wingman +Airflow plugin to enable LLMs chat in Airflow Webserver. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4431ddc --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,78 @@ + +[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", + "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-wingman" +Issues = "https://github.com/abhishekbhakat/airflow-wingman/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/src/airflow_wingman/__init__.py b/src/airflow_wingman/__init__.py new file mode 100644 index 0000000..a1c4dad --- /dev/null +++ b/src/airflow_wingman/__init__.py @@ -0,0 +1,6 @@ +from importlib.metadata import version + +from airflow_wingman.plugin import WingmanPlugin + +__version__ = version("airflow-wingman") +__all__ = ["WingmanPlugin"] diff --git a/src/airflow_wingman/llm_client.py b/src/airflow_wingman/llm_client.py new file mode 100644 index 0000000..dd58370 --- /dev/null +++ b/src/airflow_wingman/llm_client.py @@ -0,0 +1,109 @@ +""" +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/src/airflow_wingman/llms_models.py b/src/airflow_wingman/llms_models.py new file mode 100644 index 0000000..c73cd46 --- /dev/null +++ b/src/airflow_wingman/llms_models.py @@ -0,0 +1,48 @@ +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')", + }, + ], + }, +} diff --git a/src/airflow_wingman/notes.py b/src/airflow_wingman/notes.py new file mode 100644 index 0000000..5616650 --- /dev/null +++ b/src/airflow_wingman/notes.py @@ -0,0 +1,13 @@ +INTERFACE_MESSAGES = { + "model_recommendation": {"title": "Note", "content": "For best results with function/tool calling capabilities, we recommend using models like Claude-3.5 Sonnet or GPT-4."}, + "security_note": { + "title": "Security", + "content": "For your security, API keys are required for each session and are never stored. If you refresh the page or close the browser, you'll need to enter your API key again.", + }, + "context_window": { + "title": "Context Window", + "content": "Each model has a maximum context window size that determines how much text it can process. " + "For long conversations or large code snippets, consider using models with larger context windows like Claude-3 Opus (200K tokens) or GPT-4 Turbo (128K tokens). " + "For better results try to keep the context size as low as possible. Try using new chats instead of reusing the same chat.", + }, +} diff --git a/src/airflow_wingman/plugin.py b/src/airflow_wingman/plugin.py new file mode 100644 index 0000000..b1e37d6 --- /dev/null +++ b/src/airflow_wingman/plugin.py @@ -0,0 +1,32 @@ +"""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/src/airflow_wingman/prompt_engineering.py b/src/airflow_wingman/prompt_engineering.py new file mode 100644 index 0000000..31b5d8e --- /dev/null +++ b/src/airflow_wingman/prompt_engineering.py @@ -0,0 +1,34 @@ +""" +Prompt engineering for the Airflow Wingman plugin. +Contains prompts and instructions for the AI assistant. +""" + +INSTRUCTIONS = { + "default": """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: + +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/src/airflow_wingman/templates/wingman_chat.html b/src/airflow_wingman/templates/wingman_chat.html new file mode 100644 index 0000000..72657f0 --- /dev/null +++ b/src/airflow_wingman/templates/wingman_chat.html @@ -0,0 +1,384 @@ +{% 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/src/airflow_wingman/views.py b/src/airflow_wingman/views.py new file mode 100644 index 0000000..34158d7 --- /dev/null +++ b/src/airflow_wingman/views.py @@ -0,0 +1,84 @@ +"""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)