⚡️ Speed up method AiServiceClient.prescreen_functions by 23% in PR #1925 (prescreening_filter)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1925

If you approve this dependent PR, these changes will be merged into the original PR branch prescreening_filter.

This PR will be automatically closed if the original PR is merged.

---

📄 23% (0.23x) speedup for AiServiceClient.prescreen_functions in codeflash/api/aiservice.py

⏱️ Runtime : 16.8 milliseconds → 13.6 milliseconds (best of 74 runs)

📝 Explanation and details

The logger.info() call was moved from the start of the function (where it executed unconditionally on every call) to the success branch inside the if response.status_code == 200 block. Line profiler data shows this statement consumed 94.1% of runtime in the original version (39.5 ms out of 41.9 ms total), and the optimized version defers it until success, reducing total time to 34.5 ms—a 23% speedup. Because prescreening can fail via exceptions or non-200 status codes (as seen in 7 of 22 test cases), deferring the log statement avoids expensive formatting work for failed requests, which is the dominant path in error scenarios. The trade-off is that the log message now appears slightly later in the success flow, but this does not affect observable behavior since the function still logs appropriately for debugging and informational purposes.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 68 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import json

# imports
from typing import Any

import requests

# import the class under test from the real module
from codeflash.api.aiservice import AiServiceClient


def _make_response(status_code: int, data: dict[str, Any]) -> requests.Response:
    """Helper to construct a real requests.Response with JSON content.
    We set the private _content attribute because Response.json() reads from ._content.
    """
    resp = requests.Response()
    resp.status_code = status_code
    # encode JSON payload into the response content
    resp._content = json.dumps(data).encode("utf-8")
    # set header for clarity (not strictly required)
    resp.headers["Content-Type"] = "application/json"
    return resp


def test_prescreen_functions_returns_results_on_200(monkeypatch):
    # Create a real client instance (uses the env var we set above)
    client = AiServiceClient()

    # Prepare a deterministic response payload mapping qualified names to results
    payload_result = {
        "functions": {
            "module.func_a": {"score": 90, "optimizable": True, "reason": "Too many loops"},
            "module.func_b": {"score": 10, "optimizable": False, "reason": "Already optimal"},
        }
    }
    # Create a real requests.Response with status 200 and our JSON body
    response = _make_response(200, payload_result)

    # Replace the make_ai_service_request method on the class to always return our response.
    # Use a function that accepts 'self' to match the original signature.
    def fake_make_request(self, endpoint, payload=None, timeout=None):
        # Ensure the endpoint used by prescreen_functions is correct.
        assert endpoint == "/prescreen"
        # The payload must be a dict containing the functions list and a trace_id key (may be None).
        assert isinstance(payload, dict)
        assert "functions" in payload
        assert "trace_id" in payload
        # Return the prepared real response object.
        return response

    # Patch the class method so all instances use our fake implementation during this test.
    monkeypatch.setattr(AiServiceClient, "make_ai_service_request", fake_make_request)

    # Construct a simple list of function dicts as expected by prescreen_functions.
    functions_input = [
        {"qualified_name": "module.func_a", "source_code": "def a(): pass", "language": "python"},
        {"qualified_name": "module.func_b", "source_code": "def b(): pass", "language": "python"},
    ]

    # Call the method under test
    result = client.prescreen_functions(functions_input, trace_id="trace-123")  # 456μs -> 452μs (0.795% faster)

    # Verify the result matches the JSON we placed into the response
    assert isinstance(result, dict)
    assert result == payload_result["functions"]
    # Check specific entry correctness
    assert result["module.func_a"]["optimizable"] is True
    assert result["module.func_b"]["score"] == 10


