Repository: ErikBjare/gptme
File: tests/test_server_v2_tool_confirmation.py (Python)

Code before:
"""Tests for the tool confirmation flow in the V2 API."""

import logging
import unittest.mock

import pytest
import requests

from gptme.tools import ToolUse

logger = logging.getLogger(__name__)


@pytest.mark.timeout(20)
def test_tool_confirmation_flow(
    init_, setup_conversation, event_listener, mock_generation, wait_for_event
):
    """Test the tool confirmation flow."""
    port, conversation_id, session_id = setup_conversation

    # Add a user message requesting a command
    requests.post(
        f"http://localhost:{port}/api/v2/conversations/{conversation_id}",
        headers={"Authorization": "Bearer test-token-for-server-thread"},
        json={"role": "user", "content": "List files in the current directory"},
    )

    # Define the tool we expect to be used
    tool = ToolUse(
        "shell",
        args=[],
        content="ls -la",
    )

    # Create mock response with the tool
    mock_stream = mock_generation(
        [
            "I'll help you list the files. Let me run the command:\n\n"
            + tool.to_output("markdown"),
            "Done, let me know if you need anything else.",
        ]
    )

    # Start generation with mocked response
    with unittest.mock.patch("gptme.server.api_v2_sessions._stream", mock_stream):
        # Request a step
        requests.post(
            f"http://localhost:{port}/api/v2/conversations/{conversation_id}/step",
            headers={"Authorization": "Bearer test-token-for-server-thread"},

Review comment (suggestion):
I don't think we want to require this by default if the `GPTME_SERVER_TOKEN` is unset (as in tests by default).

It should be something people enable when e.g. exposing gptme-server to a local network.

Code after:
"""Tests for the tool confirmation flow in the V2 API."""

import logging
import unittest.mock

import pytest
import requests

from gptme.tools import ToolUse

logger = logging.getLogger(__name__)


@pytest.mark.timeout(20)
def test_tool_confirmation_flow(
    init_, setup_conversation, event_listener, mock_generation, wait_for_event
):
    """Test the tool confirmation flow."""
    port, conversation_id, session_id = setup_conversation

    # Add a user message requesting a command
    requests.post(
        f"http://localhost:{port}/api/v2/conversations/{conversation_id}",
        json={"role": "user", "content": "List files in the current directory"},
    )

    # Define the tool we expect to be used
    tool = ToolUse(
        "shell",
        args=[],
        content="ls -la",
    )

    # Create mock response with the tool
    mock_stream = mock_generation(
        [
            "I'll help you list the files. Let me run the command:\n\n"
            + tool.to_output("markdown"),
            "Done, let me know if you need anything else.",
        ]
    )

    # Start generation with mocked response
    with unittest.mock.patch("gptme.server.api_v2_sessions._stream", mock_stream):
        # Request a step
        requests.post(
            f"http://localhost:{port}/api/v2/conversations/{conversation_id}/step",
            json={"session_id": session_id, "model": "openai/mock-model"},
        )