From 8f8dfeb3b1fe6014df8efba5ee8a2deaa2842aee Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Mon, 24 Feb 2025 10:18:15 +0000
Subject: [PATCH] Airflow Wingman separate project

---
 airflow-wingman/__init__.py                 |   0
 airflow-wingman/llms_models.py              |  42 +++++
 airflow-wingman/plugin.py                   |  44 +++++
 airflow-wingman/pyproject.toml              |  62 +++++++
 airflow-wingman/templates/wingman_chat.html | 193 ++++++++++++++++++++
 5 files changed, 341 insertions(+)
 create mode 100644 airflow-wingman/__init__.py
 create mode 100644 airflow-wingman/llms_models.py
 create mode 100644 airflow-wingman/plugin.py
 create mode 100644 airflow-wingman/pyproject.toml
 create mode 100644 airflow-wingman/templates/wingman_chat.html

diff --git a/airflow-wingman/__init__.py b/airflow-wingman/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/airflow-wingman/llms_models.py b/airflow-wingman/llms_models.py
new file mode 100644
index 0000000..fd328ab
--- /dev/null
+++ b/airflow-wingman/llms_models.py
@@ -0,0 +1,42 @@
+MODELS = {
+        "anthropic": {
+            "name": "Anthropic",
+            "endpoint": "https://api.anthropic.com/v1/messages",
+            "models": [
+                {
+                    "id": "claude-3.5-sonnet",
+                    "name": "Claude 3.5 Sonnet",
+                    "default": True,
+                    "context_window": 200000,
+                    "description": "Input $3/M tokens, Output $15/M tokens",
+                },
+                {
+                    "id": "claude-3.5-haiku",
+                    "name": "Claude 3.5 Haiku",
+                    "default": False,
+                    "context_window": 200000,
+                    "description": "Input $0.80/M tokens, Output $4/M tokens",
+                },
+            ],
+        },
+        "openrouter": {
+            "name": "OpenRouter",
+            "endpoint": "https://openrouter.ai/api/v1/chat/completions",
+            "models": [
+                {
+                    "id": "anthropic/claude-3.5-sonnet",
+                    "name": "Claude 3.5 Sonnet",
+                    "default": False,
+                    "context_window": 200000,
+                    "description": "Input $3/M tokens, Output $15/M tokens",
+                },
+                {
+                    "id": "anthropic/claude-3.5-haiku",
+                    "name": "Claude 3.5 Haiku",
+                    "default": False,
+                    "context_window": 200000,
+                    "description": "Input $0.80/M tokens, Output $4/M tokens",
+                },
+            ],
+        },
+    }
\ No newline at end of file
diff --git a/airflow-wingman/plugin.py b/airflow-wingman/plugin.py
new file mode 100644
index 0000000..78b01fd
--- /dev/null
+++ b/airflow-wingman/plugin.py
@@ -0,0 +1,44 @@
+from airflow.plugins_manager import AirflowPlugin
+from flask_appbuilder import BaseView as AppBuilderBaseView, expose
+from flask import Blueprint
+
+from airflow_wingman.llms_models import MODELS
+
+
+bp = Blueprint(
+    "wingman",
+    __name__,
+    template_folder="templates",
+    static_folder="static",
+    static_url_path="/static/wingman",
+)
+
+
+class WingmanView(AppBuilderBaseView):
+    route_base = "/wingman"
+    default_view = "chat"
+
+    @expose("/")
+    def chat(self):
+        """
+        Chat interface for Airflow Wingman.
+        """
+        return self.render_template(
+            "wingman_chat.html", title="Airflow Wingman", models=MODELS
+        )
+
+
+# Create AppBuilder View
+v_appbuilder_view = WingmanView()
+v_appbuilder_package = {
+    "name": "Wingman",
+    "category": "AI",
+    "view": v_appbuilder_view,
+}
+
+
+# Create Plugin
+class WingmanPlugin(AirflowPlugin):
+    name = "wingman"
+    flask_blueprints = [bp]
+    appbuilder_views = [v_appbuilder_package]
diff --git a/airflow-wingman/pyproject.toml b/airflow-wingman/pyproject.toml
new file mode 100644
index 0000000..b94fbaa
--- /dev/null
+++ b/airflow-wingman/pyproject.toml
@@ -0,0 +1,62 @@
+
+[project]
+name = "airflow-wingman"
+version = "0.1.0"
+description = "Airflow plugin to enable LLMs chat"
+requires-python = ">=3.11"
+authors = [
+    {name = "Abhishek Bhakat", email = "abhishek.bhakat@hotmail.com"}
+]
+dependencies = [
+    "apache-airflow>=2.10.0",
+    "airflow-mcp-server>=0.2.0"
+]
+
+[project.urls]
+repository = "https://github.com/abhishekbhakat/airflow-mcp-server"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["airflow-wingman"]
+
+[tool.ruff]
+line-length = 200
+indent-width = 4
+fix = true
+preview = true
+
+lint.select = [
+    "E",  # pycodestyle errors
+    "F",  # pyflakes
+    "I",  # isort
+    "W",  # pycodestyle warnings
+    "C90",  # Complexity
+    "C",  # flake8-comprehensions
+    "ISC",  # flake8-implicit-str-concat
+    "T10",  # flake8-debugger
+    "A",  # flake8-builtins
+    "UP",  # pyupgrade
+]
+
+lint.ignore = [
+    "C416",  # Unnecessary list comprehension - rewrite as a generator expression
+    "C408",  # Unnecessary `dict` call - rewrite as a literal
+    "ISC001"  # Single line implicit string concatenation
+]
+
+lint.fixable = ["ALL"]
+lint.unfixable = []
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+
+[tool.ruff.lint.mccabe]
+max-complexity = 12
diff --git a/airflow-wingman/templates/wingman_chat.html b/airflow-wingman/templates/wingman_chat.html
new file mode 100644
index 0000000..dafa921
--- /dev/null
+++ b/airflow-wingman/templates/wingman_chat.html
@@ -0,0 +1,193 @@
+{% extends "appbuilder/base.html" %}
+
+{% block content %}
+<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>
+        </div>
+    </div>
+
+    <div class="row">
+        <!-- Sidebar -->
+        <div class="col-md-3">
+            <div class="panel panel-default">
+                <div class="panel-heading">
+                    <h3 class="panel-title">Model 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" data-bs-toggle="tooltip" data-bs-placement="right" title="{{ model.description }}">
+                                <input type="radio"
+                                       name="model"
+                                       value="{{ provider_id }}:{{ model.id }}"
+                                       {% if model.default %}checked{% endif %}
+                                       data-endpoint="{{ provider.endpoint }}"
+                                       data-context-window="{{ model.context_window }}">
+                                {{ model.name }}
+                            </label>
+                        </div>
+                        {% endfor %}
+                    </div>
+                    {% endfor %}
+                </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>
+
+            <!-- API Key Input -->
+            <div class="panel panel-default mt-3">
+                <div class="panel-heading">
+                    <h3 class="panel-title">API Key</h3>
+                </div>
+                <div class="panel-body">
+                    <div class="form-group">
+                        <input type="password"
+                               class="form-control"
+                               id="api-key"
+                               placeholder="Enter your API key"
+                               autocomplete="off">
+                        <small class="text-muted">
+                            Your API key will be used for the selected provider
+                        </small>
+                    </div>
+                </div>
+            </div>
+        </div>
+
+        <!-- Main Chat Window -->
+        <div class="col-md-9">
+            <div class="panel panel-default" style="height: calc(100vh - 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="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>
+
+<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() {
+    // Initialize tooltips
+    const tooltipTriggerList = [].slice.call(document.querySelectorAll('[data-bs-toggle="tooltip"]'));
+    const tooltipList = tooltipTriggerList.map(function (tooltipTriggerEl) {
+        return new bootstrap.Tooltip(tooltipTriggerEl, {
+            trigger: 'hover',
+            html: true
+        });
+    });
+
+    const messageInput = document.getElementById('message-input');
+    const sendButton = document.getElementById('send-button');
+    const chatMessages = document.getElementById('chat-messages');
+
+    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;
+    }
+
+    function sendMessage() {
+        const message = messageInput.value.trim();
+        if (message) {
+            addMessage(message, true);
+            messageInput.value = '';
+            // TODO: Add API call to send message and get response
+        }
+    }
+
+    sendButton.addEventListener('click', sendMessage);
+    messageInput.addEventListener('keypress', function(e) {
+        if (e.key === 'Enter') {
+            sendMessage();
+        }
+    });
+});
+</script>
+{% endblock %}