airflow-chat 0.1.0a1__tar.gz

airflow_chat-0.1.0a1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 hipposys-ltd
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

airflow_chat-0.1.0a1/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.3
+Name: airflow-chat
+Version: 0.1.0a1
+Summary: An Apache Airflow plugin that enables AI-powered chat interactions with your Airflow instance through MCP integration and an intuitive UI.
+License: MIT
+Author: Hipposys
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: SQLAlchemy (>=1.4.0)
+Requires-Dist: apache-airflow (>=2.4.0,<3.0.0)
+Project-URL: homepage, https://github.com/hipposys-ltd/airflow-schedule-insights
+Description-Content-Type: text/markdown
+
+# Airflow Chat Plugin

airflow_chat-0.1.0a1/airflow_chat/plugins/README.md

@@ -0,0 +1 @@
+# Airflow Chat Plugin

airflow_chat-0.1.0a1/airflow_chat/plugins/airflow_chat.py

@@ -0,0 +1,482 @@
+from airflow import settings
+from airflow.plugins_manager import AirflowPlugin
+from flask import Blueprint, request, jsonify, Response, stream_template
+from flask_login import current_user
+from flask_appbuilder import expose, BaseView as AppBuilderBaseView
+from functools import wraps
+import json
+import time
+import requests
+import os
+from datetime import datetime, timedelta
+import uuid
+from sqlalchemy import Column, String, DateTime, Text, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.exc import SQLAlchemyError
+
+# Blueprint for the chat plugin
+bp = Blueprint(
+    "airflow_chat",
+    __name__,
+    template_folder="templates",
+    static_folder="static",
+    static_url_path="/static/airflow_chat_plugin",
+)
+
+# Database model for storing chat sessions
+Base = declarative_base()
+
+
+class ChatSession(Base):
+    __tablename__ = 'airflow_chat_sessions'
+    
+    conversation_id = Column(String(255), primary_key=True)
+    cookies_data = Column(Text)  # JSON string of cookies
+    created_at = Column(DateTime, default=datetime.utcnow)
+    last_accessed = Column(DateTime, default=datetime.utcnow)
+    
+    def to_dict(self):
+        return {
+            'conversation_id': self.conversation_id,
+            'cookies_data': self.cookies_data,
+            'created_at': self.created_at.isoformat() if self.created_at else None,
+            'last_accessed': self.last_accessed.isoformat() if self.last_accessed else None
+        }
+
+
+class ChatSessionManager:
+    def __init__(self):
+        # Use Airflow's metadata database
+        self.engine = settings.engine
+        self.Session = sessionmaker(bind=self.engine)
+        
+        # Create the table if it doesn't exist
+        try:
+            Base.metadata.create_all(self.engine)
+        except Exception as e:
+            print(f"Error creating chat sessions table: {e}")
+    
+    def cleanup_old_sessions(self):
+        """Remove sessions older than 24 hours"""
+        try:
+            session = self.Session()
+            cutoff_time = datetime.utcnow() - timedelta(minutes=20)
+            
+            deleted_count = session.query(ChatSession).filter(
+                ChatSession.created_at < cutoff_time
+            ).delete()
+            
+            session.commit()
+            session.close()
+            
+            if deleted_count > 0:
+                print(f"Cleaned up {deleted_count} old chat sessions")
+                
+        except SQLAlchemyError as e:
+            print(f"Error cleaning up old sessions: {e}")
+            if session:
+                session.rollback()
+                session.close()
+    
+    def store_session(self, conversation_id, cookies_dict):
+        """Store or update a chat session"""
+        try:
+            session = self.Session()
+            
+            # Convert cookies to JSON string
+            cookies_json = json.dumps(dict(cookies_dict)) if cookies_dict else "{}"
+            
+            # Check if session exists
+            existing_session = session.query(ChatSession).filter(
+                ChatSession.conversation_id == conversation_id
+            ).first()
+            
+            if existing_session:
+                # Update existing session
+                existing_session.cookies_data = cookies_json
+                existing_session.last_accessed = datetime.utcnow()
+            else:
+                # Create new session
+                new_session = ChatSession(
+                    conversation_id=conversation_id,
+                    cookies_data=cookies_json,
+                    created_at=datetime.utcnow(),
+                    last_accessed=datetime.utcnow()
+                )
+                session.add(new_session)
+            
+            session.commit()
+            session.close()
+            return True
+            
+        except SQLAlchemyError as e:
+            print(f"Error storing session: {e}")
+            if session:
+                session.rollback()
+                session.close()
+            return False
+    
+    def get_session(self, conversation_id):
+        """Retrieve a chat session"""
+        try:
+            session = self.Session()
+            
+            chat_session = session.query(ChatSession).filter(
+                ChatSession.conversation_id == conversation_id
+            ).first()
+            
+            if chat_session:
+                # Update last accessed time
+                chat_session.last_accessed = datetime.utcnow()
+                session.commit()
+                
+                # Parse cookies from JSON
+                cookies_dict = json.loads(chat_session.cookies_data) if chat_session.cookies_data else {}
+                
+                session.close()
+                return cookies_dict
+            else:
+                session.close()
+                return None
+                
+        except SQLAlchemyError as e:
+            print(f"Error retrieving session: {e}")
+            if session:
+                session.close()
+            return None
+    
+    def delete_session(self, conversation_id):
+        """Delete a specific chat session"""
+        try:
+            session = self.Session()
+            
+            deleted_count = session.query(ChatSession).filter(
+                ChatSession.conversation_id == conversation_id
+            ).delete()
+            
+            session.commit()
+            session.close()
+            
+            return deleted_count > 0
+            
+        except SQLAlchemyError as e:
+            print(f"Error deleting session: {e}")
+            if session:
+                session.rollback()
+                session.close()
+            return False
+    
+    def get_active_sessions_count(self):
+        """Get count of active sessions (within 24 hours)"""
+        try:
+            session = self.Session()
+            cutoff_time = datetime.utcnow() - timedelta(hours=24)
+            
+            count = session.query(ChatSession).filter(
+                ChatSession.created_at >= cutoff_time
+            ).count()
+            
+            session.close()
+            return count
+            
+        except SQLAlchemyError as e:
+            print(f"Error getting session count: {e}")
+            if session:
+                session.close()
+            return 0
+
+def admin_only(f):
+    """Decorator to restrict access to Admin role."""
+    @wraps(f)
+    def decorated_function(*args, **kwargs):
+        # Check if user is authenticated
+        if not current_user or not hasattr(current_user, 'roles') or not current_user.roles:
+            return jsonify({"error": "Authentication required"}), 401
+        
+        users_roles = [role.name for role in current_user.roles]
+        approved_roles = ["Admin"]  # Adjust roles as needed
+        if not any(role in users_roles for role in approved_roles):
+            return jsonify({"error": "Access denied"}), 403
+        return f(*args, **kwargs)
+    return decorated_function
+
+def failure_tolerant(f):
+    """Decorator for error handling."""
+    @wraps(f)
+    def decorated_function(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except Exception as e:
+            print(f"CHAT PLUGIN FAILURE: {e}")
+            return jsonify({"error": "Something went wrong"}), 500
+    return decorated_function
+
+class LLMChatAgent:
+    def __init__(self):
+        self.backend_url = os.environ.get('BACKEND_URL', "http://fastapi:8080")
+        self.access_token = os.environ.get('FAST_API_ACCESS_SECRET_TOKEN', 'ThisIsATempAccessTokenForLocalEnvs.ReplaceInProd')
+        self.session_manager = ChatSessionManager()
+        
+        # Perform cleanup on initialization
+        self.session_manager.cleanup_old_sessions()
+    
+    def get_headers(self):
+        return {
+            "x-access-token": self.access_token
+        }
+    
+    def initialize_chat_session(self, conversation_id):
+        """Initialize a new chat session"""
+        try:
+            response = requests.post(
+                f"{self.backend_url}/chat/new",
+                headers=self.get_headers(),
+                timeout=10
+            )
+            response.raise_for_status()
+            
+            # Store session cookies in database
+            success = self.session_manager.store_session(conversation_id, response.cookies)
+            
+            if success:
+                print(f"Initialized and stored chat session: {conversation_id}")
+                return True, response.cookies
+            else:
+                print(f"Failed to store session in database: {conversation_id}")
+                return False, None
+            
+        except requests.RequestException as e:
+            print(f"Failed to initialize chat session: {e}")
+            return False, None
+        except Exception as e:
+            print(f"Unexpected error initializing session: {e}")
+            return False, None
+    
+    def stream_chat_response(self, message, conversation_id=None):
+        """Stream response from FastAPI backend"""
+        try:
+            # Clean up old sessions periodically (every request)
+            self.session_manager.cleanup_old_sessions()
+            
+            # Get session cookies from database
+            cookies_dict = self.session_manager.get_session(conversation_id)
+            
+            if cookies_dict is None:
+                # Initialize session if not found
+                success, cookies = self.initialize_chat_session(conversation_id)
+                if not success:
+                    yield {
+                        "content": "Failed to initialize chat session. Please try again.",
+                        "conversation_id": conversation_id,
+                        "timestamp": datetime.now().isoformat(),
+                        "error": True
+                    }
+                    return
+                cookies_dict = dict(cookies)
+            
+            # Convert dict back to requests.cookies.RequestsCookieJar if needed
+            # For requests library, we can pass cookies as dict
+            
+            # Stream chat response
+            with requests.post(
+                f"{self.backend_url}/chat/ask",
+                stream=True,
+                headers=self.get_headers(),
+                json={'message': message},
+                cookies=cookies_dict,
+                timeout=30
+            ) as response:
+                
+                response.raise_for_status()
+                
+                # Stream response chunks
+                for chunk in response.iter_content(chunk_size=1024):
+                    if chunk:
+                        try:
+                            # Decode the chunk
+                            content = str(chunk, encoding="utf-8")
+                            
+                            # Yield the content as structured response
+                            yield {
+                                "content": content,
+                                "conversation_id": conversation_id,
+                                "timestamp": datetime.now().isoformat()
+                            }
+                            
+                        except UnicodeDecodeError:
+                            # Skip chunks that can't be decoded
+                            continue
+                        except Exception as e:
+                            print(f"Error processing chunk: {e}")
+                            continue
+                
+        except requests.RequestException as e:
+            # Fallback for connection errors
+            yield {
+                "content": f"Connection error: {str(e)}. Please check if the FastAPI service is running.",
+                "conversation_id": conversation_id,
+                "timestamp": datetime.now().isoformat(),
+                "error": True
+            }
+        except Exception as e:
+            yield {
+                "content": f"An unexpected error occurred: {str(e)}",
+                "conversation_id": conversation_id,
+                "timestamp": datetime.now().isoformat(),
+                "error": True
+            }
+    
+    def clear_session(self, conversation_id):
+        """Clear session data for a conversation"""
+        print(f'Clearing session: {conversation_id}')
+        return self.session_manager.delete_session(conversation_id)
+    
+    def get_session_stats(self):
+        """Get statistics about active sessions"""
+        return {
+            'active_sessions': self.session_manager.get_active_sessions_count()
+        }
+
+class AirflowChatView(AppBuilderBaseView):
+    default_view = "chat_interface"
+    route_base = "/airflow_chat"
+    
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.llm_agent = LLMChatAgent()
+    
+    @expose("/")
+    @admin_only
+    @failure_tolerant
+    def chat_interface(self):
+        """Main chat interface"""
+        return self.render_template("chat_interface.html")
+    
+    @expose("/api/chat", methods=["POST"])
+    @admin_only
+    @failure_tolerant
+    def chat_api(self):
+        """API endpoint for chat messages"""
+        data = request.get_json()
+        message = data.get("message", "").strip()
+        conversation_id = data.get("conversation_id")
+        print(f'conversation_id: {conversation_id}')
+        
+        if not message:
+            return jsonify({"error": "Message is required"}), 400
+        
+        # Return streaming response
+        def generate():
+            try:
+                for chunk in self.llm_agent.stream_chat_response(
+                    message, 
+                    conversation_id
+                ):
+                    yield f"data: {json.dumps(chunk)}\n\n"
+                yield "data: [DONE]\n\n"
+            except Exception as e:
+                error_chunk = {
+                    "content": f"Error: {str(e)}",
+                    "conversation_id": conversation_id or str(uuid.uuid4()),
+                    "timestamp": datetime.now().isoformat(),
+                    "error": True
+                }
+                yield f"data: {json.dumps(error_chunk)}\n\n"
+                yield "data: [DONE]\n\n"
+        
+        return Response(
+            generate(),
+            mimetype="text/plain",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "Access-Control-Allow-Origin": "*"
+            }
+        )
+
+    @expose("/api/new_chat", methods=["POST"])
+    @admin_only
+    @failure_tolerant
+    def new_chat(self):
+        """Initialize a new chat session"""
+        conversation_id = str(uuid.uuid4())
+        
+        try:
+            success, cookies = self.llm_agent.initialize_chat_session(conversation_id)
+            
+            if success:
+                return jsonify({
+                    "conversation_id": conversation_id,
+                    "status": "initialized",
+                    "message": "New chat session created"
+                })
+            else:
+                return jsonify({
+                    "error": "Failed to initialize chat session"
+                }), 500
+                
+        except Exception as e:
+            return jsonify({
+                "error": f"Error creating new chat: {str(e)}"
+            }), 500
+    
+    @expose("/api/clear_session", methods=["POST"])
+    @admin_only
+    @failure_tolerant
+    def clear_session(self):
+        """Clear a specific chat session"""
+        data = request.get_json()
+        conversation_id = data.get("conversation_id")
+        
+        if not conversation_id:
+            return jsonify({"error": "Conversation ID is required"}), 400
+        
+        try:
+            success = self.llm_agent.clear_session(conversation_id)
+            
+            if success:
+                return jsonify({
+                    "message": "Session cleared successfully",
+                    "conversation_id": conversation_id
+                })
+            else:
+                return jsonify({
+                    "error": "Session not found or could not be cleared"
+                }), 404
+                
+        except Exception as e:
+            return jsonify({
+                "error": f"Error clearing session: {str(e)}"
+            }), 500
+    
+    @expose("/api/session_stats", methods=["GET"])
+    @admin_only
+    @failure_tolerant
+    def session_stats(self):
+        """Get session statistics"""
+        try:
+            stats = self.llm_agent.get_session_stats()
+            return jsonify(stats)
+        except Exception as e:
+            return jsonify({
+                "error": f"Error getting session stats: {str(e)}"
+            }), 500
+
+# Create the view instance
+chat_view = AirflowChatView()
+
+# Package for Airflow
+chat_package = {
+    "name": "AI Chat Assistant",
+    "category": "Tools",
+    "view": chat_view,
+}
+
+class AirflowChatPlugin(AirflowPlugin):
+    name = "airflow_chat_plugin"
+    hooks = []
+    macros = []
+    flask_blueprints = [bp]
+    appbuilder_views = [chat_package]
+    appbuilder_menu_items = []

