v1.1.2
Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai>
This commit is contained in:
Michel Thomazo
parent
af43abea60
commit
a340a721ea
13 changed files with 234 additions and 80 deletions
|
|
@ -17,14 +17,10 @@ from acp import (
|
|||
WriteTextFileRequest,
|
||||
)
|
||||
from acp.schema import (
|
||||
AgentCapabilities,
|
||||
AllowedOutcome,
|
||||
DeniedOutcome,
|
||||
Implementation,
|
||||
InitializeResponse,
|
||||
McpCapabilities,
|
||||
NewSessionResponse,
|
||||
PromptCapabilities,
|
||||
PromptResponse,
|
||||
SessionNotification,
|
||||
TextContentBlock,
|
||||
|
|
@ -324,48 +320,6 @@ async def initialize_session(acp_agent_process: asyncio.subprocess.Process) -> s
|
|||
return session_response_obj.result.sessionId
|
||||
|
||||
|
||||
class TestInitialization:
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_request_response(self) -> None:
|
||||
mock_env = get_mocking_env()
|
||||
async for process in get_acp_agent_process(mock_env=mock_env):
|
||||
await send_json_rpc(
|
||||
process,
|
||||
InitializeJsonRpcRequest(
|
||||
id=1, params=InitializeRequest(protocolVersion=1)
|
||||
),
|
||||
)
|
||||
|
||||
text_response = await read_response(process, timeout=10.0)
|
||||
assert text_response is not None, "No response to initialize"
|
||||
response_json = json.loads(text_response)
|
||||
response = InitializeJsonRpcResponse.model_validate(response_json)
|
||||
assert response.error is None, f"JSON-RPC error: {response.error}"
|
||||
assert response.result is not None, "No result in response"
|
||||
assert response.result.protocolVersion == 1
|
||||
assert response.result.agentCapabilities == AgentCapabilities(
|
||||
loadSession=False,
|
||||
promptCapabilities=PromptCapabilities(
|
||||
audio=False, embeddedContext=True, image=False
|
||||
),
|
||||
mcpCapabilities=McpCapabilities(http=False, sse=False),
|
||||
)
|
||||
assert response.result.agentInfo == Implementation(
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="1.1.1"
|
||||
)
|
||||
vibe_setup_method = next(
|
||||
(
|
||||
method
|
||||
for method in response.result.authMethods or []
|
||||
if method.id == "vibe-setup"
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert vibe_setup_method is not None, "vibe-setup auth not found"
|
||||
assert vibe_setup_method.field_meta is not None
|
||||
assert "terminal-auth" in vibe_setup_method.field_meta.keys()
|
||||
|
||||
|
||||
class TestSessionManagement:
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_sessions_unique_ids(self) -> None:
|
||||
|
|
|
|||
80
tests/acp/test_initialize.py
Normal file
80
tests/acp/test_initialize.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from acp import PROTOCOL_VERSION, AgentSideConnection, InitializeRequest
|
||||
from acp.schema import (
|
||||
AgentCapabilities,
|
||||
ClientCapabilities,
|
||||
Implementation,
|
||||
PromptCapabilities,
|
||||
)
|
||||
import pytest
|
||||
|
||||
from tests.stubs.fake_connection import FakeAgentSideConnection
|
||||
from vibe.acp.acp_agent import VibeAcpAgent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def acp_agent() -> VibeAcpAgent:
|
||||
vibe_acp_agent: VibeAcpAgent | None = None
|
||||
|
||||
def _create_agent(connection: AgentSideConnection) -> VibeAcpAgent:
|
||||
nonlocal vibe_acp_agent
|
||||
vibe_acp_agent = VibeAcpAgent(connection)
|
||||
return vibe_acp_agent
|
||||
|
||||
FakeAgentSideConnection(_create_agent)
|
||||
return vibe_acp_agent # pyright: ignore[reportReturnType]
|
||||
|
||||
|
||||
class TestACPInitialize:
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize(self, acp_agent: VibeAcpAgent) -> None:
|
||||
"""Test regular initialize without terminal-auth capabilities."""
|
||||
request = InitializeRequest(protocolVersion=PROTOCOL_VERSION)
|
||||
response = await acp_agent.initialize(request)
|
||||
|
||||
assert response.protocolVersion == PROTOCOL_VERSION
|
||||
assert response.agentCapabilities == AgentCapabilities(
|
||||
loadSession=False,
|
||||
promptCapabilities=PromptCapabilities(
|
||||
audio=False, embeddedContext=True, image=False
|
||||
),
|
||||
)
|
||||
assert response.agentInfo == Implementation(
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="1.1.2"
|
||||
)
|
||||
|
||||
assert response.authMethods == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_with_terminal_auth(self, acp_agent: VibeAcpAgent) -> None:
|
||||
"""Test initialize with terminal-auth capabilities to check it was included."""
|
||||
client_capabilities = ClientCapabilities(field_meta={"terminal-auth": True})
|
||||
request = InitializeRequest(
|
||||
protocolVersion=PROTOCOL_VERSION, clientCapabilities=client_capabilities
|
||||
)
|
||||
response = await acp_agent.initialize(request)
|
||||
|
||||
assert response.protocolVersion == PROTOCOL_VERSION
|
||||
assert response.agentCapabilities == AgentCapabilities(
|
||||
loadSession=False,
|
||||
promptCapabilities=PromptCapabilities(
|
||||
audio=False, embeddedContext=True, image=False
|
||||
),
|
||||
)
|
||||
assert response.agentInfo == Implementation(
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="1.1.2"
|
||||
)
|
||||
|
||||
assert response.authMethods is not None
|
||||
assert len(response.authMethods) == 1
|
||||
auth_method = response.authMethods[0]
|
||||
assert auth_method.id == "vibe-setup"
|
||||
assert auth_method.name == "Register your API Key"
|
||||
assert auth_method.description == "Register your API Key inside Mistral Vibe"
|
||||
assert auth_method.field_meta is not None
|
||||
assert "terminal-auth" in auth_method.field_meta
|
||||
terminal_auth_meta = auth_method.field_meta["terminal-auth"]
|
||||
assert "command" in terminal_auth_meta
|
||||
assert "args" in terminal_auth_meta
|
||||
assert terminal_auth_meta["label"] == "Mistral Vibe Setup"
|
||||
|
|
@ -29,12 +29,14 @@ from tests.backend.data.mistral import (
|
|||
STREAMED_TOOL_CONVERSATION_PARAMS as MISTRAL_STREAMED_TOOL_CONVERSATION_PARAMS,
|
||||
TOOL_CONVERSATION_PARAMS as MISTRAL_TOOL_CONVERSATION_PARAMS,
|
||||
)
|
||||
from vibe.core.config import ModelConfig, ProviderConfig
|
||||
from vibe.core.config import Backend, ModelConfig, ProviderConfig
|
||||
from vibe.core.llm.backend.factory import BACKEND_FACTORY
|
||||
from vibe.core.llm.backend.generic import GenericBackend
|
||||
from vibe.core.llm.backend.mistral import MistralBackend
|
||||
from vibe.core.llm.exceptions import BackendError
|
||||
from vibe.core.llm.types import BackendLike
|
||||
from vibe.core.types import LLMChunk, LLMMessage, Role, ToolCall
|
||||
from vibe.core.utils import get_user_agent
|
||||
|
||||
|
||||
class TestBackend:
|
||||
|
|
@ -246,3 +248,99 @@ class TestBackend:
|
|||
assert e.value.status == response.status_code
|
||||
assert e.value.reason == response.reason_phrase
|
||||
assert e.value.parsed_error is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend_type", [Backend.MISTRAL, Backend.GENERIC])
|
||||
async def test_backend_user_agent(self, backend_type: Backend):
|
||||
user_agent = get_user_agent(backend_type)
|
||||
base_url = "https://api.example.com"
|
||||
json_response = {
|
||||
"id": "fake_id_1234",
|
||||
"created": 1234567890,
|
||||
"model": "devstral-latest",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"total_tokens": 300,
|
||||
"completion_tokens": 200,
|
||||
},
|
||||
"object": "chat.completion",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"tool_calls": None,
|
||||
"content": "Hey",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
with respx.mock(base_url=base_url) as mock_api:
|
||||
mock_api.post("/v1/chat/completions").mock(
|
||||
return_value=httpx.Response(status_code=200, json=json_response)
|
||||
)
|
||||
|
||||
provider = ProviderConfig(
|
||||
name="provider_name",
|
||||
api_base=f"{base_url}/v1",
|
||||
api_key_env_var="API_KEY",
|
||||
)
|
||||
backend = BACKEND_FACTORY[backend_type](provider=provider)
|
||||
model = ModelConfig(
|
||||
name="model_name", provider="provider_name", alias="model_alias"
|
||||
)
|
||||
messages = [LLMMessage(role=Role.user, content="Just say hi")]
|
||||
|
||||
await backend.complete(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
tools=None,
|
||||
max_tokens=None,
|
||||
tool_choice=None,
|
||||
extra_headers={"user-agent": user_agent},
|
||||
)
|
||||
|
||||
assert mock_api.calls.last.request.headers["user-agent"] == user_agent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend_type", [Backend.MISTRAL, Backend.GENERIC])
|
||||
async def test_backend_user_agent_when_streaming(self, backend_type: Backend):
|
||||
user_agent = get_user_agent(backend_type)
|
||||
|
||||
base_url = "https://api.example.com"
|
||||
with respx.mock(base_url=base_url) as mock_api:
|
||||
chunks = [
|
||||
rb'data: {"id":"fake_id_1234","object":"chat.completion.chunk","created":1234567890,"model":"devstral-latest","choices":[{"index":0,"delta":{"role":"assistant","content":"Hey"},"finish_reason":"stop"}]}'
|
||||
]
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
stream=httpx.ByteStream(stream=b"\n\n".join(chunks)),
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
mock_api.post("/v1/chat/completions").mock(return_value=mock_response)
|
||||
|
||||
provider = ProviderConfig(
|
||||
name="provider_name",
|
||||
api_base=f"{base_url}/v1",
|
||||
api_key_env_var="API_KEY",
|
||||
)
|
||||
backend = BACKEND_FACTORY[backend_type](provider=provider)
|
||||
model = ModelConfig(
|
||||
name="model_name", provider="provider_name", alias="model_alias"
|
||||
)
|
||||
messages = [LLMMessage(role=Role.user, content="Just say hi")]
|
||||
|
||||
async for _ in backend.complete_streaming(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
tools=None,
|
||||
max_tokens=None,
|
||||
tool_choice=None,
|
||||
extra_headers={"user-agent": user_agent},
|
||||
):
|
||||
pass
|
||||
|
||||
assert mock_api.calls.last.request.headers["user-agent"] == user_agent
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue