Fix ruff linting issues and update Makefile for proper CI

Author: Pulkit-bristol

.pre-commit-config.yaml

@@ -17,8 +17,8 @@ repos:
     hooks:
       - id: pytest
         name: Run pytest tests
-        entry: bash
-        args: [-c, 'cd /Users/pulkit/Desktop/test/query_automation && python -m pytest tests/ -v --tb=short']
+        entry: uv
+        args: [run, python, -m, pytest, tests/, -v, --tb=short]
         language: system
         pass_filenames: false
         always_run: true

Makefile

@@ -3,16 +3,16 @@
 .PHONY: format requirements lint test commit push
 
 format:
-	ruff format .
+	uv run ruff format .
 
 requirements:
 	uv pip compile pyproject.toml --no-reuse-hashes --output-file=requirements.txt
 
 lint:
-	ruff check . --exclude notebooks
+	uv run ruff check . --exclude "*.ipynb"
 
 test:
-	pytest
+	uv run pytest tests/
 
 commit: format requirements lint test
 	@bash -c '{ \

scripts/cli_usage_guide.py

@@ -14,10 +14,10 @@
 
 def run_cli_command(command, description):
     """Run a CLI command and display results."""
-    print(f"\n{'='*60}")
+    print(f"\n{'=' * 60}")
     print(f" {description}")
     print(f" Command: {command}")
-    print(f"{'='*60}")
+    print(f"{'=' * 60}")
 
     try:
         # Split command for subprocess

src/agent_pipeline/agent/generate_sql.py

@@ -13,7 +13,7 @@ def generate_sql_for_subtask(subtask_text: str, prior_sql: list[str], user_reque
     """Generate SQL for a specific subtask using LLM with schema context and prior steps."""
     schema_ctx = retrieve_schema_context(user_request + "\n\n" + subtask_text)
     prior = (
-        "\n\n".join([f"-- Step {i+1}\n{q}" for i, q in enumerate(prior_sql)])
+        "\n\n".join([f"-- Step {i + 1}\n{q}" for i, q in enumerate(prior_sql)])
         if prior_sql
         else "(none)"
     )

src/agent_pipeline/agent/run_steps.py

@@ -68,7 +68,7 @@ def plan_subtasks(user_request: str, max_steps: int = None) -> list[tuple[str, s
         steps = [("Single-step query", "Directly produce the final answer in one query.")]
 
     print(f"[PLANNING] Generated {len(steps)} steps:")
-    for i, (title, desc) in enumerate(steps, 1):
+    for i, (title, _desc) in enumerate(steps, 1):
         print(f"[PLANNING]   Step {i}: {title}")
 
     return steps[:max_steps]

src/agent_pipeline/cli/main.py

@@ -5,11 +5,18 @@
 from pathlib import Path
 import sys
 
-# Add src to Python path
-project_root = Path(__file__).parent.parent.parent.parent
-sys.path.insert(0, str(project_root / "src"))
 
-from agent_pipeline.orchestration.pipeline import (
+def setup_path():
+    """Add src to Python path before importing local modules."""
+    project_root = Path(__file__).parent.parent.parent.parent
+    if str(project_root / "src") not in sys.path:
+        sys.path.insert(0, str(project_root / "src"))
+
+
+# Setup path before imports
+setup_path()
+
+from agent_pipeline.orchestration.pipeline import (  # noqa: E402
     get_final_result,
     initialize_pipeline,
     run_query_pipeline,

src/agent_pipeline/orchestration/pipeline.py

@@ -65,9 +65,8 @@ def run_query_pipeline(
     Returns:
         List of StepResult objects containing the pipeline execution results
     """
-    if auto_initialize:
-        if not initialize_pipeline():
-            raise RuntimeError("Pipeline initialization failed")
+    if auto_initialize and not initialize_pipeline():
+        raise RuntimeError("Pipeline initialization failed")
 
     if engine is None:
         engine = get_engine()

tests/test_db_execute.py

@@ -7,6 +7,7 @@
 
 import pandas as pd
 import pytest
+from sqlalchemy.exc import SQLAlchemyError
 
 # Add src to path
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
@@ -93,7 +94,7 @@ def test_execute_sql_invalid_syntax(self, temp_db):
         """Test SQL execution with invalid syntax."""
         sql = "INVALID SQL QUERY"
 
-        with pytest.raises(Exception):
+        with pytest.raises(SQLAlchemyError):
             execute_sql(temp_db, sql)
 
     def test_execute_sql_clean_formatting(self, temp_db):
@@ -174,11 +175,11 @@ def test_full_query_execution(self, temp_db):
     def test_database_error_handling(self, temp_db):
         """Test database error handling."""
         # Test table doesn't exist
-        with pytest.raises(Exception):
+        with pytest.raises(SQLAlchemyError):
             execute_sql(temp_db, "SELECT * FROM nonexistent_table")
 
         # Test invalid column
-        with pytest.raises(Exception):
+        with pytest.raises(SQLAlchemyError):
             execute_sql(temp_db, "SELECT nonexistent_column FROM users")
 
     def test_concurrent_database_access(self, temp_db):

tests/test_llm_client.py

@@ -100,9 +100,11 @@ def test_call_llm_timeout_retry(self, mock_llm):
             Mock(content="Success after retry"),
         ]
 
-        with patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm):
-            with patch("time.sleep"):  # Skip actual sleep in tests
-                response = call_llm("System prompt", "User prompt", max_retries=2)
+        with (
+            patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm),
+            patch("time.sleep"),  # Skip actual sleep in tests
+        ):
+            response = call_llm("System prompt", "User prompt", max_retries=2)
 
         assert response == "Success after retry"
         assert mock_llm.invoke.call_count == 2
@@ -111,42 +113,50 @@ def test_call_llm_timeout_exhausted(self, mock_llm):
         """Test LLM call when all retries are exhausted."""
         mock_llm.invoke.side_effect = Exception("Read timed out")
 
-        with patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm):
-            with patch("time.sleep"):  # Skip actual sleep in tests
-                with pytest.raises(TimeoutError, match="failed to respond after .* attempts"):
-                    call_llm("System prompt", "User prompt", max_retries=2)
+        with (
+            patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm),
+            patch("time.sleep"),  # Skip actual sleep in tests
+            pytest.raises(TimeoutError, match="failed to respond after .* attempts"),
+        ):
+            call_llm("System prompt", "User prompt", max_retries=2)
 
         assert mock_llm.invoke.call_count == 2
 
     def test_call_llm_non_timeout_error(self, mock_llm):
         """Test LLM call with non-timeout error (should not retry)."""
         mock_llm.invoke.side_effect = ValueError("Invalid input")
 
-        with patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm):
-            with pytest.raises(ValueError, match="Invalid input"):
-                call_llm("System prompt", "User prompt", max_retries=3)
+        with (
+            patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm),
+            pytest.raises(ValueError, match="Invalid input"),
+        ):
+            call_llm("System prompt", "User prompt", max_retries=3)
 
         assert mock_llm.invoke.call_count == 1  # Should not retry
 
     def test_call_llm_custom_retries(self, mock_llm):
         """Test LLM call with custom retry count."""
         mock_llm.invoke.side_effect = Exception("timeout")
 
-        with patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm):
-            with patch("time.sleep"):
-                with pytest.raises(TimeoutError):
-                    call_llm("System prompt", "User prompt", max_retries=5)
+        with (
+            patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm),
+            patch("time.sleep"),
+            pytest.raises(TimeoutError),
+        ):
+            call_llm("System prompt", "User prompt", max_retries=5)
 
         assert mock_llm.invoke.call_count == 5
 
     def test_call_llm_default_retries(self, mock_llm):
         """Test LLM call uses default retry count from config."""
         mock_llm.invoke.side_effect = Exception("timeout")
 
-        with patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm):
-            with patch("time.sleep"):
-                with pytest.raises(TimeoutError):
-                    call_llm("System prompt", "User prompt")  # No max_retries specified
+        with (
+            patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm),
+            patch("time.sleep"),
+            pytest.raises(TimeoutError),
+        ):
+            call_llm("System prompt", "User prompt")  # No max_retries specified
 
         assert mock_llm.invoke.call_count == LLM_MAX_RETRIES
 
@@ -235,9 +245,11 @@ def test_call_llm_retry_timing(self, mock_llm):
         # First call fails, second succeeds
         mock_llm.invoke.side_effect = [Exception("timeout"), Mock(content="Success")]
 
-        with patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm):
-            with patch("time.sleep") as mock_sleep:
-                call_llm("System", "User", max_retries=2)
+        with (
+            patch("agent_pipeline.llms.client.get_llm", return_value=mock_llm),
+            patch("time.sleep") as mock_sleep,
+        ):
+            call_llm("System", "User", max_retries=2)
 
         # Should sleep between retries
         mock_sleep.assert_called_once_with(5)