sqlsaber 0.31.0__py3-none-any.whl → 0.32.0__py3-none-any.whl

sqlsaber/cli/commands.py

@@ -1,6 +1,7 @@
 """CLI command definitions and handlers."""
 
 import asyncio
+import os
 import sys
 from typing import Annotated
 
@@ -43,6 +44,52 @@ app.command(create_threads_app(), name="threads")
 console = create_console()
 config_manager = DatabaseConfigManager()
 
+_MLFLOW_CONFIGURED = False
+
+
+def _maybe_configure_mlflow(log) -> bool:
+    """Enable mlflow autologging when environment variables are present."""
+    global _MLFLOW_CONFIGURED
+    if _MLFLOW_CONFIGURED:
+        return True
+
+    tracking_uri = os.getenv("MLFLOW_URI")
+    experiment = os.getenv("MLFLOW_EXP")
+    if not tracking_uri and not experiment:
+        return False
+
+    try:
+        import mlflow
+    except ModuleNotFoundError:
+        log.warning(
+            "mlflow.setup.skipped",
+            reason="mlflow package not installed",
+            uri=tracking_uri,
+            experiment=experiment,
+        )
+        return False
+
+    try:
+        mlflow.pydantic_ai.autolog()
+    except Exception:
+        log.warning("mlflow.autolog.failed", exc_info=True)
+    try:
+        if tracking_uri:
+            mlflow.set_tracking_uri(tracking_uri)
+        if experiment:
+            mlflow.set_experiment(experiment)
+    except Exception:
+        log.warning("mlflow.setup.failed", exc_info=True)
+        return False
+
+    _MLFLOW_CONFIGURED = True
+    log.info(
+        "mlflow.setup.enabled",
+        uri=tracking_uri,
+        experiment=experiment,
+    )
+    return True
+
 
 @app.meta.default
 def meta_handler(
@@ -175,6 +222,7 @@ def query(
 
         # Create pydantic-ai agent instance with database name for memory context
         sqlsaber_agent = SQLSaberAgent(db_conn, db_name, thinking_enabled=thinking)
+        _maybe_configure_mlflow(log)
 
         try:
             if actual_query:
@@ -188,6 +236,7 @@ def query(
                 run = await streaming_handler.execute_streaming_query(
                     actual_query, sqlsaber_agent
                 )
+
                 # Persist non-interactive run as a thread snapshot so it can be resumed later
                 try:
                     if run is not None:

{sqlsaber-0.31.0.dist-info → sqlsaber-0.32.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlsaber
-Version: 0.31.0
+Version: 0.32.0
 Summary: SQLsaber - Open-source agentic SQL assistant
 License-File: LICENSE
 Requires-Python: >=3.12

{sqlsaber-0.31.0.dist-info → sqlsaber-0.32.0.dist-info}/RECORD

@@ -10,7 +10,7 @@ sqlsaber/application/model_selection.py,sha256=fSC06MZNKinHDR-csMFVYYJFyK8MydKf6
 sqlsaber/application/prompts.py,sha256=4rMGcWpYJbNWPMzqVWseUMx0nwvXOkWS6GaTAJ5mhfc,3473
 sqlsaber/cli/__init__.py,sha256=qVSLVJLLJYzoC6aj6y9MFrzZvAwc4_OgxU9DlkQnZ4M,86
 sqlsaber/cli/auth.py,sha256=8BsAkOYarLSpGipYi2CCUqyAv7EbOCqVXKLmdrX_g2c,6966
-sqlsaber/cli/commands.py,sha256=eH6jC2coaH6QhJAlOySBxEHfETWW283fBr6BQG60Bl8,9732
+sqlsaber/cli/commands.py,sha256=WocWlLrxA5kM8URfvIvWFtc0ocfgKWAwoYTxVNZhmM4,10962
 sqlsaber/cli/completers.py,sha256=g-hLDq5fiBx7gg8Bte1Lq8GU-ZxCYVs4dcPsmHPIcK4,6574
 sqlsaber/cli/database.py,sha256=BBGj0eyduh5DDXNLZLDtWfY9kWpeT_ZX0J9R9INZyyU,12421
 sqlsaber/cli/display.py,sha256=WB5JCumhXadziDEX1EZHG3vN1Chol5FNAaTXHieqFK0,17892
@@ -51,8 +51,8 @@ sqlsaber/tools/base.py,sha256=NKEEooliPKTJj_Pomwte_wW0Xd9Z5kXNfVdCRfTppuw,883
 sqlsaber/tools/registry.py,sha256=XmBzERq0LJXtg3BZ-r8cEyt8J54NUekgUlTJ_EdSYMk,2204
 sqlsaber/tools/sql_guard.py,sha256=dTDwcZP-N4xPGzcr7MQtKUxKrlDzlc1irr9aH5a4wvk,6182
 sqlsaber/tools/sql_tools.py,sha256=eo-NTxiXGHMopAjujvDDjmv9hf5bQNbiy3nTpxoJ_E8,7369
-sqlsaber-0.31.0.dist-info/METADATA,sha256=B5_gtNMDPZ9pEIMaAHvuvx9XvyaTS3ENl3A37d1-A_o,5915
-sqlsaber-0.31.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sqlsaber-0.31.0.dist-info/entry_points.txt,sha256=tw1mB0fjlkXQiOsC0434X6nE-o1cFCuQwt2ZYHv_WAE,91
-sqlsaber-0.31.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sqlsaber-0.31.0.dist-info/RECORD,,
+sqlsaber-0.32.0.dist-info/METADATA,sha256=Ai3-BNq7KNZ1oA4W-39jp6YTHIE_9ZApz_z0vCkVPyI,5915
+sqlsaber-0.32.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sqlsaber-0.32.0.dist-info/entry_points.txt,sha256=tw1mB0fjlkXQiOsC0434X6nE-o1cFCuQwt2ZYHv_WAE,91
+sqlsaber-0.32.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sqlsaber-0.32.0.dist-info/RECORD,,