b8e9a80d93
* wait I forgot to comit locally * cp the entire core directory and then rm the .git subdir
125 lines
3.8 KiB
Python
125 lines
3.8 KiB
Python
"""
|
|
Integration module for SQLAlchemy synchronous operation instrumentation.
|
|
|
|
This module provides easy integration with the existing Letta application,
|
|
including automatic discovery of database engines and integration with
|
|
the existing OpenTelemetry setup.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Any, Dict, Optional
|
|
|
|
from letta.otel.sqlalchemy_instrumentation import (
|
|
configure_instrumentation,
|
|
get_instrumentation_config,
|
|
is_instrumentation_active,
|
|
setup_sqlalchemy_sync_instrumentation,
|
|
teardown_sqlalchemy_sync_instrumentation,
|
|
)
|
|
from letta.server.db import db_registry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def setup_letta_db_instrumentation(
|
|
enable_joined_monitoring: bool = True,
|
|
sql_truncate_length: int = 1000,
|
|
additional_config: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
"""
|
|
Set up SQLAlchemy instrumentation for Letta application.
|
|
|
|
Args:
|
|
enable_joined_monitoring: Whether to monitor joined loading operations
|
|
sql_truncate_length: Maximum length of SQL statements in traces
|
|
additional_config: Additional configuration options
|
|
"""
|
|
if is_instrumentation_active():
|
|
logger.info("SQLAlchemy instrumentation already active")
|
|
return
|
|
|
|
# Build configuration
|
|
config = {
|
|
"enabled": True,
|
|
"monitor_joined_loading": enable_joined_monitoring,
|
|
"sql_truncate_length": sql_truncate_length,
|
|
"log_instrumentation_errors": True,
|
|
}
|
|
|
|
if additional_config:
|
|
config.update(additional_config)
|
|
|
|
# Get engines from db_registry
|
|
engines = []
|
|
try:
|
|
if hasattr(db_registry, "_async_engines"):
|
|
engines.extend(db_registry._async_engines.values())
|
|
if hasattr(db_registry, "_sync_engines"):
|
|
engines.extend(db_registry._sync_engines.values())
|
|
except Exception as e:
|
|
logger.warning(f"Could not discover engines from db_registry: {e}")
|
|
|
|
if not engines:
|
|
logger.warning("No SQLAlchemy engines found for instrumentation")
|
|
return
|
|
|
|
try:
|
|
setup_sqlalchemy_sync_instrumentation(
|
|
engines=engines,
|
|
config_overrides=config,
|
|
)
|
|
logger.info(f"SQLAlchemy instrumentation setup complete for {len(engines)} engines")
|
|
|
|
# Log configuration
|
|
logger.info("Instrumentation configuration:")
|
|
for key, value in get_instrumentation_config().items():
|
|
logger.info(f" {key}: {value}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to setup SQLAlchemy instrumentation: {e}")
|
|
raise
|
|
|
|
|
|
def teardown_letta_db_instrumentation() -> None:
|
|
"""Tear down SQLAlchemy instrumentation for Letta application."""
|
|
if not is_instrumentation_active():
|
|
logger.info("SQLAlchemy instrumentation not active")
|
|
return
|
|
|
|
try:
|
|
teardown_sqlalchemy_sync_instrumentation()
|
|
logger.info("SQLAlchemy instrumentation teardown complete")
|
|
except Exception as e:
|
|
logger.error(f"Failed to teardown SQLAlchemy instrumentation: {e}")
|
|
raise
|
|
|
|
|
|
def configure_letta_db_instrumentation(**kwargs) -> None:
|
|
"""
|
|
Configure SQLAlchemy instrumentation for Letta application.
|
|
|
|
Args:
|
|
**kwargs: Configuration options to update
|
|
"""
|
|
configure_instrumentation(**kwargs)
|
|
logger.info(f"SQLAlchemy instrumentation configuration updated: {kwargs}")
|
|
|
|
|
|
# FastAPI integration
|
|
def setup_fastapi_db_instrumentation(app, **config_kwargs):
|
|
"""
|
|
Set up SQLAlchemy instrumentation for FastAPI application.
|
|
|
|
Args:
|
|
app: FastAPI application instance
|
|
**config_kwargs: Configuration options for instrumentation
|
|
"""
|
|
|
|
@app.on_event("startup")
|
|
async def startup_db_instrumentation():
|
|
setup_letta_db_instrumentation(**config_kwargs)
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_db_instrumentation():
|
|
teardown_letta_db_instrumentation()
|