airflow_chat-0.1.0a1/airflow_chat/plugins/templates/chat_interface.html

@@ -0,0 +1,833 @@
+{% extends base_template %}
+{% block content %}
+<div class="chat-container">
+  <div class="chat-header">
+    <h1><i class="fa fa-comments"></i> AI Assistant</h1>
+    <div class="chat-controls">
+      <button id="newChat" class="btn btn-sm btn-primary">
+        <i class="fa fa-plus"></i> New Chat
+      </button>
+    </div>
+  </div>
+
+  <div class="chat-messages" id="chatMessages">
+    <div class="welcome-message">
+      <div class="message assistant-message">
+        <div class="message-avatar">
+          <i class="fa fa-robot"></i>
+        </div>
+        <div class="message-content">
+          <div class="message-text">
+            Hello! I'm your Airflow AI assistant. I can help you with:
+            <ul>
+              <li>DAG management and troubleshooting</li>
+              <li>Task configuration and dependencies</li>
+              <li>Airflow best practices</li>
+              <li>Connection and variable management</li>
+              <li>General workflow questions</li>
+            </ul>
+            What would you like to know?
+          </div>
+          <div class="message-time">Just now</div>
+        </div>
+      </div>
+    </div>
+  </div>
+
+  <div class="chat-input-container">
+    <div class="chat-input-wrapper">
+      <textarea 
+        id="chatInput" 
+        placeholder="Ask me anything about Airflow..." 
+        rows="1"
+        maxlength="2000"></textarea>
+      <button id="sendButton" class="btn btn-primary" disabled>
+        <i class="fa fa-paper-plane"></i>
+      </button>
+    </div>
+    <div class="input-footer">
+      <span class="char-count">0/2000</span>
+      <span class="status" id="connectionStatus">Ready</span>
+    </div>
+  </div>
+</div>
+
+<style>
+.chat-container {
+  max-width: 1000px;
+  margin: 0 auto;
+  height: calc(100vh - 120px);
+  display: flex;
+  flex-direction: column;
+  background: #f8f9fa;
+  border-radius: 10px;
+  overflow: hidden;
+  box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+}
+
+.chat-header {
+  background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+  color: white;
+  padding: 20px;
+  display: flex;
+  justify-content: space-between;
+  align-items: center;
+}
+
+.chat-header h1 {
+  margin: 0;
+  font-size: 24px;
+  font-weight: 600;
+}
+
+.chat-controls {
+  display: flex;
+  gap: 10px;
+}
+
+.chat-messages {
+  flex: 1;
+  overflow-y: auto;
+  padding: 20px;
+  background: white;
+}
+
+.message {
+  display: flex;
+  margin-bottom: 20px;
+  max-width: 80%;
+}
+
+.user-message {
+  margin-left: auto;
+  flex-direction: row-reverse;
+}
+
+.assistant-message {
+  margin-right: auto;
+}
+
+.message-avatar {
+  width: 40px;
+  height: 40px;
+  border-radius: 50%;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  margin: 0 10px;
+  font-size: 18px;
+}
+
+.user-message .message-avatar {
+  background: #007bff;
+  color: white;
+}
+
+.assistant-message .message-avatar {
+  background: #6f42c1;
+  color: white;
+}
+
+.message-content {
+  flex: 1;
+}
+
+.message-text {
+  background: #f1f3f4;
+  padding: 12px 16px;
+  border-radius: 18px;
+  word-wrap: break-word;
+  line-height: 1.4;
+}
+
+.user-message .message-text {
+  background: #007bff;
+  color: white;
+}
+
+.assistant-message .message-text {
+  background: #f1f3f4;
+  color: #333;
+}
+
+.message-time {
+  font-size: 11px;
+  color: #666;
+  margin-top: 4px;
+  text-align: right;
+}
+
+.user-message .message-time {
+  text-align: left;
+}
+
+.chat-input-container {
+  background: white;
+  border-top: 1px solid #e0e0e0;
+  padding: 20px;
+}
+
+.chat-input-wrapper {
+  display: flex;
+  align-items: flex-end;
+  gap: 10px;
+  max-width: 100%;
+}
+
+#chatInput {
+  flex: 1;
+  border: 2px solid #e0e0e0;
+  border-radius: 20px;
+  padding: 12px 16px;
+  resize: none;
+  font-family: inherit;
+  font-size: 14px;
+  line-height: 1.4;
+  max-height: 120px;
+  transition: border-color 0.3s;
+}
+
+#chatInput:focus {
+  outline: none;
+  border-color: #007bff;
+}
+
+#sendButton {
+  width: 44px;
+  height: 44px;
+  border-radius: 50%;
+  border: none;
+  background: #007bff;
+  color: white;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  transition: all 0.3s;
+}
+
+#sendButton:hover:not(:disabled) {
+  background: #0056b3;
+  transform: scale(1.05);
+}
+
+#sendButton:disabled {
+  background: #ccc;
+  cursor: not-allowed;
+}
+
+.input-footer {
+  display: flex;
+  justify-content: space-between;
+  align-items: center;
+  margin-top: 8px;
+  font-size: 12px;
+  color: #666;
+}
+
+.char-count {
+  font-size: 11px;
+}
+
+.status {
+  font-weight: 500;
+}
+
+.status.connecting {
+  color: #ffc107;
+}
+
+.status.ready {
+  color: #28a745;
+}
+
+.status.error {
+  color: #dc3545;
+}
+
+.typing-indicator {
+  display: flex;
+  align-items: center;
+  gap: 4px;
+  padding: 8px 16px;
+  background: #f1f3f4;
+  border-radius: 18px;
+  margin: 8px 0;
+}
+
+.typing-dots {
+  display: flex;
+  gap: 2px;
+}
+
+.typing-dot {
+  width: 6px;
+  height: 6px;
+  background: #666;
+  border-radius: 50%;
+  animation: typing 1.4s infinite;
+}
+
+.typing-dot:nth-child(2) {
+  animation-delay: 0.2s;
+}
+
+.typing-dot:nth-child(3) {
+  animation-delay: 0.4s;
+}
+
+@keyframes typing {
+  0%, 60%, 100% {
+    transform: translateY(0);
+    opacity: 0.4;
+  }
+  30% {
+    transform: translateY(-10px);
+    opacity: 1;
+  }
+}
+
+.welcome-message ul {
+  margin: 10px 0;
+  padding-left: 20px;
+}
+
+.welcome-message li {
+  margin-bottom: 5px;
+}
+
+/* Tool execution styling */
+.tool-execution {
+  margin: 10px 0;
+  border: 1px solid #e0e0e0;
+  border-radius: 8px;
+  background: #f8f9fa;
+}
+
+.tool-execution[open] {
+  border-color: #007bff;
+  background: #f0f8ff;
+}
+
+.tool-summary {
+  padding: 12px 16px;
+  cursor: pointer;
+  font-weight: 600;
+  color: #007bff;
+  background: #f8f9fa;
+  border-radius: 8px;
+  user-select: none;
+  transition: background-color 0.2s;
+}
+
+.tool-summary:hover {
+  background: #e9ecef;
+}
+
+.tool-execution[open] .tool-summary {
+  border-bottom: 1px solid #e0e0e0;
+  border-radius: 8px 8px 0 0;
+  background: #e3f2fd;
+}
+
+.tool-execution > div:not(.tool-summary) {
+  padding: 16px;
+}
+
+.tool-execution pre {
+  background: #2d3748;
+  color: #e2e8f0;
+  padding: 12px;
+  border-radius: 6px;
+  overflow-x: auto;
+  margin: 8px 0;
+  font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
+  font-size: 13px;
+  line-height: 1.4;
+}
+
+.tool-execution code {
+  background: #f1f3f4;
+  padding: 2px 6px;
+  border-radius: 4px;
+  font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
+  font-size: 13px;
+}
+
+.tool-execution h1,
+.tool-execution h2,
+.tool-execution h3 {
+  margin: 16px 0 8px 0;
+  color: #333;
+}
+
+.tool-execution ul {
+  margin: 8px 0;
+  padding-left: 20px;
+}
+
+.tool-execution li {
+  margin: 4px 0;
+  list-style-type: disc;
+}
+
+/* Formatting styles */
+.message-text h1,
+.message-text h2,
+.message-text h3 {
+  margin: 16px 0 8px 0;
+  color: #333;
+  line-height: 1.3;
+}
+
+.message-text h1 { font-size: 1.5em; font-weight: 700; }
+.message-text h2 { font-size: 1.3em; font-weight: 600; }
+.message-text h3 { font-size: 1.1em; font-weight: 600; }
+
+.message-text strong {
+  font-weight: 600;
+  color: #333;
+}
+
+.message-text code {
+  background: #f1f3f4;
+  padding: 2px 6px;
+  border-radius: 4px;
+  font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
+  font-size: 13px;
+  color: #d63384;
+}
+
+.message-text pre {
+  background: #2d3748;
+  color: #e2e8f0;
+  padding: 12px;
+  border-radius: 6px;
+  overflow-x: auto;
+  margin: 12px 0;
+  font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace;
+  font-size: 13px;
+  line-height: 1.4;
+}
+
+.message-text pre code {
+  background: none;
+  padding: 0;
+  color: inherit;
+  border-radius: 0;
+}
+
+.message-text ul,
+.message-text ol {
+  margin: 8px 0;
+  padding-left: 20px;
+}
+
+.message-text li {
+  margin: 4px 0;
+  line-height: 1.4;
+}
+
+.message-text ul li {
+  list-style-type: disc;
+}
+
+.message-text ol li {
+  list-style-type: decimal;
+}
+
+@media (max-width: 768px) {
+  .chat-container {
+    height: calc(100vh - 80px);
+    margin: 0;
+    border-radius: 0;
+  }
+  
+  .message {
+    max-width: 90%;
+  }
+  
+  .chat-header {
+    padding: 15px;
+  }
+  
+  .chat-header h1 {
+    font-size: 20px;
+  }
+}
+</style>
+
+<script>
+class AirflowChatClient {
+  constructor() {
+    this.conversationId = null;
+    this.isConnected = true;
+    this.currentStream = null;
+    this.isStreaming = false; // Track streaming state
+    
+    this.initializeElements();
+    this.attachEventListeners();
+    this.autoResize();
+  }
+  
+  initializeElements() {
+    this.chatInput = document.getElementById('chatInput');
+    this.sendButton = document.getElementById('sendButton');
+    this.chatMessages = document.getElementById('chatMessages');
+    this.newChatBtn = document.getElementById('newChat');
+    this.charCount = document.querySelector('.char-count');
+    this.status = document.getElementById('connectionStatus');
+  }
+  
+  attachEventListeners() {
+    // Send message events
+    this.sendButton.addEventListener('click', () => this.sendMessage());
+    this.chatInput.addEventListener('keydown', (e) => {
+      if (e.key === 'Enter' && !e.shiftKey) {
+        e.preventDefault();
+        this.sendMessage();
+      }
+    });
+    
+    // Input events
+    this.chatInput.addEventListener('input', () => {
+      this.updateCharCount();
+      this.updateSendButton();
+      this.autoResize();
+    });
+    
+    this.newChatBtn.addEventListener('click', () => this.newChat());
+  }
+  
+  updateCharCount() {
+    const length = this.chatInput.value.length;
+    this.charCount.textContent = `${length}/2000`;
+    
+    if (length > 1800) {
+      this.charCount.style.color = '#dc3545';
+    } else if (length > 1500) {
+      this.charCount.style.color = '#ffc107';
+    } else {
+      this.charCount.style.color = '#666';
+    }
+  }
+  
+  updateSendButton() {
+    const hasText = this.chatInput.value.trim().length > 0;
+    // Disable button if no text, not connected, or currently streaming
+    this.sendButton.disabled = !hasText || !this.isConnected || this.isStreaming;
+  }
+  
+  autoResize() {
+    this.chatInput.style.height = 'auto';
+    this.chatInput.style.height = Math.min(this.chatInput.scrollHeight, 120) + 'px';
+  }
+  
+  setStatus(text, className = 'ready') {
+    this.status.textContent = text;
+    this.status.className = `status ${className}`;
+  }
+  
+  setStreaming(streaming) {
+    this.isStreaming = streaming;
+    this.updateSendButton();
+  }
+  
+  async sendMessage() {
+    const message = this.chatInput.value.trim();
+    if (!message || !this.isConnected || this.isStreaming) return;
+    
+    // Set streaming state to true
+    this.setStreaming(true);
+    
+    // Initialize conversation if not exists
+    if (!this.conversationId) {
+      this.setStatus('Initializing chat...', 'connecting');
+      try {
+        const initResponse = await fetch("{{ url_for('AirflowChatView.new_chat') }}", {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+            'X-CSRFToken': '{{ csrf_token() }}'
+          }
+        });
+        
+        const initData = await initResponse.json();
+        if (initData.conversation_id) {
+          this.conversationId = initData.conversation_id;
+        } else {
+          throw new Error('Failed to initialize chat session');
+        }
+      } catch (error) {
+        console.error('Failed to initialize chat:', error);
+        this.setStatus('Failed to initialize chat', 'error');
+        this.setStreaming(false); // Re-enable button on error
+        return;
+      }
+    }
+    
+    // Add user message to chat
+    this.addMessage(message, 'user');
+    
+    // Clear input
+    this.chatInput.value = '';
+    this.updateCharCount();
+    this.updateSendButton();
+    this.autoResize();
+    
+    // Show typing indicator
+    const typingElement = this.showTypingIndicator();
+    
+    try {
+      this.setStatus('Sending...', 'connecting');
+      console.log('this.conversationId')
+      console.log(this.conversationId)
+      const response = await fetch("{{ url_for('AirflowChatView.chat_api') }}", {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          'X-CSRFToken': '{{ csrf_token() }}'
+        },
+        body: JSON.stringify({
+          message: message,
+          conversation_id: this.conversationId
+        })
+      });
+      
+      if (!response.ok) {
+        throw new Error(`HTTP ${response.status}: ${response.statusText}`);
+      }
+      
+      // Remove typing indicator
+      this.removeTypingIndicator(typingElement);
+      
+      // Handle streaming response
+      await this.handleStreamingResponse(response);
+      
+    } catch (error) {
+      console.error('Chat error:', error);
+      this.removeTypingIndicator(typingElement);
+      this.addMessage(`Sorry, I encountered an error: ${error.message}`, 'assistant', true);
+      this.setStatus('Error', 'error');
+    } finally {
+      // Always re-enable button when done
+      this.setStreaming(false);
+    }
+  }
+  
+  async handleStreamingResponse(response) {
+    const reader = response.body.getReader();
+    const decoder = new TextDecoder();
+    let assistantMessage = '';
+    let messageElement = null;
+    
+    this.setStatus('Receiving...', 'connecting');
+    
+    try {
+      while (true) {
+        const { done, value } = 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 data = line.slice(6);
+            
+            if (data === '[DONE]') {
+              this.setStatus('Ready', 'ready');
+              return;
+            }
+            
+            try {
+              const parsed = JSON.parse(data);
+              
+            //   if (parsed.conversation_id && !this.conversationId) {
+            //     this.conversationId = parsed.conversation_id;
+            //   }
+              
+              if (parsed.content) {
+                assistantMessage += parsed.content;
+                
+                if (!messageElement) {
+                  messageElement = this.addMessage('', 'assistant');
+                }
+                
+                this.updateMessage(messageElement, assistantMessage);
+              }
+              
+              if (parsed.error) {
+                this.setStatus('Error', 'error');
+              }
+              
+            } catch (e) {
+              console.warn('Failed to parse chunk:', data);
+            }
+          }
+        }
+      }
+    } catch (error) {
+      console.error('Streaming error:', error);
+      this.addMessage('Sorry, the connection was interrupted.', 'assistant', true);
+    } finally {
+      this.setStatus('Ready', 'ready');
+    }
+  }
+  
+  addMessage(text, sender, isError = false) {
+    const messageDiv = document.createElement('div');
+    messageDiv.className = `message ${sender}-message`;
+    
+    const now = new Date();
+    const timeStr = now.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' });
+    
+    // Parse markdown for better formatting
+    const formattedText = this.parseMarkdown(text);
+    
+    messageDiv.innerHTML = `
+      <div class="message-avatar">
+        <i class="fa fa-${sender === 'user' ? 'user' : 'robot'}"></i>
+      </div>
+      <div class="message-content">
+        <div class="message-text${isError ? ' error' : ''}">${formattedText}</div>
+        <div class="message-time">${timeStr}</div>
+      </div>
+    `;
+    
+    this.chatMessages.appendChild(messageDiv);
+    this.scrollToBottom();
+    
+    return messageDiv;
+  }
+  
+  updateMessage(messageElement, text) {
+    const textElement = messageElement.querySelector('.message-text');
+    
+    // Parse markdown-like formatting
+    const formattedText = this.parseMarkdown(text);
+    textElement.innerHTML = formattedText;
+    
+    this.scrollToBottom();
+  }
+  
+  parseMarkdown(text) {
+    // Convert basic markdown to HTML
+    let html = text
+      // Headers
+      .replace(/^### (.*$)/gm, '<h3>$1</h3>')
+      .replace(/^## (.*$)/gm, '<h2>$1</h2>')
+      .replace(/^# (.*$)/gm, '<h1>$1</h1>')
+      
+      // Bold
+      .replace(/\*\*(.*?)\*\*/g, '<strong>$1</strong>')
+      
+      // Code blocks
+      .replace(/```(\w+)?\n([\s\S]*?)```/g, '<pre><code class="language-$1">$2</code></pre>')
+      .replace(/```([\s\S]*?)```/g, '<pre><code>$1</code></pre>')
+      
+      // Inline code
+      .replace(/`([^`]+)`/g, '<code>$1</code>')
+      
+      // Lists
+      .replace(/^\d+\.\s+(.*)$/gm, '<li>$1</li>')
+      .replace(/^[\*\-]\s+(.*)$/gm, '<li>$1</li>')
+      
+      // Line breaks
+      .replace(/\n/g, '<br>');
+    
+    // Wrap lists
+    html = html.replace(/(<li>.*<\/li>)/g, '<ul>$1</ul>');
+    
+    // Handle details/summary (collapsible sections)
+    html = html.replace(/<details>/g, '<details class="tool-execution">')
+               .replace(/<summary>/g, '<summary class="tool-summary">');
+    
+    return html;
+  }
+  
+  showTypingIndicator() {
+    const typingDiv = document.createElement('div');
+    typingDiv.className = 'message assistant-message typing-message';
+    typingDiv.innerHTML = `
+      <div class="message-avatar">
+        <i class="fa fa-robot"></i>
+      </div>
+      <div class="message-content">
+        <div class="typing-indicator">
+          <span>AI is thinking</span>
+          <div class="typing-dots">
+            <div class="typing-dot"></div>
+            <div class="typing-dot"></div>
+            <div class="typing-dot"></div>
+          </div>
+        </div>
+      </div>
+    `;
+    
+    this.chatMessages.appendChild(typingDiv);
+    this.scrollToBottom();
+    return typingDiv;
+  }
+  
+  removeTypingIndicator(element) {
+    if (element && element.parentNode) {
+      element.parentNode.removeChild(element);
+    }
+  }
+  
+  clearChat() {
+    if (confirm('Are you sure you want to clear the chat history?')) {
+      // Keep the welcome message
+      const welcomeMessage = this.chatMessages.querySelector('.welcome-message');
+      this.chatMessages.innerHTML = '';
+      if (welcomeMessage) {
+        this.chatMessages.appendChild(welcomeMessage);
+      }
+      this.conversationId = null;
+    }
+  }
+  
+  newChat() {
+    // Call the new chat API endpoint
+    fetch("{{ url_for('AirflowChatView.new_chat') }}", {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'X-CSRFToken': '{{ csrf_token() }}'
+      }
+    })
+    .then(response => response.json())
+    .then(data => {
+      if (data.conversation_id) {
+        // this.conversationId = data.conversation_id;
+        this.clearChat();
+        this.setStatus('New chat initialized', 'ready');
+      } else {
+        console.error('Failed to initialize new chat:', data);
+        this.setStatus('Failed to create new chat', 'error');
+      }
+    })
+    .catch(error => {
+      console.error('Error creating new chat:', error);
+      this.setStatus('Error creating new chat', 'error');
+    });
+  }
+  
+  scrollToBottom() {
+    this.chatMessages.scrollTop = this.chatMessages.scrollHeight;
+  }
+  
+  escapeHtml(text) {
+    const div = document.createElement('div');
+    div.textContent = text;
+    return div.innerHTML;
+  }
+}
+
+// Initialize chat when page loads
+document.addEventListener('DOMContentLoaded', () => {
+  new AirflowChatClient();
+});
+</script>
+{% endblock %}

airflow_chat-0.1.0a1/pyproject.toml

@@ -0,0 +1,22 @@
+[tool.poetry]
+name = "airflow-chat"
+version = "0.1.0a1"
+description = "An Apache Airflow plugin that enables AI-powered chat interactions with your Airflow instance through MCP integration and an intuitive UI."
+license = "MIT"
+authors = ["Hipposys"]
+readme = "airflow_chat/plugins/README.md"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+apache-airflow = "^2.4.0"
+SQLAlchemy = ">=1.4.0"
+
+[tool.poetry.plugins."airflow.plugins"]
+airflow-chat-plugin = "airflow_chat.plugins.airflow_chat:AirflowChatPlugin"
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.urls]
+homepage = "https://github.com/hipposys-ltd/airflow-schedule-insights"