86 lines
2.5 KiB
Python
86 lines
2.5 KiB
Python
from airflow.plugins_manager import AirflowPlugin
|
|
from flask_appbuilder import BaseView as AppBuilderBaseView, expose
|
|
from flask import Blueprint
|
|
|
|
|
|
bp = Blueprint(
|
|
"wingman",
|
|
__name__,
|
|
template_folder="templates",
|
|
static_folder="static",
|
|
static_url_path="/static/wingman",
|
|
)
|
|
|
|
|
|
class WingmanView(AppBuilderBaseView):
|
|
route_base = "/wingman"
|
|
default_view = "chat"
|
|
|
|
AVAILABLE_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",
|
|
},
|
|
],
|
|
},
|
|
}
|
|
|
|
@expose("/")
|
|
def chat(self):
|
|
"""
|
|
Chat interface for Airflow Wingman.
|
|
"""
|
|
return self.render_template(
|
|
"wingman_chat.html", title="Airflow Wingman", models=self.AVAILABLE_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]
|