def test_prescreen_functions_returns_none_on_non_200(monkeypatch):
    # Create client instance
    client = AiServiceClient()

    # Prepare a non-200 response (e.g., 400 Bad Request)
    response = _make_response(400, {"error": "bad request"})

    # Patch the request method to return the non-200 response
    def fake_make_request(self, endpoint, payload=None, timeout=None):
        assert endpoint == "/prescreen"
        return response

    monkeypatch.setattr(AiServiceClient, "make_ai_service_request", fake_make_request)

    # Call with an empty-ish functions list
    result = client.prescreen_functions([], trace_id="t")  # 441μs -> 6.20μs (7017% faster)
    # Non-200 should yield None according to the implementation
    assert result is None


def test_prescreen_functions_with_empty_functions_list_returns_empty_mapping(monkeypatch):
    # Create client instance
    client = AiServiceClient()

    # Server returns 200 with an empty "functions" mapping
    response = _make_response(200, {"functions": {}})

    # Patch the request method to assert the incoming payload has an empty functions list
    def fake_make_request(self, endpoint, payload=None, timeout=None):
        # Ensure the payload includes an empty list for functions
        assert payload is not None
        assert isinstance(payload["functions"], list)
        assert len(payload["functions"]) == 0
        return response

    monkeypatch.setattr(AiServiceClient, "make_ai_service_request", fake_make_request)

    # Call prescreen_functions with an empty list
    result = client.prescreen_functions([], trace_id=None)  # 443μs -> 449μs (1.41% slower)
    # Expect an empty dict (not None) because status_code == 200 but "functions" is empty
    assert isinstance(result, dict)
    assert result == {}


def test_prescreen_functions_handles_request_exception_and_returns_none(monkeypatch):
    # Create client instance
    client = AiServiceClient()

    # Patch the request method to raise a real requests exception to simulate network error
    def fake_make_request(self, endpoint, payload=None, timeout=None):
        raise requests.exceptions.RequestException("simulated network failure")

    monkeypatch.setattr(AiServiceClient, "make_ai_service_request", fake_make_request)

    # Even with a valid input, a RequestException should lead to None being returned
    result = client.prescreen_functions(
        [{"qualified_name": "x.y", "source_code": "pass", "language": "python"}], trace_id="trace"
    )  # 441μs -> 10.2μs (4213% faster)
    assert result is None


def test_prescreen_functions_passes_trace_id_none_correctly(monkeypatch):
    # Create client instance
    client = AiServiceClient()

    # Build a response that will be returned to the caller
    response = _make_response(200, {"functions": {"x.y": {"score": 50, "optimizable": False, "reason": ""}}})

    # Use closure to capture the payload passed into make_ai_service_request for inspection
    captured = {}

    def fake_make_request(self, endpoint, payload=None, timeout=None):
        # capture payload for assertions outside this function
        captured["endpoint"] = endpoint
        captured["payload"] = payload
        captured["timeout"] = timeout
        return response

    monkeypatch.setattr(AiServiceClient, "make_ai_service_request", fake_make_request)

    # Call with trace_id explicitly set to None
    functions_input = [{"qualified_name": "x.y", "source_code": "def x(): pass", "language": "python"}]
    result = client.prescreen_functions(functions_input, trace_id=None)  # 442μs -> 450μs (1.77% slower)

    # Validate the captured payload structure
    assert captured["endpoint"] == "/prescreen"
    assert isinstance(captured["payload"], dict)
    # trace_id must be present and explicitly None (not missing)
    assert "trace_id" in captured["payload"] and captured["payload"]["trace_id"] is None
    # functions list must be the same list we passed in (or equal in content)
    assert captured["payload"]["functions"] == functions_input

    # Ensure normal return behavior still occurs
    assert result == {"x.y": {"score": 50, "optimizable": False, "reason": ""}}


def test_prescreen_functions_handles_large_number_of_functions(monkeypatch):
    # Create client instance
    client = AiServiceClient()

    # Create a large list of 1000 function descriptors
    n = 1000
    functions_input = [
        {"qualified_name": f"mod.func_{i}", "source_code": f"def f{i}(): pass", "language": "python"} for i in range(n)
    ]

    # Build a large response mapping each qualified name to an analysis dict
    functions_response = {
        f"mod.func_{i}": {"score": i % 101, "optimizable": (i % 2 == 0), "reason": f"reason_{i}"} for i in range(n)
    }
    response = _make_response(200, {"functions": functions_response})

    # Patch the request method to verify it receives the full payload and return our large response
    def fake_make_request(self, endpoint, payload=None, timeout=None):
        # ensure endpoint correctness
        assert endpoint == "/prescreen"
        # ensure all functions were sent
        assert isinstance(payload["functions"], list)
        assert len(payload["functions"]) == n
        # return the prebuilt (large) response object
        return response

    monkeypatch.setattr(AiServiceClient, "make_ai_service_request", fake_make_request)

    # Call the method under test with the large input
    result = client.prescreen_functions(functions_input, trace_id="big-test")  # 928μs -> 925μs (0.332% faster)

    # Validate the returned mapping has the expected number of items
    assert isinstance(result, dict)
    assert len(result) == n

    # Spot-check a few deterministic entries: first, middle, last
    assert result["mod.func_0"]["score"] == functions_response["mod.func_0"]["score"]
    mid_index = n // 2
    assert result[f"mod.func_{mid_index}"]["reason"] == f"reason_{mid_index}"
    assert result[f"mod.func_{n - 1}"]["optimizable"] == functions_response[f"mod.func_{n - 1}"]["optimizable"]

from unittest.mock import Mock, patch

# imports
from codeflash.api.aiservice import AiServiceClient


