Refactor main.py for improved organization and clarity
This commit reorganizes the import statements in main.py, enhancing the structure and readability of the code. It also includes minor formatting adjustments for consistency. The changes contribute to a cleaner codebase, facilitating easier maintenance and future development.
This commit is contained in:
parent
a10b35f631
commit
4be48270b7
1 changed files with 107 additions and 106 deletions
213
src/main.py
213
src/main.py
|
|
@ -31,15 +31,16 @@ from auth_middleware import optional_auth, require_auth
|
||||||
|
|
||||||
# Configuration and setup
|
# Configuration and setup
|
||||||
from config.settings import INDEX_BODY, INDEX_NAME, SESSION_SECRET, clients
|
from config.settings import INDEX_BODY, INDEX_NAME, SESSION_SECRET, clients
|
||||||
|
|
||||||
# Existing services
|
# Existing services
|
||||||
from connectors.service import ConnectorService
|
from connectors.service import ConnectorService
|
||||||
from services.auth_service import AuthService
|
from services.auth_service import AuthService
|
||||||
from services.chat_service import ChatService
|
from services.chat_service import ChatService
|
||||||
|
from services.document_service import DocumentService
|
||||||
|
from services.knowledge_filter_service import KnowledgeFilterService
|
||||||
|
|
||||||
# Services
|
# Services
|
||||||
from services.langflow_file_service import LangflowFileService
|
from services.langflow_file_service import LangflowFileService
|
||||||
from services.document_service import DocumentService
|
|
||||||
from services.knowledge_filter_service import KnowledgeFilterService
|
|
||||||
from services.monitor_service import MonitorService
|
from services.monitor_service import MonitorService
|
||||||
from services.search_service import SearchService
|
from services.search_service import SearchService
|
||||||
from services.task_service import TaskService
|
from services.task_service import TaskService
|
||||||
|
|
@ -53,7 +54,7 @@ async def wait_for_opensearch():
|
||||||
"""Wait for OpenSearch to be ready with retries"""
|
"""Wait for OpenSearch to be ready with retries"""
|
||||||
max_retries = 30
|
max_retries = 30
|
||||||
retry_delay = 2
|
retry_delay = 2
|
||||||
|
|
||||||
for attempt in range(max_retries):
|
for attempt in range(max_retries):
|
||||||
try:
|
try:
|
||||||
await clients.opensearch.info()
|
await clients.opensearch.info()
|
||||||
|
|
@ -74,11 +75,11 @@ async def configure_alerting_security():
|
||||||
alerting_settings = {
|
alerting_settings = {
|
||||||
"persistent": {
|
"persistent": {
|
||||||
"plugins.alerting.filter_by_backend_roles": "false",
|
"plugins.alerting.filter_by_backend_roles": "false",
|
||||||
"opendistro.alerting.filter_by_backend_roles": "false",
|
"opendistro.alerting.filter_by_backend_roles": "false",
|
||||||
"opensearch.notifications.general.filter_by_backend_roles": "false"
|
"opensearch.notifications.general.filter_by_backend_roles": "false"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
# Use admin client (clients.opensearch uses admin credentials)
|
# Use admin client (clients.opensearch uses admin credentials)
|
||||||
response = await clients.opensearch.cluster.put_settings(body=alerting_settings)
|
response = await clients.opensearch.cluster.put_settings(body=alerting_settings)
|
||||||
print("Alerting security settings configured successfully")
|
print("Alerting security settings configured successfully")
|
||||||
|
|
@ -90,14 +91,14 @@ async def configure_alerting_security():
|
||||||
async def init_index():
|
async def init_index():
|
||||||
"""Initialize OpenSearch index and security roles"""
|
"""Initialize OpenSearch index and security roles"""
|
||||||
await wait_for_opensearch()
|
await wait_for_opensearch()
|
||||||
|
|
||||||
# Create documents index
|
# Create documents index
|
||||||
if not await clients.opensearch.indices.exists(index=INDEX_NAME):
|
if not await clients.opensearch.indices.exists(index=INDEX_NAME):
|
||||||
await clients.opensearch.indices.create(index=INDEX_NAME, body=INDEX_BODY)
|
await clients.opensearch.indices.create(index=INDEX_NAME, body=INDEX_BODY)
|
||||||
print(f"Created index '{INDEX_NAME}'")
|
print(f"Created index '{INDEX_NAME}'")
|
||||||
else:
|
else:
|
||||||
print(f"Index '{INDEX_NAME}' already exists, skipping creation.")
|
print(f"Index '{INDEX_NAME}' already exists, skipping creation.")
|
||||||
|
|
||||||
# Create knowledge filters index
|
# Create knowledge filters index
|
||||||
knowledge_filter_index_name = "knowledge_filters"
|
knowledge_filter_index_name = "knowledge_filters"
|
||||||
knowledge_filter_index_body = {
|
knowledge_filter_index_body = {
|
||||||
|
|
@ -116,13 +117,13 @@ async def init_index():
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if not await clients.opensearch.indices.exists(index=knowledge_filter_index_name):
|
if not await clients.opensearch.indices.exists(index=knowledge_filter_index_name):
|
||||||
await clients.opensearch.indices.create(index=knowledge_filter_index_name, body=knowledge_filter_index_body)
|
await clients.opensearch.indices.create(index=knowledge_filter_index_name, body=knowledge_filter_index_body)
|
||||||
print(f"Created index '{knowledge_filter_index_name}'")
|
print(f"Created index '{knowledge_filter_index_name}'")
|
||||||
else:
|
else:
|
||||||
print(f"Index '{knowledge_filter_index_name}' already exists, skipping creation.")
|
print(f"Index '{knowledge_filter_index_name}' already exists, skipping creation.")
|
||||||
|
|
||||||
# Configure alerting plugin security settings
|
# Configure alerting plugin security settings
|
||||||
await configure_alerting_security()
|
await configure_alerting_security()
|
||||||
|
|
||||||
|
|
@ -131,10 +132,10 @@ def generate_jwt_keys():
|
||||||
keys_dir = "keys"
|
keys_dir = "keys"
|
||||||
private_key_path = os.path.join(keys_dir, "private_key.pem")
|
private_key_path = os.path.join(keys_dir, "private_key.pem")
|
||||||
public_key_path = os.path.join(keys_dir, "public_key.pem")
|
public_key_path = os.path.join(keys_dir, "public_key.pem")
|
||||||
|
|
||||||
# Create keys directory if it doesn't exist
|
# Create keys directory if it doesn't exist
|
||||||
os.makedirs(keys_dir, exist_ok=True)
|
os.makedirs(keys_dir, exist_ok=True)
|
||||||
|
|
||||||
# Generate keys if they don't exist
|
# Generate keys if they don't exist
|
||||||
if not os.path.exists(private_key_path):
|
if not os.path.exists(private_key_path):
|
||||||
try:
|
try:
|
||||||
|
|
@ -142,12 +143,12 @@ def generate_jwt_keys():
|
||||||
subprocess.run([
|
subprocess.run([
|
||||||
"openssl", "genrsa", "-out", private_key_path, "2048"
|
"openssl", "genrsa", "-out", private_key_path, "2048"
|
||||||
], check=True, capture_output=True)
|
], check=True, capture_output=True)
|
||||||
|
|
||||||
# Generate public key
|
# Generate public key
|
||||||
subprocess.run([
|
subprocess.run([
|
||||||
"openssl", "rsa", "-in", private_key_path, "-pubout", "-out", public_key_path
|
"openssl", "rsa", "-in", private_key_path, "-pubout", "-out", public_key_path
|
||||||
], check=True, capture_output=True)
|
], check=True, capture_output=True)
|
||||||
|
|
||||||
print("Generated RSA keys for JWT signing")
|
print("Generated RSA keys for JWT signing")
|
||||||
except subprocess.CalledProcessError as e:
|
except subprocess.CalledProcessError as e:
|
||||||
print(f"Failed to generate RSA keys: {e}")
|
print(f"Failed to generate RSA keys: {e}")
|
||||||
|
|
@ -163,19 +164,19 @@ async def init_index_when_ready():
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"OpenSearch index initialization failed: {e}")
|
print(f"OpenSearch index initialization failed: {e}")
|
||||||
print("OIDC endpoints will still work, but document operations may fail until OpenSearch is ready")
|
print("OIDC endpoints will still work, but document operations may fail until OpenSearch is ready")
|
||||||
|
|
||||||
|
|
||||||
async def initialize_services():
|
async def initialize_services():
|
||||||
"""Initialize all services and their dependencies"""
|
"""Initialize all services and their dependencies"""
|
||||||
# Generate JWT keys if they don't exist
|
# Generate JWT keys if they don't exist
|
||||||
generate_jwt_keys()
|
generate_jwt_keys()
|
||||||
|
|
||||||
# Initialize clients (now async to generate Langflow API key)
|
# Initialize clients (now async to generate Langflow API key)
|
||||||
await clients.initialize()
|
await clients.initialize()
|
||||||
|
|
||||||
# Initialize session manager
|
# Initialize session manager
|
||||||
session_manager = SessionManager(SESSION_SECRET)
|
session_manager = SessionManager(SESSION_SECRET)
|
||||||
|
|
||||||
# Initialize services
|
# Initialize services
|
||||||
document_service = DocumentService(session_manager=session_manager)
|
document_service = DocumentService(session_manager=session_manager)
|
||||||
search_service = SearchService(session_manager)
|
search_service = SearchService(session_manager)
|
||||||
|
|
@ -183,10 +184,10 @@ async def initialize_services():
|
||||||
chat_service = ChatService()
|
chat_service = ChatService()
|
||||||
knowledge_filter_service = KnowledgeFilterService(session_manager)
|
knowledge_filter_service = KnowledgeFilterService(session_manager)
|
||||||
monitor_service = MonitorService(session_manager)
|
monitor_service = MonitorService(session_manager)
|
||||||
|
|
||||||
# Set process pool for document service
|
# Set process pool for document service
|
||||||
document_service.process_pool = process_pool
|
document_service.process_pool = process_pool
|
||||||
|
|
||||||
# Initialize connector service
|
# Initialize connector service
|
||||||
connector_service = ConnectorService(
|
connector_service = ConnectorService(
|
||||||
patched_async_client=clients.patched_async_client,
|
patched_async_client=clients.patched_async_client,
|
||||||
|
|
@ -196,10 +197,10 @@ async def initialize_services():
|
||||||
task_service=task_service,
|
task_service=task_service,
|
||||||
session_manager=session_manager
|
session_manager=session_manager
|
||||||
)
|
)
|
||||||
|
|
||||||
# Initialize auth service
|
# Initialize auth service
|
||||||
auth_service = AuthService(session_manager, connector_service)
|
auth_service = AuthService(session_manager, connector_service)
|
||||||
|
|
||||||
# Load persisted connector connections at startup so webhooks and syncs
|
# Load persisted connector connections at startup so webhooks and syncs
|
||||||
# can resolve existing subscriptions immediately after server boot
|
# can resolve existing subscriptions immediately after server boot
|
||||||
# Skip in no-auth mode since connectors require OAuth
|
# Skip in no-auth mode since connectors require OAuth
|
||||||
|
|
@ -213,7 +214,7 @@ async def initialize_services():
|
||||||
print(f"[WARNING] Failed to load persisted connections on startup: {e}")
|
print(f"[WARNING] Failed to load persisted connections on startup: {e}")
|
||||||
else:
|
else:
|
||||||
print("[CONNECTORS] Skipping connection loading in no-auth mode")
|
print("[CONNECTORS] Skipping connection loading in no-auth mode")
|
||||||
|
|
||||||
# New: Langflow file service
|
# New: Langflow file service
|
||||||
|
|
||||||
langflow_file_service = LangflowFileService()
|
langflow_file_service = LangflowFileService()
|
||||||
|
|
@ -234,13 +235,13 @@ async def initialize_services():
|
||||||
async def create_app():
|
async def create_app():
|
||||||
"""Create and configure the Starlette application"""
|
"""Create and configure the Starlette application"""
|
||||||
services = await initialize_services()
|
services = await initialize_services()
|
||||||
|
|
||||||
# Create route handlers with service dependencies injected
|
# Create route handlers with service dependencies injected
|
||||||
routes = [
|
routes = [
|
||||||
# Upload endpoints
|
# Upload endpoints
|
||||||
Route("/upload",
|
Route("/upload",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(upload.upload,
|
partial(upload.upload,
|
||||||
document_service=services['document_service'],
|
document_service=services['document_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
@ -266,15 +267,15 @@ async def create_app():
|
||||||
langflow_file_service=services['langflow_file_service'],
|
langflow_file_service=services['langflow_file_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["DELETE"]),
|
), methods=["DELETE"]),
|
||||||
|
|
||||||
Route("/upload_context",
|
Route("/upload_context",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(upload.upload_context,
|
partial(upload.upload_context,
|
||||||
document_service=services['document_service'],
|
document_service=services['document_service'],
|
||||||
chat_service=services['chat_service'],
|
chat_service=services['chat_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/upload_path",
|
Route("/upload_path",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(upload.upload_path,
|
partial(upload.upload_path,
|
||||||
|
|
@ -294,227 +295,227 @@ async def create_app():
|
||||||
task_service=services['task_service'],
|
task_service=services['task_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/tasks/{task_id}",
|
Route("/tasks/{task_id}",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(tasks.task_status,
|
partial(tasks.task_status,
|
||||||
task_service=services['task_service'],
|
task_service=services['task_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/tasks",
|
Route("/tasks",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(tasks.all_tasks,
|
partial(tasks.all_tasks,
|
||||||
task_service=services['task_service'],
|
task_service=services['task_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/tasks/{task_id}/cancel",
|
Route("/tasks/{task_id}/cancel",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(tasks.cancel_task,
|
partial(tasks.cancel_task,
|
||||||
task_service=services['task_service'],
|
task_service=services['task_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
# Search endpoint
|
# Search endpoint
|
||||||
Route("/search",
|
Route("/search",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(search.search,
|
partial(search.search,
|
||||||
search_service=services['search_service'],
|
search_service=services['search_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
# Knowledge Filter endpoints
|
# Knowledge Filter endpoints
|
||||||
Route("/knowledge-filter",
|
Route("/knowledge-filter",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.create_knowledge_filter,
|
partial(knowledge_filter.create_knowledge_filter,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/knowledge-filter/search",
|
Route("/knowledge-filter/search",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.search_knowledge_filters,
|
partial(knowledge_filter.search_knowledge_filters,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/knowledge-filter/{filter_id}",
|
Route("/knowledge-filter/{filter_id}",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.get_knowledge_filter,
|
partial(knowledge_filter.get_knowledge_filter,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/knowledge-filter/{filter_id}",
|
Route("/knowledge-filter/{filter_id}",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.update_knowledge_filter,
|
partial(knowledge_filter.update_knowledge_filter,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["PUT"]),
|
), methods=["PUT"]),
|
||||||
|
|
||||||
Route("/knowledge-filter/{filter_id}",
|
Route("/knowledge-filter/{filter_id}",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.delete_knowledge_filter,
|
partial(knowledge_filter.delete_knowledge_filter,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["DELETE"]),
|
), methods=["DELETE"]),
|
||||||
|
|
||||||
# Knowledge Filter Subscription endpoints
|
# Knowledge Filter Subscription endpoints
|
||||||
Route("/knowledge-filter/{filter_id}/subscribe",
|
Route("/knowledge-filter/{filter_id}/subscribe",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.subscribe_to_knowledge_filter,
|
partial(knowledge_filter.subscribe_to_knowledge_filter,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
monitor_service=services['monitor_service'],
|
monitor_service=services['monitor_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/knowledge-filter/{filter_id}/subscriptions",
|
Route("/knowledge-filter/{filter_id}/subscriptions",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.list_knowledge_filter_subscriptions,
|
partial(knowledge_filter.list_knowledge_filter_subscriptions,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/knowledge-filter/{filter_id}/subscribe/{subscription_id}",
|
Route("/knowledge-filter/{filter_id}/subscribe/{subscription_id}",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(knowledge_filter.cancel_knowledge_filter_subscription,
|
partial(knowledge_filter.cancel_knowledge_filter_subscription,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
monitor_service=services['monitor_service'],
|
monitor_service=services['monitor_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["DELETE"]),
|
), methods=["DELETE"]),
|
||||||
|
|
||||||
# Knowledge Filter Webhook endpoint (no auth required - called by OpenSearch)
|
# Knowledge Filter Webhook endpoint (no auth required - called by OpenSearch)
|
||||||
Route("/knowledge-filter/{filter_id}/webhook/{subscription_id}",
|
Route("/knowledge-filter/{filter_id}/webhook/{subscription_id}",
|
||||||
partial(knowledge_filter.knowledge_filter_webhook,
|
partial(knowledge_filter.knowledge_filter_webhook,
|
||||||
knowledge_filter_service=services['knowledge_filter_service'],
|
knowledge_filter_service=services['knowledge_filter_service'],
|
||||||
session_manager=services['session_manager']),
|
session_manager=services['session_manager']),
|
||||||
methods=["POST"]),
|
methods=["POST"]),
|
||||||
|
|
||||||
# Chat endpoints
|
# Chat endpoints
|
||||||
Route("/chat",
|
Route("/chat",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(chat.chat_endpoint,
|
partial(chat.chat_endpoint,
|
||||||
chat_service=services['chat_service'],
|
chat_service=services['chat_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/langflow",
|
Route("/langflow",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(chat.langflow_endpoint,
|
partial(chat.langflow_endpoint,
|
||||||
chat_service=services['chat_service'],
|
chat_service=services['chat_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
# Chat history endpoints
|
# Chat history endpoints
|
||||||
Route("/chat/history",
|
Route("/chat/history",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(chat.chat_history_endpoint,
|
partial(chat.chat_history_endpoint,
|
||||||
chat_service=services['chat_service'],
|
chat_service=services['chat_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/langflow/history",
|
Route("/langflow/history",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(chat.langflow_history_endpoint,
|
partial(chat.langflow_history_endpoint,
|
||||||
chat_service=services['chat_service'],
|
chat_service=services['chat_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
# Authentication endpoints
|
# Authentication endpoints
|
||||||
Route("/auth/init",
|
Route("/auth/init",
|
||||||
optional_auth(services['session_manager'])(
|
optional_auth(services['session_manager'])(
|
||||||
partial(auth.auth_init,
|
partial(auth.auth_init,
|
||||||
auth_service=services['auth_service'],
|
auth_service=services['auth_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/auth/callback",
|
Route("/auth/callback",
|
||||||
partial(auth.auth_callback,
|
partial(auth.auth_callback,
|
||||||
auth_service=services['auth_service'],
|
auth_service=services['auth_service'],
|
||||||
session_manager=services['session_manager']),
|
session_manager=services['session_manager']),
|
||||||
methods=["POST"]),
|
methods=["POST"]),
|
||||||
|
|
||||||
Route("/auth/me",
|
Route("/auth/me",
|
||||||
optional_auth(services['session_manager'])(
|
optional_auth(services['session_manager'])(
|
||||||
partial(auth.auth_me,
|
partial(auth.auth_me,
|
||||||
auth_service=services['auth_service'],
|
auth_service=services['auth_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/auth/logout",
|
Route("/auth/logout",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(auth.auth_logout,
|
partial(auth.auth_logout,
|
||||||
auth_service=services['auth_service'],
|
auth_service=services['auth_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
# Connector endpoints
|
# Connector endpoints
|
||||||
Route("/connectors",
|
Route("/connectors",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(connectors.list_connectors,
|
partial(connectors.list_connectors,
|
||||||
connector_service=services['connector_service'],
|
connector_service=services['connector_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/connectors/{connector_type}/sync",
|
Route("/connectors/{connector_type}/sync",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(connectors.connector_sync,
|
partial(connectors.connector_sync,
|
||||||
connector_service=services['connector_service'],
|
connector_service=services['connector_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["POST"]),
|
), methods=["POST"]),
|
||||||
|
|
||||||
Route("/connectors/{connector_type}/status",
|
Route("/connectors/{connector_type}/status",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(connectors.connector_status,
|
partial(connectors.connector_status,
|
||||||
connector_service=services['connector_service'],
|
connector_service=services['connector_service'],
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
|
|
||||||
Route("/connectors/{connector_type}/webhook",
|
Route("/connectors/{connector_type}/webhook",
|
||||||
partial(connectors.connector_webhook,
|
partial(connectors.connector_webhook,
|
||||||
connector_service=services['connector_service'],
|
connector_service=services['connector_service'],
|
||||||
session_manager=services['session_manager']),
|
session_manager=services['session_manager']),
|
||||||
methods=["POST", "GET"]),
|
methods=["POST", "GET"]),
|
||||||
|
|
||||||
# OIDC endpoints
|
# OIDC endpoints
|
||||||
Route("/.well-known/openid-configuration",
|
Route("/.well-known/openid-configuration",
|
||||||
partial(oidc.oidc_discovery,
|
partial(oidc.oidc_discovery,
|
||||||
session_manager=services['session_manager']),
|
session_manager=services['session_manager']),
|
||||||
methods=["GET"]),
|
methods=["GET"]),
|
||||||
|
|
||||||
Route("/auth/jwks",
|
Route("/auth/jwks",
|
||||||
partial(oidc.jwks_endpoint,
|
partial(oidc.jwks_endpoint,
|
||||||
session_manager=services['session_manager']),
|
session_manager=services['session_manager']),
|
||||||
methods=["GET"]),
|
methods=["GET"]),
|
||||||
|
|
||||||
Route("/auth/introspect",
|
Route("/auth/introspect",
|
||||||
partial(oidc.token_introspection,
|
partial(oidc.token_introspection,
|
||||||
session_manager=services['session_manager']),
|
session_manager=services['session_manager']),
|
||||||
methods=["POST"]),
|
methods=["POST"]),
|
||||||
|
|
||||||
# Settings endpoint
|
# Settings endpoint
|
||||||
Route("/settings",
|
Route("/settings",
|
||||||
require_auth(services['session_manager'])(
|
require_auth(services['session_manager'])(
|
||||||
partial(settings.get_settings,
|
partial(settings.get_settings,
|
||||||
session_manager=services['session_manager'])
|
session_manager=services['session_manager'])
|
||||||
), methods=["GET"]),
|
), methods=["GET"]),
|
||||||
]
|
]
|
||||||
|
|
||||||
app = Starlette(debug=True, routes=routes)
|
app = Starlette(debug=True, routes=routes)
|
||||||
app.state.services = services # Store services for cleanup
|
app.state.services = services # Store services for cleanup
|
||||||
|
|
||||||
# Add startup event handler
|
# Add startup event handler
|
||||||
@app.on_event("startup")
|
@app.on_event("startup")
|
||||||
async def startup_event():
|
async def startup_event():
|
||||||
# Start index initialization in background to avoid blocking OIDC endpoints
|
# Start index initialization in background to avoid blocking OIDC endpoints
|
||||||
asyncio.create_task(init_index_when_ready())
|
asyncio.create_task(init_index_when_ready())
|
||||||
|
|
||||||
# Add shutdown event handler
|
# Add shutdown event handler
|
||||||
@app.on_event("shutdown")
|
@app.on_event("shutdown")
|
||||||
async def shutdown_event():
|
async def shutdown_event():
|
||||||
await cleanup_subscriptions_proper(services)
|
await cleanup_subscriptions_proper(services)
|
||||||
|
|
||||||
return app
|
return app
|
||||||
|
|
||||||
async def startup():
|
async def startup():
|
||||||
|
|
@ -533,15 +534,15 @@ def cleanup():
|
||||||
async def cleanup_subscriptions_proper(services):
|
async def cleanup_subscriptions_proper(services):
|
||||||
"""Cancel all active webhook subscriptions"""
|
"""Cancel all active webhook subscriptions"""
|
||||||
print("[CLEANUP] Cancelling active webhook subscriptions...")
|
print("[CLEANUP] Cancelling active webhook subscriptions...")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
connector_service = services['connector_service']
|
connector_service = services['connector_service']
|
||||||
await connector_service.connection_manager.load_connections()
|
await connector_service.connection_manager.load_connections()
|
||||||
|
|
||||||
# Get all active connections with webhook subscriptions
|
# Get all active connections with webhook subscriptions
|
||||||
all_connections = await connector_service.connection_manager.list_connections()
|
all_connections = await connector_service.connection_manager.list_connections()
|
||||||
active_connections = [c for c in all_connections if c.is_active and c.config.get('webhook_channel_id')]
|
active_connections = [c for c in all_connections if c.is_active and c.config.get('webhook_channel_id')]
|
||||||
|
|
||||||
for connection in active_connections:
|
for connection in active_connections:
|
||||||
try:
|
try:
|
||||||
print(f"[CLEANUP] Cancelling subscription for connection {connection.connection_id}")
|
print(f"[CLEANUP] Cancelling subscription for connection {connection.connection_id}")
|
||||||
|
|
@ -552,22 +553,22 @@ async def cleanup_subscriptions_proper(services):
|
||||||
print(f"[CLEANUP] Cancelled subscription {subscription_id}")
|
print(f"[CLEANUP] Cancelled subscription {subscription_id}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"[ERROR] Failed to cancel subscription for {connection.connection_id}: {e}")
|
print(f"[ERROR] Failed to cancel subscription for {connection.connection_id}: {e}")
|
||||||
|
|
||||||
print(f"[CLEANUP] Finished cancelling {len(active_connections)} subscriptions")
|
print(f"[CLEANUP] Finished cancelling {len(active_connections)} subscriptions")
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"[ERROR] Failed to cleanup subscriptions: {e}")
|
print(f"[ERROR] Failed to cleanup subscriptions: {e}")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
import uvicorn
|
import uvicorn
|
||||||
|
|
||||||
# Register cleanup function
|
# Register cleanup function
|
||||||
atexit.register(cleanup)
|
atexit.register(cleanup)
|
||||||
|
|
||||||
# Create app asynchronously
|
# Create app asynchronously
|
||||||
app = asyncio.run(create_app())
|
app = asyncio.run(create_app())
|
||||||
|
|
||||||
# Run the server (startup tasks now handled by Starlette startup event)
|
# Run the server (startup tasks now handled by Starlette startup event)
|
||||||
uvicorn.run(
|
uvicorn.run(
|
||||||
app,
|
app,
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue