2.0.0
Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai> Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai> Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai> Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-Authored-By: Clément Siriex <clement.sirieix@mistral.ai> Co-Authored-By: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai> Co-Authored-By: Thaddee Tyl <thaddee.tyl@gmail.com> Co-Authored-By: David Brochart <david.brochart@gmail.com> Co-Authored-By: Joseph Guhlin <joseph.guhlin@gmail.com> Co-Authored-By: Thomas Kenbeek <thomaskenbeek@gmail.com> Co-Authored-By: Remenby31 <baptiste.cruvellier31@gmail.com>
This commit is contained in:
Mathias Gesbert
Mathias Gesbert
parent
79f215d91c
commit
d33db9fff8
217 changed files with 16911 additions and 4305 deletions
|
|
@ -2,101 +2,52 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
from uuid import uuid4
|
||||
|
||||
from acp import (
|
||||
PROTOCOL_VERSION,
|
||||
InitializeRequest,
|
||||
NewSessionRequest,
|
||||
PromptRequest,
|
||||
RequestError,
|
||||
)
|
||||
from acp import PROTOCOL_VERSION, RequestError
|
||||
from acp.schema import TextContentBlock
|
||||
import pytest
|
||||
from pytest import raises
|
||||
|
||||
from tests.mock.utils import mock_llm_chunk
|
||||
from tests.stubs.fake_backend import FakeBackend
|
||||
from tests.stubs.fake_connection import FakeAgentSideConnection
|
||||
from vibe.acp.acp_agent import VibeAcpAgent
|
||||
from vibe.core.agent import Agent
|
||||
from vibe.core.config import ModelConfig, VibeConfig
|
||||
from vibe.acp.acp_agent_loop import VibeAcpAgentLoop
|
||||
from vibe.core.types import Role
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def backend() -> FakeBackend:
|
||||
backend = FakeBackend()
|
||||
return backend
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def acp_agent(backend: FakeBackend) -> VibeAcpAgent:
|
||||
config = VibeConfig(
|
||||
active_model="devstral-latest",
|
||||
models=[
|
||||
ModelConfig(
|
||||
name="devstral-latest", provider="mistral", alias="devstral-latest"
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
class PatchedAgent(Agent):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.backend = backend
|
||||
self.config = config
|
||||
|
||||
patch("vibe.acp.acp_agent.VibeAgent", side_effect=PatchedAgent).start()
|
||||
|
||||
vibe_acp_agent: VibeAcpAgent | None = None
|
||||
|
||||
def _create_agent(connection: Any) -> 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 TestMultiSessionCore:
|
||||
@pytest.mark.asyncio
|
||||
async def test_different_sessions_use_different_agents(
|
||||
self, acp_agent: VibeAcpAgent
|
||||
self, acp_agent_loop: VibeAcpAgentLoop
|
||||
) -> None:
|
||||
await acp_agent.initialize(InitializeRequest(protocolVersion=PROTOCOL_VERSION))
|
||||
session1_response = await acp_agent.newSession(
|
||||
NewSessionRequest(cwd=str(Path.cwd()), mcpServers=[])
|
||||
await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION)
|
||||
session1_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
session1 = acp_agent.sessions[session1_response.sessionId]
|
||||
session2_response = await acp_agent.newSession(
|
||||
NewSessionRequest(cwd=str(Path.cwd()), mcpServers=[])
|
||||
session1 = acp_agent_loop.sessions[session1_response.session_id]
|
||||
session2_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
session2 = acp_agent.sessions[session2_response.sessionId]
|
||||
session2 = acp_agent_loop.sessions[session2_response.session_id]
|
||||
|
||||
assert session1.id != session2.id
|
||||
# Each agent should be independent
|
||||
assert session1.agent is not session2.agent
|
||||
assert id(session1.agent) != id(session2.agent)
|
||||
# Each agent loop should be independent
|
||||
assert session1.agent_loop is not session2.agent_loop
|
||||
assert id(session1.agent_loop) != id(session2.agent_loop)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_on_nonexistent_session(self, acp_agent: VibeAcpAgent) -> None:
|
||||
await acp_agent.initialize(InitializeRequest(protocolVersion=PROTOCOL_VERSION))
|
||||
await acp_agent.newSession(
|
||||
NewSessionRequest(cwd=str(Path.cwd()), mcpServers=[])
|
||||
)
|
||||
async def test_error_on_nonexistent_session(
|
||||
self, acp_agent_loop: VibeAcpAgentLoop
|
||||
) -> None:
|
||||
await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION)
|
||||
await acp_agent_loop.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
fake_session_id = "fake-session-id-" + str(uuid4())
|
||||
|
||||
with raises(RequestError) as exc_info:
|
||||
await acp_agent.prompt(
|
||||
PromptRequest(
|
||||
sessionId=fake_session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello, world!")],
|
||||
)
|
||||
await acp_agent_loop.prompt(
|
||||
session_id=fake_session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello, world!")],
|
||||
)
|
||||
|
||||
assert isinstance(exc_info.value, RequestError)
|
||||
|
|
@ -104,17 +55,17 @@ class TestMultiSessionCore:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simultaneous_message_processing(
|
||||
self, acp_agent: VibeAcpAgent, backend: FakeBackend
|
||||
self, acp_agent_loop: VibeAcpAgentLoop, backend: FakeBackend
|
||||
) -> None:
|
||||
await acp_agent.initialize(InitializeRequest(protocolVersion=PROTOCOL_VERSION))
|
||||
session1_response = await acp_agent.newSession(
|
||||
NewSessionRequest(cwd=str(Path.cwd()), mcpServers=[])
|
||||
await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION)
|
||||
session1_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
session1 = acp_agent.sessions[session1_response.sessionId]
|
||||
session2_response = await acp_agent.newSession(
|
||||
NewSessionRequest(cwd=str(Path.cwd()), mcpServers=[])
|
||||
session1 = acp_agent_loop.sessions[session1_response.session_id]
|
||||
session2_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
session2 = acp_agent.sessions[session2_response.sessionId]
|
||||
session2 = acp_agent_loop.sessions[session2_response.session_id]
|
||||
|
||||
backend._streams = [
|
||||
[mock_llm_chunk(content="Response 1")],
|
||||
|
|
@ -122,40 +73,38 @@ class TestMultiSessionCore:
|
|||
]
|
||||
|
||||
async def run_session1():
|
||||
await acp_agent.prompt(
|
||||
PromptRequest(
|
||||
sessionId=session1.id,
|
||||
prompt=[TextContentBlock(type="text", text="Prompt for session 1")],
|
||||
)
|
||||
await acp_agent_loop.prompt(
|
||||
session_id=session1.id,
|
||||
prompt=[TextContentBlock(type="text", text="Prompt for session 1")],
|
||||
)
|
||||
|
||||
async def run_session2():
|
||||
await acp_agent.prompt(
|
||||
PromptRequest(
|
||||
sessionId=session2.id,
|
||||
prompt=[TextContentBlock(type="text", text="Prompt for session 2")],
|
||||
)
|
||||
await acp_agent_loop.prompt(
|
||||
session_id=session2.id,
|
||||
prompt=[TextContentBlock(type="text", text="Prompt for session 2")],
|
||||
)
|
||||
|
||||
await asyncio.gather(run_session1(), run_session2())
|
||||
|
||||
user_message1 = next(
|
||||
(msg for msg in session1.agent.messages if msg.role == Role.user), None
|
||||
(msg for msg in session1.agent_loop.messages if msg.role == Role.user), None
|
||||
)
|
||||
assert user_message1 is not None
|
||||
assert user_message1.content == "Prompt for session 1"
|
||||
assistant_message1 = next(
|
||||
(msg for msg in session1.agent.messages if msg.role == Role.assistant), None
|
||||
(msg for msg in session1.agent_loop.messages if msg.role == Role.assistant),
|
||||
None,
|
||||
)
|
||||
assert assistant_message1 is not None
|
||||
assert assistant_message1.content == "Response 1"
|
||||
user_message2 = next(
|
||||
(msg for msg in session2.agent.messages if msg.role == Role.user), None
|
||||
(msg for msg in session2.agent_loop.messages if msg.role == Role.user), None
|
||||
)
|
||||
assert user_message2 is not None
|
||||
assert user_message2.content == "Prompt for session 2"
|
||||
assistant_message2 = next(
|
||||
(msg for msg in session2.agent.messages if msg.role == Role.assistant), None
|
||||
(msg for msg in session2.agent_loop.messages if msg.role == Role.assistant),
|
||||
None,
|
||||
)
|
||||
assert assistant_message2 is not None
|
||||
assert assistant_message2.content == "Response 2"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue