airflow-chat 0.1.0a1__py3-none-any.whl → 0.1.0a2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,7 @@
1
1
  from airflow import settings
2
2
  from airflow.plugins_manager import AirflowPlugin
3
- from flask import Blueprint, request, jsonify, Response, stream_template
3
+ from flask import Blueprint, request, jsonify, Response, stream_template, \
4
+ flash, url_for, redirect
4
5
  from flask_login import current_user
5
6
  from flask_appbuilder import expose, BaseView as AppBuilderBaseView
6
7
  from functools import wraps
@@ -14,6 +15,7 @@ from sqlalchemy import Column, String, DateTime, Text, create_engine
14
15
  from sqlalchemy.ext.declarative import declarative_base
15
16
  from sqlalchemy.orm import sessionmaker
16
17
  from sqlalchemy.exc import SQLAlchemyError
18
+ import asyncio
17
19
 
18
20
  # Blueprint for the chat plugin
19
21
  bp = Blueprint(
@@ -195,9 +197,10 @@ def admin_only(f):
195
197
  return jsonify({"error": "Authentication required"}), 401
196
198
 
197
199
  users_roles = [role.name for role in current_user.roles]
198
- approved_roles = ["Admin"] # Adjust roles as needed
200
+ approved_roles = ["Admin", "AIRFLOW_AI"] # Adjust roles as needed
199
201
  if not any(role in users_roles for role in approved_roles):
200
- return jsonify({"error": "Access denied"}), 403
202
+ flash("You do not have permission to access this page.", "danger")
203
+ return redirect(url_for("Airflow.index"))
201
204
  return f(*args, **kwargs)
202
205
  return decorated_function
203
206
 
@@ -369,12 +372,49 @@ class AirflowChatView(AppBuilderBaseView):
369
372
  # Return streaming response
370
373
  def generate():
371
374
  try:
372
- for chunk in self.llm_agent.stream_chat_response(
373
- message,
374
- conversation_id
375
- ):
376
- yield f"data: {json.dumps(chunk)}\n\n"
377
- yield "data: [DONE]\n\n"
375
+ if os.environ.get('INTERNAL_AI_ASSISTANT_SERVER', True):
376
+ from app.server.llm import get_stream_agent_responce
377
+ md_uri = str(settings.Session().bind.url).replace('postgresql+psycopg2', 'postgres')
378
+ loop = asyncio.new_event_loop()
379
+ asyncio.set_event_loop(loop)
380
+ try:
381
+ # Get the async generator function
382
+ stream_agent_response = loop.run_until_complete(
383
+ get_stream_agent_responce(conversation_id, message, md_uri)
384
+ )
385
+
386
+ # Stream each chunk as it comes
387
+ async_gen = stream_agent_response()
388
+ while True:
389
+ try:
390
+ chunk = loop.run_until_complete(async_gen.__anext__())
391
+ # Ensure chunk has the same structure as the else branch
392
+ if isinstance(chunk, dict):
393
+ # If chunk is already a dict, use it as is
394
+ formatted_chunk = chunk
395
+ else:
396
+ # If chunk is raw content, wrap it in the expected format
397
+ formatted_chunk = {
398
+ "content": str(chunk),
399
+ "conversation_id": conversation_id or str(uuid.uuid4()),
400
+ "timestamp": datetime.now().isoformat(),
401
+ "error": False
402
+ }
403
+ yield f"data: {json.dumps(formatted_chunk)}\n\n"
404
+ except StopAsyncIteration:
405
+ break
406
+
407
+ yield "data: [DONE]\n\n"
408
+
409
+ finally:
410
+ loop.close()
411
+ else:
412
+ for chunk in self.llm_agent.stream_chat_response(
413
+ message,
414
+ conversation_id
415
+ ):
416
+ yield f"data: {json.dumps(chunk)}\n\n"
417
+ yield "data: [DONE]\n\n"
378
418
  except Exception as e:
379
419
  error_chunk = {
380
420
  "content": f"Error: {str(e)}",
@@ -403,7 +443,14 @@ class AirflowChatView(AppBuilderBaseView):
403
443
  conversation_id = str(uuid.uuid4())
404
444
 
405
445
  try:
406
- success, cookies = self.llm_agent.initialize_chat_session(conversation_id)
446
+ if not os.environ.get('INTERNAL_AI_ASSISTANT_SERVER', True):
447
+ success, cookies = self.llm_agent.initialize_chat_session(conversation_id)
448
+ else:
449
+ from app.databases.postgres import Database
450
+ md_uri = str(settings.Session().bind.url).replace('postgresql+psycopg2', 'postgres')
451
+ asyncio.run(Database.setup(md_uri))
452
+ print('Database setup complete')
453
+ success = True
407
454
 
408
455
  if success:
409
456
  return jsonify({
File without changes
File without changes
@@ -0,0 +1,38 @@
1
+ import os
2
+
3
+ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
4
+
5
+ from app.utils.singleton import Singleton
6
+
7
+
8
+ class Database(metaclass=Singleton):
9
+ """Represents the main database.
10
+ Currently, provides only the connection string to the database.
11
+ """
12
+
13
+ def __init__(self, md_uri: str = None):
14
+ """Initialize the database connection."""
15
+
16
+ super().__init__()
17
+
18
+ self.user = os.environ.get('POSTGRES_USER', 'postgres')
19
+ self.password = os.environ.get('POSTGRES_PASSWORD')
20
+ self.host = os.environ.get('POSTGRES_HOSTNAME', 'postgres')
21
+ self.port = os.environ.get('POSTGRES_PORT', 5432)
22
+ self.database = os.environ.get('POSTGRES_DB', 'chat_db')
23
+ auth = f'{self.user}:{self.password}'
24
+ self.uri = \
25
+ f'postgres://{auth}@{self.host}:{self.port}/{self.database}'
26
+ if md_uri:
27
+ self.uri = md_uri
28
+
29
+ def get_connection_string(self) -> str:
30
+ """Get a URI representation of the database connection params."""
31
+ return self.uri
32
+
33
+ @staticmethod
34
+ async def setup(md_uri: str = None):
35
+ """Setup the database."""
36
+ async with AsyncPostgresSaver.from_conn_string(
37
+ Database(md_uri).get_connection_string()) as saver:
38
+ await saver.setup()
@@ -0,0 +1,14 @@
1
+ import os
2
+ model_type, model_id = os.environ.get('LLM_MODEL_ID', 'none:none').split(':', 1)
3
+
4
+ if model_type == 'bedrock':
5
+ from app.models.inference.bedrock_model import ChatBedrock
6
+ ChatModel = ChatBedrock
7
+ elif model_type == 'antropic':
8
+ from app.models.inference.antropic_model import ChatAnthropic
9
+ ChatModel = ChatAnthropic
10
+ elif model_type == 'openai':
11
+ from app.models.inference.openai_model import ChatOpenAI
12
+ ChatModel = ChatOpenAI
13
+ else:
14
+ ChatModel = None
File without changes
@@ -0,0 +1,17 @@
1
+ import os
2
+
3
+ from langchain_anthropic import ChatAnthropic as BaseChatAnthropic
4
+
5
+
6
+ class ChatAnthropic(BaseChatAnthropic):
7
+ """A wrapper for the `langchain_aws.ChatBedrock`."""
8
+
9
+ def __init__(self, **kwargs):
10
+ """Initialize the `ChatBedrock` with specific configuration."""
11
+ model_type, model_id = os.environ['LLM_MODEL_ID'].split(':', 1)
12
+ default_kwargs = {
13
+ 'model': model_id,
14
+ 'temperature': 0,
15
+ }
16
+
17
+ super().__init__(**(default_kwargs | kwargs))
@@ -0,0 +1,18 @@
1
+ import os
2
+
3
+ from langchain_aws import ChatBedrock as BaseChatBedrock
4
+
5
+
6
+ class ChatBedrock(BaseChatBedrock):
7
+ """A wrapper for the `langchain_aws.ChatBedrock`."""
8
+
9
+ def __init__(self, **kwargs):
10
+ """Initialize the `ChatBedrock` with specific configuration."""
11
+ model_type, model_id = os.environ['LLM_MODEL_ID'].split(':', 1)
12
+ default_kwargs = {
13
+ 'model_id': model_id,
14
+ 'region_name': os.environ.get('AWS_DEFAULT_REGION', 'us-east-1'),
15
+ 'model_kwargs': dict(temperature=0),
16
+ }
17
+
18
+ super().__init__(**(default_kwargs | kwargs))
@@ -0,0 +1,17 @@
1
+ import os
2
+
3
+ from langchain_openai import ChatOpenAI as BaseChatOpenAI
4
+
5
+
6
+ class ChatOpenAI(BaseChatOpenAI):
7
+ """A wrapper for the `langchain_aws.ChatOpenAI`."""
8
+
9
+ def __init__(self, model_id: str = None, **kwargs):
10
+ """Initialize the `ChatOpenAI` with specific configuration."""
11
+ model_type, model_id_env = os.environ['LLM_MODEL_ID'].split(':', 1)
12
+ default_kwargs = {
13
+ 'model': model_id or model_id_env,
14
+ 'temperature': 0,
15
+ }
16
+
17
+ super().__init__(**(default_kwargs | kwargs))
File without changes
@@ -0,0 +1,37 @@
1
+ import uuid
2
+
3
+ from fastapi import APIRouter, Request
4
+ from fastapi.responses import StreamingResponse
5
+ from pydantic import BaseModel
6
+ from app.server.llm import get_stream_agent_responce
7
+
8
+
9
+ chat_router = APIRouter()
10
+
11
+
12
+ class ChatRequest(BaseModel):
13
+ message: str
14
+
15
+
16
+ @chat_router.post("/new")
17
+ async def new_chat(request: Request):
18
+ """Create a new chat session."""
19
+ request.session['chat_session_id'] = f'user_{uuid.uuid4()}'
20
+ return {'results': 'ok'}
21
+
22
+
23
+ @chat_router.post("/ask")
24
+ async def chat(
25
+ request: Request,
26
+ chat_request: ChatRequest,
27
+ ):
28
+ if 'chat_session_id' not in request.session:
29
+ await new_chat(request)
30
+ # Get the user chat configuration and the LLM agent.
31
+ stream_agent_response = await get_stream_agent_responce(request.session[
32
+ 'chat_session_id'],
33
+ chat_request.message)
34
+
35
+ # Return the agent's response as a stream of JSON objects.
36
+ return StreamingResponse(stream_agent_response(),
37
+ media_type='application/json')
@@ -0,0 +1,245 @@
1
+ from enum import Enum
2
+ from typing import AsyncGenerator
3
+
4
+ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
5
+ from langgraph.prebuilt import create_react_agent
6
+ from langchain_core.messages import HumanMessage, BaseMessage, \
7
+ SystemMessage, ToolMessage, AIMessage, AIMessageChunk
8
+
9
+ from app.databases.postgres import Database
10
+ from app.models import ChatModel
11
+ from app.utils.logger import Logger
12
+
13
+ import os
14
+
15
+ from langchain_mcp_adapters.client import MultiServerMCPClient
16
+ from datetime import datetime
17
+ from langchain.tools import Tool
18
+
19
+
20
+ PROMPT_MESSAGE = """Be a chatbot."""
21
+
22
+
23
+ class LLMEventType(Enum):
24
+ """Event types for the LLM agent."""
25
+
26
+ STORED_MESSAGE = 'stored_message'
27
+ RETRIEVER_START = 'on_retriever_start'
28
+ RETRIEVER_END = 'on_retriever_end'
29
+ CHAT_CHUNK = 'on_chat_model_stream'
30
+ DONE = 'done'
31
+
32
+
33
+ class ChatMessage:
34
+ class Sender(Enum):
35
+ """The sender of the message."""
36
+ SYSTEM = 'system'
37
+ AI = 'ai'
38
+ HUMAN = 'human'
39
+ TOOL = 'tool'
40
+
41
+ def __init__(self, type: LLMEventType, sender: Sender,
42
+ content: str, payload: dict = None):
43
+ self.type = type
44
+ self.sender: str = sender.value
45
+ self.content = content
46
+ self.payload = payload or {}
47
+
48
+ @classmethod
49
+ def from_base_message(cls, message: BaseMessage) -> 'ChatMessage':
50
+ message_type_lookup = {
51
+ HumanMessage: cls.Sender.HUMAN,
52
+ SystemMessage: cls.Sender.SYSTEM,
53
+ ToolMessage: cls.Sender.TOOL,
54
+ AIMessage: cls.Sender.AI,
55
+ AIMessageChunk: cls.Sender.AI,
56
+ }
57
+
58
+ # Different message types have different structures.
59
+ try:
60
+ content = message.content[0]['text']
61
+ except (KeyError, TypeError, IndexError):
62
+ content = message.content
63
+
64
+ return ChatMessage(
65
+ LLMEventType.STORED_MESSAGE,
66
+ sender=message_type_lookup[type(message)],
67
+ content=content,
68
+ )
69
+
70
+ @classmethod
71
+ def from_event(cls, event: dict) -> 'ChatMessage':
72
+ """Convert an event from the LLM agent to a `ChatMessage` object."""
73
+ if event['event'] in ('on_tool_start', 'on_tool_end'):
74
+ print(event)
75
+ print('--------------------')
76
+ match event['event']:
77
+ case 'on_chat_model_stream':
78
+ if event['data']['chunk'].content:
79
+ return cls._handle_on_chat_model_stream(event)
80
+ case 'on_tool_start':
81
+ return ChatMessage(LLMEventType.CHAT_CHUNK, cls.Sender.AI,
82
+ f'''\n\nStart Running Tool:
83
+ ```
84
+ Data: {event['data']['input']}
85
+ Function: {event['name']}
86
+ ```
87
+ ''')
88
+ case 'on_tool_end':
89
+ return ChatMessage(LLMEventType.CHAT_CHUNK, cls.Sender.AI,
90
+ f'''\n\nTool Output: \n```\n
91
+ {event['data']['output'].content}
92
+ \n```\n''')
93
+ # The conversation is done.
94
+ case 'done':
95
+ return ChatMessage(
96
+ LLMEventType.DONE,
97
+ cls.Sender.SYSTEM,
98
+ 'Done',
99
+ )
100
+ # Known events that we ignore.
101
+ case 'on_chat_model_start' | 'on_chain_start' | 'on_chain_end' \
102
+ | 'on_chat_model_stream' | 'on_chat_model_end' | \
103
+ 'on_chain_stream' | 'on_tool_start' | 'on_tool_end':
104
+ Logger().get_logger().debug('Ignoring message', event['event'])
105
+ return ''
106
+ # Unknown events.
107
+ case _:
108
+ raise ValueError('Unknown event', event)
109
+
110
+ @classmethod
111
+ def _handle_on_chat_model_stream(cls, event: dict) -> 'ChatMessage':
112
+ content = event['data']['chunk'].content
113
+ content_type = ''
114
+ if not isinstance(content, str):
115
+ content_type = content[0]['type']
116
+ content = content[0].get('text')
117
+
118
+ # If the message is a tool call, just print a debug message.
119
+ if content_type in ('tool_use', 'tool_call'):
120
+ Logger().get_logger().debug('Stream.tool_calls:',
121
+ event['data']['chunk'].tool_calls,
122
+ flush=True)
123
+ return ''
124
+ else:
125
+ return ChatMessage(LLMEventType.CHAT_CHUNK, cls.Sender.AI, content
126
+ if content is not None else '')
127
+
128
+ def to_dict(self) -> dict:
129
+ """Returns a dictionary representation of the message."""
130
+ return {
131
+ 'sender': self.sender,
132
+ 'content': self.content,
133
+ 'payload': self.payload,
134
+ }
135
+
136
+
137
+ class LLMAgent:
138
+
139
+ def __init__(self, tools, md_uri=None):
140
+ self._agent = None
141
+ self._llm = None
142
+ self.retriever_tool_name = 'Internal_Company_Info_Retriever'
143
+ self._checkpointer_ctx = None
144
+ self.tools = tools
145
+ self.md_uri = md_uri
146
+
147
+ async def __aenter__(self) -> 'LLMAgent':
148
+
149
+ tools = self.tools
150
+ self._llm = ChatModel()
151
+
152
+ # Checkpointer for the agent.
153
+ self._checkpointer_ctx = AsyncPostgresSaver\
154
+ .from_conn_string(
155
+ Database(self.md_uri).get_connection_string())
156
+ checkpointer = await self._checkpointer_ctx.__aenter__()
157
+
158
+ # Create the agent itself.
159
+ self._agent = create_react_agent(
160
+ self._llm,
161
+ tools,
162
+ checkpointer=checkpointer,
163
+ # state_modifier=SystemMessage(PROMPT_MESSAGE),
164
+ prompt=PROMPT_MESSAGE
165
+ )
166
+
167
+ return self
168
+
169
+ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
170
+ """Close the agent and the checkpointer."""
171
+ await self._checkpointer_ctx.__aexit__(exc_type, exc_val, exc_tb)
172
+ self._llm = None
173
+ self._agent = None
174
+ self._checkpointer_ctx = None
175
+
176
+ async def astream_events(self, message: str,
177
+ chat_session: dict) -> AsyncGenerator[ChatMessage,
178
+ None]:
179
+ async for event in self._agent.astream_events(
180
+ {"messages": [HumanMessage(content=message)]},
181
+ config=chat_session,
182
+ version='v2',
183
+ ):
184
+ message = ChatMessage.from_event(event)
185
+ if message:
186
+ yield message
187
+
188
+ # Let the client know that the conversation is done.
189
+ # yield ChatMessage.from_event({'event': 'done'})
190
+
191
+
192
+ def get_user_chat_config(session_id: str) -> dict:
193
+ return {'configurable': {'thread_id': session_id},
194
+ "recursion_limit": 100}
195
+
196
+
197
+ async def get_stream_agent_responce(session_id, message,
198
+ md_uri: str = None):
199
+ user_config = get_user_chat_config(session_id)
200
+ mcp_host = os.environ.get('mcp_host', 'mcp_sse_server:8000')
201
+ TRANSPORT_TYPE = os.environ.get('TRANSPORT_TYPE', 'stdio')
202
+ if TRANSPORT_TYPE == 'stdio':
203
+ mcps = {
204
+ "AirflowMCP":
205
+ {
206
+ 'command': "python",
207
+ 'args': ["-m", "airflow_mcp_hipposys.mcp_airflow"],
208
+ "transport": "stdio",
209
+ 'env': {k: v for k, v in {
210
+ 'AIRFLOW_ASSISTENT_AI_CONN': os.getenv(
211
+ 'AIRFLOW_ASSISTENT_AI_CONN'),
212
+ 'airflow_api_url': os.getenv('airflow_api_url'),
213
+ 'airflow_username': os.getenv('airflow_username'),
214
+ 'airflow_password': os.getenv('airflow_password'),
215
+ 'AIRFLOW_INSIGHTS_MODE':
216
+ os.getenv('AIRFLOW_INSIGHTS_MODE'),
217
+ 'POST_MODE': os.getenv('POST_MODE'),
218
+ 'TRANSPORT_TYPE': 'stdio'
219
+ }.items() if v is not None}
220
+ }
221
+ }
222
+ elif TRANSPORT_TYPE == 'sse':
223
+ mcps = {
224
+ "AirflowMCP": {
225
+ "url": f"http://{mcp_host}/sse",
226
+ "transport": "sse",
227
+ "headers": {"Authorization": f"""Bearer {
228
+ os.environ.get('MCP_TOKEN')}"""}
229
+ }
230
+ }
231
+ datetime_tool = Tool(
232
+ name="Datetime",
233
+ func=lambda x: datetime.now().isoformat(),
234
+ description="Returns the current datetime",
235
+ )
236
+
237
+ client = MultiServerMCPClient(mcps)
238
+ tools = await client.get_tools() + [datetime_tool]
239
+
240
+ async def stream_agent_response():
241
+ async with LLMAgent(tools=tools, md_uri=md_uri) as llm_agent:
242
+ async for chat_msg in llm_agent.astream_events(
243
+ message, user_config):
244
+ yield chat_msg.content
245
+ return stream_agent_response
@@ -0,0 +1,45 @@
1
+ import os
2
+
3
+ from contextlib import asynccontextmanager
4
+ from fastapi import FastAPI, Request
5
+ from fastapi.responses import JSONResponse
6
+ from starlette.middleware.sessions import SessionMiddleware
7
+
8
+ from app.server.chat import chat_router
9
+ from app.databases.postgres import Database
10
+ from app.utils.config import Config
11
+ from app.utils.logger import Logger
12
+
13
+
14
+ @asynccontextmanager
15
+ async def lifespan(app: FastAPI):
16
+ """Run the database setup and teardown."""
17
+ await Database.setup()
18
+ Logger().get_logger().info('Database setup complete')
19
+
20
+ yield
21
+ # Optionally add teardown code here
22
+
23
+
24
+ app = FastAPI(lifespan=lifespan)
25
+
26
+ app.add_middleware(
27
+ SessionMiddleware,
28
+ secret_key=os.environ.get('SECRET_KEY'),
29
+ https_only=Config.get_deploy_env() == 'PROD',
30
+ )
31
+
32
+
33
+ @app.middleware("http")
34
+ async def check_token_middleware(request: Request, call_next):
35
+ """Allow only requests with the correct token."""
36
+ token = request.headers.get("x-access-token")
37
+ if (Config.get_deploy_env() != 'LOCAL') \
38
+ and (token != os.environ['FAST_API_ACCESS_SECRET_TOKEN']):
39
+ return JSONResponse(
40
+ status_code=403, content={'reason': 'Invalid or missing token'})
41
+ response = await call_next(request)
42
+ return response
43
+
44
+
45
+ app.include_router(chat_router, prefix='/chat')
File without changes
@@ -0,0 +1,11 @@
1
+ import os
2
+
3
+ from app.utils.singleton import Singleton
4
+
5
+
6
+ class Config(metaclass=Singleton):
7
+ """Holds "global" configuration for the application."""
8
+ @staticmethod
9
+ def get_deploy_env() -> str:
10
+ """Get the current deployment environment."""
11
+ return os.environ.get('DEPLOY_ENV', 'prod').upper()
@@ -0,0 +1,34 @@
1
+ import sys
2
+ import logging
3
+
4
+ from app.utils.singleton import Singleton
5
+
6
+
7
+ class Logger(metaclass=Singleton):
8
+ default_config = {
9
+ 'name': 'RAG-App',
10
+ 'level': 'INFO',
11
+ 'format': '[%(asctime)s|%(name)s|%(levelname)s|%(processName)s:%(threadName)s|%(filename)s, '
12
+ 'line %(lineno)s in %(funcName)s] %(message)s',
13
+ }
14
+
15
+ def __init__(self,
16
+ config_to_use: dict = None):
17
+
18
+ config = {**Logger.default_config, **(config_to_use or {})}
19
+ self.logger = logging.getLogger(config['name'])
20
+
21
+ # Set the root logging level.
22
+ logging_lvl = getattr(logging, config['level'])
23
+ logging.root.setLevel(logging_lvl)
24
+
25
+ # Create an handler and configure it.
26
+ handler = logging.StreamHandler(sys.stdout)
27
+ handler.setLevel(logging_lvl)
28
+ handler.setFormatter(logging.Formatter(config['format']))
29
+
30
+ # Remove all other handlers in the process.
31
+ self.logger.addHandler(handler)
32
+
33
+ def get_logger(self) -> logging.Logger:
34
+ return self.logger
@@ -0,0 +1,52 @@
1
+ from datetime import datetime, timedelta
2
+
3
+
4
+ class Singleton(type):
5
+ """A metaclass that implements the singleton pattern.
6
+
7
+ This metaclass ensures that only one(*) instance of a class is created and that the instance is reused for all subsequent calls.
8
+
9
+ The instance is refreshed every `MAX_INSTANCE_TTL` seconds.
10
+
11
+ Usage:
12
+ ```python
13
+ >>> class MyClass(metaclass=Singleton):
14
+ ... pass
15
+
16
+ >>> a = MyClass()
17
+ >>> b = MyClass()
18
+ >>> a is b
19
+ True
20
+ >>> sleep(MyClass.MAX_INSTANCE_TTL + 1)
21
+ >>> c = MyClass()
22
+ >>> a is c
23
+ False
24
+ ```
25
+ """
26
+
27
+ _instances = {}
28
+ _creation_time = {}
29
+
30
+ MAX_INSTANCE_TTL = timedelta(seconds=5 * 60) # 5 minutes
31
+
32
+ def __create_instance(cls, *args, **kwargs):
33
+ """Create a new instance of this class."""
34
+
35
+ cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
36
+ cls._creation_time[cls] = datetime.now()
37
+
38
+ def __call__(cls, *args, force_recreate=False, **kwargs):
39
+ """Create a new instance of this class if it does not exist or if it is no longer valid.
40
+
41
+ If the `force_recreate` parameter is set to `True`, a new instance will be created regardless of the validity of the existing instance.
42
+ """
43
+
44
+ # Create a new instance if it does not exist or if `force_recreate` is set.
45
+ if force_recreate or (cls not in cls._instances):
46
+ cls.__create_instance(*args, **kwargs)
47
+ # Create a new instance if the existing instance is no longer valid.
48
+ elif (cls not in cls._creation_time) or (cls._creation_time[cls] + cls.MAX_INSTANCE_TTL < datetime.now()):
49
+ cls.__create_instance(*args, **kwargs)
50
+
51
+ # Return the existing instance.
52
+ return cls._instances[cls]
@@ -18,15 +18,7 @@
18
18
  </div>
19
19
  <div class="message-content">
20
20
  <div class="message-text">
21
- Hello! I'm your Airflow AI assistant. I can help you with:
22
- <ul>
23
- <li>DAG management and troubleshooting</li>
24
- <li>Task configuration and dependencies</li>
25
- <li>Airflow best practices</li>
26
- <li>Connection and variable management</li>
27
- <li>General workflow questions</li>
28
- </ul>
29
- What would you like to know?
21
+ Hello! I'm your Airflow AI assistant.
30
22
  </div>
31
23
  <div class="message-time">Just now</div>
32
24
  </div>
@@ -0,0 +1,37 @@
1
+ Metadata-Version: 2.3
2
+ Name: airflow-chat
3
+ Version: 0.1.0a2
4
+ Summary: An Apache Airflow plugin that enables AI-powered chat interactions with your Airflow instance through MCP integration and an intuitive UI.
5
+ License: MIT
6
+ Author: Hipposys
7
+ Requires-Python: >=3.10,<3.13
8
+ Classifier: License :: OSI Approved :: MIT License
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Programming Language :: Python :: 3.10
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Requires-Dist: SQLAlchemy (>=1.4.0)
14
+ Requires-Dist: aiosqlite (>=0.21,<0.22)
15
+ Requires-Dist: airflow-mcp-hipposys (==0.1.0a7)
16
+ Requires-Dist: apache-airflow (>=2.4.0,<3.0.0)
17
+ Requires-Dist: fastapi[standard] (>=0.115,<0.116)
18
+ Requires-Dist: httpx (>=0.27,<0.28)
19
+ Requires-Dist: itsdangerous (>=2.2,<3.0)
20
+ Requires-Dist: langchain (>=0.3,<0.4)
21
+ Requires-Dist: langchain-anthropic (>=0.3,<0.4)
22
+ Requires-Dist: langchain-aws (>=0.2,<0.3)
23
+ Requires-Dist: langchain-community (>=0.3,<0.4)
24
+ Requires-Dist: langchain-core (>=0.3,<0.4)
25
+ Requires-Dist: langchain-mcp-adapters (==0.1.1)
26
+ Requires-Dist: langchain-openai (>=0.3,<0.4)
27
+ Requires-Dist: langchain-text-splitters (>=0.3,<0.4)
28
+ Requires-Dist: langgraph (>=0.5,<0.6)
29
+ Requires-Dist: langgraph-checkpoint-postgres (>=2.0,<3.0)
30
+ Requires-Dist: langgraph-checkpoint-sqlite (>=2.0,<3.0)
31
+ Requires-Dist: mcp[cli] (>=1.9,<2.0)
32
+ Requires-Dist: psycopg[binary] (>=3.2,<4.0)
33
+ Requires-Dist: pydantic-settings (>=2.10,<3.0)
34
+ Project-URL: homepage, https://github.com/hipposys-ltd/airflow-schedule-insights
35
+ Description-Content-Type: text/markdown
36
+
37
+ # Airflow Chat Plugin
@@ -0,0 +1,26 @@
1
+ airflow_chat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
2
+ airflow_chat/plugins/README.md,sha256=Hwmn7ohaPBIoeww2xQ_aSeB6Fy_EgkHXA3LYFi8beIk,21
3
+ airflow_chat/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
+ airflow_chat/plugins/airflow_chat.py,sha256=aGtciZJigmLn0rdG27027x8MpgzRrTfvwobMALm4eX8,19636
5
+ airflow_chat/plugins/app/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
6
+ airflow_chat/plugins/app/databases/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
7
+ airflow_chat/plugins/app/databases/postgres.py,sha256=0UYnV8Z4F33CubrFwkzuXoOhPi11IvhIuu1L1WsZKW8,1287
8
+ airflow_chat/plugins/app/models/__init__.py,sha256=KmN9_vswQd53wm7I38eg5ps0E23YRDtgh_Y25rYLMCM,482
9
+ airflow_chat/plugins/app/models/inference/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
10
+ airflow_chat/plugins/app/models/inference/antropic_model.py,sha256=8kC7cXjzGiy77eDnetq7h9SEfUfIEid4T6_2Hd8E5Hc,509
11
+ airflow_chat/plugins/app/models/inference/bedrock_model.py,sha256=8ebK4oNwqIqhEIBSkVqzNcfQpU2-RtW9KJ_Lxd7yYqA,595
12
+ airflow_chat/plugins/app/models/inference/openai_model.py,sha256=8iNDgJRXBDP9HR9rMFwYKg-paUUQK2NBJw25ZY0EFWk,534
13
+ airflow_chat/plugins/app/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
14
+ airflow_chat/plugins/app/server/chat.py,sha256=Fvgq_OiILqOApCIFqEH7HaXMrnxMoQ99wohz3Y2fK5M,1054
15
+ airflow_chat/plugins/app/server/llm.py,sha256=QH2q8hvrEHf-EyPIcdCtwQO2EzlQx2fsLKLjAX2fNeE,8725
16
+ airflow_chat/plugins/app/server/main.py,sha256=x3FRFWA8TB4ShD-XIo6FQLT-AW_e5H8MSq_Ncc3FkkE,1288
17
+ airflow_chat/plugins/app/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
18
+ airflow_chat/plugins/app/utils/config.py,sha256=l0FfBlneeTdI69X9tGXFblMH8EQPVVrJdCoH9Ukxe7Y,315
19
+ airflow_chat/plugins/app/utils/logger.py,sha256=EC_bekFeVJHVKwGABffDhXSrpHqgm8r00Q27WETuo90,1056
20
+ airflow_chat/plugins/app/utils/singleton.py,sha256=bQCsVc8XPpM6etUbbg2_bKKReCjLm0bFh6OgXmoBO0I,1789
21
+ airflow_chat/plugins/templates/chat_interface.html,sha256=ZnMeIeaknXf7ZQI7PRV6fyxyzMsOMUz8KmMaX4k_mJo,18719
22
+ airflow_chat-0.1.0a2.dist-info/LICENSE,sha256=XFYCNJc3ykWUpIIuB6uHwgnWKWUm3iez5vxgFF352as,1069
23
+ airflow_chat-0.1.0a2.dist-info/METADATA,sha256=P_Mr3zE-lquNCJGW718GgxNa7QI88Q6lpQErf452bVo,1605
24
+ airflow_chat-0.1.0a2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
25
+ airflow_chat-0.1.0a2.dist-info/entry_points.txt,sha256=KzmeXDfhihgaLTQgh4vOz2ts7obTzZCg8CNgkHX1zaU,91
26
+ airflow_chat-0.1.0a2.dist-info/RECORD,,
@@ -1,19 +0,0 @@
1
- Metadata-Version: 2.3
2
- Name: airflow-chat
3
- Version: 0.1.0a1
4
- Summary: An Apache Airflow plugin that enables AI-powered chat interactions with your Airflow instance through MCP integration and an intuitive UI.
5
- License: MIT
6
- Author: Hipposys
7
- Requires-Python: >=3.10,<4.0
8
- Classifier: License :: OSI Approved :: MIT License
9
- Classifier: Programming Language :: Python :: 3
10
- Classifier: Programming Language :: Python :: 3.10
11
- Classifier: Programming Language :: Python :: 3.11
12
- Classifier: Programming Language :: Python :: 3.12
13
- Classifier: Programming Language :: Python :: 3.13
14
- Requires-Dist: SQLAlchemy (>=1.4.0)
15
- Requires-Dist: apache-airflow (>=2.4.0,<3.0.0)
16
- Project-URL: homepage, https://github.com/hipposys-ltd/airflow-schedule-insights
17
- Description-Content-Type: text/markdown
18
-
19
- # Airflow Chat Plugin
@@ -1,10 +0,0 @@
1
- airflow_chat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
2
- airflow_chat/plugins/README.md,sha256=Hwmn7ohaPBIoeww2xQ_aSeB6Fy_EgkHXA3LYFi8beIk,21
3
- airflow_chat/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
- airflow_chat/plugins/airflow_chat.py,sha256=bpygj75yfRYU7a6NNndnfSFqpYOzHWdwk0JlosdbJ-Q,17000
5
- airflow_chat/plugins/templates/chat_interface.html,sha256=cV-_dDSKHUo1T8IQJ0iDCLDS-Dt2oJk7rHcecC8CeIc,19087
6
- airflow_chat-0.1.0a1.dist-info/LICENSE,sha256=XFYCNJc3ykWUpIIuB6uHwgnWKWUm3iez5vxgFF352as,1069
7
- airflow_chat-0.1.0a1.dist-info/METADATA,sha256=SaMdd-O9JaTHBhJMxQKqbKwc9ctZDAd8_fPJ_7QeT8Y,795
8
- airflow_chat-0.1.0a1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
9
- airflow_chat-0.1.0a1.dist-info/entry_points.txt,sha256=KzmeXDfhihgaLTQgh4vOz2ts7obTzZCg8CNgkHX1zaU,91
10
- airflow_chat-0.1.0a1.dist-info/RECORD,,