mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Literal, overload
|
|
|
|
import pytest
|
|
import tomlkit
|
|
|
|
from plyngent.agent import ChatAgent
|
|
from plyngent.cli.repl import handle_slash
|
|
from plyngent.cli.state import ReplState
|
|
from plyngent.config.models import DatabaseConfig, OpenAIProvider
|
|
from plyngent.config.store import ConfigStore
|
|
from plyngent.lmproto.openai_compatible.model import (
|
|
AssistantChatMessage,
|
|
ChatCompletionChoice,
|
|
ChatCompletionChunk,
|
|
ChatCompletionResponse,
|
|
ChatCompletionsParam,
|
|
)
|
|
from plyngent.memory import MemoryStore
|
|
from plyngent.tools import set_workspace_root
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
from pathlib import Path
|
|
|
|
|
|
class DummyClient:
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param
|
|
if stream:
|
|
|
|
async def empty() -> AsyncIterator[ChatCompletionChunk]:
|
|
empty_chunks: list[ChatCompletionChunk] = []
|
|
for chunk in empty_chunks:
|
|
yield chunk
|
|
|
|
return empty()
|
|
return ChatCompletionResponse(
|
|
id="1",
|
|
object="chat.completion",
|
|
created=0,
|
|
model="m",
|
|
choices=[
|
|
ChatCompletionChoice(
|
|
index=0,
|
|
message=AssistantChatMessage(content="ok"),
|
|
logprobs={},
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
system_fingerprint="",
|
|
usage={},
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
async def state(tmp_path: Path) -> AsyncIterator[ReplState]:
|
|
_ = set_workspace_root(tmp_path)
|
|
memory = await MemoryStore.open(DatabaseConfig())
|
|
provider = OpenAIProvider(access_key_or_token="sk-test")
|
|
config = ConfigStore(path=tmp_path / "plyngent.toml", document=tomlkit.document())
|
|
config.providers = {"local": provider}
|
|
st = ReplState(
|
|
config=config,
|
|
memory=memory,
|
|
workspace=tmp_path,
|
|
provider_name="local",
|
|
provider=provider,
|
|
model="gpt-test",
|
|
tools_enabled=False,
|
|
)
|
|
st.client = DummyClient()
|
|
st.agent = ChatAgent(st.client, model=st.model, memory=st.memory, session_id=None)
|
|
await st.new_session("t")
|
|
yield st
|
|
await memory.close()
|
|
|
|
|
|
async def test_help_and_clear(state: ReplState) -> None:
|
|
assert await handle_slash(state, "/help") is True
|
|
state.agent.messages.append(AssistantChatMessage(content="x"))
|
|
assert await handle_slash(state, "/clear") is True
|
|
assert state.agent.messages == []
|
|
|
|
|
|
async def test_quit(state: ReplState) -> None:
|
|
assert await handle_slash(state, "/quit") is False
|
|
|
|
|
|
async def test_new_and_sessions(state: ReplState, capsys: pytest.CaptureFixture[str]) -> None:
|
|
first = state.session_id
|
|
assert await handle_slash(state, "/new other") is True
|
|
assert state.session_id != first
|
|
assert await handle_slash(state, "/sessions") is True
|
|
out = capsys.readouterr().out
|
|
assert str(state.session_id) in out
|
|
|
|
|
|
async def test_tools_toggle(state: ReplState) -> None:
|
|
assert await handle_slash(state, "/tools on") is True
|
|
assert state.tools_enabled is True
|
|
assert await handle_slash(state, "/tools off") is True
|
|
assert state.tools_enabled is False
|
|
|
|
|
|
async def test_resume(state: ReplState) -> None:
|
|
sid = state.session_id
|
|
assert sid is not None
|
|
state.agent.messages.clear()
|
|
assert await handle_slash(state, f"/resume {sid}") is True
|