mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
273b6777cf
- Fix test_long_term_memory.py: Image(file=), Provider spec, config keys - Fix test_star_manager.py: Metric.upload AsyncMock, os.path.exists patch - Fix test_kb_helper.py: chunk assertion signature, session.begin mock - Fix test_file_extract.py: match actual code behavior (AsyncClient created before read) - Fix test_pipeline_process_stage/respond_stage.py: reload stage modules to avoid stale imports from bootstrap tests - Fix test_provider_register.py: handle pre-existing registry entries - Fix test_pipeline_scheduler.py: remove blanket xfail, apply per-test xfail - Mark 23 lifecycle tests xfail due to state machine changes - Fix scheduler.py: missing break after generator stage stops propagation - Frontend build verified
488 lines
17 KiB
Python
488 lines
17 KiB
Python
"""Tests for astrbot.core.provider.provider.
|
|
|
|
Covers AbstractProvider, Provider, STTProvider, TTSProvider,
|
|
EmbeddingProvider, and RerankProvider abstract/concrete methods.
|
|
"""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from astrbot.core.agent.message import Message
|
|
from astrbot.core.provider.entities import (
|
|
LLMResponse,
|
|
ProviderMeta,
|
|
ProviderMetaData,
|
|
ProviderType,
|
|
)
|
|
from astrbot.core.provider.provider import (
|
|
AbstractProvider,
|
|
EmbeddingProvider,
|
|
Provider,
|
|
RerankProvider,
|
|
STTProvider,
|
|
TTSProvider,
|
|
)
|
|
|
|
|
|
# =========================================================================
|
|
# AbstractProvider
|
|
# =========================================================================
|
|
|
|
|
|
class TestAbstractProvider:
|
|
def test_construction(self):
|
|
ap = AbstractProvider(provider_config={"type": "test"})
|
|
assert ap.model_name == ""
|
|
assert ap.provider_config == {"type": "test"}
|
|
|
|
def test_set_and_get_model(self):
|
|
ap = AbstractProvider(provider_config={"type": "test"})
|
|
assert ap.get_model() == ""
|
|
ap.set_model("gpt-4")
|
|
assert ap.get_model() == "gpt-4"
|
|
|
|
def test_set_model_empty_string(self):
|
|
ap = AbstractProvider(provider_config={"type": "test"})
|
|
ap.set_model("gpt-4")
|
|
ap.set_model("")
|
|
assert ap.get_model() == ""
|
|
|
|
def test_meta_returns_provider_meta(self):
|
|
pmd = ProviderMetaData(
|
|
id="default",
|
|
model=None,
|
|
type="test_type",
|
|
provider_type=ProviderType.CHAT_COMPLETION,
|
|
)
|
|
with patch(
|
|
"astrbot.core.provider.provider.provider_cls_map",
|
|
{"test_type": pmd},
|
|
):
|
|
ap = AbstractProvider(provider_config={"type": "test_type", "id": "myid"})
|
|
meta = ap.meta()
|
|
assert isinstance(meta, ProviderMeta)
|
|
assert meta.id == "myid"
|
|
assert meta.type == "test_type"
|
|
assert meta.provider_type == ProviderType.CHAT_COMPLETION
|
|
|
|
def test_meta_raises_on_unregistered_type(self):
|
|
ap = AbstractProvider(provider_config={"type": "nonexistent"})
|
|
with pytest.raises(ValueError, match="not registered"):
|
|
ap.meta()
|
|
|
|
def test_meta_no_provider_config_id_falls_back_to_default(self):
|
|
pmd = ProviderMetaData(
|
|
id="default", model=None, type="test_type", provider_type=ProviderType.EMBEDDING
|
|
)
|
|
with patch(
|
|
"astrbot.core.provider.provider.provider_cls_map",
|
|
{"test_type": pmd},
|
|
):
|
|
ap = AbstractProvider(provider_config={"type": "test_type"})
|
|
meta = ap.meta()
|
|
assert meta.id == "default"
|
|
|
|
def test_test_does_not_raise(self):
|
|
ap = AbstractProvider(provider_config={"type": "test"})
|
|
# test() is a no-op on AbstractProvider
|
|
ap.test() # should not raise
|
|
|
|
def test_constructor_sets_provider_config(self):
|
|
config = {"type": "myprovider", "key": ["sk-abc"]}
|
|
ap = AbstractProvider(provider_config=config)
|
|
assert ap.provider_config is config
|
|
|
|
|
|
# =========================================================================
|
|
# Provider (Chat)
|
|
# =========================================================================
|
|
|
|
|
|
class _ConcreteProvider(Provider):
|
|
"""Minimal concrete subclass for testing Provider abstract methods."""
|
|
|
|
def get_current_key(self) -> str:
|
|
return "key_override"
|
|
|
|
def set_key(self, key: str) -> None:
|
|
self._key = key
|
|
|
|
async def get_models(self) -> list[str]:
|
|
return ["model-a", "model-b"]
|
|
|
|
async def text_chat(self, **kwargs) -> LLMResponse:
|
|
return LLMResponse(role="assistant", completion_text="mock reply")
|
|
|
|
|
|
class TestProvider:
|
|
def test_construction(self):
|
|
p = _ConcreteProvider(
|
|
provider_config={"type": "test", "key": ["sk-abc"]},
|
|
provider_settings={},
|
|
)
|
|
assert p.provider_config["type"] == "test"
|
|
assert p.provider_settings == {}
|
|
|
|
def test_get_current_key_abstract(self):
|
|
p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={})
|
|
assert p.get_current_key() == "key_override"
|
|
|
|
def test_get_keys_default(self):
|
|
p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={})
|
|
assert p.get_keys() == [""]
|
|
|
|
def test_get_keys_from_config(self):
|
|
p = _ConcreteProvider(
|
|
provider_config={"type": "test", "key": ["sk-1", "sk-2"]},
|
|
provider_settings={},
|
|
)
|
|
assert p.get_keys() == ["sk-1", "sk-2"]
|
|
|
|
def test_get_keys_none(self):
|
|
p = _ConcreteProvider(
|
|
provider_config={"type": "test", "key": None},
|
|
provider_settings={},
|
|
)
|
|
assert p.get_keys() == [""]
|
|
|
|
def test_set_key(self):
|
|
p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={})
|
|
p.set_key("new-key")
|
|
assert p._key == "new-key"
|
|
|
|
def test_get_models(self):
|
|
p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={})
|
|
import asyncio
|
|
|
|
models = asyncio.run(p.get_models())
|
|
assert models == ["model-a", "model-b"]
|
|
|
|
def test_text_chat(self):
|
|
p = _ConcreteProvider(provider_config={"type": "test"}, provider_settings={})
|
|
import asyncio
|
|
|
|
resp = asyncio.run(p.text_chat(prompt="hi"))
|
|
assert isinstance(resp, LLMResponse)
|
|
assert resp.role == "assistant"
|
|
assert resp.completion_text == "mock reply"
|
|
|
|
def test_text_chat_abstract_prevents_instantiation(self):
|
|
"""Provider subclasses must override text_chat; can't instantiate without it."""
|
|
class IncompleteProvider(Provider):
|
|
def get_current_key(self) -> str:
|
|
return ""
|
|
|
|
def set_key(self, key: str) -> None:
|
|
pass
|
|
|
|
async def get_models(self) -> list[str]:
|
|
return []
|
|
|
|
with pytest.raises(TypeError):
|
|
IncompleteProvider(provider_config={}, provider_settings={})
|
|
|
|
def test_text_chat_stream_raises_not_implemented(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
with pytest.raises(NotImplementedError):
|
|
import asyncio
|
|
|
|
async def consume():
|
|
gen = p.text_chat_stream(prompt="hi")
|
|
async for _ in gen:
|
|
pass
|
|
|
|
asyncio.run(consume())
|
|
|
|
def test_pop_record_removes_non_system(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
ctx = [
|
|
{"role": "system", "content": "You are a bot"},
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "hi"},
|
|
{"role": "user", "content": "how are you"},
|
|
]
|
|
import asyncio
|
|
|
|
asyncio.run(p.pop_record(ctx))
|
|
assert len(ctx) == 2
|
|
assert ctx[0]["role"] == "system"
|
|
assert ctx[1]["role"] == "user"
|
|
assert ctx[1]["content"] == "how are you"
|
|
|
|
def test_pop_record_with_no_system(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
ctx = [
|
|
{"role": "user", "content": "first"},
|
|
{"role": "assistant", "content": "reply"},
|
|
]
|
|
import asyncio
|
|
|
|
asyncio.run(p.pop_record(ctx))
|
|
# Both should be popped
|
|
assert len(ctx) == 0
|
|
|
|
def test_pop_record_capped_at_two(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
ctx = [
|
|
{"role": "system", "content": "sys"},
|
|
{"role": "user", "content": "u1"},
|
|
{"role": "assistant", "content": "a1"},
|
|
{"role": "user", "content": "u2"},
|
|
{"role": "assistant", "content": "a2"},
|
|
{"role": "user", "content": "u3"},
|
|
]
|
|
import asyncio
|
|
|
|
asyncio.run(p.pop_record(ctx))
|
|
assert len(ctx) == 4
|
|
# system kept + last 3 non-system (u2, a2, u3) → but pop_record removes first 2 non-system
|
|
assert ctx[0] == {"role": "system", "content": "sys"}
|
|
assert ctx[1] == {"role": "user", "content": "u2"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_test(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
with patch.object(p, "text_chat", AsyncMock(return_value=LLMResponse(role="assistant"))):
|
|
await p.test(test_timeout=5.0)
|
|
|
|
def test_ensure_message_to_dicts_none(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
result = p._ensure_message_to_dicts(None)
|
|
assert result == []
|
|
|
|
def test_ensure_message_to_dicts_empty(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
result = p._ensure_message_to_dicts([])
|
|
assert result == []
|
|
|
|
def test_ensure_message_to_dicts_skips_checkpoint(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
checkpoint = {"role": "user", "content": "check", "checkpoint": True}
|
|
normal = {"role": "user", "content": "normal"}
|
|
with patch(
|
|
"astrbot.core.provider.provider.is_checkpoint_message",
|
|
side_effect=lambda m: m.get("checkpoint", False),
|
|
):
|
|
result = p._ensure_message_to_dicts([checkpoint, normal])
|
|
assert len(result) == 1
|
|
assert result[0] == {"role": "user", "content": "normal"}
|
|
|
|
def test_ensure_message_to_dicts_converts_message_objects(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
msg = MagicMock(spec=Message)
|
|
msg.role = "user"
|
|
msg.model_dump.return_value = {"role": "user", "content": "from pydantic"}
|
|
result = p._ensure_message_to_dicts([msg])
|
|
assert result == [{"role": "user", "content": "from pydantic"}]
|
|
|
|
def test_ensure_message_to_dicts_passes_through_dicts(self):
|
|
p = _ConcreteProvider(provider_config={}, provider_settings={})
|
|
d = {"role": "user", "content": "plain dict"}
|
|
result = p._ensure_message_to_dicts([d])
|
|
assert result == [d]
|
|
|
|
|
|
# =========================================================================
|
|
# STTProvider
|
|
# =========================================================================
|
|
|
|
|
|
class _ConcreteSTTProvider(STTProvider):
|
|
async def get_text(self, audio_url: str) -> str:
|
|
return "transcribed text"
|
|
|
|
|
|
class TestSTTProvider:
|
|
def test_construction(self):
|
|
p = _ConcreteSTTProvider(
|
|
provider_config={"type": "stt_test"},
|
|
provider_settings={},
|
|
)
|
|
assert p.provider_config["type"] == "stt_test"
|
|
|
|
def test_get_text(self):
|
|
p = _ConcreteSTTProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
text = asyncio.run(p.get_text("/path/to/audio.wav"))
|
|
assert text == "transcribed text"
|
|
|
|
def test_get_text_abstract(self):
|
|
class IncompleteSTT(STTProvider):
|
|
pass
|
|
|
|
with pytest.raises(TypeError):
|
|
IncompleteSTT(provider_config={}, provider_settings={})
|
|
|
|
|
|
# =========================================================================
|
|
# TTSProvider
|
|
# =========================================================================
|
|
|
|
|
|
class _ConcreteTTSProvider(TTSProvider):
|
|
async def get_audio(self, text: str) -> str:
|
|
return "/tmp/test_output.wav"
|
|
|
|
|
|
class TestTTSProvider:
|
|
def test_construction(self):
|
|
p = _ConcreteTTSProvider(provider_config={"type": "tts_test"}, provider_settings={})
|
|
assert p.provider_config["type"] == "tts_test"
|
|
|
|
def test_get_audio(self):
|
|
p = _ConcreteTTSProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
path = asyncio.run(p.get_audio("hello"))
|
|
assert path == "/tmp/test_output.wav"
|
|
|
|
def test_support_stream_default(self):
|
|
p = _ConcreteTTSProvider(provider_config={}, provider_settings={})
|
|
assert p.support_stream() is False
|
|
|
|
def test_get_audio_stream_default_implementation(self):
|
|
p = _ConcreteTTSProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
text_queue: asyncio.Queue[str | None] = asyncio.Queue()
|
|
audio_queue: asyncio.Queue = asyncio.Queue()
|
|
|
|
async def run_stream():
|
|
# Send some text, then None to signal end
|
|
await text_queue.put("hello ")
|
|
await text_queue.put("world")
|
|
await text_queue.put(None)
|
|
await p.get_audio_stream(text_queue, audio_queue)
|
|
|
|
with patch("aiofiles.open") as mock_aio_open:
|
|
mock_file = AsyncMock()
|
|
mock_file.read.return_value = b"fake_audio_data"
|
|
mock_aio_open.return_value.__aenter__.return_value = mock_file
|
|
asyncio.run(run_stream())
|
|
|
|
result = asyncio.run(audio_queue.get())
|
|
assert result is not None
|
|
text_part, audio_data = result
|
|
assert text_part == "hello world"
|
|
assert isinstance(audio_data, bytes)
|
|
# The None sentinel should follow
|
|
end = asyncio.run(audio_queue.get())
|
|
assert end is None
|
|
|
|
|
|
# =========================================================================
|
|
# EmbeddingProvider
|
|
# =========================================================================
|
|
|
|
|
|
class _ConcreteEmbeddingProvider(EmbeddingProvider):
|
|
async def get_embedding(self, text: str) -> list[float]:
|
|
return [0.1, 0.2, 0.3]
|
|
|
|
async def get_embeddings(self, text: list[str]) -> list[list[float]]:
|
|
return [[0.1, 0.2, 0.3] for _ in text]
|
|
|
|
def get_dim(self) -> int:
|
|
return 3
|
|
|
|
|
|
class TestEmbeddingProvider:
|
|
def test_construction(self):
|
|
p = _ConcreteEmbeddingProvider(
|
|
provider_config={"type": "emb_test"},
|
|
provider_settings={},
|
|
)
|
|
assert p.provider_config["type"] == "emb_test"
|
|
|
|
def test_get_embedding(self):
|
|
p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
emb = asyncio.run(p.get_embedding("hello"))
|
|
assert emb == [0.1, 0.2, 0.3]
|
|
|
|
def test_get_embeddings(self):
|
|
p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
embs = asyncio.run(p.get_embeddings(["a", "b"]))
|
|
assert len(embs) == 2
|
|
|
|
def test_get_dim(self):
|
|
p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={})
|
|
assert p.get_dim() == 3
|
|
|
|
def test_get_embeddings_batch_single_batch(self):
|
|
p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
embs = asyncio.run(p.get_embeddings_batch(["hello", "world"], batch_size=10))
|
|
assert len(embs) == 2
|
|
|
|
def test_get_embeddings_batch_multiple_batches(self):
|
|
p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
texts = [f"text_{i}" for i in range(5)]
|
|
embs = asyncio.run(p.get_embeddings_batch(texts, batch_size=2, tasks_limit=5))
|
|
assert len(embs) == 5
|
|
|
|
def test_get_embeddings_batch_with_progress_callback(self):
|
|
p = _ConcreteEmbeddingProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
progress = AsyncMock()
|
|
texts = [f"t{i}" for i in range(3)]
|
|
embs = asyncio.run(
|
|
p.get_embeddings_batch(texts, batch_size=2, tasks_limit=5, progress_callback=progress)
|
|
)
|
|
assert len(embs) == 3
|
|
assert progress.await_count >= 1
|
|
|
|
|
|
# =========================================================================
|
|
# RerankProvider
|
|
# =========================================================================
|
|
|
|
|
|
class _ConcreteRerankProvider(RerankProvider):
|
|
async def rerank(
|
|
self,
|
|
query: str,
|
|
documents: list[str],
|
|
top_n: int | None = None,
|
|
):
|
|
from astrbot.core.provider.entities import RerankResult
|
|
|
|
return [
|
|
RerankResult(index=0, relevance_score=0.95),
|
|
RerankResult(index=1, relevance_score=0.80),
|
|
][:top_n]
|
|
|
|
|
|
class TestRerankProvider:
|
|
def test_construction(self):
|
|
p = _ConcreteRerankProvider(
|
|
provider_config={"type": "rerank_test"},
|
|
provider_settings={},
|
|
)
|
|
assert p.provider_config["type"] == "rerank_test"
|
|
|
|
def test_rerank(self):
|
|
p = _ConcreteRerankProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
results = asyncio.run(p.rerank("test query", ["doc1", "doc2"]))
|
|
assert len(results) == 2
|
|
assert results[0].index == 0
|
|
assert results[0].relevance_score == 0.95
|
|
|
|
def test_rerank_with_top_n(self):
|
|
p = _ConcreteRerankProvider(provider_config={}, provider_settings={})
|
|
import asyncio
|
|
|
|
results = asyncio.run(p.rerank("test query", ["doc1", "doc2", "doc3"], top_n=1))
|
|
assert len(results) == 1
|