airflow-chat 0.1.0a1__tar.gz → 0.1.0a3__tar.gz

airflow_chat-0.1.0a3/PKG-INFO

@@ -0,0 +1,37 @@
+Metadata-Version: 2.3
+Name: airflow-chat
+Version: 0.1.0a3
+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,<3.13
+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
+Requires-Dist: SQLAlchemy (>=1.4.0)
+Requires-Dist: aiosqlite (>=0.21,<0.22)
+Requires-Dist: airflow-mcp-hipposys (==0.1.0a7)
+Requires-Dist: apache-airflow (>=2.4.0,<3.0.0)
+Requires-Dist: fastapi[standard] (>=0.115,<0.116)
+Requires-Dist: httpx (>=0.27,<0.28)
+Requires-Dist: itsdangerous (>=2.2,<3.0)
+Requires-Dist: langchain (>=0.3,<0.4)
+Requires-Dist: langchain-anthropic (>=0.3,<0.4)
+Requires-Dist: langchain-aws (>=0.2,<0.3)
+Requires-Dist: langchain-community (>=0.3,<0.4)
+Requires-Dist: langchain-core (>=0.3,<0.4)
+Requires-Dist: langchain-mcp-adapters (==0.1.1)
+Requires-Dist: langchain-openai (>=0.3,<0.4)
+Requires-Dist: langchain-text-splitters (>=0.3,<0.4)
+Requires-Dist: langgraph (>=0.5,<0.6)
+Requires-Dist: langgraph-checkpoint-postgres (>=2.0,<3.0)
+Requires-Dist: langgraph-checkpoint-sqlite (>=2.0,<3.0)
+Requires-Dist: mcp[cli] (>=1.9,<2.0)
+Requires-Dist: psycopg[binary] (>=3.2,<4.0)
+Requires-Dist: pydantic-settings (>=2.10,<3.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-0.1.0a3}/airflow_chat/plugins/airflow_chat.py

@@ -1,6 +1,7 @@
 from airflow import settings
 from airflow.plugins_manager import AirflowPlugin
-from flask import Blueprint, request, jsonify, Response, stream_template
+from flask import Blueprint, request, jsonify, Response, stream_template, \
+    flash, url_for, redirect
 from flask_login import current_user
 from flask_appbuilder import expose, BaseView as AppBuilderBaseView
 from functools import wraps
@@ -14,6 +15,7 @@ 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
+import asyncio
 
 # Blueprint for the chat plugin
 bp = Blueprint(
@@ -195,9 +197,10 @@ def admin_only(f):
             return jsonify({"error": "Authentication required"}), 401
         
         users_roles = [role.name for role in current_user.roles]
-        approved_roles = ["Admin"]  # Adjust roles as needed
+        approved_roles = ["Admin", "AIRFLOW_AI"]  # Adjust roles as needed
         if not any(role in users_roles for role in approved_roles):
-            return jsonify({"error": "Access denied"}), 403
+            flash("You do not have permission to access this page.", "danger")
+            return redirect(url_for("Airflow.index"))
         return f(*args, **kwargs)
     return decorated_function
 
@@ -369,12 +372,49 @@ class AirflowChatView(AppBuilderBaseView):
         # 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"
+                if os.environ.get('INTERNAL_AI_ASSISTANT_SERVER', True):
+                    from app.server.llm import get_stream_agent_responce
+                    md_uri = str(settings.Session().bind.url).replace('postgresql+psycopg2', 'postgres')
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                    try:
+                        # Get the async generator function
+                        stream_agent_response = loop.run_until_complete(
+                            get_stream_agent_responce(conversation_id, message, md_uri)
+                        )
+                        
+                        # Stream each chunk as it comes
+                        async_gen = stream_agent_response()
+                        while True:
+                            try:
+                                chunk = loop.run_until_complete(async_gen.__anext__())
+                                # Ensure chunk has the same structure as the else branch
+                                if isinstance(chunk, dict):
+                                    # If chunk is already a dict, use it as is
+                                    formatted_chunk = chunk
+                                else:
+                                    # If chunk is raw content, wrap it in the expected format
+                                    formatted_chunk = {
+                                        "content": str(chunk),
+                                        "conversation_id": conversation_id or str(uuid.uuid4()),
+                                        "timestamp": datetime.now().isoformat(),
+                                        "error": False
+                                    }
+                                yield f"data: {json.dumps(formatted_chunk)}\n\n"
+                            except StopAsyncIteration:
+                                break
+                        
+                        yield "data: [DONE]\n\n"
+                        
+                    finally:
+                        loop.close()
+                else:
+                    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)}",
@@ -403,7 +443,14 @@ class AirflowChatView(AppBuilderBaseView):
         conversation_id = str(uuid.uuid4())
         
         try:
-            success, cookies = self.llm_agent.initialize_chat_session(conversation_id)
+            if not os.environ.get('INTERNAL_AI_ASSISTANT_SERVER', True):
+                success, cookies = self.llm_agent.initialize_chat_session(conversation_id)
+            else:
+                from app.databases.postgres import Database
+                md_uri = str(settings.Session().bind.url).replace('postgresql+psycopg2', 'postgres')
+                asyncio.run(Database.setup(md_uri))
+                print('Database setup complete')
+                success = True
             
             if success:
                 return jsonify({

airflow_chat-0.1.0a3/airflow_chat/plugins/app/databases/postgres.py

@@ -0,0 +1,38 @@
+import os
+
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+from app.utils.singleton import Singleton
+
+
+class Database(metaclass=Singleton):
+    """Represents the main database.
+    Currently, provides only the connection string to the database.
+    """
+
+    def __init__(self, md_uri: str = None):
+        """Initialize the database connection."""
+
+        super().__init__()
+
+        self.user = os.environ.get('POSTGRES_USER', 'postgres')
+        self.password = os.environ.get('POSTGRES_PASSWORD')
+        self.host = os.environ.get('POSTGRES_HOSTNAME', 'postgres')
+        self.port = os.environ.get('POSTGRES_PORT', 5432)
+        self.database = os.environ.get('POSTGRES_DB', 'chat_db')
+        auth = f'{self.user}:{self.password}'
+        self.uri = \
+            f'postgres://{auth}@{self.host}:{self.port}/{self.database}'
+        if md_uri:
+            self.uri = md_uri
+
+    def get_connection_string(self) -> str:
+        """Get a URI representation of the database connection params."""
+        return self.uri
+
+    @staticmethod
+    async def setup(md_uri: str = None):
+        """Setup the database."""
+        async with AsyncPostgresSaver.from_conn_string(
+           Database(md_uri).get_connection_string()) as saver:
+            await saver.setup()

airflow_chat-0.1.0a3/airflow_chat/plugins/app/models/__init__.py

@@ -0,0 +1,14 @@
+import os
+model_type, model_id = os.environ.get('LLM_MODEL_ID', 'none:none').split(':', 1)
+
+if model_type == 'bedrock':
+    from app.models.inference.bedrock_model import ChatBedrock
+    ChatModel = ChatBedrock
+elif model_type == 'antropic':
+    from app.models.inference.antropic_model import ChatAnthropic
+    ChatModel = ChatAnthropic
+elif model_type == 'openai':
+    from app.models.inference.openai_model import ChatOpenAI
+    ChatModel = ChatOpenAI
+else:
+    ChatModel = None

airflow_chat-0.1.0a3/airflow_chat/plugins/app/models/inference/antropic_model.py

@@ -0,0 +1,17 @@
+import os
+
+from langchain_anthropic import ChatAnthropic as BaseChatAnthropic
+
+
+class ChatAnthropic(BaseChatAnthropic):
+    """A wrapper for the `langchain_aws.ChatBedrock`."""
+
+    def __init__(self, **kwargs):
+        """Initialize the `ChatBedrock` with specific configuration."""
+        model_type, model_id = os.environ['LLM_MODEL_ID'].split(':', 1)
+        default_kwargs = {
+            'model': model_id,
+            'temperature': 0,
+        }
+
+        super().__init__(**(default_kwargs | kwargs))

airflow_chat-0.1.0a3/airflow_chat/plugins/app/models/inference/bedrock_model.py

@@ -0,0 +1,18 @@
+import os
+
+from langchain_aws import ChatBedrock as BaseChatBedrock
+
+
+class ChatBedrock(BaseChatBedrock):
+    """A wrapper for the `langchain_aws.ChatBedrock`."""
+
+    def __init__(self, **kwargs):
+        """Initialize the `ChatBedrock` with specific configuration."""
+        model_type, model_id = os.environ['LLM_MODEL_ID'].split(':', 1)
+        default_kwargs = {
+            'model_id': model_id,
+            'region_name': os.environ.get('AWS_DEFAULT_REGION', 'us-east-1'),
+            'model_kwargs': dict(temperature=0),
+        }
+
+        super().__init__(**(default_kwargs | kwargs))

airflow_chat-0.1.0a3/airflow_chat/plugins/app/models/inference/openai_model.py

@@ -0,0 +1,17 @@
+import os
+
+from langchain_openai import ChatOpenAI as BaseChatOpenAI
+
+
+class ChatOpenAI(BaseChatOpenAI):
+    """A wrapper for the `langchain_aws.ChatOpenAI`."""
+
+    def __init__(self, model_id: str = None, **kwargs):
+        """Initialize the `ChatOpenAI` with specific configuration."""
+        model_type, model_id_env = os.environ['LLM_MODEL_ID'].split(':', 1)
+        default_kwargs = {
+            'model': model_id or model_id_env,
+            'temperature': 0,
+        }
+
+        super().__init__(**(default_kwargs | kwargs))

airflow_chat-0.1.0a3/airflow_chat/plugins/app/server/chat.py

@@ -0,0 +1,37 @@
+import uuid
+
+from fastapi import APIRouter, Request
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from app.server.llm import get_stream_agent_responce
+
+
+chat_router = APIRouter()
+
+
+class ChatRequest(BaseModel):
+    message: str
+
+
+@chat_router.post("/new")
+async def new_chat(request: Request):
+    """Create a new chat session."""
+    request.session['chat_session_id'] = f'user_{uuid.uuid4()}'
+    return {'results': 'ok'}
+
+
+@chat_router.post("/ask")
+async def chat(
+    request: Request,
+    chat_request: ChatRequest,
+):
+    if 'chat_session_id' not in request.session:
+        await new_chat(request)
+    # Get the user chat configuration and the LLM agent.
+    stream_agent_response = await get_stream_agent_responce(request.session[
+        'chat_session_id'],
+                                                      chat_request.message)
+
+    # Return the agent's response as a stream of JSON objects.
+    return StreamingResponse(stream_agent_response(),
+                             media_type='application/json')

airflow_chat-0.1.0a3/airflow_chat/plugins/app/server/llm.py

@@ -0,0 +1,245 @@
+from enum import Enum
+from typing import AsyncGenerator
+
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.prebuilt import create_react_agent
+from langchain_core.messages import HumanMessage, BaseMessage, \
+    SystemMessage, ToolMessage, AIMessage, AIMessageChunk
+
+from app.databases.postgres import Database
+from app.models import ChatModel
+from app.utils.logger import Logger
+
+import os
+
+from langchain_mcp_adapters.client import MultiServerMCPClient
+from datetime import datetime
+from langchain.tools import Tool
+
+
+PROMPT_MESSAGE = """Be a chatbot."""
+
+
+class LLMEventType(Enum):
+    """Event types for the LLM agent."""
+
+    STORED_MESSAGE = 'stored_message'
+    RETRIEVER_START = 'on_retriever_start'
+    RETRIEVER_END = 'on_retriever_end'
+    CHAT_CHUNK = 'on_chat_model_stream'
+    DONE = 'done'
+
+
+class ChatMessage:
+    class Sender(Enum):
+        """The sender of the message."""
+        SYSTEM = 'system'
+        AI = 'ai'
+        HUMAN = 'human'
+        TOOL = 'tool'
+
+    def __init__(self, type: LLMEventType, sender: Sender,
+                 content: str, payload: dict = None):
+        self.type = type
+        self.sender: str = sender.value
+        self.content = content
+        self.payload = payload or {}
+
+    @classmethod
+    def from_base_message(cls, message: BaseMessage) -> 'ChatMessage':
+        message_type_lookup = {
+            HumanMessage: cls.Sender.HUMAN,
+            SystemMessage: cls.Sender.SYSTEM,
+            ToolMessage: cls.Sender.TOOL,
+            AIMessage: cls.Sender.AI,
+            AIMessageChunk: cls.Sender.AI,
+        }
+
+        # Different message types have different structures.
+        try:
+            content = message.content[0]['text']
+        except (KeyError, TypeError, IndexError):
+            content = message.content
+
+        return ChatMessage(
+            LLMEventType.STORED_MESSAGE,
+            sender=message_type_lookup[type(message)],
+            content=content,
+        )
+
+    @classmethod
+    def from_event(cls, event: dict) -> 'ChatMessage':
+        """Convert an event from the LLM agent to a `ChatMessage` object."""
+        if event['event'] in ('on_tool_start', 'on_tool_end'):
+            print(event)
+            print('--------------------')
+        match event['event']:
+            case 'on_chat_model_stream':
+                if event['data']['chunk'].content:
+                    return cls._handle_on_chat_model_stream(event)
+            case 'on_tool_start':
+                return ChatMessage(LLMEventType.CHAT_CHUNK, cls.Sender.AI,
+                                   f'''\n\nStart Running Tool:
+```
+Data: {event['data']['input']}
+Function: {event['name']}
+```
+''')
+            case 'on_tool_end':
+                return ChatMessage(LLMEventType.CHAT_CHUNK, cls.Sender.AI,
+                                   f'''\n\nTool Output: \n```\n
+{event['data']['output'].content}
+\n```\n''')
+            # The conversation is done.
+            case 'done':
+                return ChatMessage(
+                    LLMEventType.DONE,
+                    cls.Sender.SYSTEM,
+                    'Done',
+                )
+            # Known events that we ignore.
+            case 'on_chat_model_start' | 'on_chain_start' | 'on_chain_end' \
+                | 'on_chat_model_stream' | 'on_chat_model_end' | \
+                    'on_chain_stream' | 'on_tool_start' | 'on_tool_end':
+                Logger().get_logger().debug('Ignoring message', event['event'])
+                return ''
+            # Unknown events.
+            case _:
+                raise ValueError('Unknown event', event)
+
+    @classmethod
+    def _handle_on_chat_model_stream(cls, event: dict) -> 'ChatMessage':
+        content = event['data']['chunk'].content
+        content_type = ''
+        if not isinstance(content, str):
+            content_type = content[0]['type']
+            content = content[0].get('text')
+
+        # If the message is a tool call, just print a debug message.
+        if content_type in ('tool_use', 'tool_call'):
+            Logger().get_logger().debug('Stream.tool_calls:',
+                                        event['data']['chunk'].tool_calls,
+                                        flush=True)
+            return ''
+        else:
+            return ChatMessage(LLMEventType.CHAT_CHUNK, cls.Sender.AI, content
+                               if content is not None else '')
+
+    def to_dict(self) -> dict:
+        """Returns a dictionary representation of the message."""
+        return {
+            'sender': self.sender,
+            'content': self.content,
+            'payload': self.payload,
+        }
+
+
+class LLMAgent:
+
+    def __init__(self, tools, md_uri=None):
+        self._agent = None
+        self._llm = None
+        self.retriever_tool_name = 'Internal_Company_Info_Retriever'
+        self._checkpointer_ctx = None
+        self.tools = tools
+        self.md_uri = md_uri
+
+    async def __aenter__(self) -> 'LLMAgent':
+
+        tools = self.tools
+        self._llm = ChatModel()
+
+        # Checkpointer for the agent.
+        self._checkpointer_ctx = AsyncPostgresSaver\
+            .from_conn_string(
+                Database(self.md_uri).get_connection_string())
+        checkpointer = await self._checkpointer_ctx.__aenter__()
+
+        # Create the agent itself.
+        self._agent = create_react_agent(
+            self._llm,
+            tools,
+            checkpointer=checkpointer,
+            # state_modifier=SystemMessage(PROMPT_MESSAGE),
+            prompt=PROMPT_MESSAGE
+        )
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Close the agent and the checkpointer."""
+        await self._checkpointer_ctx.__aexit__(exc_type, exc_val, exc_tb)
+        self._llm = None
+        self._agent = None
+        self._checkpointer_ctx = None
+
+    async def astream_events(self, message: str,
+                             chat_session: dict) -> AsyncGenerator[ChatMessage,
+                                                                   None]:
+        async for event in self._agent.astream_events(
+            {"messages": [HumanMessage(content=message)]},
+            config=chat_session,
+            version='v2',
+        ):
+            message = ChatMessage.from_event(event)
+            if message:
+                yield message
+
+        # Let the client know that the conversation is done.
+        # yield ChatMessage.from_event({'event': 'done'})
+
+
+def get_user_chat_config(session_id: str) -> dict:
+    return {'configurable': {'thread_id': session_id},
+            "recursion_limit": 100}
+
+
+async def get_stream_agent_responce(session_id, message,
+                                    md_uri: str = None):
+    user_config = get_user_chat_config(session_id)
+    mcp_host = os.environ.get('mcp_host', 'mcp_sse_server:8000')
+    TRANSPORT_TYPE = os.environ.get('TRANSPORT_TYPE', 'stdio')
+    if TRANSPORT_TYPE == 'stdio':
+        mcps = {
+                "AirflowMCP":
+                {
+                    'command': "python",
+                    'args': ["-m", "airflow_mcp_hipposys.mcp_airflow"],
+                    "transport": "stdio",
+                    'env': {k: v for k, v in {
+                        'AIRFLOW_ASSISTENT_AI_CONN': os.getenv(
+                            'AIRFLOW_ASSISTENT_AI_CONN'),
+                        'airflow_api_url': os.getenv('airflow_api_url'),
+                        'airflow_username': os.getenv('airflow_username'),
+                        'airflow_password': os.getenv('airflow_password'),
+                        'AIRFLOW_INSIGHTS_MODE':
+                            os.getenv('AIRFLOW_INSIGHTS_MODE'),
+                        'POST_MODE': os.getenv('POST_MODE'),
+                        'TRANSPORT_TYPE': 'stdio'
+                    }.items() if v is not None}
+                }
+            }
+    elif TRANSPORT_TYPE == 'sse':
+        mcps = {
+                    "AirflowMCP": {
+                        "url": f"http://{mcp_host}/sse",
+                        "transport": "sse",
+                        "headers": {"Authorization": f"""Bearer {
+                            os.environ.get('MCP_TOKEN')}"""}
+                    }
+                }
+    datetime_tool = Tool(
+        name="Datetime",
+        func=lambda x: datetime.now().isoformat(),
+        description="Returns the current datetime",
+    )
+
+    client = MultiServerMCPClient(mcps)
+    tools = await client.get_tools() + [datetime_tool]
+
+    async def stream_agent_response():
+        async with LLMAgent(tools=tools, md_uri=md_uri) as llm_agent:
+            async for chat_msg in llm_agent.astream_events(
+                 message, user_config):
+                yield chat_msg.content
+    return stream_agent_response

airflow_chat-0.1.0a3/airflow_chat/plugins/app/server/main.py

@@ -0,0 +1,45 @@
+import os
+
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from starlette.middleware.sessions import SessionMiddleware
+
+from app.server.chat import chat_router
+from app.databases.postgres import Database
+from app.utils.config import Config
+from app.utils.logger import Logger
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Run the database setup and teardown."""
+    await Database.setup()
+    Logger().get_logger().info('Database setup complete')
+
+    yield
+    # Optionally add teardown code here
+
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    SessionMiddleware,
+    secret_key=os.environ.get('SECRET_KEY'),
+    https_only=Config.get_deploy_env() == 'PROD',
+)
+
+
+@app.middleware("http")
+async def check_token_middleware(request: Request, call_next):
+    """Allow only requests with the correct token."""
+    token = request.headers.get("x-access-token")
+    if (Config.get_deploy_env() != 'LOCAL') \
+       and (token != os.environ['FAST_API_ACCESS_SECRET_TOKEN']):
+        return JSONResponse(
+            status_code=403, content={'reason': 'Invalid or missing token'})
+    response = await call_next(request)
+    return response
+
+
+app.include_router(chat_router, prefix='/chat')

airflow_chat-0.1.0a3/airflow_chat/plugins/app/utils/config.py

@@ -0,0 +1,11 @@
+import os
+
+from app.utils.singleton import Singleton
+
+
+class Config(metaclass=Singleton):
+    """Holds "global" configuration for the application."""
+    @staticmethod
+    def get_deploy_env() -> str:
+        """Get the current deployment environment."""
+        return os.environ.get('DEPLOY_ENV', 'prod').upper()

airflow_chat-0.1.0a3/airflow_chat/plugins/app/utils/logger.py

@@ -0,0 +1,34 @@
+import sys
+import logging
+
+from app.utils.singleton import Singleton
+
+
+class Logger(metaclass=Singleton):
+    default_config = {
+        'name': 'RAG-App',
+        'level': 'INFO',
+        'format': '[%(asctime)s|%(name)s|%(levelname)s|%(processName)s:%(threadName)s|%(filename)s, '
+                  'line %(lineno)s in %(funcName)s] %(message)s',
+    }
+
+    def __init__(self,
+                 config_to_use: dict = None):
+
+        config = {**Logger.default_config, **(config_to_use or {})}
+        self.logger = logging.getLogger(config['name'])
+
+        # Set the root logging level.
+        logging_lvl = getattr(logging, config['level'])
+        logging.root.setLevel(logging_lvl)
+
+        # Create an handler and configure it.
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(logging_lvl)
+        handler.setFormatter(logging.Formatter(config['format']))
+
+        # Remove all other handlers in the process.
+        self.logger.addHandler(handler)
+
+    def get_logger(self) -> logging.Logger:
+        return self.logger

airflow_chat-0.1.0a3/airflow_chat/plugins/app/utils/singleton.py

@@ -0,0 +1,52 @@
+from datetime import datetime, timedelta
+
+
+class Singleton(type):
+    """A metaclass that implements the singleton pattern.
+
+    This metaclass ensures that only one(*) instance of a class is created and that the instance is reused for all subsequent calls.
+    
+    The instance is refreshed every `MAX_INSTANCE_TTL` seconds.
+    
+    Usage:
+    ```python
+    >>> class MyClass(metaclass=Singleton):
+    ...     pass
+        
+    >>> a = MyClass()
+    >>> b = MyClass()
+    >>> a is b
+        True
+    >>> sleep(MyClass.MAX_INSTANCE_TTL + 1)
+    >>> c = MyClass()
+    >>> a is c
+        False
+    ```
+    """
+
+    _instances = {}
+    _creation_time = {}
+    
+    MAX_INSTANCE_TTL = timedelta(seconds=5 * 60)  # 5 minutes
+
+    def __create_instance(cls, *args, **kwargs):
+        """Create a new instance of this class."""
+
+        cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        cls._creation_time[cls] = datetime.now()
+
+    def __call__(cls, *args, force_recreate=False, **kwargs):
+        """Create a new instance of this class if it does not exist or if it is no longer valid.
+        
+        If the `force_recreate` parameter is set to `True`, a new instance will be created regardless of the validity of the existing instance.
+        """
+
+        # Create a new instance if it does not exist or if `force_recreate` is set.
+        if force_recreate or (cls not in cls._instances):
+            cls.__create_instance(*args, **kwargs)
+        # Create a new instance if the existing instance is no longer valid.
+        elif (cls not in cls._creation_time) or (cls._creation_time[cls] + cls.MAX_INSTANCE_TTL < datetime.now()):
+            cls.__create_instance(*args, **kwargs)
+
+        # Return the existing instance.
+        return cls._instances[cls]

{airflow_chat-0.1.0a1 → airflow_chat-0.1.0a3}/airflow_chat/plugins/templates/chat_interface.html

@@ -18,15 +18,7 @@
         </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?
+            Hello! I'm your Airflow AI assistant.
           </div>
           <div class="message-time">Just now</div>
         </div>

airflow_chat-0.1.0a3/pyproject.toml

@@ -0,0 +1,42 @@
+[tool.poetry]
+name = "airflow-chat"
+version = "0.1.0a3"
+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,<3.13"
+apache-airflow = "^2.4.0"
+SQLAlchemy = ">=1.4.0"
+fastapi = {extras = ["standard"], version = "^0.115"}
+pydantic-settings = "^2.10"
+aiosqlite = "^0.21"
+itsdangerous = "^2.2"
+langchain = "^0.3"
+langchain-community = "^0.3"
+langchain-core = "^0.3"
+langchain-text-splitters = "^0.3"
+langgraph = "^0.5"
+langgraph-checkpoint-postgres = "^2.0"
+langgraph-checkpoint-sqlite = "^2.0"
+langchain-aws = "^0.2"
+langchain-anthropic = "^0.3"
+langchain-openai = "^0.3"
+mcp = {extras = ["cli"], version = "^1.9"}
+httpx = "^0.27"
+psycopg = {extras = ["binary"], version = "^3.2"}
+langchain-mcp-adapters = "0.1.1"
+airflow-mcp-hipposys = "0.1.0a7"
+
+[tool.poetry.plugins."airflow.plugins"]
+airflow_chat = "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"
+

airflow_chat-0.1.0a1/PKG-INFO

@@ -1,19 +0,0 @@
-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/pyproject.toml

@@ -1,22 +0,0 @@
-[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"