@@ -180,7 +177,7 @@ document.addEventListener('DOMContentLoaded', function() {
const modelName = this.getAttribute('data-model-name');
console.log('Selected provider:', provider);
console.log('Model name:', modelName);
-
+
if (provider === 'openrouter') {
console.log('Enabling model name input');
modelNameInput.disabled = false;
@@ -201,7 +198,7 @@ document.addEventListener('DOMContentLoaded', function() {
const modelName = defaultSelected.getAttribute('data-model-name');
console.log('Initial provider:', provider);
console.log('Initial model name:', modelName);
-
+
if (provider === 'openrouter') {
console.log('Initially enabling model name input');
modelNameInput.disabled = false;
@@ -218,20 +215,105 @@ document.addEventListener('DOMContentLoaded', function() {
const sendButton = document.getElementById('send-button');
const chatMessages = document.getElementById('chat-messages');
+ let currentMessageDiv = null;
+
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;
+ return messageDiv;
}
- function sendMessage() {
+ async function sendMessage() {
const message = messageInput.value.trim();
- if (message) {
- addMessage(message, true);
- messageInput.value = '';
- // TODO: Add API call to send message and get response
+ if (!message) return;
+
+ // Get selected model
+ const selectedModel = document.querySelector('input[name="model"]:checked');
+ if (!selectedModel) {
+ alert('Please select a model');
+ return;
+ }
+
+ const provider = selectedModel.getAttribute('data-provider');
+ const modelId = selectedModel.value.split(':')[1];
+ const modelName = provider === 'openrouter' ? modelNameInput.value : modelId;
+
+ // Clear input and add user message
+ messageInput.value = '';
+ addMessage(message, true);
+
+ try {
+ // Create messages array with system message
+ const messages = [
+ {
+ role: 'system',
+ content: 'You are a helpful AI assistant integrated into Apache Airflow.'
+ },
+ {
+ role: 'user',
+ content: message
+ }
+ ];
+
+ // Create assistant message div
+ currentMessageDiv = addMessage('', false);
+
+ // Get API key
+ const apiKey = document.getElementById('api-key').value.trim();
+ if (!apiKey) {
+ alert('Please enter an API key');
+ return;
+ }
+
+ // Send request
+ const response = await fetch('/wingman/chat', {
+ method: 'POST',
+ headers: {
+ 'Content-Type': 'application/json',
+ },
+ body: JSON.stringify({
+ provider: provider,
+ model: modelName,
+ messages: messages,
+ api_key: apiKey,
+ stream: true,
+ temperature: 0.7
+ })
+ });
+
+ if (!response.ok) {
+ const error = await response.json();
+ throw new Error(error.error || 'Failed to get response');
+ }
+
+ // Handle streaming response
+ const reader = response.body.getReader();
+ const decoder = new TextDecoder();
+
+ while (true) {
+ const { value, done } = await reader.read();
+ if (done) break;
+
+ const chunk = decoder.decode(value);
+ const lines = chunk.split('\n');
+
+ for (const line of lines) {
+ if (line.startsWith('data: ')) {
+ const content = line.slice(6);
+ if (content) {
+ currentMessageDiv.textContent += content;
+ chatMessages.scrollTop = chatMessages.scrollHeight;
+ }
+ }
+ }
+ }
+ } catch (error) {
+ console.error('Error:', error);
+ currentMessageDiv.textContent = `Error: ${error.message}`;
+ currentMessageDiv.style.color = 'red';
}
}
diff --git a/airflow-wingman/src/airflow_wingman/views.py b/airflow-wingman/src/airflow_wingman/views.py
new file mode 100644
index 0000000..0dc6934
--- /dev/null
+++ b/airflow-wingman/src/airflow_wingman/views.py
@@ -0,0 +1,76 @@
+"""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
+
+
+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)
+
+ @expose("/chat", methods=["POST"])
+ async 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 await 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)}")
+
+ return {
+ "provider": data["provider"],
+ "model": data["model"],
+ "messages": data["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."""
+
+ async def generate():
+ async for chunk in await 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"
+
+ return Response(stream_with_context(generate()), mimetype="text/event-stream")
+
+ async def _handle_regular_response(self, client: LLMClient, data: dict) -> Response:
+ """Handle regular response."""
+ response = await 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)