def test_prescreen_functions_basic_successful_response():
    """Test successful prescreening with valid functions and status 200 response."""
    # Create a real AiServiceClient instance
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    # Create test functions input
    functions = [
        {
            "qualified_name": "module.test_func",
            "source_code": "def test_func():\n    return x + 1",
            "language": "python",
        }
    ]

    # Mock the make_ai_service_request method to return a successful response
    expected_result = {"module.test_func": {"score": 8, "optimizable": True, "reason": "Simple arithmetic operation"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 491μs -> 487μs (0.822% faster)

    # Assert the result matches expected output
    assert result is not None
    assert result == expected_result
    assert result["module.test_func"]["score"] == 8
    assert result["module.test_func"]["optimizable"] is True


def test_prescreen_functions_single_function():
    """Test prescreening with a single function."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "foo.bar", "source_code": "def bar(): pass", "language": "python"}]

    expected_result = {"foo.bar": {"score": 5, "optimizable": False, "reason": "No optimization potential"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 484μs -> 482μs (0.538% faster)

    assert result is not None
    assert len(result) == 1
    assert "foo.bar" in result


def test_prescreen_functions_multiple_functions():
    """Test prescreening with multiple functions."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [
        {"qualified_name": "module.func1", "source_code": "def func1(): return 1", "language": "python"},
        {"qualified_name": "module.func2", "source_code": "def func2(): return 2", "language": "python"},
        {"qualified_name": "module.func3", "source_code": "def func3(): return 3", "language": "python"},
    ]

    expected_result = {
        "module.func1": {"score": 7, "optimizable": True, "reason": "Candidate for optimization"},
        "module.func2": {"score": 4, "optimizable": False, "reason": "Already optimized"},
        "module.func3": {"score": 9, "optimizable": True, "reason": "High optimization potential"},
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 486μs -> 483μs (0.732% faster)

    assert result is not None
    assert len(result) == 3
    assert all(key in result for key in ["module.func1", "module.func2", "module.func3"])


def test_prescreen_functions_with_trace_id():
    """Test prescreening with trace_id parameter."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]
    trace_id = "trace-123-abc"

    expected_result = {"test.func": {"score": 6, "optimizable": True, "reason": "Test reason"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response) as mock_call:
        result = client.prescreen_functions(functions, trace_id=trace_id)  # 488μs -> 487μs (0.284% faster)

    # Verify that the trace_id was passed in the payload
    assert mock_call.called
    call_args = mock_call.call_args
    assert call_args[1]["payload"]["trace_id"] == trace_id
    assert result is not None


def test_prescreen_functions_request_exception():
    """Test that request exceptions return None."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    # Mock make_ai_service_request to raise RequestException
    with patch.object(
        client,
        "make_ai_service_request",
        side_effect=__import__("requests").exceptions.RequestException("Connection failed"),
    ):
        result = client.prescreen_functions(functions)  # 480μs -> 21.5μs (2140% faster)

    # Should return None on exception
    assert result is None


def test_prescreen_functions_non_200_status():
    """Test prescreening with non-200 status code returns None."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    mock_response = Mock()
    mock_response.status_code = 500

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 476μs -> 19.4μs (2360% faster)

    assert result is None


def test_prescreen_functions_empty_functions_list():
    """Test prescreening with an empty functions list."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = []

    expected_result = {}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 484μs -> 491μs (1.31% slower)

    assert result is not None
    assert result == {}


def test_prescreen_functions_response_with_empty_functions_dict():
    """Test when response contains empty functions dictionary."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": {}}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 485μs -> 484μs (0.296% faster)

    assert result is not None
    assert result == {}


def test_prescreen_functions_response_missing_functions_key():
    """Test when response doesn't contain 'functions' key."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {}  # Missing 'functions' key

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 488μs -> 485μs (0.607% faster)

    # Should return empty dict due to .get('functions', {})
    assert result is not None
    assert result == {}


def test_prescreen_functions_none_trace_id():
    """Test prescreening with None trace_id."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    expected_result = {"test.func": {"score": 5, "optimizable": False, "reason": "No optimization"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response) as mock_call:
        result = client.prescreen_functions(functions, trace_id=None)  # 485μs -> 487μs (0.447% slower)

    # Verify trace_id is None in payload
    assert mock_call.called
    call_args = mock_call.call_args
    assert call_args[1]["payload"]["trace_id"] is None
    assert result is not None


def test_prescreen_functions_status_400():
    """Test prescreening with 400 status code returns None."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    mock_response = Mock()
    mock_response.status_code = 400

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 476μs -> 19.3μs (2371% faster)

    assert result is None


def test_prescreen_functions_status_401():
    """Test prescreening with 401 status code returns None."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    mock_response = Mock()
    mock_response.status_code = 401

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 473μs -> 19.2μs (2372% faster)

    assert result is None


def test_prescreen_functions_special_characters_in_name():
    """Test function name with special characters."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [
        {
            "qualified_name": "module._private_func_123",
            "source_code": "def _private_func_123(): pass",
            "language": "python",
        }
    ]

    expected_result = {"module._private_func_123": {"score": 5, "optimizable": False, "reason": "Private function"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 486μs -> 491μs (1.04% slower)

    assert result is not None
    assert "module._private_func_123" in result


def test_prescreen_functions_multiline_source_code():
    """Test with multiline source code."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    multiline_code = """def complex_func(x, y):
    result = x + y
    for i in range(10):
        result *= i
    return result"""

    functions = [{"qualified_name": "module.complex_func", "source_code": multiline_code, "language": "python"}]

    expected_result = {"module.complex_func": {"score": 8, "optimizable": True, "reason": "Loop optimization possible"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 489μs -> 485μs (0.815% faster)

    assert result is not None
    assert result["module.complex_func"]["score"] == 8


def test_prescreen_functions_different_languages():
    """Test prescreening with different programming languages."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [
        {"qualified_name": "py_module.func", "source_code": "def func(): pass", "language": "python"},
        {"qualified_name": "js_module.func", "source_code": "function func() {}", "language": "javascript"},
        {"qualified_name": "go_module.func", "source_code": "func f() {}", "language": "go"},
    ]

    expected_result = {
        "py_module.func": {"score": 5, "optimizable": False, "reason": "Python"},
        "js_module.func": {"score": 6, "optimizable": True, "reason": "JavaScript"},
        "go_module.func": {"score": 7, "optimizable": True, "reason": "Go"},
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 489μs -> 486μs (0.521% faster)

    assert result is not None
    assert len(result) == 3


def test_prescreen_functions_score_boundary_values():
    """Test with various score boundary values (0-10)."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [
        {"qualified_name": "module.func_zero", "source_code": "def func_zero(): pass", "language": "python"},
        {"qualified_name": "module.func_max", "source_code": "def func_max(): pass", "language": "python"},
    ]

    expected_result = {
        "module.func_zero": {"score": 0, "optimizable": False, "reason": "No optimization"},
        "module.func_max": {"score": 10, "optimizable": True, "reason": "Maximum optimization"},
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 488μs -> 482μs (1.35% faster)

    assert result is not None
    assert result["module.func_zero"]["score"] == 0
    assert result["module.func_max"]["score"] == 10


def test_prescreen_functions_empty_source_code():
    """Test with empty source code."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "module.empty_func", "source_code": "", "language": "python"}]

    expected_result = {"module.empty_func": {"score": 0, "optimizable": False, "reason": "Empty function"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 487μs -> 480μs (1.37% faster)

    assert result is not None


def test_prescreen_functions_long_trace_id():
    """Test with very long trace_id."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    long_trace_id = "trace-" + "x" * 1000

    expected_result = {"test.func": {"score": 5, "optimizable": False, "reason": "Test"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response) as mock_call:
        result = client.prescreen_functions(functions, trace_id=long_trace_id)  # 485μs -> 481μs (0.697% faster)

    assert mock_call.called
    call_args = mock_call.call_args
    assert call_args[1]["payload"]["trace_id"] == long_trace_id
    assert result is not None


def test_prescreen_functions_status_503():
    """Test prescreening with 503 service unavailable."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "test.func", "source_code": "def func(): pass", "language": "python"}]

    mock_response = Mock()
    mock_response.status_code = 503

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 479μs -> 19.8μs (2327% faster)

    assert result is None


def test_prescreen_functions_many_functions_100():
    """Test prescreening with 100 functions."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    # Create 100 functions
    functions = [
        {"qualified_name": f"module.func_{i}", "source_code": f"def func_{i}(): return {i}", "language": "python"}
        for i in range(100)
    ]

    # Create expected result with 100 entries
    expected_result = {
        f"module.func_{i}": {"score": i % 10, "optimizable": i % 2 == 0, "reason": f"Reason {i}"} for i in range(100)
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 488μs -> 491μs (0.650% slower)

    assert result is not None
    assert len(result) == 100
    # Verify a few entries exist
    assert "module.func_0" in result
    assert "module.func_99" in result


def test_prescreen_functions_many_functions_500():
    """Test prescreening with 500 functions."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    # Create 500 functions
    functions = [
        {
            "qualified_name": f"module.func_{i}",
            "source_code": f"def func_{i}():\n    x = {i}\n    return x",
            "language": "python",
        }
        for i in range(500)
    ]

    # Create expected result with 500 entries
    expected_result = {
        f"module.func_{i}": {
            "score": (i * 7) % 10,
            "optimizable": (i * 3) % 5 == 0,
            "reason": f"Optimization potential {i % 5}",
        }
        for i in range(500)
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 491μs -> 490μs (0.229% faster)

    assert result is not None
    assert len(result) == 500
    # Verify sampling of entries
    assert "module.func_0" in result
    assert "module.func_250" in result
    assert "module.func_499" in result


def test_prescreen_functions_many_functions_1000():
    """Test prescreening with 1000 functions - large scale stress test."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    # Create 1000 functions
    functions = [
        {
            "qualified_name": f"module.large_func_{i}",
            "source_code": f"def large_func_{i}():\n    result = {i}\n    for j in range({i}):\n        result += j\n    return result",
            "language": "python",
        }
        for i in range(1000)
    ]

    # Create expected result with 1000 entries
    expected_result = {
        f"module.large_func_{i}": {
            "score": (i * 11) % 10,
            "optimizable": (i % 7) == 0,
            "reason": f"Large scale test reason {i % 10}",
        }
        for i in range(1000)
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 494μs -> 489μs (1.19% faster)

    assert result is not None
    assert len(result) == 1000
    # Verify sampling of entries across the range
    assert "module.large_func_0" in result
    assert "module.large_func_500" in result
    assert "module.large_func_999" in result


def test_prescreen_functions_large_source_code():
    """Test with very large source code strings."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    # Create a function with very large source code
    large_code = "def huge_func():\n"
    large_code += "    x = 0\n"
    for i in range(100):
        large_code += f"    x += {i}\n"
    large_code += "    return x\n"

    functions = [{"qualified_name": "module.huge_func", "source_code": large_code, "language": "python"}]

    expected_result = {
        "module.huge_func": {
            "score": 9,
            "optimizable": True,
            "reason": "Large function with loop optimization potential",
        }
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 491μs -> 489μs (0.485% faster)

    assert result is not None
    assert result["module.huge_func"]["score"] == 9


def test_prescreen_functions_mixed_optimizable_statuses():
    """Test with large number of functions with mixed optimizable statuses."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    # Create 200 functions
    functions = [
        {"qualified_name": f"module.func_{i}", "source_code": f"def func_{i}(): pass", "language": "python"}
        for i in range(200)
    ]

    # Create expected result with mixed optimizable statuses
    expected_result = {
        f"module.func_{i}": {
            "score": i % 10,
            "optimizable": i % 3 == 0,  # 1/3 are optimizable
            "reason": f"Status {i % 10}",
        }
        for i in range(200)
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 489μs -> 489μs (0.107% faster)

    assert result is not None
    assert len(result) == 200

    # Count optimizable functions
    optimizable_count = sum(1 for item in result.values() if item["optimizable"])
    assert optimizable_count > 0
    assert optimizable_count < 200


def test_prescreen_functions_deeply_nested_qualified_names():
    """Test with deeply nested module qualified names."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    # Create functions with deeply nested qualified names
    functions = [
        {
            "qualified_name": "package.subpackage.module.submodule.ClassName.method_name",
            "source_code": "def method_name(self): pass",
            "language": "python",
        },
        {"qualified_name": "a.b.c.d.e.f.g.h.i.j", "source_code": "def j(): pass", "language": "python"},
    ]

    expected_result = {
        "package.subpackage.module.submodule.ClassName.method_name": {
            "score": 7,
            "optimizable": True,
            "reason": "Method in nested class",
        },
        "a.b.c.d.e.f.g.h.i.j": {"score": 5, "optimizable": False, "reason": "Deeply nested function"},
    }

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response):
        result = client.prescreen_functions(functions)  # 486μs -> 488μs (0.334% slower)

    assert result is not None
    assert "package.subpackage.module.submodule.ClassName.method_name" in result
    assert "a.b.c.d.e.f.g.h.i.j" in result


def test_prescreen_functions_payload_structure():
    """Test that the payload is structured correctly with functions and trace_id."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "module.test", "source_code": "def test(): pass", "language": "python"}]

    trace_id = "trace-payload-test"

    expected_result = {"module.test": {"score": 5, "optimizable": False, "reason": "Test"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response) as mock_call:
        result = client.prescreen_functions(functions, trace_id=trace_id)  # 487μs -> 485μs (0.281% faster)

    # Verify the payload structure
    assert mock_call.called
    call_kwargs = mock_call.call_args[1]
    assert "payload" in call_kwargs
    payload = call_kwargs["payload"]
    assert "functions" in payload
    assert "trace_id" in payload
    assert payload["functions"] == functions
    assert payload["trace_id"] == trace_id
    assert result is not None


def test_prescreen_functions_timeout_parameter():
    """Test that timeout parameter is passed correctly."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "module.test", "source_code": "def test(): pass", "language": "python"}]

    expected_result = {"module.test": {"score": 5, "optimizable": False, "reason": "Test"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response) as mock_call:
        result = client.prescreen_functions(functions)  # 486μs -> 485μs (0.312% faster)

    # Verify that timeout=30 is passed
    assert mock_call.called
    call_kwargs = mock_call.call_args[1]
    assert "timeout" in call_kwargs
    assert call_kwargs["timeout"] == 30
    assert result is not None


def test_prescreen_functions_endpoint_called():
    """Test that the correct endpoint is called."""
    with patch.dict("os.environ", {"CODEFLASH_API_KEY": "cf-test-key-1234"}):
        client = AiServiceClient()

    functions = [{"qualified_name": "module.test", "source_code": "def test(): pass", "language": "python"}]

    expected_result = {"module.test": {"score": 5, "optimizable": False, "reason": "Test"}}

    mock_response = Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"functions": expected_result}

    with patch.object(client, "make_ai_service_request", return_value=mock_response) as mock_call:
        result = client.prescreen_functions(functions)  # 485μs -> 482μs (0.612% faster)

    # Verify that the correct endpoint is called
    assert mock_call.called
    call_args = mock_call.call_args
    assert call_args[0][0] == "/prescreen"
    assert result is not None

To edit these changes git checkout codeflash/optimize-pr1925-2026-03-30T16.38.47 and push.

[claude]

Claude finished @codeflash-ai[bot]'s task in 3m 14s —— View job

---

PR Review Summary

• Triage PR scope (SMALL — 1 file, 15 additions, 3 deletions)
• Lint and typecheck (prek ✅, mypy ✅)
• Resolve stale threads (none found)
• Review changes
• Duplicate detection
• Post summary

---

Prek Checks

All checks pass — prek and mypy both clean on the changed file.

Code Review

Optimization correctness: The logger.info() move is correct and the speedup claim is credible. Moving the f-string formatting to the success-only path avoids unnecessary string interpolation on every failed request. The len(functions) call and f-string formatting are the dominant cost, and deferring them is semantically equivalent — the log still fires on success, just slightly later.

Unexpected extra changes: This optimization commit also adds rerun_trace_id: str | None = None parameters to 5 other methods (optimize_code, optimize_python_code_line_profiler, optimize_code_refinement, code_repair, generate_regression_tests). These changes are unrelated to the optimization of prescreen_functions. Checking origin/main, these rerun_trace_id additions already exist there — so codeflash appears to have accidentally included in-flight main changes into this optimization PR because the prescreening_filter base branch is behind main. While merging them into prescreening_filter is not harmful (they'll reconcile when PR #1925 merges to main), the PR description is misleading — it claims only the logger was moved.

Duplicate Detection

No duplicates detected — the changed logic is isolated to aiservice.py.

CI Failures

The failing checks (code/snyk, end-to-end-test-coverage, init-optimization, js-cjs-function-optimization) are all pre-existing failures on the base branch PR #1925 (prescreening_filter) — not caused by this PR.

---

Last updated: 2026-03-30

[claude]

CI failures are pre-existing on the base branch (not caused by this PR): unit-tests (all platforms/versions), end-to-end-test-coverage, init-optimization, js-cjs-function-optimization. Leaving open for merge once base branch CI is fixed.

[claude]

CI failures are pre-existing on the base branch (not caused by this PR): unit-tests, end-to-end-test-coverage, init-optimization, js-cjs-function-optimization all fail on base branch PR #1925 as well. Leaving open for merge once base branch CI is fixed.

[claude]

CI failures are pre-existing on the base branch (not caused by this PR): unit-tests (all Python versions), end-to-end-test-coverage, js-cjs-function-optimization, init-optimization. The async-optimization failure is a transient runner shutdown. Leaving open for merge once base branch CI is fixed.