Airflow Wingman separate project
This commit is contained in:
0
airflow-wingman/__init__.py
Normal file
0
airflow-wingman/__init__.py
Normal file
42
airflow-wingman/llms_models.py
Normal file
42
airflow-wingman/llms_models.py
Normal file
@@ -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",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
}
|
||||||
44
airflow-wingman/plugin.py
Normal file
44
airflow-wingman/plugin.py
Normal file
@@ -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]
|
||||||
62
airflow-wingman/pyproject.toml
Normal file
62
airflow-wingman/pyproject.toml
Normal file
@@ -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
|
||||||
193
airflow-wingman/templates/wingman_chat.html
Normal file
193
airflow-wingman/templates/wingman_chat.html
Normal file
@@ -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 %}
|
||||||
Reference in New Issue
Block a user