openwebui-cli/tests/test_chat_rag.py
2025-12-01 04:24:51 +01:00

762 lines
27 KiB
Python

"""Tests for RAG context features in chat commands."""
import json
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch
import pytest
from typer.testing import CliRunner
from openwebui_cli.main import app
runner = CliRunner()
@pytest.fixture
def mock_config(tmp_path, monkeypatch):
"""Mock configuration for testing."""
config_dir = tmp_path / "openwebui"
config_path = config_dir / "config.yaml"
monkeypatch.setattr("openwebui_cli.config.get_config_dir", lambda: config_dir)
monkeypatch.setattr("openwebui_cli.config.get_config_path", lambda: config_path)
# Create default config
from openwebui_cli.config import Config, save_config
config = Config()
save_config(config)
return config_path
@pytest.fixture
def mock_keyring(monkeypatch):
"""Mock keyring for testing."""
token_store = {}
def get_password(service, key):
return token_store.get(f"{service}:{key}")
def set_password(service, key, password):
token_store[f"{service}:{key}"] = password
monkeypatch.setattr("keyring.get_password", get_password)
monkeypatch.setattr("keyring.set_password", set_password)
class TestRAGContextFeatures:
"""Test suite for RAG context features."""
def test_file_and_collection_together(self, mock_config, mock_keyring):
"""Test --file and --collection populate body['files'] correctly."""
response_data = {
"choices": [
{
"message": {
"content": "Response with RAG context"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Search my docs",
"--file", "file-id-123",
"--collection", "collection-xyz",
"--no-stream"
],
)
assert result.exit_code == 0
# Verify the request body structure
call_args = mock_http_client.post.call_args
assert call_args is not None
body = call_args.kwargs["json"]
# Assert 'files' key exists in body
assert "files" in body, "body should contain 'files' key"
# Assert correct number of entries
assert len(body["files"]) == 2, "should have 2 entries (1 file, 1 collection)"
# Check types are present
types = [f["type"] for f in body["files"]]
assert "file" in types, "should have 'file' type"
assert "collection" in types, "should have 'collection' type"
# Verify correct IDs
file_entry = next((f for f in body["files"] if f["type"] == "file"), None)
collection_entry = next((f for f in body["files"] if f["type"] == "collection"), None)
assert file_entry is not None, "should have file entry"
assert collection_entry is not None, "should have collection entry"
assert file_entry["id"] == "file-id-123", "file ID should match"
assert collection_entry["id"] == "collection-xyz", "collection ID should match"
def test_file_only(self, mock_config, mock_keyring):
"""Test --file alone populates body['files'] with only file entry."""
response_data = {
"choices": [
{
"message": {
"content": "Response with file context"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Use this file",
"--file", "file-456",
"--no-stream"
],
)
assert result.exit_code == 0
# Verify the request body
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "files" in body
assert len(body["files"]) == 1
assert body["files"][0]["type"] == "file"
assert body["files"][0]["id"] == "file-456"
def test_collection_only(self, mock_config, mock_keyring):
"""Test --collection alone populates body['files'] with only collection entry."""
response_data = {
"choices": [
{
"message": {
"content": "Response with collection context"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Search the collection",
"--collection", "docs-789",
"--no-stream"
],
)
assert result.exit_code == 0
# Verify the request body
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "files" in body
assert len(body["files"]) == 1
assert body["files"][0]["type"] == "collection"
assert body["files"][0]["id"] == "docs-789"
def test_multiple_files(self, mock_config, mock_keyring):
"""Test multiple --file options work correctly."""
response_data = {
"choices": [
{
"message": {
"content": "Response"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Search multiple files",
"--file", "file-1",
"--file", "file-2",
"--file", "file-3",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "files" in body
assert len(body["files"]) == 3
# All should be of type 'file'
for entry in body["files"]:
assert entry["type"] == "file"
# Check all IDs are present
ids = [f["id"] for f in body["files"]]
assert "file-1" in ids
assert "file-2" in ids
assert "file-3" in ids
def test_multiple_collections(self, mock_config, mock_keyring):
"""Test multiple --collection options work correctly."""
response_data = {
"choices": [
{
"message": {
"content": "Response"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Search multiple collections",
"--collection", "coll-a",
"--collection", "coll-b",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "files" in body
assert len(body["files"]) == 2
# All should be of type 'collection'
for entry in body["files"]:
assert entry["type"] == "collection"
# Check all IDs are present
ids = [f["id"] for f in body["files"]]
assert "coll-a" in ids
assert "coll-b" in ids
def test_mixed_files_and_collections(self, mock_config, mock_keyring):
"""Test combination of multiple files and collections."""
response_data = {
"choices": [
{
"message": {
"content": "Response"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Search mixed context",
"--file", "file-1",
"--file", "file-2",
"--collection", "coll-x",
"--collection", "coll-y",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "files" in body
assert len(body["files"]) == 4
# Verify structure
file_entries = [f for f in body["files"] if f["type"] == "file"]
collection_entries = [f for f in body["files"] if f["type"] == "collection"]
assert len(file_entries) == 2
assert len(collection_entries) == 2
file_ids = [f["id"] for f in file_entries]
collection_ids = [f["id"] for f in collection_entries]
assert "file-1" in file_ids
assert "file-2" in file_ids
assert "coll-x" in collection_ids
assert "coll-y" in collection_ids
def test_no_rag_context(self, mock_config, mock_keyring):
"""Test that files key is not present when no RAG context specified."""
response_data = {
"choices": [
{
"message": {
"content": "Response without RAG"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Hello without RAG",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
# files key should not be present
assert "files" not in body
def test_rag_with_system_prompt(self, mock_config, mock_keyring):
"""Test RAG context works alongside system prompt."""
response_data = {
"choices": [
{
"message": {
"content": "Response with system and RAG"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Question about docs",
"-s", "You are a helpful assistant",
"--file", "file-doc",
"--collection", "coll-main",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
# Should have both system message and RAG files
assert "messages" in body
assert any(msg.get("role") == "system" for msg in body["messages"])
assert "files" in body
assert len(body["files"]) == 2
def test_rag_with_chat_id(self, mock_config, mock_keyring):
"""Test RAG context works with chat_id for conversation continuation."""
response_data = {
"choices": [
{
"message": {
"content": "Continued response with RAG"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Continue with docs",
"--chat-id", "chat-xyz-123",
"--file", "file-continuing",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "chat_id" in body
assert body["chat_id"] == "chat-xyz-123"
assert "files" in body
assert len(body["files"]) == 1
def test_rag_with_temperature_and_tokens(self, mock_config, mock_keyring):
"""Test RAG context works with temperature and max_tokens."""
response_data = {
"choices": [
{
"message": {
"content": "Response with temperature"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Creative response",
"-T", "1.5",
"--max-tokens", "500",
"--file", "file-creative",
"--collection", "coll-creative",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert body["temperature"] == 1.5
assert body["max_tokens"] == 500
assert "files" in body
assert len(body["files"]) == 2
def test_rag_streaming_with_context(self, mock_config, mock_keyring):
"""Test RAG context works with streaming responses."""
streaming_lines = [
'data: {"choices": [{"delta": {"content": "Response"}}]}',
'data: {"choices": [{"delta": {"content": " with"}}]}',
'data: {"choices": [{"delta": {"content": " RAG"}}]}',
"data: [DONE]",
]
class MockStreamResponse:
def __init__(self, lines):
self.lines = lines
self.status_code = 200
def __enter__(self):
return self
def __exit__(self, *args):
pass
def iter_lines(self):
for line in self.lines:
yield line
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_stream = MockStreamResponse(streaming_lines)
mock_http_client.stream.return_value = mock_stream
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Stream with RAG",
"--file", "file-stream",
"--collection", "coll-stream",
],
)
assert result.exit_code == 0
assert "Response with RAG" in result.stdout
# Verify streaming request was made with RAG context
call_args = mock_http_client.stream.call_args
body = call_args.kwargs["json"]
assert "files" in body
assert len(body["files"]) == 2
def test_rag_context_structure_validation(self, mock_config, mock_keyring):
"""Test that RAG context entries have correct structure."""
response_data = {
"choices": [
{
"message": {
"content": "Response"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Test structure",
"--file", "f1",
"--collection", "c1",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
# Validate structure of each entry
for entry in body["files"]:
assert "type" in entry, "Each entry must have 'type' field"
assert "id" in entry, "Each entry must have 'id' field"
assert entry["type"] in ["file", "collection"], "type must be 'file' or 'collection'"
assert isinstance(entry["id"], str), "id must be a string"
assert len(entry) == 2, "Entry should only have 'type' and 'id' fields"
class TestRAGEdgeCases:
"""Test edge cases and error handling for RAG context."""
def test_empty_file_id_handling(self, mock_config, mock_keyring):
"""Test handling of empty file ID."""
response_data = {
"choices": [
{
"message": {
"content": "Response"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Test",
"--file", "",
"--no-stream"
],
)
# Should still execute but with empty ID
call_args = mock_http_client.post.call_args
if call_args:
body = call_args.kwargs["json"]
# Even empty IDs should be passed through
if "files" in body:
assert any(f["id"] == "" for f in body["files"] if f["type"] == "file")
def test_special_characters_in_ids(self, mock_config, mock_keyring):
"""Test IDs with special characters."""
response_data = {
"choices": [
{
"message": {
"content": "Response"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
result = runner.invoke(
app,
[
"chat", "send",
"-m", "test-model",
"-p", "Test",
"--file", "file-with-dashes-123_special.chars",
"--collection", "coll/with/slashes",
"--no-stream"
],
)
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "files" in body
ids = [f["id"] for f in body["files"]]
assert "file-with-dashes-123_special.chars" in ids
assert "coll/with/slashes" in ids
def test_large_number_of_files(self, mock_config, mock_keyring):
"""Test handling many files in context."""
response_data = {
"choices": [
{
"message": {
"content": "Response"
}
}
]
}
with patch("openwebui_cli.commands.chat.create_client") as mock_client_factory:
mock_http_client = MagicMock()
mock_http_client.__enter__.return_value = mock_http_client
mock_http_client.__exit__.return_value = None
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = response_data
mock_http_client.post.return_value = mock_response
mock_client_factory.return_value = mock_http_client
# Build command with many files
cmd = ["chat", "send", "-m", "test-model", "-p", "Test"]
for i in range(10):
cmd.extend(["--file", f"file-{i}"])
result = runner.invoke(app, cmd + ["--no-stream"])
assert result.exit_code == 0
call_args = mock_http_client.post.call_args
body = call_args.kwargs["json"]
assert "files" in body
assert len(body["files"]) == 10
assert all(f["type"] == "file" for f in body["files"])