chore: import upstream snapshot with attribution
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,224 @@
|
||||
"""Tests for application/api/answer/routes/answer.py"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Static IDs
|
||||
_CONV_ID = "507f1f77bcf86cd799439011"
|
||||
_AGENT_ID = "507f1f77bcf86cd799439012"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_stream_processor():
|
||||
"""Create a mock StreamProcessor."""
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.StreamProcessor"
|
||||
) as MockProcessor:
|
||||
processor = MagicMock()
|
||||
processor.decoded_token = {"sub": "test_user"}
|
||||
processor.conversation_id = _CONV_ID
|
||||
processor.agent_config = {}
|
||||
processor.agent_id = _AGENT_ID
|
||||
processor.is_shared_usage = False
|
||||
processor.shared_token = None
|
||||
processor.model_id = "gpt-4"
|
||||
processor.build_agent.return_value = MagicMock()
|
||||
MockProcessor.return_value = processor
|
||||
yield processor
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def answer_client(mock_mongo_db, flask_app):
|
||||
"""Create a test client with the answer route registered."""
|
||||
from flask_restx import Api
|
||||
|
||||
from application.api.answer.routes.answer import answer_ns
|
||||
|
||||
api = Api(flask_app)
|
||||
api.add_namespace(answer_ns)
|
||||
flask_app.config["TESTING"] = True
|
||||
return flask_app.test_client()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAnswerResourcePost:
|
||||
def test_missing_question_returns_400(self, answer_client, mock_stream_processor):
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_successful_answer(self, answer_client, mock_stream_processor):
|
||||
conv_id = str(uuid.uuid4())
|
||||
with patch.object(
|
||||
mock_stream_processor.build_agent.return_value,
|
||||
"gen",
|
||||
return_value=iter([]),
|
||||
):
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.check_usage",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.complete_stream",
|
||||
return_value=iter(
|
||||
[
|
||||
f'data: {json.dumps({"type": "answer", "answer": "Hello"})}\n\n',
|
||||
f'data: {json.dumps({"type": "id", "id": conv_id})}\n\n',
|
||||
f'data: {json.dumps({"type": "end"})}\n\n',
|
||||
]
|
||||
),
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.process_response_stream",
|
||||
return_value={"conversation_id": conv_id, "answer": "Hello", "sources": [], "tool_calls": [], "thought": "", "error": None},
|
||||
):
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({"question": "What is Python?"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.get_json()
|
||||
assert data["answer"] == "Hello"
|
||||
assert data["conversation_id"] == conv_id
|
||||
|
||||
def test_unauthorized_returns_401(self, answer_client, mock_stream_processor):
|
||||
mock_stream_processor.decoded_token = None
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.validate_request",
|
||||
return_value=None,
|
||||
):
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
assert resp.get_json()["error"] == "Unauthorized"
|
||||
|
||||
def test_usage_exceeded_returns_error(self, answer_client, mock_stream_processor):
|
||||
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.check_usage",
|
||||
) as mock_check:
|
||||
with flask_app_context(answer_client):
|
||||
mock_check.return_value = ({"error": "Usage limit exceeded"}, 429)
|
||||
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 429
|
||||
|
||||
def test_stream_error_returns_400(self, answer_client, mock_stream_processor):
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.check_usage",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.complete_stream",
|
||||
return_value=iter([]),
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.process_response_stream",
|
||||
return_value={"conversation_id": None, "answer": None, "sources": None, "tool_calls": None, "thought": None, "error": "Stream error"},
|
||||
):
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
assert resp.get_json()["error"] == "Stream error"
|
||||
|
||||
def test_exception_returns_500(self, answer_client, mock_stream_processor):
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.check_usage",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.complete_stream",
|
||||
side_effect=RuntimeError("unexpected"),
|
||||
):
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 500
|
||||
assert "error" in resp.get_json()
|
||||
|
||||
def test_structured_info_merged_into_result(
|
||||
self, answer_client, mock_stream_processor
|
||||
):
|
||||
conv_id = str(uuid.uuid4())
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.check_usage",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.complete_stream",
|
||||
return_value=iter([]),
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.process_response_stream",
|
||||
return_value={"conversation_id": conv_id, "answer": '{"key": "val"}', "sources": [], "tool_calls": [], "thought": "", "error": None, "extra": {"structured": True, "schema": {"type": "object"}}},
|
||||
):
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.get_json()
|
||||
assert data["structured"] is True
|
||||
assert data["schema"] == {"type": "object"}
|
||||
|
||||
def test_result_contains_all_expected_fields(
|
||||
self, answer_client, mock_stream_processor
|
||||
):
|
||||
conv_id = str(uuid.uuid4())
|
||||
with patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.check_usage",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.complete_stream",
|
||||
return_value=iter([]),
|
||||
), patch(
|
||||
"application.api.answer.routes.answer.AnswerResource.process_response_stream",
|
||||
return_value={"conversation_id": conv_id, "answer": "answer text", "sources": [{"title": "src"}], "tool_calls": [{"tool": "t"}], "thought": "thinking...", "error": None},
|
||||
):
|
||||
resp = answer_client.post(
|
||||
"/api/answer",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
data = resp.get_json()
|
||||
assert data["conversation_id"] == conv_id
|
||||
assert data["answer"] == "answer text"
|
||||
assert data["sources"] == [{"title": "src"}]
|
||||
assert data["tool_calls"] == [{"tool": "t"}]
|
||||
assert data["thought"] == "thinking..."
|
||||
|
||||
|
||||
def flask_app_context(client):
|
||||
"""Helper to get app context from test client."""
|
||||
return client.application.app_context()
|
||||
@@ -0,0 +1,700 @@
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestBaseAnswerValidation:
|
||||
pass
|
||||
|
||||
def test_validate_request_passes_with_required_fields(
|
||||
self, mock_mongo_db, flask_app
|
||||
):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
data = {"question": "What is Python?"}
|
||||
|
||||
result = resource.validate_request(data)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_validate_request_fails_without_question(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
data = {}
|
||||
|
||||
result = resource.validate_request(data)
|
||||
|
||||
assert result is not None
|
||||
assert result.status_code == 400
|
||||
assert "question" in result.json["message"].lower()
|
||||
|
||||
def test_validate_with_conversation_id_required(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
data = {"question": "Test"}
|
||||
|
||||
result = resource.validate_request(data, require_conversation_id=True)
|
||||
|
||||
assert result is not None
|
||||
assert result.status_code == 400
|
||||
assert "conversation_id" in result.json["message"].lower()
|
||||
|
||||
def test_validate_passes_with_all_required_fields(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
data = {"question": "Test", "conversation_id": str(uuid.uuid4())}
|
||||
|
||||
result = resource.validate_request(data, require_conversation_id=True)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestUsageChecking:
|
||||
pass
|
||||
|
||||
def test_returns_none_when_no_api_key(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
agent_config = {}
|
||||
|
||||
result = resource.check_usage(agent_config)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGPTModelRetrieval:
|
||||
pass
|
||||
|
||||
def test_initializes_gpt_model(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
assert hasattr(resource, "default_model_id")
|
||||
assert resource.default_model_id is not None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestConversationServiceIntegration:
|
||||
pass
|
||||
|
||||
def test_initializes_conversation_service(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
assert hasattr(resource, "conversation_service")
|
||||
assert resource.conversation_service is not None
|
||||
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompleteStreamMethod:
|
||||
pass
|
||||
|
||||
def test_streams_answer_chunks(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.return_value = iter(
|
||||
[
|
||||
{"answer": "Hello "},
|
||||
{"answer": "world!"},
|
||||
]
|
||||
)
|
||||
|
||||
decoded_token = {"sub": "user123"}
|
||||
|
||||
stream = list(
|
||||
resource.complete_stream(
|
||||
question="Test question",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token=decoded_token,
|
||||
should_persist=False,
|
||||
)
|
||||
)
|
||||
|
||||
answer_chunks = [s for s in stream if '"type": "answer"' in s]
|
||||
assert len(answer_chunks) == 2
|
||||
assert '"answer": "Hello "' in answer_chunks[0]
|
||||
assert '"answer": "world!"' in answer_chunks[1]
|
||||
|
||||
def test_streams_sources(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.return_value = iter(
|
||||
[
|
||||
{"answer": "Test answer"},
|
||||
{"sources": [{"title": "doc1.txt", "text": "x" * 200}]},
|
||||
]
|
||||
)
|
||||
|
||||
decoded_token = {"sub": "user123"}
|
||||
|
||||
stream = list(
|
||||
resource.complete_stream(
|
||||
question="Test?",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token=decoded_token,
|
||||
should_persist=False,
|
||||
)
|
||||
)
|
||||
|
||||
source_chunks = [s for s in stream if '"type": "source"' in s]
|
||||
assert len(source_chunks) == 1
|
||||
assert '"title": "doc1.txt"' in source_chunks[0]
|
||||
|
||||
def test_handles_error_during_streaming(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.side_effect = Exception("Test error")
|
||||
|
||||
decoded_token = {"sub": "user123"}
|
||||
|
||||
stream = list(
|
||||
resource.complete_stream(
|
||||
question="Test?",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token=decoded_token,
|
||||
should_persist=False,
|
||||
)
|
||||
)
|
||||
|
||||
assert any('"type": "error"' in s for s in stream)
|
||||
|
||||
def test_user_facing_error_is_not_sanitized(self, mock_mongo_db, flask_app):
|
||||
"""A user_facing error (e.g. an artifact-quota notice) streams verbatim.
|
||||
|
||||
Without the flag, sanitize_api_error substring-matches "quota" and rewrites the
|
||||
message into a misleading rate-limit notice.
|
||||
"""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.return_value = iter(
|
||||
[
|
||||
{
|
||||
"type": "error",
|
||||
"user_facing": True,
|
||||
"error": "This run's input documents exceed your artifact storage quota.",
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
stream = list(
|
||||
resource.complete_stream(
|
||||
question="Test?",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": "user123"},
|
||||
should_persist=False,
|
||||
)
|
||||
)
|
||||
|
||||
error_chunks = [s for s in stream if '"type": "error"' in s]
|
||||
assert error_chunks
|
||||
assert "artifact storage quota" in error_chunks[0]
|
||||
assert "Rate limit exceeded" not in error_chunks[0]
|
||||
|
||||
def test_notice_is_forwarded_verbatim_and_not_an_error(self, mock_mongo_db, flask_app):
|
||||
"""A non-fatal ``notice`` streams through as a notice, never as an error.
|
||||
|
||||
A ``notice`` (e.g. some workflow input documents were dropped) must not be
|
||||
emitted as ``type: error`` -- the client treats an error event as terminal and
|
||||
disables reconnect -- and its text must not be run through sanitize_api_error.
|
||||
"""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.return_value = iter(
|
||||
[{"type": "notice", "notice": "big.txt exceeds the per-file size limit"}]
|
||||
)
|
||||
|
||||
stream = list(
|
||||
resource.complete_stream(
|
||||
question="Test?",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": "user123"},
|
||||
should_persist=False,
|
||||
)
|
||||
)
|
||||
|
||||
notice_chunks = [s for s in stream if '"type": "notice"' in s]
|
||||
assert notice_chunks
|
||||
assert "big.txt exceeds the per-file size limit" in notice_chunks[0]
|
||||
# Crucially, it is not surfaced as an error event.
|
||||
assert not [s for s in stream if '"type": "error"' in s]
|
||||
|
||||
def test_non_user_facing_error_is_sanitized(self, mock_mongo_db, flask_app):
|
||||
"""A raw error without the flag is still routed through sanitize_api_error."""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.return_value = iter(
|
||||
[{"type": "error", "error": "OpenAI 429: quota exceeded for this key"}]
|
||||
)
|
||||
|
||||
stream = list(
|
||||
resource.complete_stream(
|
||||
question="Test?",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": "user123"},
|
||||
should_persist=False,
|
||||
)
|
||||
)
|
||||
|
||||
error_chunks = [s for s in stream if '"type": "error"' in s]
|
||||
assert error_chunks
|
||||
assert "Rate limit exceeded" in error_chunks[0]
|
||||
|
||||
def test_saves_conversation_when_enabled(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.gen.return_value = iter(
|
||||
[
|
||||
{"answer": "Test answer"},
|
||||
]
|
||||
)
|
||||
|
||||
decoded_token = {"sub": "user123"}
|
||||
|
||||
# The fresh-question path now reserves a row before agent.gen()
|
||||
# and calls finalize_message at end of stream — assert both fire.
|
||||
with patch.object(
|
||||
resource.conversation_service, "save_user_question"
|
||||
) as mock_reserve, patch.object(
|
||||
resource.conversation_service, "finalize_message"
|
||||
) as mock_finalize:
|
||||
mock_reserve.return_value = {
|
||||
"conversation_id": str(uuid.uuid4()),
|
||||
"message_id": str(uuid.uuid4()),
|
||||
"request_id": "req-1",
|
||||
}
|
||||
mock_finalize.return_value = True
|
||||
|
||||
list(
|
||||
resource.complete_stream(
|
||||
question="Test?",
|
||||
agent=mock_agent,
|
||||
conversation_id=None,
|
||||
user_api_key=None,
|
||||
decoded_token=decoded_token,
|
||||
should_persist=True,
|
||||
)
|
||||
)
|
||||
|
||||
mock_reserve.assert_called_once()
|
||||
mock_finalize.assert_called_once()
|
||||
|
||||
def test_tool_executor_conversation_id_set_after_reserve(
|
||||
self, mock_mongo_db, flask_app,
|
||||
):
|
||||
"""Regression: ``save_user_question`` may mint a fresh
|
||||
``conversation_id`` (first turn). The propagation MUST land on
|
||||
``agent.tool_executor.conversation_id`` BEFORE ``agent.gen`` runs,
|
||||
so tools needing a conversation home (``scheduler`` in an agentless
|
||||
chat) see it on the very first call.
|
||||
"""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
fresh_conv_id = str(uuid.uuid4())
|
||||
seen_conv_id_on_gen: dict = {}
|
||||
|
||||
mock_agent = MagicMock()
|
||||
tool_executor = MagicMock()
|
||||
# Start with no conversation_id — the propagation must set it.
|
||||
tool_executor.conversation_id = None
|
||||
mock_agent.tool_executor = tool_executor
|
||||
|
||||
def _gen(**_kwargs):
|
||||
# Capture the executor's id at the exact moment gen runs;
|
||||
# this is what tools see when called from the agent loop.
|
||||
seen_conv_id_on_gen["value"] = (
|
||||
mock_agent.tool_executor.conversation_id
|
||||
)
|
||||
yield {"answer": "ok"}
|
||||
|
||||
mock_agent.gen.side_effect = _gen
|
||||
mock_agent.gen.return_value = None # use side_effect instead
|
||||
|
||||
with patch.object(
|
||||
resource.conversation_service, "save_user_question"
|
||||
) as mock_reserve, patch.object(
|
||||
resource.conversation_service, "finalize_message",
|
||||
return_value=True,
|
||||
):
|
||||
mock_reserve.return_value = {
|
||||
"conversation_id": fresh_conv_id,
|
||||
"message_id": str(uuid.uuid4()),
|
||||
"request_id": "req-prop",
|
||||
}
|
||||
|
||||
list(
|
||||
resource.complete_stream(
|
||||
question="schedule something",
|
||||
agent=mock_agent,
|
||||
conversation_id=None, # caller had no conv yet
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": "user-prop"},
|
||||
should_persist=True,
|
||||
)
|
||||
)
|
||||
|
||||
# The fresh id reserved by save_user_question must reach the
|
||||
# tool_executor before agent.gen consumes it.
|
||||
assert seen_conv_id_on_gen["value"] == fresh_conv_id
|
||||
assert tool_executor.conversation_id == fresh_conv_id
|
||||
|
||||
def _run_paused(self, resource, pending_calls):
|
||||
"""Drive complete_stream into its paused branch with the given pending
|
||||
tool calls, mocking out the WAL row and continuation save."""
|
||||
agent = MagicMock()
|
||||
agent.gen.return_value = iter(
|
||||
[{"type": "tool_calls_pending", "data": {"pending_tool_calls": pending_calls}}]
|
||||
)
|
||||
agent._pending_continuation = {
|
||||
"messages": [],
|
||||
"pending_tool_calls": pending_calls,
|
||||
"tools_dict": {},
|
||||
}
|
||||
# Make the WAL reservation no-op so we stay off the journal/DB.
|
||||
resource.conversation_service = MagicMock()
|
||||
resource.conversation_service.save_user_question.side_effect = Exception("skip")
|
||||
list(
|
||||
resource.complete_stream(
|
||||
question="Do the test",
|
||||
agent=agent,
|
||||
conversation_id="conv-1",
|
||||
user_api_key=None,
|
||||
decoded_token={"sub": "user123"},
|
||||
should_persist=True,
|
||||
)
|
||||
)
|
||||
|
||||
def test_paused_skips_notification_for_client_execution(
|
||||
self, mock_mongo_db, flask_app
|
||||
):
|
||||
"""A pure ``requires_client_execution`` pause must NOT publish a
|
||||
``tool.approval.required`` event — the client resolves it, so the
|
||||
notification would be non-actionable noise."""
|
||||
from application.api.answer.routes import base as base_mod
|
||||
|
||||
with flask_app.app_context(), patch.object(
|
||||
base_mod, "publish_user_event"
|
||||
) as published, patch.object(
|
||||
base_mod, "ContinuationService", MagicMock
|
||||
):
|
||||
self._run_paused(
|
||||
base_mod.BaseAnswerResource(),
|
||||
[
|
||||
{
|
||||
"call_id": "c1",
|
||||
"name": "create_file",
|
||||
"tool_name": "create_file",
|
||||
"action_name": "create_file",
|
||||
"pause_type": "requires_client_execution",
|
||||
}
|
||||
],
|
||||
)
|
||||
published.assert_not_called()
|
||||
|
||||
def test_paused_publishes_notification_only_for_awaiting_approval(
|
||||
self, mock_mongo_db, flask_app
|
||||
):
|
||||
"""A pause with an ``awaiting_approval`` call publishes once, and the
|
||||
payload surfaces only the approval call (not the client-side one)."""
|
||||
from application.api.answer.routes import base as base_mod
|
||||
|
||||
with flask_app.app_context(), patch.object(
|
||||
base_mod, "publish_user_event"
|
||||
) as published, patch.object(
|
||||
base_mod, "ContinuationService", MagicMock
|
||||
):
|
||||
self._run_paused(
|
||||
base_mod.BaseAnswerResource(),
|
||||
[
|
||||
{
|
||||
"call_id": "a1",
|
||||
"name": "delete_thing",
|
||||
"tool_name": "api_tool",
|
||||
"action_name": "delete_thing",
|
||||
"pause_type": "awaiting_approval",
|
||||
},
|
||||
{
|
||||
"call_id": "c1",
|
||||
"name": "create_file",
|
||||
"tool_name": "create_file",
|
||||
"action_name": "create_file",
|
||||
"pause_type": "requires_client_execution",
|
||||
},
|
||||
],
|
||||
)
|
||||
published.assert_called_once()
|
||||
args, _ = published.call_args
|
||||
assert args[1] == "tool.approval.required"
|
||||
summaries = args[2]["pending_tool_calls"]
|
||||
assert [s["call_id"] for s in summaries] == ["a1"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestProcessResponseStream:
|
||||
pass
|
||||
|
||||
def test_processes_complete_stream(self, mock_mongo_db, flask_app):
|
||||
import json
|
||||
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
conv_id = str(uuid.uuid4())
|
||||
stream = [
|
||||
f'data: {json.dumps({"type": "answer", "answer": "Hello "})}\n\n',
|
||||
f'data: {json.dumps({"type": "answer", "answer": "world"})}\n\n',
|
||||
f'data: {json.dumps({"type": "source", "source": [{"title": "doc1"}]})}\n\n',
|
||||
f'data: {json.dumps({"type": "id", "id": conv_id})}\n\n',
|
||||
f'data: {json.dumps({"type": "end"})}\n\n',
|
||||
]
|
||||
|
||||
result = resource.process_response_stream(iter(stream))
|
||||
|
||||
assert result["conversation_id"] == conv_id
|
||||
assert result["answer"] == "Hello world"
|
||||
assert result["sources"] == [{"title": "doc1"}]
|
||||
assert result["error"] is None
|
||||
|
||||
def test_handles_stream_error(self, mock_mongo_db, flask_app):
|
||||
import json
|
||||
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
stream = [
|
||||
f'data: {json.dumps({"type": "error", "error": "Test error"})}\n\n',
|
||||
]
|
||||
|
||||
result = resource.process_response_stream(iter(stream))
|
||||
|
||||
assert result["conversation_id"] is None
|
||||
assert result["error"] == "Test error"
|
||||
|
||||
def test_handles_malformed_stream_data(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
stream = [
|
||||
"data: invalid json\n\n",
|
||||
'data: {"type": "end"}\n\n',
|
||||
]
|
||||
|
||||
result = resource.process_response_stream(iter(stream))
|
||||
|
||||
assert result is not None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestErrorStreamGenerate:
|
||||
pass
|
||||
|
||||
def test_generates_error_stream(self, mock_mongo_db, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
|
||||
error_stream = list(resource.error_stream_generate("Test error message"))
|
||||
|
||||
assert len(error_stream) == 1
|
||||
assert '"type": "error"' in error_stream[0]
|
||||
assert '"error": "Test error message"' in error_stream[0]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Real-PG tests for check_usage against seeded agents + token usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_base_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.answer.routes.base.db_readonly", _yield
|
||||
), patch(
|
||||
"application.api.answer.routes.base.db_session", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCheckUsagePgConn:
|
||||
def test_invalid_api_key_returns_401(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
|
||||
with _patch_base_db(pg_conn), flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
result = resource.check_usage({"user_api_key": "does-not-exist"})
|
||||
assert result is not None
|
||||
assert result.status_code == 401
|
||||
|
||||
def test_no_limits_returns_none(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
"owner", "a", "published", key="k1",
|
||||
limited_token_mode=False, limited_request_mode=False,
|
||||
)
|
||||
with _patch_base_db(pg_conn), flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
result = resource.check_usage({"user_api_key": "k1"})
|
||||
assert result is None
|
||||
|
||||
def test_within_limit_returns_none(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
"owner", "a", "published", key="k2",
|
||||
limited_token_mode=True, token_limit=10000,
|
||||
)
|
||||
with _patch_base_db(pg_conn), flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
result = resource.check_usage({"user_api_key": "k2"})
|
||||
assert result is None
|
||||
|
||||
def test_token_limit_exceeded_returns_429(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.token_usage import (
|
||||
TokenUsageRepository,
|
||||
)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
"owner", "a", "published", key="k3",
|
||||
limited_token_mode=True, token_limit=100,
|
||||
)
|
||||
# Seed token usage exceeding the limit
|
||||
TokenUsageRepository(pg_conn).insert(
|
||||
api_key="k3", prompt_tokens=500, generated_tokens=0,
|
||||
)
|
||||
|
||||
with _patch_base_db(pg_conn), flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
result = resource.check_usage({"user_api_key": "k3"})
|
||||
assert result is not None
|
||||
assert result.status_code == 429
|
||||
|
||||
def test_request_limit_exceeded_returns_429(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.token_usage import (
|
||||
TokenUsageRepository,
|
||||
)
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
"owner", "a", "published", key="k4",
|
||||
limited_request_mode=True, request_limit=1,
|
||||
)
|
||||
# Two request entries exceed limit=1
|
||||
TokenUsageRepository(pg_conn).insert(api_key="k4", prompt_tokens=10, generated_tokens=10)
|
||||
TokenUsageRepository(pg_conn).insert(api_key="k4", prompt_tokens=10, generated_tokens=10)
|
||||
|
||||
with _patch_base_db(pg_conn), flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
result = resource.check_usage({"user_api_key": "k4"})
|
||||
assert result is not None
|
||||
assert result.status_code == 429
|
||||
|
||||
def test_string_True_limited_token_mode_parsed(self, pg_conn, flask_app):
|
||||
"""Legacy Mongo sometimes stored ``limited_token_mode`` as the
|
||||
string 'True'; verify the parse branch."""
|
||||
from application.api.answer.routes.base import BaseAnswerResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
# Store bool=False in DB (limited_token_mode default). Test uses
|
||||
# string 'True' by mutating the row directly.
|
||||
from sqlalchemy import text
|
||||
AgentsRepository(pg_conn).create(
|
||||
"owner", "a", "published", key="k5",
|
||||
)
|
||||
pg_conn.execute(
|
||||
text(
|
||||
"UPDATE agents SET limited_token_mode = :v WHERE key = :k"
|
||||
),
|
||||
{"v": True, "k": "k5"},
|
||||
)
|
||||
with _patch_base_db(pg_conn), flask_app.app_context():
|
||||
resource = BaseAnswerResource()
|
||||
result = resource.check_usage({"user_api_key": "k5"})
|
||||
# With default limit and no token usage, should pass
|
||||
assert result is None
|
||||
@@ -0,0 +1,208 @@
|
||||
"""Tests for /api/search route (application/api/answer/routes/search.py).
|
||||
|
||||
Retrieval logic lives in ``application/services/search_service.py`` and
|
||||
has its own unit tests in ``tests/services/test_search_service.py``. The
|
||||
tests below focus on what the route specifically owns:
|
||||
|
||||
* Request validation (400 for missing fields).
|
||||
* Translation of the service's ``InvalidAPIKey`` / ``SearchFailed``
|
||||
exceptions to HTTP status codes (401 / 500).
|
||||
* End-to-end happy path against a real ephemeral Postgres via
|
||||
``pg_conn``, to catch regressions in the route's wiring to the
|
||||
service and repositories.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSearchResourceValidation:
|
||||
def test_returns_400_when_question_missing(self, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
|
||||
with flask_app.app_context():
|
||||
with flask_app.test_request_context(json={"api_key": "test_key"}):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 400
|
||||
assert "question" in result.json["error"]
|
||||
|
||||
def test_returns_400_when_api_key_missing(self, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
|
||||
with flask_app.app_context():
|
||||
with flask_app.test_request_context(json={"question": "test query"}):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 400
|
||||
assert "api_key" in result.json["error"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSearchResourceExceptionMapping:
|
||||
"""Verify the route maps service exceptions to HTTP status codes.
|
||||
|
||||
The service function itself is patched; these tests do not care about
|
||||
the search logic — only that 401/500/200 are produced correctly from
|
||||
the three possible service outcomes.
|
||||
"""
|
||||
|
||||
def test_invalid_api_key_returns_401(self, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
from application.services.search_service import InvalidAPIKey
|
||||
|
||||
with flask_app.app_context(), flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "bad"}
|
||||
), patch(
|
||||
"application.api.answer.routes.search.search",
|
||||
side_effect=InvalidAPIKey(),
|
||||
):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 401
|
||||
assert result.json == {"error": "Invalid API key"}
|
||||
|
||||
def test_search_failed_returns_500(self, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
from application.services.search_service import SearchFailed
|
||||
|
||||
with flask_app.app_context(), flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "k"}
|
||||
), patch(
|
||||
"application.api.answer.routes.search.search",
|
||||
side_effect=SearchFailed("boom"),
|
||||
):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 500
|
||||
assert result.json == {"error": "Search failed"}
|
||||
|
||||
def test_happy_path_passes_service_result_through(self, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
|
||||
hits = [{"text": "t", "title": "T", "source": "s"}]
|
||||
with flask_app.app_context(), flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "k", "chunks": 7}
|
||||
), patch(
|
||||
"application.api.answer.routes.search.search",
|
||||
return_value=hits,
|
||||
) as mock_search:
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 200
|
||||
assert result.json == hits
|
||||
mock_search.assert_called_once_with("k", "q", 7)
|
||||
|
||||
def test_default_chunks_is_5(self, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
|
||||
with flask_app.app_context(), flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "k"} # no chunks field
|
||||
), patch(
|
||||
"application.api.answer.routes.search.search",
|
||||
return_value=[],
|
||||
) as mock_search:
|
||||
SearchResource().post()
|
||||
mock_search.assert_called_once_with("k", "q", 5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end against a real ephemeral Postgres.
|
||||
#
|
||||
# These exercise the full route → service → repository → DB path, patching
|
||||
# only ``VectorCreator.create_vectorstore`` (so we don't need real embeddings
|
||||
# or a vector index). ``db_readonly`` is redirected at the *service* module
|
||||
# since that's where the import now lives.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_search_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.services.search_service.db_readonly", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
class TestSearchResourcePgConn:
|
||||
def test_invalid_api_key_returns_401(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
|
||||
with _patch_search_db(pg_conn), flask_app.app_context():
|
||||
with flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "does-not-exist"},
|
||||
):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 401
|
||||
|
||||
def test_no_sources_returns_empty(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
AgentsRepository(pg_conn).create(
|
||||
"u", "a", "published", key="no-src-key",
|
||||
)
|
||||
with _patch_search_db(pg_conn), flask_app.app_context():
|
||||
with flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "no-src-key"},
|
||||
):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 200
|
||||
assert result.json == []
|
||||
|
||||
def test_search_returns_results(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
src = SourcesRepository(pg_conn).create("src", user_id="u")
|
||||
AgentsRepository(pg_conn).create(
|
||||
"u", "a", "published",
|
||||
key="search-key",
|
||||
source_id=str(src["id"]),
|
||||
)
|
||||
|
||||
fake_vs = MagicMock()
|
||||
fake_vs.search.return_value = [
|
||||
{"text": "answer text", "metadata": {"title": "Doc"}},
|
||||
]
|
||||
|
||||
with _patch_search_db(pg_conn), patch(
|
||||
"application.services.search_service.VectorCreator.create_vectorstore",
|
||||
return_value=fake_vs,
|
||||
), flask_app.app_context():
|
||||
with flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "search-key"},
|
||||
):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 200
|
||||
assert len(result.json) == 1
|
||||
|
||||
def test_search_uses_extra_source_ids(self, pg_conn, flask_app):
|
||||
from application.api.answer.routes.search import SearchResource
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
src1 = SourcesRepository(pg_conn).create("s1", user_id="u")
|
||||
src2 = SourcesRepository(pg_conn).create("s2", user_id="u")
|
||||
AgentsRepository(pg_conn).create(
|
||||
"u", "a", "published",
|
||||
key="extra-key",
|
||||
extra_source_ids=[str(src1["id"]), str(src2["id"])],
|
||||
)
|
||||
|
||||
fake_vs = MagicMock()
|
||||
fake_vs.search.return_value = [
|
||||
{"text": "one", "metadata": {"title": "A"}},
|
||||
]
|
||||
with _patch_search_db(pg_conn), patch(
|
||||
"application.services.search_service.VectorCreator.create_vectorstore",
|
||||
return_value=fake_vs,
|
||||
), flask_app.app_context():
|
||||
with flask_app.test_request_context(
|
||||
json={"question": "q", "api_key": "extra-key", "chunks": 4},
|
||||
):
|
||||
result = SearchResource().post()
|
||||
assert result.status_code == 200
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Tests for application/api/answer/routes/stream.py"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Static IDs — valid 24-hex-char strings
|
||||
_CONV_ID = "507f1f77bcf86cd799439011"
|
||||
_AGENT_ID = "507f1f77bcf86cd799439012"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_stream_processor():
|
||||
"""Create a mock StreamProcessor for stream tests."""
|
||||
with patch(
|
||||
"application.api.answer.routes.stream.StreamProcessor"
|
||||
) as MockProcessor:
|
||||
processor = MagicMock()
|
||||
processor.decoded_token = {"sub": "test_user"}
|
||||
processor.conversation_id = _CONV_ID
|
||||
processor.agent_config = {}
|
||||
processor.agent_id = _AGENT_ID
|
||||
processor.is_shared_usage = False
|
||||
processor.shared_token = None
|
||||
processor.model_id = "gpt-4"
|
||||
processor.build_agent.return_value = MagicMock()
|
||||
MockProcessor.return_value = processor
|
||||
yield processor
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stream_client(mock_mongo_db, flask_app):
|
||||
"""Create a test client with the stream route registered."""
|
||||
from flask_restx import Api
|
||||
|
||||
from application.api.answer.routes.stream import answer_ns
|
||||
|
||||
api = Api(flask_app)
|
||||
api.add_namespace(answer_ns)
|
||||
flask_app.config["TESTING"] = True
|
||||
return flask_app.test_client()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestStreamResourcePost:
|
||||
def test_missing_question_returns_400(self, stream_client, mock_stream_processor):
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps({}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_successful_stream(self, stream_client, mock_stream_processor):
|
||||
def fake_stream(*args, **kwargs):
|
||||
yield f'data: {json.dumps({"type": "answer", "answer": "Hi"})}\n\n'
|
||||
yield f'data: {json.dumps({"type": "end"})}\n\n'
|
||||
|
||||
with patch(
|
||||
"application.api.answer.routes.stream.StreamResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.stream.StreamResource.check_usage",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.stream.StreamResource.complete_stream",
|
||||
side_effect=fake_stream,
|
||||
):
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps({"question": "What is Python?"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert "text/event-stream" in resp.content_type
|
||||
data = resp.get_data(as_text=True)
|
||||
assert '"type": "answer"' in data
|
||||
assert '"answer": "Hi"' in data
|
||||
|
||||
def test_unauthorized_returns_401_stream(
|
||||
self, stream_client, mock_stream_processor
|
||||
):
|
||||
mock_stream_processor.decoded_token = None
|
||||
with patch(
|
||||
"application.api.answer.routes.stream.StreamResource.validate_request",
|
||||
return_value=None,
|
||||
):
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
assert "text/event-stream" in resp.content_type
|
||||
data = resp.get_data(as_text=True)
|
||||
assert "Unauthorized" in data
|
||||
|
||||
def test_usage_exceeded_returns_error(
|
||||
self, stream_client, mock_stream_processor
|
||||
):
|
||||
with patch(
|
||||
"application.api.answer.routes.stream.StreamResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.stream.StreamResource.check_usage",
|
||||
) as mock_check:
|
||||
mock_check.return_value = ({"error": "Usage limit exceeded"}, 429)
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 429
|
||||
|
||||
def test_value_error_returns_400_stream(
|
||||
self, stream_client, mock_stream_processor
|
||||
):
|
||||
mock_stream_processor.build_agent.side_effect = ValueError("bad data")
|
||||
with patch(
|
||||
"application.api.answer.routes.stream.StreamResource.validate_request",
|
||||
return_value=None,
|
||||
):
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
assert "text/event-stream" in resp.content_type
|
||||
data = resp.get_data(as_text=True)
|
||||
assert "Malformed request body" in data
|
||||
|
||||
def test_general_exception_returns_400_stream(
|
||||
self, stream_client, mock_stream_processor
|
||||
):
|
||||
mock_stream_processor.build_agent.side_effect = RuntimeError("crash")
|
||||
with patch(
|
||||
"application.api.answer.routes.stream.StreamResource.validate_request",
|
||||
return_value=None,
|
||||
):
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps({"question": "test"}),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
assert "text/event-stream" in resp.content_type
|
||||
data = resp.get_data(as_text=True)
|
||||
assert "Unknown error occurred" in data
|
||||
|
||||
def test_index_in_data_requires_conversation_id(
|
||||
self, stream_client, mock_stream_processor
|
||||
):
|
||||
"""When 'index' is present, validate_request is called with require_conversation_id=True."""
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps({"question": "test", "index": 0}),
|
||||
content_type="application/json",
|
||||
)
|
||||
# Should get 400 since conversation_id is missing
|
||||
assert resp.status_code == 400
|
||||
|
||||
def test_stream_passes_attachments_and_index(
|
||||
self, stream_client, mock_stream_processor
|
||||
):
|
||||
"""Verify attachments and index params are forwarded to complete_stream."""
|
||||
|
||||
def fake_stream(*args, **kwargs):
|
||||
yield f'data: {json.dumps({"type": "end"})}\n\n'
|
||||
|
||||
conv_id = str(uuid.uuid4())
|
||||
with patch(
|
||||
"application.api.answer.routes.stream.StreamResource.validate_request",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.stream.StreamResource.check_usage",
|
||||
return_value=None,
|
||||
), patch(
|
||||
"application.api.answer.routes.stream.StreamResource.complete_stream",
|
||||
side_effect=fake_stream,
|
||||
) as mock_complete:
|
||||
resp = stream_client.post(
|
||||
"/stream",
|
||||
data=json.dumps(
|
||||
{
|
||||
"question": "test",
|
||||
"conversation_id": conv_id,
|
||||
"index": 3,
|
||||
"attachments": ["att1", "att2"],
|
||||
}
|
||||
),
|
||||
content_type="application/json",
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = mock_complete.call_args
|
||||
assert call_kwargs.kwargs.get("index") == 3
|
||||
assert call_kwargs.kwargs.get("attachment_ids") == ["att1", "att2"]
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Tests for application/api/answer/services/compression/message_builder.py"""
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.compression.message_builder import MessageBuilder
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestBuildFromCompressedContext:
|
||||
def test_no_compression_returns_system_only(self):
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="You are helpful.",
|
||||
compressed_summary=None,
|
||||
recent_queries=[],
|
||||
)
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "system"
|
||||
assert messages[0]["content"] == "You are helpful."
|
||||
|
||||
def test_with_recent_queries_no_compression(self):
|
||||
queries = [
|
||||
{"prompt": "Hello", "response": "Hi there!"},
|
||||
{"prompt": "How are you?", "response": "I'm fine."},
|
||||
]
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="System prompt",
|
||||
compressed_summary=None,
|
||||
recent_queries=queries,
|
||||
)
|
||||
# system + 2 * (user + assistant) = 5
|
||||
assert len(messages) == 5
|
||||
assert messages[1] == {"role": "user", "content": "Hello"}
|
||||
assert messages[2] == {"role": "assistant", "content": "Hi there!"}
|
||||
assert messages[3] == {"role": "user", "content": "How are you?"}
|
||||
assert messages[4] == {"role": "assistant", "content": "I'm fine."}
|
||||
|
||||
def test_with_compressed_summary_appended_to_system(self):
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="You are helpful.",
|
||||
compressed_summary="Previous: user asked about Python.",
|
||||
recent_queries=[{"prompt": "More?", "response": "Sure."}],
|
||||
)
|
||||
system_content = messages[0]["content"]
|
||||
assert "This session is being continued" in system_content
|
||||
assert "Previous: user asked about Python." in system_content
|
||||
|
||||
def test_mid_execution_context_type(self):
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="System",
|
||||
compressed_summary="Summary here",
|
||||
recent_queries=[{"prompt": "q", "response": "r"}],
|
||||
context_type="mid_execution",
|
||||
)
|
||||
system_content = messages[0]["content"]
|
||||
assert "Context window limit reached" in system_content
|
||||
|
||||
def test_include_tool_calls(self):
|
||||
queries = [
|
||||
{
|
||||
"prompt": "Search for X",
|
||||
"response": "Found X",
|
||||
"tool_calls": [
|
||||
{
|
||||
"call_id": "call-1",
|
||||
"action_name": "search",
|
||||
"arguments": {"q": "X"},
|
||||
"result": "X found",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="System",
|
||||
compressed_summary=None,
|
||||
recent_queries=queries,
|
||||
include_tool_calls=True,
|
||||
)
|
||||
# system + user + assistant + tool_call_assistant + tool_response = 5
|
||||
assert len(messages) == 5
|
||||
assert messages[3]["role"] == "assistant"
|
||||
assert messages[3].get("tool_calls") is not None
|
||||
assert messages[3]["tool_calls"][0]["function"]["name"] == "search"
|
||||
assert messages[4]["role"] == "tool"
|
||||
assert messages[4].get("tool_call_id") == "call-1"
|
||||
|
||||
def test_tool_calls_not_included_by_default(self):
|
||||
queries = [
|
||||
{
|
||||
"prompt": "Search",
|
||||
"response": "Found",
|
||||
"tool_calls": [
|
||||
{
|
||||
"call_id": "c1",
|
||||
"action_name": "search",
|
||||
"arguments": {},
|
||||
"result": "ok",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="System",
|
||||
compressed_summary=None,
|
||||
recent_queries=queries,
|
||||
include_tool_calls=False,
|
||||
)
|
||||
# system + user + assistant = 3 (no tool messages)
|
||||
assert len(messages) == 3
|
||||
|
||||
def test_tool_call_without_call_id_generates_uuid(self):
|
||||
queries = [
|
||||
{
|
||||
"prompt": "q",
|
||||
"response": "r",
|
||||
"tool_calls": [
|
||||
{
|
||||
"action_name": "act",
|
||||
"arguments": {},
|
||||
"result": "res",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="S",
|
||||
compressed_summary=None,
|
||||
recent_queries=queries,
|
||||
include_tool_calls=True,
|
||||
)
|
||||
assistant_msg = messages[3]
|
||||
call_id = assistant_msg["tool_calls"][0]["id"]
|
||||
assert call_id is not None
|
||||
assert len(call_id) > 0
|
||||
|
||||
def test_continuation_message_when_no_recent_queries_but_has_summary(self):
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="System",
|
||||
compressed_summary="Everything was compressed",
|
||||
recent_queries=[],
|
||||
)
|
||||
# system + continuation user message = 2
|
||||
assert len(messages) == 2
|
||||
assert messages[1]["role"] == "user"
|
||||
assert "continue" in messages[1]["content"].lower()
|
||||
|
||||
def test_no_continuation_when_no_summary(self):
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="System",
|
||||
compressed_summary=None,
|
||||
recent_queries=[],
|
||||
)
|
||||
assert len(messages) == 1
|
||||
|
||||
def test_queries_without_prompt_or_response_skipped(self):
|
||||
queries = [
|
||||
{"other_field": "value"},
|
||||
{"prompt": "real", "response": "answer"},
|
||||
]
|
||||
messages = MessageBuilder.build_from_compressed_context(
|
||||
system_prompt="S",
|
||||
compressed_summary=None,
|
||||
recent_queries=queries,
|
||||
)
|
||||
# system + 1 valid query (user + assistant) = 3
|
||||
assert len(messages) == 3
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAppendCompressionContext:
|
||||
def test_pre_request_context(self):
|
||||
result = MessageBuilder._append_compression_context(
|
||||
"Original prompt", "Summary text", "pre_request"
|
||||
)
|
||||
assert "This session is being continued" in result
|
||||
assert "Summary text" in result
|
||||
assert result.startswith("Original prompt")
|
||||
|
||||
def test_mid_execution_context(self):
|
||||
result = MessageBuilder._append_compression_context(
|
||||
"Original prompt", "Summary text", "mid_execution"
|
||||
)
|
||||
assert "Context window limit reached" in result
|
||||
assert "Summary text" in result
|
||||
|
||||
def test_removes_existing_compression_context(self):
|
||||
prompt_with_existing = (
|
||||
"Original prompt\n\n---\n\nThis session is being continued from old"
|
||||
)
|
||||
result = MessageBuilder._append_compression_context(
|
||||
prompt_with_existing, "New summary", "pre_request"
|
||||
)
|
||||
# Should not contain old context twice
|
||||
assert result.count("This session is being continued") == 1
|
||||
assert "New summary" in result
|
||||
|
||||
def test_removes_mid_execution_context(self):
|
||||
prompt_with_existing = (
|
||||
"Original\n\n---\n\nContext window limit reached during execution. Old."
|
||||
)
|
||||
result = MessageBuilder._append_compression_context(
|
||||
prompt_with_existing, "New", "mid_execution"
|
||||
)
|
||||
assert result.count("Context window limit reached") == 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRebuildMessagesAfterCompression:
|
||||
def test_basic_rebuild(self):
|
||||
messages = [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "old message"},
|
||||
{"role": "assistant", "content": "old reply"},
|
||||
]
|
||||
recent = [{"prompt": "new q", "response": "new r"}]
|
||||
|
||||
result = MessageBuilder.rebuild_messages_after_compression(
|
||||
messages=messages,
|
||||
compressed_summary="Everything was compressed.",
|
||||
recent_queries=recent,
|
||||
)
|
||||
assert result is not None
|
||||
# system + user + assistant = 3
|
||||
assert len(result) == 3
|
||||
assert "Context window limit reached" in result[0]["content"]
|
||||
assert result[1] == {"role": "user", "content": "new q"}
|
||||
assert result[2] == {"role": "assistant", "content": "new r"}
|
||||
|
||||
def test_returns_none_without_system_message(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "hello"},
|
||||
]
|
||||
result = MessageBuilder.rebuild_messages_after_compression(
|
||||
messages=messages,
|
||||
compressed_summary="summary",
|
||||
recent_queries=[],
|
||||
)
|
||||
assert result is None
|
||||
|
||||
def test_no_summary_keeps_system_unchanged(self):
|
||||
messages = [{"role": "system", "content": "Be helpful."}]
|
||||
result = MessageBuilder.rebuild_messages_after_compression(
|
||||
messages=messages,
|
||||
compressed_summary=None,
|
||||
recent_queries=[],
|
||||
)
|
||||
assert result is not None
|
||||
assert result[0]["content"] == "Be helpful."
|
||||
|
||||
def test_include_tool_calls_in_rebuild(self):
|
||||
messages = [{"role": "system", "content": "S"}]
|
||||
recent = [
|
||||
{
|
||||
"prompt": "q",
|
||||
"response": "r",
|
||||
"tool_calls": [
|
||||
{
|
||||
"call_id": "c1",
|
||||
"action_name": "act",
|
||||
"arguments": {"a": 1},
|
||||
"result": "done",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
result = MessageBuilder.rebuild_messages_after_compression(
|
||||
messages=messages,
|
||||
compressed_summary="s",
|
||||
recent_queries=recent,
|
||||
include_tool_calls=True,
|
||||
)
|
||||
# system + user + assistant + tool_call + tool_response = 5
|
||||
assert len(result) == 5
|
||||
|
||||
def test_continuation_added_when_no_recent_queries(self):
|
||||
messages = [{"role": "system", "content": "S"}]
|
||||
result = MessageBuilder.rebuild_messages_after_compression(
|
||||
messages=messages,
|
||||
compressed_summary="All compressed",
|
||||
recent_queries=[],
|
||||
)
|
||||
assert len(result) == 2
|
||||
assert result[1]["role"] == "user"
|
||||
assert "continue" in result[1]["content"].lower()
|
||||
|
||||
def test_include_current_execution_preserves_extra_messages(self):
|
||||
messages = [
|
||||
{"role": "system", "content": "S"},
|
||||
{"role": "user", "content": "q1"},
|
||||
{"role": "assistant", "content": "r1"},
|
||||
{"role": "user", "content": "current execution msg"},
|
||||
]
|
||||
recent = [{"prompt": "q1", "response": "r1"}]
|
||||
|
||||
result = MessageBuilder.rebuild_messages_after_compression(
|
||||
messages=messages,
|
||||
compressed_summary="summary",
|
||||
recent_queries=recent,
|
||||
include_current_execution=True,
|
||||
)
|
||||
assert result is not None
|
||||
# Should include the current execution message
|
||||
contents = [m.get("content") for m in result]
|
||||
assert "current execution msg" in contents
|
||||
@@ -0,0 +1,458 @@
|
||||
"""Tests for application/api/answer/services/compression/orchestrator.py"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.compression.orchestrator import (
|
||||
CompressionOrchestrator,
|
||||
)
|
||||
from application.api.answer.services.compression.types import (
|
||||
CompressionMetadata,
|
||||
CompressionResult,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_conversation_service():
|
||||
svc = MagicMock()
|
||||
return svc
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_threshold_checker():
|
||||
checker = MagicMock()
|
||||
return checker
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orchestrator(mock_conversation_service, mock_threshold_checker):
|
||||
return CompressionOrchestrator(
|
||||
conversation_service=mock_conversation_service,
|
||||
threshold_checker=mock_threshold_checker,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_conversation():
|
||||
return {
|
||||
"queries": [
|
||||
{"prompt": "q0", "response": "r0"},
|
||||
{"prompt": "q1", "response": "r1"},
|
||||
{"prompt": "q2", "response": "r2"},
|
||||
],
|
||||
"compression_metadata": {},
|
||||
"agent_id": "agent-1",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def decoded_token():
|
||||
return {"sub": "user123"}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressIfNeeded:
|
||||
def test_conversation_not_found_returns_failure(
|
||||
self, orchestrator, mock_conversation_service
|
||||
):
|
||||
mock_conversation_service.get_conversation.return_value = None
|
||||
|
||||
result = orchestrator.compress_if_needed(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token={"sub": "user1"},
|
||||
)
|
||||
|
||||
assert result.success is False
|
||||
assert "not found" in result.error
|
||||
|
||||
def test_no_compression_needed(
|
||||
self,
|
||||
orchestrator,
|
||||
mock_conversation_service,
|
||||
mock_threshold_checker,
|
||||
sample_conversation,
|
||||
):
|
||||
mock_conversation_service.get_conversation.return_value = sample_conversation
|
||||
mock_threshold_checker.should_compress.return_value = False
|
||||
|
||||
result = orchestrator.compress_if_needed(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token={"sub": "user1"},
|
||||
)
|
||||
|
||||
assert result.success is True
|
||||
assert result.compression_performed is False
|
||||
assert len(result.recent_queries) == 3
|
||||
|
||||
def test_compression_performed_successfully(
|
||||
self,
|
||||
orchestrator,
|
||||
mock_conversation_service,
|
||||
mock_threshold_checker,
|
||||
sample_conversation,
|
||||
decoded_token,
|
||||
):
|
||||
mock_conversation_service.get_conversation.return_value = sample_conversation
|
||||
mock_threshold_checker.should_compress.return_value = True
|
||||
|
||||
mock_metadata = MagicMock(spec=CompressionMetadata)
|
||||
mock_metadata.compression_ratio = 5.0
|
||||
mock_metadata.original_token_count = 1000
|
||||
mock_metadata.compressed_token_count = 200
|
||||
mock_metadata.to_dict.return_value = {"query_index": 2}
|
||||
|
||||
with patch.object(
|
||||
orchestrator, "_perform_compression"
|
||||
) as mock_perform:
|
||||
mock_perform.return_value = CompressionResult.success_with_compression(
|
||||
"compressed summary",
|
||||
[{"prompt": "q2", "response": "r2"}],
|
||||
mock_metadata,
|
||||
)
|
||||
|
||||
result = orchestrator.compress_if_needed(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token=decoded_token,
|
||||
)
|
||||
|
||||
assert result.success is True
|
||||
assert result.compression_performed is True
|
||||
assert result.compressed_summary == "compressed summary"
|
||||
mock_perform.assert_called_once()
|
||||
|
||||
def test_exception_returns_failure(
|
||||
self,
|
||||
orchestrator,
|
||||
mock_conversation_service,
|
||||
):
|
||||
mock_conversation_service.get_conversation.side_effect = RuntimeError("DB down")
|
||||
|
||||
result = orchestrator.compress_if_needed(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token={"sub": "user1"},
|
||||
)
|
||||
|
||||
assert result.success is False
|
||||
assert "DB down" in result.error
|
||||
|
||||
def test_custom_query_tokens(
|
||||
self,
|
||||
orchestrator,
|
||||
mock_conversation_service,
|
||||
mock_threshold_checker,
|
||||
sample_conversation,
|
||||
):
|
||||
mock_conversation_service.get_conversation.return_value = sample_conversation
|
||||
mock_threshold_checker.should_compress.return_value = False
|
||||
|
||||
orchestrator.compress_if_needed(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token={"sub": "user1"},
|
||||
current_query_tokens=1000,
|
||||
)
|
||||
|
||||
# user_id flows through so BYOM custom-model UUIDs resolve to
|
||||
# the user's declared context window in the threshold check.
|
||||
mock_threshold_checker.should_compress.assert_called_once_with(
|
||||
sample_conversation, "gpt-4", 1000, user_id="user1"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPerformCompression:
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_provider_from_model_id"
|
||||
)
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_api_key_for_provider"
|
||||
)
|
||||
@patch("application.api.answer.services.compression.orchestrator.LLMCreator")
|
||||
@patch("application.api.answer.services.compression.orchestrator.CompressionService")
|
||||
@patch("application.api.answer.services.compression.orchestrator.settings")
|
||||
def test_successful_compression(
|
||||
self,
|
||||
mock_settings,
|
||||
MockCompressionService,
|
||||
MockLLMCreator,
|
||||
mock_get_api_key,
|
||||
mock_get_provider,
|
||||
mock_conversation_service,
|
||||
mock_threshold_checker,
|
||||
sample_conversation,
|
||||
decoded_token,
|
||||
):
|
||||
mock_settings.COMPRESSION_MODEL_OVERRIDE = None
|
||||
mock_get_provider.return_value = "openai"
|
||||
mock_get_api_key.return_value = "sk-test"
|
||||
MockLLMCreator.create_llm.return_value = MagicMock()
|
||||
|
||||
mock_metadata = MagicMock(spec=CompressionMetadata)
|
||||
mock_metadata.compression_ratio = 5.0
|
||||
mock_metadata.original_token_count = 500
|
||||
mock_metadata.compressed_token_count = 100
|
||||
|
||||
mock_svc_instance = MagicMock()
|
||||
mock_svc_instance.compress_and_save.return_value = mock_metadata
|
||||
mock_svc_instance.get_compressed_context.return_value = (
|
||||
"compressed text",
|
||||
[{"prompt": "q2", "response": "r2"}],
|
||||
)
|
||||
MockCompressionService.return_value = mock_svc_instance
|
||||
|
||||
# After compression, reload conversation
|
||||
mock_conversation_service.get_conversation.return_value = sample_conversation
|
||||
|
||||
orch = CompressionOrchestrator(
|
||||
conversation_service=mock_conversation_service,
|
||||
threshold_checker=mock_threshold_checker,
|
||||
)
|
||||
|
||||
result = orch._perform_compression(
|
||||
"conv1", sample_conversation, "gpt-4", decoded_token
|
||||
)
|
||||
|
||||
assert result.success is True
|
||||
assert result.compression_performed is True
|
||||
assert result.compressed_summary == "compressed text"
|
||||
mock_svc_instance.compress_and_save.assert_called_once()
|
||||
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_provider_from_model_id"
|
||||
)
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_api_key_for_provider"
|
||||
)
|
||||
@patch("application.api.answer.services.compression.orchestrator.LLMCreator")
|
||||
@patch("application.api.answer.services.compression.orchestrator.settings")
|
||||
def test_uses_compression_model_override(
|
||||
self,
|
||||
mock_settings,
|
||||
MockLLMCreator,
|
||||
mock_get_api_key,
|
||||
mock_get_provider,
|
||||
mock_conversation_service,
|
||||
mock_threshold_checker,
|
||||
decoded_token,
|
||||
):
|
||||
mock_settings.COMPRESSION_MODEL_OVERRIDE = "gpt-3.5-turbo"
|
||||
mock_get_provider.return_value = "openai"
|
||||
mock_get_api_key.return_value = "sk-test"
|
||||
MockLLMCreator.create_llm.return_value = MagicMock()
|
||||
|
||||
conversation = {"queries": [{"prompt": "q", "response": "r"}], "agent_id": "a"}
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.compression.orchestrator.CompressionService"
|
||||
) as MockCS:
|
||||
mock_svc = MagicMock()
|
||||
mock_svc.compress_and_save.return_value = MagicMock(
|
||||
compression_ratio=3.0,
|
||||
original_token_count=300,
|
||||
compressed_token_count=100,
|
||||
)
|
||||
mock_svc.get_compressed_context.return_value = ("s", [])
|
||||
MockCS.return_value = mock_svc
|
||||
|
||||
mock_conversation_service.get_conversation.return_value = conversation
|
||||
|
||||
orch = CompressionOrchestrator(
|
||||
conversation_service=mock_conversation_service,
|
||||
threshold_checker=mock_threshold_checker,
|
||||
)
|
||||
orch._perform_compression("c1", conversation, "gpt-4", decoded_token)
|
||||
|
||||
# Verify the override model was used. user_id flows from
|
||||
# decoded_token['sub'] so per-user BYOM custom-model UUIDs
|
||||
# resolve.
|
||||
mock_get_provider.assert_called_with(
|
||||
"gpt-3.5-turbo", user_id=decoded_token["sub"]
|
||||
)
|
||||
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_provider_from_model_id"
|
||||
)
|
||||
@patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_api_key_for_provider"
|
||||
)
|
||||
@patch("application.api.answer.services.compression.orchestrator.LLMCreator")
|
||||
@patch("application.api.answer.services.compression.orchestrator.CompressionService")
|
||||
@patch("application.api.answer.services.compression.orchestrator.settings")
|
||||
def test_no_queries_returns_no_compression(
|
||||
self,
|
||||
mock_settings,
|
||||
MockCompressionService,
|
||||
MockLLMCreator,
|
||||
mock_get_api_key,
|
||||
mock_get_provider,
|
||||
mock_conversation_service,
|
||||
mock_threshold_checker,
|
||||
decoded_token,
|
||||
):
|
||||
mock_settings.COMPRESSION_MODEL_OVERRIDE = None
|
||||
mock_get_provider.return_value = "openai"
|
||||
mock_get_api_key.return_value = "sk-test"
|
||||
MockLLMCreator.create_llm.return_value = MagicMock()
|
||||
|
||||
conversation = {"queries": [], "agent_id": "a"}
|
||||
|
||||
orch = CompressionOrchestrator(
|
||||
conversation_service=mock_conversation_service,
|
||||
threshold_checker=mock_threshold_checker,
|
||||
)
|
||||
result = orch._perform_compression("c1", conversation, "gpt-4", decoded_token)
|
||||
|
||||
assert result.success is True
|
||||
assert result.compression_performed is False
|
||||
|
||||
def test_exception_returns_failure(
|
||||
self,
|
||||
mock_conversation_service,
|
||||
mock_threshold_checker,
|
||||
decoded_token,
|
||||
):
|
||||
conversation = {
|
||||
"queries": [{"prompt": "q", "response": "r"}],
|
||||
"agent_id": "a",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.compression.orchestrator.settings"
|
||||
) as mock_settings, patch(
|
||||
"application.api.answer.services.compression.orchestrator.get_provider_from_model_id",
|
||||
side_effect=RuntimeError("provider error"),
|
||||
):
|
||||
mock_settings.COMPRESSION_MODEL_OVERRIDE = None
|
||||
|
||||
orch = CompressionOrchestrator(
|
||||
conversation_service=mock_conversation_service,
|
||||
threshold_checker=mock_threshold_checker,
|
||||
)
|
||||
result = orch._perform_compression(
|
||||
"c1", conversation, "gpt-4", decoded_token
|
||||
)
|
||||
|
||||
assert result.success is False
|
||||
assert "provider error" in result.error
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressMidExecution:
|
||||
def test_with_provided_conversation(
|
||||
self,
|
||||
orchestrator,
|
||||
sample_conversation,
|
||||
decoded_token,
|
||||
):
|
||||
with patch.object(
|
||||
orchestrator, "_perform_compression"
|
||||
) as mock_perform:
|
||||
mock_perform.return_value = CompressionResult.success_no_compression([])
|
||||
|
||||
orchestrator.compress_mid_execution(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token=decoded_token,
|
||||
current_conversation=sample_conversation,
|
||||
)
|
||||
|
||||
mock_perform.assert_called_once_with(
|
||||
"conv1",
|
||||
sample_conversation,
|
||||
"gpt-4",
|
||||
decoded_token,
|
||||
user_id="user1",
|
||||
model_user_id=None,
|
||||
)
|
||||
|
||||
def test_loads_conversation_when_not_provided(
|
||||
self,
|
||||
orchestrator,
|
||||
mock_conversation_service,
|
||||
sample_conversation,
|
||||
decoded_token,
|
||||
):
|
||||
mock_conversation_service.get_conversation.return_value = sample_conversation
|
||||
|
||||
with patch.object(
|
||||
orchestrator, "_perform_compression"
|
||||
) as mock_perform:
|
||||
mock_perform.return_value = CompressionResult.success_no_compression([])
|
||||
|
||||
orchestrator.compress_mid_execution(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token=decoded_token,
|
||||
)
|
||||
|
||||
mock_conversation_service.get_conversation.assert_called_once_with(
|
||||
"conv1", "user1"
|
||||
)
|
||||
mock_perform.assert_called_once()
|
||||
|
||||
def test_conversation_not_found_returns_failure(
|
||||
self,
|
||||
orchestrator,
|
||||
mock_conversation_service,
|
||||
decoded_token,
|
||||
):
|
||||
mock_conversation_service.get_conversation.return_value = None
|
||||
|
||||
result = orchestrator.compress_mid_execution(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token=decoded_token,
|
||||
)
|
||||
|
||||
assert result.success is False
|
||||
assert "not found" in result.error
|
||||
|
||||
def test_exception_returns_failure(
|
||||
self,
|
||||
orchestrator,
|
||||
mock_conversation_service,
|
||||
decoded_token,
|
||||
):
|
||||
mock_conversation_service.get_conversation.side_effect = RuntimeError("fail")
|
||||
|
||||
result = orchestrator.compress_mid_execution(
|
||||
conversation_id="conv1",
|
||||
user_id="user1",
|
||||
model_id="gpt-4",
|
||||
decoded_token=decoded_token,
|
||||
)
|
||||
|
||||
assert result.success is False
|
||||
assert "fail" in result.error
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestOrchestratorInit:
|
||||
def test_default_threshold_checker(self, mock_conversation_service):
|
||||
orch = CompressionOrchestrator(
|
||||
conversation_service=mock_conversation_service
|
||||
)
|
||||
assert orch.threshold_checker is not None
|
||||
assert orch.conversation_service is mock_conversation_service
|
||||
|
||||
def test_custom_threshold_checker(
|
||||
self, mock_conversation_service, mock_threshold_checker
|
||||
):
|
||||
orch = CompressionOrchestrator(
|
||||
conversation_service=mock_conversation_service,
|
||||
threshold_checker=mock_threshold_checker,
|
||||
)
|
||||
assert orch.threshold_checker is mock_threshold_checker
|
||||
@@ -0,0 +1,423 @@
|
||||
"""Tests for application/api/answer/services/compression/service.py"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.compression.service import CompressionService
|
||||
from application.api.answer.services.compression.types import CompressionMetadata
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm():
|
||||
llm = MagicMock()
|
||||
llm.gen.return_value = "<summary>Compressed summary content</summary>"
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_conversation_service():
|
||||
svc = MagicMock()
|
||||
svc.update_compression_metadata = MagicMock()
|
||||
return svc
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_conversation():
|
||||
return {
|
||||
"queries": [
|
||||
{"prompt": "What is Python?", "response": "A programming language."},
|
||||
{"prompt": "Tell me more.", "response": "It's versatile and popular."},
|
||||
{
|
||||
"prompt": "What about tools?",
|
||||
"response": "Python has many tools.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"tool_name": "search",
|
||||
"action_name": "web_search",
|
||||
"arguments": {"q": "python tools"},
|
||||
"result": "Found 10 results",
|
||||
"status": "success",
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
"compression_metadata": {},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressionServiceInit:
|
||||
@patch("application.api.answer.services.compression.service.settings")
|
||||
def test_default_prompt_builder(self, mock_settings, mock_llm):
|
||||
mock_settings.COMPRESSION_PROMPT_VERSION = "v1.0"
|
||||
with patch(
|
||||
"application.api.answer.services.compression.service.CompressionPromptBuilder"
|
||||
):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
assert svc.llm is mock_llm
|
||||
assert svc.model_id == "gpt-4"
|
||||
|
||||
def test_custom_prompt_builder(self, mock_llm):
|
||||
custom_builder = MagicMock()
|
||||
svc = CompressionService(
|
||||
llm=mock_llm, model_id="gpt-4", prompt_builder=custom_builder
|
||||
)
|
||||
assert svc.prompt_builder is custom_builder
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressConversation:
|
||||
def test_successful_compression(self, mock_llm, sample_conversation):
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.build_prompt.return_value = [
|
||||
{"role": "system", "content": "Compress"},
|
||||
{"role": "user", "content": "Conversation..."},
|
||||
]
|
||||
mock_builder.version = "v1.0"
|
||||
|
||||
svc = CompressionService(
|
||||
llm=mock_llm, model_id="gpt-4", prompt_builder=mock_builder
|
||||
)
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.compression.service.TokenCounter"
|
||||
) as MockTC:
|
||||
MockTC.count_query_tokens.return_value = 1000
|
||||
MockTC.count_message_tokens.return_value = 100
|
||||
|
||||
result = svc.compress_conversation(sample_conversation, 2)
|
||||
|
||||
assert isinstance(result, CompressionMetadata)
|
||||
assert result.query_index == 2
|
||||
assert result.compressed_summary == "Compressed summary content"
|
||||
assert result.original_token_count == 1000
|
||||
assert result.compressed_token_count == 100
|
||||
assert result.compression_ratio == 10.0
|
||||
assert result.model_used == "gpt-4"
|
||||
assert result.compression_prompt_version == "v1.0"
|
||||
|
||||
def test_invalid_index_negative(self, mock_llm, sample_conversation):
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.version = "v1.0"
|
||||
svc = CompressionService(
|
||||
llm=mock_llm, model_id="gpt-4", prompt_builder=mock_builder
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid compress_up_to_index"):
|
||||
svc.compress_conversation(sample_conversation, -1)
|
||||
|
||||
def test_invalid_index_too_large(self, mock_llm, sample_conversation):
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.version = "v1.0"
|
||||
svc = CompressionService(
|
||||
llm=mock_llm, model_id="gpt-4", prompt_builder=mock_builder
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid compress_up_to_index"):
|
||||
svc.compress_conversation(sample_conversation, 10)
|
||||
|
||||
def test_with_existing_compressions(self, mock_llm):
|
||||
conversation = {
|
||||
"queries": [
|
||||
{"prompt": "q1", "response": "r1"},
|
||||
{"prompt": "q2", "response": "r2"},
|
||||
],
|
||||
"compression_metadata": {
|
||||
"compression_points": [
|
||||
{
|
||||
"query_index": 0,
|
||||
"compressed_summary": "Previous summary",
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.build_prompt.return_value = [
|
||||
{"role": "system", "content": "Compress"},
|
||||
{"role": "user", "content": "..."},
|
||||
]
|
||||
mock_builder.version = "v1.0"
|
||||
|
||||
svc = CompressionService(
|
||||
llm=mock_llm, model_id="gpt-4", prompt_builder=mock_builder
|
||||
)
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.compression.service.TokenCounter"
|
||||
) as MockTC:
|
||||
MockTC.count_query_tokens.return_value = 500
|
||||
MockTC.count_message_tokens.return_value = 50
|
||||
|
||||
result = svc.compress_conversation(conversation, 1)
|
||||
assert isinstance(result, CompressionMetadata)
|
||||
# Verify existing compressions were passed to prompt builder
|
||||
call_args = mock_builder.build_prompt.call_args
|
||||
assert call_args[0][1] == [
|
||||
{"query_index": 0, "compressed_summary": "Previous summary"}
|
||||
]
|
||||
|
||||
def test_zero_compressed_tokens_ratio(self, mock_llm, sample_conversation):
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.build_prompt.return_value = [
|
||||
{"role": "system", "content": "C"},
|
||||
{"role": "user", "content": "..."},
|
||||
]
|
||||
mock_builder.version = "v1.0"
|
||||
|
||||
svc = CompressionService(
|
||||
llm=mock_llm, model_id="gpt-4", prompt_builder=mock_builder
|
||||
)
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.compression.service.TokenCounter"
|
||||
) as MockTC:
|
||||
MockTC.count_query_tokens.return_value = 1000
|
||||
MockTC.count_message_tokens.return_value = 0
|
||||
|
||||
result = svc.compress_conversation(sample_conversation, 2)
|
||||
assert result.compression_ratio == 0
|
||||
|
||||
def test_llm_error_propagates(self, sample_conversation):
|
||||
llm = MagicMock()
|
||||
llm.gen.side_effect = RuntimeError("LLM error")
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.build_prompt.return_value = [
|
||||
{"role": "system", "content": "C"},
|
||||
{"role": "user", "content": "..."},
|
||||
]
|
||||
mock_builder.version = "v1.0"
|
||||
|
||||
svc = CompressionService(
|
||||
llm=llm, model_id="gpt-4", prompt_builder=mock_builder
|
||||
)
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.compression.service.TokenCounter"
|
||||
) as MockTC:
|
||||
MockTC.count_query_tokens.return_value = 100
|
||||
with pytest.raises(RuntimeError, match="LLM error"):
|
||||
svc.compress_conversation(sample_conversation, 2)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressAndSave:
|
||||
def test_saves_metadata_to_db(
|
||||
self, mock_llm, mock_conversation_service, sample_conversation
|
||||
):
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.build_prompt.return_value = [
|
||||
{"role": "system", "content": "C"},
|
||||
{"role": "user", "content": "..."},
|
||||
]
|
||||
mock_builder.version = "v1.0"
|
||||
|
||||
svc = CompressionService(
|
||||
llm=mock_llm,
|
||||
model_id="gpt-4",
|
||||
conversation_service=mock_conversation_service,
|
||||
prompt_builder=mock_builder,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.compression.service.TokenCounter"
|
||||
) as MockTC:
|
||||
MockTC.count_query_tokens.return_value = 500
|
||||
MockTC.count_message_tokens.return_value = 50
|
||||
|
||||
result = svc.compress_and_save("conv_123", sample_conversation, 2)
|
||||
|
||||
assert isinstance(result, CompressionMetadata)
|
||||
mock_conversation_service.update_compression_metadata.assert_called_once_with(
|
||||
"conv_123", result.to_dict()
|
||||
)
|
||||
|
||||
def test_raises_without_conversation_service(self, mock_llm, sample_conversation):
|
||||
mock_builder = MagicMock()
|
||||
mock_builder.version = "v1.0"
|
||||
svc = CompressionService(
|
||||
llm=mock_llm, model_id="gpt-4", prompt_builder=mock_builder
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="conversation_service required"):
|
||||
svc.compress_and_save("conv_123", sample_conversation, 2)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGetCompressedContext:
|
||||
def test_no_compression_returns_full_history(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
conversation = {
|
||||
"queries": [{"prompt": "q1", "response": "r1"}],
|
||||
"compression_metadata": {},
|
||||
}
|
||||
|
||||
summary, queries = svc.get_compressed_context(conversation)
|
||||
|
||||
assert summary is None
|
||||
assert queries == [{"prompt": "q1", "response": "r1"}]
|
||||
|
||||
def test_no_compression_points_returns_full_history(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
conversation = {
|
||||
"queries": [{"prompt": "q1", "response": "r1"}],
|
||||
"compression_metadata": {"is_compressed": True, "compression_points": []},
|
||||
}
|
||||
|
||||
summary, queries = svc.get_compressed_context(conversation)
|
||||
assert summary is None
|
||||
assert len(queries) == 1
|
||||
|
||||
def test_with_compression_returns_summary_and_recent(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
conversation = {
|
||||
"queries": [
|
||||
{"prompt": "q0", "response": "r0"},
|
||||
{"prompt": "q1", "response": "r1"},
|
||||
{"prompt": "q2", "response": "r2"},
|
||||
],
|
||||
"compression_metadata": {
|
||||
"is_compressed": True,
|
||||
"compression_points": [
|
||||
{
|
||||
"query_index": 1,
|
||||
"compressed_summary": "Summary of q0 and q1",
|
||||
"compressed_token_count": 50,
|
||||
"original_token_count": 500,
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
summary, queries = svc.get_compressed_context(conversation)
|
||||
|
||||
assert summary == "Summary of q0 and q1"
|
||||
assert len(queries) == 1
|
||||
assert queries[0]["prompt"] == "q2"
|
||||
|
||||
def test_none_queries_returns_empty(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
conversation = {
|
||||
"queries": None,
|
||||
"compression_metadata": {},
|
||||
}
|
||||
|
||||
summary, queries = svc.get_compressed_context(conversation)
|
||||
assert summary is None
|
||||
assert queries == []
|
||||
|
||||
def test_exception_falls_back_to_full_history(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
conversation = {
|
||||
"queries": [{"prompt": "q", "response": "r"}],
|
||||
"compression_metadata": {
|
||||
"is_compressed": True,
|
||||
"compression_points": "invalid", # This will cause an error
|
||||
},
|
||||
}
|
||||
|
||||
summary, queries = svc.get_compressed_context(conversation)
|
||||
assert summary is None
|
||||
assert queries == [{"prompt": "q", "response": "r"}]
|
||||
|
||||
def test_exception_with_none_queries_returns_empty(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
# Force exception by making compression_points non-iterable
|
||||
conversation = {
|
||||
"queries": None,
|
||||
"compression_metadata": {
|
||||
"is_compressed": True,
|
||||
"compression_points": "bad",
|
||||
},
|
||||
}
|
||||
|
||||
summary, queries = svc.get_compressed_context(conversation)
|
||||
assert summary is None
|
||||
assert queries == []
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestExtractSummary:
|
||||
def test_extracts_from_summary_tags(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
response = "<analysis>Some analysis</analysis><summary>The actual summary</summary>"
|
||||
result = svc._extract_summary(response)
|
||||
assert result == "The actual summary"
|
||||
|
||||
def test_removes_analysis_tags_when_no_summary(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
response = "<analysis>analysis text</analysis>Raw summary text here"
|
||||
result = svc._extract_summary(response)
|
||||
assert result == "Raw summary text here"
|
||||
|
||||
def test_returns_full_response_when_no_tags(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
response = "Just a plain text response"
|
||||
result = svc._extract_summary(response)
|
||||
assert result == "Just a plain text response"
|
||||
|
||||
def test_multiline_summary(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
response = "<summary>Line 1\nLine 2\nLine 3</summary>"
|
||||
result = svc._extract_summary(response)
|
||||
assert "Line 1" in result
|
||||
assert "Line 3" in result
|
||||
|
||||
def test_strips_whitespace(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
response = "<summary> Trimmed </summary>"
|
||||
result = svc._extract_summary(response)
|
||||
assert result == "Trimmed"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLogToolCallStats:
|
||||
def test_no_tool_calls(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
queries = [{"prompt": "q", "response": "r"}]
|
||||
# Should not raise
|
||||
svc._log_tool_call_stats(queries)
|
||||
|
||||
def test_with_tool_calls(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
queries = [
|
||||
{
|
||||
"prompt": "q",
|
||||
"response": "r",
|
||||
"tool_calls": [
|
||||
{
|
||||
"tool_name": "search",
|
||||
"action_name": "web",
|
||||
"result": "result text",
|
||||
},
|
||||
{
|
||||
"tool_name": "search",
|
||||
"action_name": "web",
|
||||
"result": "more text",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
# Should not raise - just logs
|
||||
svc._log_tool_call_stats(queries)
|
||||
|
||||
def test_empty_queries(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
svc._log_tool_call_stats([])
|
||||
|
||||
def test_tool_call_with_none_result(self, mock_llm):
|
||||
svc = CompressionService(llm=mock_llm, model_id="gpt-4")
|
||||
queries = [
|
||||
{
|
||||
"prompt": "q",
|
||||
"response": "r",
|
||||
"tool_calls": [
|
||||
{
|
||||
"tool_name": "t",
|
||||
"action_name": "a",
|
||||
"result": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
svc._log_tool_call_stats(queries)
|
||||
@@ -0,0 +1,46 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressionThresholdChecker:
|
||||
|
||||
def _make_checker(self, pct=0.7):
|
||||
from application.api.answer.services.compression.threshold_checker import (
|
||||
CompressionThresholdChecker,
|
||||
)
|
||||
|
||||
return CompressionThresholdChecker(threshold_percentage=pct)
|
||||
|
||||
@patch(
|
||||
"application.api.answer.services.compression.threshold_checker.get_token_limit",
|
||||
return_value=8000,
|
||||
)
|
||||
@patch(
|
||||
"application.api.answer.services.compression.threshold_checker.TokenCounter.count_message_tokens",
|
||||
return_value=6000,
|
||||
)
|
||||
def test_check_message_tokens_above_threshold(self, mock_count, mock_limit):
|
||||
checker = self._make_checker(0.7)
|
||||
assert checker.check_message_tokens([{"role": "user"}], "gpt-4") is True
|
||||
|
||||
@patch(
|
||||
"application.api.answer.services.compression.threshold_checker.get_token_limit",
|
||||
return_value=8000,
|
||||
)
|
||||
@patch(
|
||||
"application.api.answer.services.compression.threshold_checker.TokenCounter.count_message_tokens",
|
||||
return_value=1000,
|
||||
)
|
||||
def test_check_message_tokens_below_threshold(self, mock_count, mock_limit):
|
||||
checker = self._make_checker(0.7)
|
||||
assert checker.check_message_tokens([{"role": "user"}], "gpt-4") is False
|
||||
|
||||
@patch(
|
||||
"application.api.answer.services.compression.threshold_checker.TokenCounter.count_message_tokens",
|
||||
side_effect=Exception("Token error"),
|
||||
)
|
||||
def test_check_message_tokens_exception_returns_false(self, mock_count):
|
||||
checker = self._make_checker(0.7)
|
||||
assert checker.check_message_tokens([], "gpt-4") is False
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Tests for application/api/answer/services/compression/types.py"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.compression.types import (
|
||||
CompressionMetadata,
|
||||
CompressionResult,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressionMetadata:
|
||||
def _make_metadata(self, **overrides):
|
||||
defaults = dict(
|
||||
timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc),
|
||||
query_index=5,
|
||||
compressed_summary="Summary of conversation",
|
||||
original_token_count=5000,
|
||||
compressed_token_count=500,
|
||||
compression_ratio=10.0,
|
||||
model_used="gpt-4",
|
||||
compression_prompt_version="v1.0",
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return CompressionMetadata(**defaults)
|
||||
|
||||
def test_to_dict_contains_all_fields(self):
|
||||
meta = self._make_metadata()
|
||||
d = meta.to_dict()
|
||||
|
||||
assert d["timestamp"] == datetime(2025, 1, 1, tzinfo=timezone.utc)
|
||||
assert d["query_index"] == 5
|
||||
assert d["compressed_summary"] == "Summary of conversation"
|
||||
assert d["original_token_count"] == 5000
|
||||
assert d["compressed_token_count"] == 500
|
||||
assert d["compression_ratio"] == 10.0
|
||||
assert d["model_used"] == "gpt-4"
|
||||
assert d["compression_prompt_version"] == "v1.0"
|
||||
|
||||
def test_to_dict_returns_dict_type(self):
|
||||
meta = self._make_metadata()
|
||||
assert isinstance(meta.to_dict(), dict)
|
||||
|
||||
def test_to_dict_field_count(self):
|
||||
meta = self._make_metadata()
|
||||
d = meta.to_dict()
|
||||
assert len(d) == 8
|
||||
|
||||
def test_attributes_accessible(self):
|
||||
meta = self._make_metadata(query_index=10, compression_ratio=5.5)
|
||||
assert meta.query_index == 10
|
||||
assert meta.compression_ratio == 5.5
|
||||
|
||||
def test_zero_compressed_tokens(self):
|
||||
meta = self._make_metadata(compressed_token_count=0, compression_ratio=0)
|
||||
d = meta.to_dict()
|
||||
assert d["compressed_token_count"] == 0
|
||||
assert d["compression_ratio"] == 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCompressionResult:
|
||||
def test_success_with_compression(self):
|
||||
meta = CompressionMetadata(
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
query_index=3,
|
||||
compressed_summary="summary",
|
||||
original_token_count=1000,
|
||||
compressed_token_count=100,
|
||||
compression_ratio=10.0,
|
||||
model_used="gpt-4",
|
||||
compression_prompt_version="v1.0",
|
||||
)
|
||||
queries = [{"prompt": "q1", "response": "r1"}]
|
||||
result = CompressionResult.success_with_compression("summary", queries, meta)
|
||||
|
||||
assert result.success is True
|
||||
assert result.compressed_summary == "summary"
|
||||
assert result.recent_queries == queries
|
||||
assert result.metadata is meta
|
||||
assert result.compression_performed is True
|
||||
assert result.error is None
|
||||
|
||||
def test_success_no_compression(self):
|
||||
queries = [{"prompt": "q1", "response": "r1"}]
|
||||
result = CompressionResult.success_no_compression(queries)
|
||||
|
||||
assert result.success is True
|
||||
assert result.compressed_summary is None
|
||||
assert result.recent_queries == queries
|
||||
assert result.metadata is None
|
||||
assert result.compression_performed is False
|
||||
assert result.error is None
|
||||
|
||||
def test_failure(self):
|
||||
result = CompressionResult.failure("something went wrong")
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "something went wrong"
|
||||
assert result.compression_performed is False
|
||||
assert result.compressed_summary is None
|
||||
assert result.recent_queries == []
|
||||
assert result.metadata is None
|
||||
|
||||
def test_as_history_extracts_prompt_response(self):
|
||||
queries = [
|
||||
{"prompt": "Hello", "response": "Hi", "extra": "ignored"},
|
||||
{"prompt": "How?", "response": "Fine"},
|
||||
]
|
||||
result = CompressionResult.success_no_compression(queries)
|
||||
history = result.as_history()
|
||||
|
||||
assert len(history) == 2
|
||||
assert history[0] == {"prompt": "Hello", "response": "Hi"}
|
||||
assert history[1] == {"prompt": "How?", "response": "Fine"}
|
||||
|
||||
def test_as_history_empty_queries(self):
|
||||
result = CompressionResult.success_no_compression([])
|
||||
assert result.as_history() == []
|
||||
|
||||
def test_default_recent_queries_is_empty_list(self):
|
||||
result = CompressionResult(success=True)
|
||||
assert result.recent_queries == []
|
||||
assert result.as_history() == []
|
||||
|
||||
def test_success_no_compression_with_empty_list(self):
|
||||
result = CompressionResult.success_no_compression([])
|
||||
assert result.success is True
|
||||
assert result.recent_queries == []
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Tests for application/api/answer/services/continuation_service.py using pg_conn."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import patch
|
||||
from uuid import uuid4
|
||||
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.continuation_service.db_readonly",
|
||||
_yield,
|
||||
), patch(
|
||||
"application.api.answer.services.continuation_service.db_session",
|
||||
_yield,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
class TestMakeSerializable:
|
||||
def test_uuid_becomes_string(self):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
u = uuid4()
|
||||
assert _make_serializable(u) == str(u)
|
||||
|
||||
def test_dict_keys_stringified(self):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
got = _make_serializable({42: "a", "b": 2})
|
||||
assert got == {"42": "a", "b": 2}
|
||||
|
||||
def test_list_elements_recursively_serialized(self):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
u = uuid4()
|
||||
got = _make_serializable([u, {"x": u}, 1])
|
||||
assert got == [str(u), {"x": str(u)}, 1]
|
||||
|
||||
def test_bytes_base64_encoded(self):
|
||||
# Migrated from UTF-8-replace to base64 once the helper moved to
|
||||
# the shared serialization module — base64 is lossless and round-
|
||||
# trippable (UTF-8-replace silently corrupted binary payloads).
|
||||
import base64
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
got = _make_serializable(b"hello")
|
||||
assert got == base64.b64encode(b"hello").decode("ascii")
|
||||
|
||||
def test_bytes_arbitrary_binary_roundtrips(self):
|
||||
import base64
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
raw = b"\xff\xfe\x00\x10"
|
||||
got = _make_serializable(raw)
|
||||
assert isinstance(got, str)
|
||||
assert base64.b64decode(got) == raw
|
||||
|
||||
def test_passes_through_primitives(self):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
assert _make_serializable("hello") == "hello"
|
||||
assert _make_serializable(42) == 42
|
||||
assert _make_serializable(None) is None
|
||||
assert _make_serializable(True) is True
|
||||
|
||||
def test_datetime_becomes_iso_string(self):
|
||||
# PG SELECT * pulls timestamptz columns through as datetime —
|
||||
# tools_dict carries ``created_at``/``updated_at`` from user_tools
|
||||
# rows, which would otherwise blow up json.dumps in pending_tool_state.
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
|
||||
ts = datetime(2026, 5, 2, 12, 14, 32, tzinfo=timezone.utc)
|
||||
got = _make_serializable(ts)
|
||||
assert got == "2026-05-02T12:14:32+00:00"
|
||||
json.dumps(got) # would raise on raw datetime
|
||||
|
||||
def test_datetime_nested_in_tools_dict(self):
|
||||
# Mirrors the production failure: tools_dict is a dict-of-dicts
|
||||
# where each tool row has timestamp fields buried under string keys.
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
|
||||
ts = datetime(2026, 5, 2, 12, 14, 32, tzinfo=timezone.utc)
|
||||
tools_dict = {
|
||||
"0": {
|
||||
"name": "mcp_tool",
|
||||
"actions": [{"name": "search", "active": True}],
|
||||
"created_at": ts,
|
||||
"updated_at": ts,
|
||||
}
|
||||
}
|
||||
got = _make_serializable(tools_dict)
|
||||
json.dumps(got)
|
||||
assert got["0"]["created_at"] == "2026-05-02T12:14:32+00:00"
|
||||
|
||||
def test_date_becomes_iso_string(self):
|
||||
from datetime import date
|
||||
from application.api.answer.services.continuation_service import (
|
||||
_make_serializable,
|
||||
)
|
||||
assert _make_serializable(date(2026, 5, 2)) == "2026-05-02"
|
||||
|
||||
|
||||
class TestContinuationServiceSaveLoad:
|
||||
def test_save_and_load_state(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-cont"
|
||||
conv = ConversationsRepository(pg_conn).create(user, name="c")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
service.save_state(
|
||||
conversation_id=conv_id,
|
||||
user=user,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
pending_tool_calls=[{"id": "call-1", "name": "search"}],
|
||||
tools_dict={"search": {"arg": "value"}},
|
||||
tool_schemas=[{"name": "search", "params": {}}],
|
||||
agent_config={"model_id": "gpt-4"},
|
||||
client_tools=[{"name": "client-tool"}],
|
||||
)
|
||||
loaded = service.load_state(conv_id, user)
|
||||
assert loaded is not None
|
||||
assert loaded["messages"] == [{"role": "user", "content": "hi"}]
|
||||
|
||||
def test_load_state_returns_none_when_not_found(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
got = service.load_state(
|
||||
"00000000-0000-0000-0000-000000000000", "u",
|
||||
)
|
||||
assert got is None
|
||||
|
||||
def test_save_state_no_client_tools(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-no-client"
|
||||
conv = ConversationsRepository(pg_conn).create(user, name="c")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
service.save_state(
|
||||
conversation_id=conv_id,
|
||||
user=user,
|
||||
messages=[],
|
||||
pending_tool_calls=[{"id": "c"}],
|
||||
tools_dict={},
|
||||
tool_schemas=[],
|
||||
agent_config={},
|
||||
client_tools=None,
|
||||
)
|
||||
loaded = service.load_state(conv_id, user)
|
||||
assert loaded is not None
|
||||
|
||||
def test_delete_state_returns_true_when_deleted(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-del"
|
||||
conv = ConversationsRepository(pg_conn).create(user, name="c")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
service.save_state(
|
||||
conversation_id=conv_id,
|
||||
user=user,
|
||||
messages=[],
|
||||
pending_tool_calls=[{"id": "c"}],
|
||||
tools_dict={},
|
||||
tool_schemas=[],
|
||||
agent_config={},
|
||||
)
|
||||
got = service.delete_state(conv_id, user)
|
||||
assert got is True
|
||||
|
||||
def test_delete_state_returns_false_when_missing(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
got = service.delete_state(
|
||||
"00000000-0000-0000-0000-000000000000", "u",
|
||||
)
|
||||
assert got is False
|
||||
|
||||
|
||||
class TestContinuationServiceLegacyIdHandling:
|
||||
"""An unresolvable Mongo ObjectId must not reach ``CAST(:conv_id AS uuid)``
|
||||
in ``PendingToolStateRepository`` — that cast raises and aborts the
|
||||
enclosing transaction. This can happen when an OpenAI-compatible client
|
||||
echoes back a ``chatcmpl-<legacy_objectid>`` id from before the PG
|
||||
cutover and the conversation hasn't been backfilled."""
|
||||
|
||||
LEGACY_OBJECTID = "507f1f77bcf86cd799439011"
|
||||
|
||||
def test_load_state_unresolvable_legacy_id_returns_none(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
got = service.load_state(self.LEGACY_OBJECTID, "u")
|
||||
assert got is None
|
||||
|
||||
def test_delete_state_unresolvable_legacy_id_returns_false(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
got = service.delete_state(self.LEGACY_OBJECTID, "u")
|
||||
assert got is False
|
||||
|
||||
def test_save_state_unresolvable_legacy_id_raises(self, pg_conn):
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
with pytest.raises(ValueError):
|
||||
service.save_state(
|
||||
conversation_id=self.LEGACY_OBJECTID,
|
||||
user="u",
|
||||
messages=[],
|
||||
pending_tool_calls=[{"id": "c"}],
|
||||
tools_dict={},
|
||||
tool_schemas=[],
|
||||
agent_config={},
|
||||
)
|
||||
|
||||
def test_load_state_resolves_backfilled_legacy_id(self, pg_conn):
|
||||
from application.api.answer.services.continuation_service import (
|
||||
ContinuationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "legacy-cont"
|
||||
ConversationsRepository(pg_conn).create(
|
||||
user, name="c", legacy_mongo_id=self.LEGACY_OBJECTID,
|
||||
)
|
||||
|
||||
service = ContinuationService()
|
||||
with _patch_db(pg_conn):
|
||||
# Save via the legacy id — should resolve to the PG UUID internally.
|
||||
service.save_state(
|
||||
conversation_id=self.LEGACY_OBJECTID,
|
||||
user=user,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
pending_tool_calls=[{"id": "c"}],
|
||||
tools_dict={},
|
||||
tool_schemas=[],
|
||||
agent_config={},
|
||||
)
|
||||
# And load via the legacy id — should round-trip.
|
||||
loaded = service.load_state(self.LEGACY_OBJECTID, user)
|
||||
assert loaded is not None
|
||||
assert loaded["messages"] == [{"role": "user", "content": "hi"}]
|
||||
@@ -0,0 +1,873 @@
|
||||
"""Tests for application/api/answer/services/conversation_service.py."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.conversation_service.db_session",
|
||||
_yield,
|
||||
), patch(
|
||||
"application.api.answer.services.conversation_service.db_readonly",
|
||||
_yield,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
class TestConversationServiceGet:
|
||||
def test_returns_none_when_no_conversation_id(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
assert ConversationService().get_conversation("", "u") is None
|
||||
|
||||
def test_returns_none_when_no_user_id(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
assert (
|
||||
ConversationService().get_conversation(
|
||||
"00000000-0000-0000-0000-000000000001", ""
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
def test_returns_none_when_not_found(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
with _patch_db(pg_conn):
|
||||
got = ConversationService().get_conversation(
|
||||
"00000000-0000-0000-0000-000000000000", "u",
|
||||
)
|
||||
assert got is None
|
||||
|
||||
def test_returns_conversation_with_messages(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-get"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="hi")
|
||||
conv_id = str(conv["id"])
|
||||
repo.append_message(
|
||||
conv_id, {"prompt": "q1", "response": "r1"}
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
got = ConversationService().get_conversation(conv_id, user)
|
||||
assert got is not None
|
||||
assert got["_id"] == conv_id
|
||||
assert len(got["queries"]) == 1
|
||||
|
||||
def test_handles_exception_returns_none(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _broken():
|
||||
raise RuntimeError("boom")
|
||||
yield
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.conversation_service.db_readonly",
|
||||
_broken,
|
||||
):
|
||||
got = ConversationService().get_conversation("abc", "u")
|
||||
assert got is None
|
||||
|
||||
|
||||
class TestConversationServiceSave:
|
||||
def test_raises_for_none_token(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
ConversationService().save_conversation(
|
||||
conversation_id=None,
|
||||
question="q", response="r", thought="", sources=[],
|
||||
tool_calls=[], llm=MagicMock(), model_id="gpt",
|
||||
decoded_token=None,
|
||||
)
|
||||
|
||||
def test_raises_when_no_user_in_token(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
ConversationService().save_conversation(
|
||||
conversation_id=None,
|
||||
question="q", response="r", thought="", sources=[],
|
||||
tool_calls=[], llm=MagicMock(), model_id="gpt",
|
||||
decoded_token={},
|
||||
)
|
||||
|
||||
def test_creates_new_conversation(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-save-new"
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.gen.return_value = "Title"
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
conv_id = ConversationService().save_conversation(
|
||||
conversation_id=None,
|
||||
question="q1", response="r1", thought="",
|
||||
sources=[{"text": "x" * 2000}], tool_calls=[],
|
||||
llm=mock_llm, model_id="gpt-4",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
got = ConversationsRepository(pg_conn).get_any(conv_id, user)
|
||||
assert got is not None
|
||||
assert got["name"] == "Title"
|
||||
|
||||
def test_appends_message_to_existing(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-append"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="existing")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
got = ConversationService().save_conversation(
|
||||
conversation_id=conv_id,
|
||||
question="q-new", response="r-new", thought="",
|
||||
sources=[], tool_calls=[],
|
||||
llm=MagicMock(), model_id="gpt-4",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
assert got == conv_id
|
||||
messages = repo.get_messages(conv_id)
|
||||
assert any(m["prompt"] == "q-new" for m in messages)
|
||||
|
||||
def test_updates_message_at_index(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-update-at"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="edit")
|
||||
conv_id = str(conv["id"])
|
||||
repo.append_message(conv_id, {"prompt": "old", "response": "old"})
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
got = ConversationService().save_conversation(
|
||||
conversation_id=conv_id,
|
||||
question="updated", response="reply", thought="",
|
||||
sources=[], tool_calls=[],
|
||||
llm=MagicMock(), model_id="gpt-4",
|
||||
decoded_token={"sub": user},
|
||||
index=0,
|
||||
)
|
||||
assert got == conv_id
|
||||
messages = repo.get_messages(conv_id)
|
||||
assert messages[0]["prompt"] == "updated"
|
||||
|
||||
def test_raises_when_conversation_missing(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
with _patch_db(pg_conn), pytest.raises(ValueError):
|
||||
ConversationService().save_conversation(
|
||||
conversation_id="00000000-0000-0000-0000-000000000000",
|
||||
question="q", response="r", thought="",
|
||||
sources=[], tool_calls=[],
|
||||
llm=MagicMock(), model_id="gpt",
|
||||
decoded_token={"sub": "u"},
|
||||
)
|
||||
|
||||
def test_save_with_empty_llm_title_falls_back(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-empty-title"
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.gen.return_value = ""
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
conv_id = ConversationService().save_conversation(
|
||||
conversation_id=None,
|
||||
question="q-fallback", response="r", thought="",
|
||||
sources=[], tool_calls=[],
|
||||
llm=mock_llm, model_id="gpt-4",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
got = ConversationsRepository(pg_conn).get_any(conv_id, user)
|
||||
assert got["name"] == "q-fallback"
|
||||
|
||||
|
||||
class TestSaveUserQuestion:
|
||||
def test_creates_conversation_and_reserves_message(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
TERMINATED_RESPONSE_PLACEHOLDER,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-wal-new"
|
||||
with _patch_db(pg_conn):
|
||||
result = ConversationService().save_user_question(
|
||||
conversation_id=None,
|
||||
question="what is python?",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
assert result["conversation_id"]
|
||||
assert result["message_id"]
|
||||
assert result["request_id"]
|
||||
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.get_any(result["conversation_id"], user)
|
||||
assert conv is not None
|
||||
messages = repo.get_messages(result["conversation_id"])
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["status"] == "pending"
|
||||
assert messages[0]["prompt"] == "what is python?"
|
||||
assert messages[0]["response"] == TERMINATED_RESPONSE_PLACEHOLDER
|
||||
assert messages[0]["request_id"] == result["request_id"]
|
||||
|
||||
def test_appends_to_existing_conversation(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-wal-existing"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="hi")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
result = ConversationService().save_user_question(
|
||||
conversation_id=conv_id,
|
||||
question="follow-up",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
assert result["conversation_id"] == conv_id
|
||||
msgs = repo.get_messages(conv_id)
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["prompt"] == "follow-up"
|
||||
|
||||
def test_raises_when_token_missing(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
ConversationService().save_user_question(
|
||||
conversation_id=None,
|
||||
question="q",
|
||||
decoded_token=None,
|
||||
)
|
||||
|
||||
def test_regenerate_at_index_replaces_old_message(self, pg_conn):
|
||||
"""Regenerate at ``index`` truncates the old message *and
|
||||
everything after* before reserving the placeholder, so the new
|
||||
WAL row lands at ``position=index`` rather than appending at
|
||||
the end. Pre-fix the WAL path appended unconditionally and the
|
||||
old answer survived alongside the regenerated one.
|
||||
"""
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-wal-regen"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="regen-test")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
# Seed five completed messages at positions 0..4.
|
||||
for i in range(5):
|
||||
repo.append_message(
|
||||
conv_id,
|
||||
{
|
||||
"prompt": f"q{i}",
|
||||
"response": f"a{i}",
|
||||
"thought": "",
|
||||
"sources": [],
|
||||
"tool_calls": [],
|
||||
"metadata": {},
|
||||
},
|
||||
)
|
||||
seeded = repo.get_messages(conv_id)
|
||||
assert len(seeded) == 5
|
||||
assert [m["position"] for m in seeded] == [0, 1, 2, 3, 4]
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
result = ConversationService().save_user_question(
|
||||
conversation_id=conv_id,
|
||||
question="q3-regen",
|
||||
decoded_token={"sub": user},
|
||||
index=3,
|
||||
)
|
||||
|
||||
msgs = repo.get_messages(conv_id)
|
||||
# Positions 0,1,2 from the seed plus the new placeholder at 3.
|
||||
assert [m["position"] for m in msgs] == [0, 1, 2, 3]
|
||||
# The placeholder carries the regenerated prompt.
|
||||
regen = next(m for m in msgs if m["position"] == 3)
|
||||
assert regen["prompt"] == "q3-regen"
|
||||
assert regen["status"] == "pending"
|
||||
assert str(regen["id"]) == result["message_id"]
|
||||
# The old answer at index 3 is gone.
|
||||
assert not any(m["response"] == "a3" for m in msgs)
|
||||
# And anything after index 3 was truncated.
|
||||
assert not any(m["prompt"] == "q4" for m in msgs)
|
||||
|
||||
def test_regenerate_at_index_zero_truncates_everything(self, pg_conn):
|
||||
"""``index=0`` is a valid edge: it should drop every prior
|
||||
message and reseat the placeholder at position 0.
|
||||
"""
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-wal-regen-zero"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="regen-zero")
|
||||
conv_id = str(conv["id"])
|
||||
for i in range(3):
|
||||
repo.append_message(
|
||||
conv_id,
|
||||
{
|
||||
"prompt": f"old-{i}",
|
||||
"response": f"old-a-{i}",
|
||||
"thought": "",
|
||||
"sources": [],
|
||||
"tool_calls": [],
|
||||
"metadata": {},
|
||||
},
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
ConversationService().save_user_question(
|
||||
conversation_id=conv_id,
|
||||
question="fresh-from-start",
|
||||
decoded_token={"sub": user},
|
||||
index=0,
|
||||
)
|
||||
|
||||
msgs = repo.get_messages(conv_id)
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["position"] == 0
|
||||
assert msgs[0]["prompt"] == "fresh-from-start"
|
||||
|
||||
def test_regenerate_index_ignored_without_conversation_id(self, pg_conn):
|
||||
"""``index`` only makes sense against an existing conversation;
|
||||
the create-then-reserve path silently treats it as a no-op
|
||||
rather than truncating a freshly-created conversation.
|
||||
"""
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-wal-regen-no-conv"
|
||||
with _patch_db(pg_conn):
|
||||
result = ConversationService().save_user_question(
|
||||
conversation_id=None,
|
||||
question="brand new q",
|
||||
decoded_token={"sub": user},
|
||||
index=2,
|
||||
)
|
||||
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
msgs = repo.get_messages(result["conversation_id"])
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["position"] == 0
|
||||
assert msgs[0]["prompt"] == "brand new q"
|
||||
|
||||
def test_raises_when_conversation_unauthorized(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
with _patch_db(pg_conn), pytest.raises(ValueError):
|
||||
ConversationService().save_user_question(
|
||||
conversation_id="00000000-0000-0000-0000-000000000000",
|
||||
question="q",
|
||||
decoded_token={"sub": "u"},
|
||||
)
|
||||
|
||||
|
||||
class TestFinalizeMessage:
|
||||
def test_finalizes_complete(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-fin-ok"
|
||||
with _patch_db(pg_conn):
|
||||
svc = ConversationService()
|
||||
res = svc.save_user_question(
|
||||
conversation_id=None,
|
||||
question="q",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
MessageUpdateOutcome,
|
||||
)
|
||||
assert svc.finalize_message(
|
||||
res["message_id"],
|
||||
"real answer",
|
||||
thought="thinking",
|
||||
sources=[{"text": "x" * 2000, "title": "doc"}],
|
||||
tool_calls=[{"name": "search"}],
|
||||
model_id="gpt-4",
|
||||
metadata={"foo": "bar"},
|
||||
status="complete",
|
||||
) is MessageUpdateOutcome.UPDATED
|
||||
|
||||
msgs = ConversationsRepository(pg_conn).get_messages(
|
||||
res["conversation_id"],
|
||||
)
|
||||
assert msgs[0]["response"] == "real answer"
|
||||
assert msgs[0]["status"] == "complete"
|
||||
assert msgs[0]["thought"] == "thinking"
|
||||
assert msgs[0]["model_id"] == "gpt-4"
|
||||
# source text trimmed to 1000 chars at finalize time
|
||||
assert len(msgs[0]["sources"][0]["text"]) == 1000
|
||||
assert msgs[0]["metadata"]["foo"] == "bar"
|
||||
|
||||
def test_finalizes_failed_records_error(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-fin-fail"
|
||||
with _patch_db(pg_conn):
|
||||
svc = ConversationService()
|
||||
res = svc.save_user_question(
|
||||
conversation_id=None,
|
||||
question="q",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
err = RuntimeError("provider down")
|
||||
from application.storage.db.repositories.conversations import (
|
||||
MessageUpdateOutcome,
|
||||
)
|
||||
assert svc.finalize_message(
|
||||
res["message_id"],
|
||||
"fallback text",
|
||||
status="failed",
|
||||
error=err,
|
||||
) is MessageUpdateOutcome.UPDATED
|
||||
|
||||
msgs = ConversationsRepository(pg_conn).get_messages(
|
||||
res["conversation_id"],
|
||||
)
|
||||
assert msgs[0]["status"] == "failed"
|
||||
assert msgs[0]["metadata"]["error"] == "RuntimeError: provider down"
|
||||
|
||||
def test_finalize_flips_executed_tool_calls(self, pg_conn):
|
||||
"""finalize_message must mark tool_call_attempts.status='executed'
|
||||
rows as 'confirmed' for the same message_id."""
|
||||
from sqlalchemy import text as sql_text
|
||||
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
user = "u-fin-tools"
|
||||
with _patch_db(pg_conn):
|
||||
svc = ConversationService()
|
||||
res = svc.save_user_question(
|
||||
conversation_id=None,
|
||||
question="q",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
pg_conn.execute(
|
||||
sql_text(
|
||||
"INSERT INTO tool_call_attempts "
|
||||
"(call_id, message_id, tool_name, action_name, arguments, status) "
|
||||
"VALUES (:cid, CAST(:mid AS uuid), 't', 'a', '{}'::jsonb, 'executed')"
|
||||
),
|
||||
{"cid": "c1", "mid": res["message_id"]},
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
MessageUpdateOutcome,
|
||||
)
|
||||
assert svc.finalize_message(
|
||||
res["message_id"], "ans", status="complete",
|
||||
) is MessageUpdateOutcome.UPDATED
|
||||
|
||||
status = pg_conn.execute(
|
||||
sql_text("SELECT status FROM tool_call_attempts WHERE call_id = :cid"),
|
||||
{"cid": "c1"},
|
||||
).scalar()
|
||||
assert status == "confirmed"
|
||||
|
||||
def test_finalize_returns_not_found_for_unknown_message(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
MessageUpdateOutcome,
|
||||
)
|
||||
with _patch_db(pg_conn):
|
||||
assert ConversationService().finalize_message(
|
||||
"00000000-0000-0000-0000-000000000000",
|
||||
"x",
|
||||
status="complete",
|
||||
) is MessageUpdateOutcome.NOT_FOUND
|
||||
|
||||
def test_finalize_rolls_back_tool_call_confirm_on_message_update_failure(
|
||||
self, pg_conn
|
||||
):
|
||||
"""Atomicity: if ``update_message_by_id`` raises after the
|
||||
tool_call_attempts confirm ran on the same connection, the
|
||||
confirm rolls back with the rest of the transaction. The
|
||||
``pg_conn`` fixture pins one connection inside an outer
|
||||
rolled-back transaction; we patch ``db_session`` to wrap each
|
||||
call in a SAVEPOINT so the production-code ``with`` block
|
||||
actually rolls back when the message-update raises.
|
||||
"""
|
||||
from contextlib import contextmanager
|
||||
|
||||
from sqlalchemy import text as sql_text
|
||||
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories import conversations as conv_module
|
||||
|
||||
user = "u-fin-rollback"
|
||||
|
||||
@contextmanager
|
||||
def _savepoint_session():
|
||||
nested = pg_conn.begin_nested()
|
||||
try:
|
||||
yield pg_conn
|
||||
nested.commit()
|
||||
except Exception:
|
||||
nested.rollback()
|
||||
raise
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.conversation_service.db_session",
|
||||
_savepoint_session,
|
||||
), patch(
|
||||
"application.api.answer.services.conversation_service.db_readonly",
|
||||
_savepoint_session,
|
||||
):
|
||||
svc = ConversationService()
|
||||
res = svc.save_user_question(
|
||||
conversation_id=None,
|
||||
question="q",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
pg_conn.execute(
|
||||
sql_text(
|
||||
"INSERT INTO tool_call_attempts "
|
||||
"(call_id, message_id, tool_name, action_name, "
|
||||
"arguments, status) VALUES (:cid, CAST(:mid AS uuid), "
|
||||
"'t', 'a', '{}'::jsonb, 'executed')"
|
||||
),
|
||||
{"cid": "rb-1", "mid": res["message_id"]},
|
||||
)
|
||||
original = conv_module.ConversationsRepository.update_message_by_id
|
||||
|
||||
def boom(self, *args, **kwargs):
|
||||
_ = (args, kwargs)
|
||||
raise RuntimeError("simulated message-update failure")
|
||||
|
||||
conv_module.ConversationsRepository.update_message_by_id = boom
|
||||
try:
|
||||
with pytest.raises(RuntimeError):
|
||||
svc.finalize_message(
|
||||
res["message_id"], "answer", status="complete",
|
||||
)
|
||||
finally:
|
||||
conv_module.ConversationsRepository.update_message_by_id = original
|
||||
|
||||
# The tool_call confirm rolled back: row stays at ``executed``.
|
||||
status = pg_conn.execute(
|
||||
sql_text(
|
||||
"SELECT status FROM tool_call_attempts WHERE call_id = :cid"
|
||||
),
|
||||
{"cid": "rb-1"},
|
||||
).scalar()
|
||||
assert status == "executed"
|
||||
msg_status = pg_conn.execute(
|
||||
sql_text(
|
||||
"SELECT status FROM conversation_messages "
|
||||
"WHERE id = CAST(:mid AS uuid)"
|
||||
),
|
||||
{"mid": res["message_id"]},
|
||||
).scalar()
|
||||
assert msg_status == "pending"
|
||||
|
||||
def test_finalize_generates_title_when_provided(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-fin-title"
|
||||
mock_llm = MagicMock()
|
||||
mock_llm.gen.return_value = "Short Title"
|
||||
with _patch_db(pg_conn):
|
||||
svc = ConversationService()
|
||||
res = svc.save_user_question(
|
||||
conversation_id=None,
|
||||
question="long question that becomes the fallback name",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
MessageUpdateOutcome,
|
||||
)
|
||||
assert svc.finalize_message(
|
||||
res["message_id"],
|
||||
"answer",
|
||||
status="complete",
|
||||
title_inputs={
|
||||
"llm": mock_llm,
|
||||
"question": "long question that becomes the fallback name",
|
||||
"response": "answer",
|
||||
"model_id": "gpt-4",
|
||||
"fallback_name": (
|
||||
"long question that becomes the fallback name"[:50]
|
||||
),
|
||||
},
|
||||
) is MessageUpdateOutcome.UPDATED
|
||||
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.get_any(res["conversation_id"], user)
|
||||
assert conv["name"] == "Short Title"
|
||||
|
||||
|
||||
class TestSaveUserQuestionFinalizeFailedFlow:
|
||||
"""LLM fails immediately; question stays queryable with status='failed' + error metadata."""
|
||||
|
||||
def test_failed_llm_leaves_question_persisted(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-acceptance"
|
||||
with _patch_db(pg_conn):
|
||||
svc = ConversationService()
|
||||
# Simulates the WAL pre-persist before LLM call.
|
||||
res = svc.save_user_question(
|
||||
conversation_id=None,
|
||||
question="why did this fail?",
|
||||
decoded_token={"sub": user},
|
||||
)
|
||||
# Simulates the LLM raising immediately, caught by complete_stream.
|
||||
try:
|
||||
raise RuntimeError("upstream 503")
|
||||
except RuntimeError as e:
|
||||
svc.finalize_message(
|
||||
res["message_id"],
|
||||
"",
|
||||
status="failed",
|
||||
error=e,
|
||||
)
|
||||
|
||||
msgs = ConversationsRepository(pg_conn).get_messages(
|
||||
res["conversation_id"],
|
||||
)
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["prompt"] == "why did this fail?"
|
||||
assert msgs[0]["status"] == "failed"
|
||||
assert "RuntimeError" in msgs[0]["metadata"]["error"]
|
||||
assert "upstream 503" in msgs[0]["metadata"]["error"]
|
||||
|
||||
|
||||
class TestCompressionMetadata:
|
||||
def test_update_compression_metadata(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-compress"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="c")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
ConversationService().update_compression_metadata(
|
||||
conv_id,
|
||||
{
|
||||
"timestamp": datetime.now(timezone.utc),
|
||||
"compressed_summary": "summary",
|
||||
"model_used": "gpt",
|
||||
},
|
||||
)
|
||||
|
||||
def test_update_compression_raises_on_error(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _broken():
|
||||
raise RuntimeError("boom")
|
||||
yield
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.conversation_service.db_session",
|
||||
_broken,
|
||||
), pytest.raises(RuntimeError):
|
||||
ConversationService().update_compression_metadata("abc", {})
|
||||
|
||||
def test_append_compression_message_skips_empty_summary(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
ConversationService().append_compression_message(
|
||||
"any-id", {"compressed_summary": ""}
|
||||
)
|
||||
|
||||
def test_append_compression_message_appends_summary(self, pg_conn):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
user = "u-append-cs"
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create(user, name="c")
|
||||
conv_id = str(conv["id"])
|
||||
|
||||
with _patch_db(pg_conn):
|
||||
ConversationService().append_compression_message(
|
||||
conv_id,
|
||||
{
|
||||
"compressed_summary": "A summary",
|
||||
"timestamp": datetime.now(timezone.utc),
|
||||
"model_used": "gpt-4",
|
||||
},
|
||||
)
|
||||
|
||||
messages = repo.get_messages(conv_id)
|
||||
assert any(m["response"] == "A summary" for m in messages)
|
||||
|
||||
def test_append_compression_message_swallows_error(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _broken():
|
||||
raise RuntimeError("boom")
|
||||
yield
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.conversation_service.db_session",
|
||||
_broken,
|
||||
):
|
||||
ConversationService().append_compression_message(
|
||||
"abc", {"compressed_summary": "x"}
|
||||
)
|
||||
|
||||
def test_get_compression_metadata_returns_none_for_missing(
|
||||
self, pg_conn,
|
||||
):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
with _patch_db(pg_conn):
|
||||
got = ConversationService().get_compression_metadata(
|
||||
"00000000-0000-0000-0000-000000000000"
|
||||
)
|
||||
assert got is None
|
||||
|
||||
def test_get_compression_metadata_non_uuid_does_not_raise(
|
||||
self, pg_conn, caplog,
|
||||
):
|
||||
"""Non-UUID ids (legacy Mongo ObjectIds with no legacy_mongo_id row)
|
||||
must return None without hitting ``CAST(:id AS uuid)`` — which
|
||||
would raise and pollute logs with a stack trace every call."""
|
||||
import logging
|
||||
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
with _patch_db(pg_conn):
|
||||
got = ConversationService().get_compression_metadata(
|
||||
"507f1f77bcf86cd799439011"
|
||||
)
|
||||
assert got is None
|
||||
assert not any(
|
||||
"Error getting compression metadata" in r.message for r in caplog.records
|
||||
)
|
||||
|
||||
def test_get_compression_metadata_handles_exception(self):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _broken():
|
||||
raise RuntimeError("boom")
|
||||
yield
|
||||
|
||||
with patch(
|
||||
"application.api.answer.services.conversation_service.db_readonly",
|
||||
_broken,
|
||||
):
|
||||
got = ConversationService().get_compression_metadata("abc")
|
||||
assert got is None
|
||||
@@ -0,0 +1,44 @@
|
||||
"""Unit tests for ``resolve_persistence`` (persist + sidebar visibility)."""
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.persistence_policy import (
|
||||
VISIBILITY_HIDDEN,
|
||||
VISIBILITY_LISTED,
|
||||
resolve_persistence,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestResolvePersistence:
|
||||
def test_default_is_hidden(self):
|
||||
persist, visibility = resolve_persistence()
|
||||
assert persist is True
|
||||
assert visibility == VISIBILITY_HIDDEN
|
||||
|
||||
def test_explicit_listed_lists(self):
|
||||
persist, visibility = resolve_persistence(visibility_flag=VISIBILITY_LISTED)
|
||||
assert persist is True
|
||||
assert visibility == VISIBILITY_LISTED
|
||||
|
||||
def test_explicit_hidden_stays_hidden(self):
|
||||
_, visibility = resolve_persistence(visibility_flag=VISIBILITY_HIDDEN)
|
||||
assert visibility == VISIBILITY_HIDDEN
|
||||
|
||||
def test_unknown_visibility_value_falls_back_to_hidden(self):
|
||||
_, visibility = resolve_persistence(visibility_flag="sidebar-please")
|
||||
assert visibility == VISIBILITY_HIDDEN
|
||||
|
||||
def test_non_string_visibility_value_falls_back_to_hidden(self):
|
||||
# Clients sending booleans (e.g. a stray ``visibility: true``) must
|
||||
# not accidentally list — only the exact string "listed" opts in.
|
||||
_, visibility = resolve_persistence(visibility_flag=True)
|
||||
assert visibility == VISIBILITY_HIDDEN
|
||||
|
||||
def test_persist_opt_out(self):
|
||||
persist, _ = resolve_persistence(persist_flag=False)
|
||||
assert persist is False
|
||||
|
||||
def test_persist_defaults_true_when_listed(self):
|
||||
persist, _ = resolve_persistence(visibility_flag=VISIBILITY_LISTED)
|
||||
assert persist is True
|
||||
@@ -0,0 +1,903 @@
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestTemplateEngine:
|
||||
|
||||
def test_render_simple_template(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
result = engine.render("Hello {{ name }}", {"name": "World"})
|
||||
|
||||
assert result == "Hello World"
|
||||
|
||||
def test_render_with_namespace(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
context = {
|
||||
"user": {"name": "Alice", "role": "admin"},
|
||||
"system": {"date": "2025-10-22"},
|
||||
}
|
||||
result = engine.render(
|
||||
"{{ user.name }} is a {{ user.role }} on {{ system.date }}", context
|
||||
)
|
||||
|
||||
assert result == "Alice is a admin on 2025-10-22"
|
||||
|
||||
def test_render_empty_template(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
result = engine.render("", {"key": "value"})
|
||||
|
||||
assert result == ""
|
||||
|
||||
def test_render_template_without_variables(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
result = engine.render("Just plain text", {})
|
||||
|
||||
assert result == "Just plain text"
|
||||
|
||||
def test_render_undefined_variable_returns_empty_string(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
|
||||
result = engine.render("Hello {{ undefined_var }}", {})
|
||||
assert result == "Hello "
|
||||
|
||||
def test_render_syntax_error_raises_error(self):
|
||||
from application.templates.template_engine import (
|
||||
TemplateEngine,
|
||||
TemplateRenderError,
|
||||
)
|
||||
|
||||
engine = TemplateEngine()
|
||||
|
||||
with pytest.raises(TemplateRenderError, match="Template syntax error"):
|
||||
engine.render("Hello {{ name", {"name": "World"})
|
||||
|
||||
def test_render_blocks_unsafe_attribute_access(self):
|
||||
from application.templates.template_engine import (
|
||||
TemplateEngine,
|
||||
TemplateRenderError,
|
||||
)
|
||||
|
||||
engine = TemplateEngine()
|
||||
payload = "{{ cycler.__init__.__globals__.os.popen('id').read() }}"
|
||||
|
||||
with pytest.raises(TemplateRenderError, match="Template security violation"):
|
||||
engine.render(payload, {})
|
||||
|
||||
def test_validate_template_valid(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
assert engine.validate_template("Valid {{ variable }}") is True
|
||||
|
||||
def test_validate_template_invalid(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
assert engine.validate_template("Invalid {{ variable") is False
|
||||
|
||||
def test_validate_empty_template(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
assert engine.validate_template("") is True
|
||||
|
||||
def test_extract_variables(self):
|
||||
from application.templates.template_engine import TemplateEngine
|
||||
|
||||
engine = TemplateEngine()
|
||||
template = "{{ user.name }} and {{ user.email }}"
|
||||
|
||||
result = engine.extract_variables(template)
|
||||
|
||||
assert isinstance(result, set)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSystemNamespace:
|
||||
|
||||
def test_system_namespace_build(self):
|
||||
from application.templates.namespaces import SystemNamespace
|
||||
|
||||
builder = SystemNamespace()
|
||||
context = builder.build(
|
||||
request_id="req_123", user_id="user_456", extra_param="ignored"
|
||||
)
|
||||
|
||||
assert context["request_id"] == "req_123"
|
||||
assert context["user_id"] == "user_456"
|
||||
assert "date" in context
|
||||
assert "time" in context
|
||||
assert "timestamp" in context
|
||||
|
||||
def test_system_namespace_generates_request_id(self):
|
||||
from application.templates.namespaces import SystemNamespace
|
||||
|
||||
builder = SystemNamespace()
|
||||
context = builder.build(user_id="user_123")
|
||||
|
||||
assert context["request_id"] is not None
|
||||
assert len(context["request_id"]) > 0
|
||||
|
||||
def test_system_namespace_name(self):
|
||||
from application.templates.namespaces import SystemNamespace
|
||||
|
||||
builder = SystemNamespace()
|
||||
assert builder.namespace_name == "system"
|
||||
|
||||
def test_system_namespace_date_format(self):
|
||||
from application.templates.namespaces import SystemNamespace
|
||||
|
||||
builder = SystemNamespace()
|
||||
context = builder.build()
|
||||
|
||||
import re
|
||||
|
||||
assert re.match(r"\d{4}-\d{2}-\d{2}", context["date"])
|
||||
assert re.match(r"\d{2}:\d{2}:\d{2}", context["time"])
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPassthroughNamespace:
|
||||
|
||||
def test_passthrough_namespace_build(self):
|
||||
from application.templates.namespaces import PassthroughNamespace
|
||||
|
||||
builder = PassthroughNamespace()
|
||||
passthrough_data = {"company": "Acme", "user_name": "John", "count": 42}
|
||||
|
||||
context = builder.build(passthrough_data=passthrough_data)
|
||||
|
||||
assert context["company"] == "Acme"
|
||||
assert context["user_name"] == "John"
|
||||
assert context["count"] == 42
|
||||
|
||||
def test_passthrough_namespace_empty(self):
|
||||
from application.templates.namespaces import PassthroughNamespace
|
||||
|
||||
builder = PassthroughNamespace()
|
||||
context = builder.build(passthrough_data=None)
|
||||
|
||||
assert context == {}
|
||||
|
||||
def test_passthrough_namespace_filters_unsafe_values(self):
|
||||
from application.templates.namespaces import PassthroughNamespace
|
||||
|
||||
builder = PassthroughNamespace()
|
||||
passthrough_data = {
|
||||
"safe_string": "value",
|
||||
"unsafe_object": {"key": "value"},
|
||||
"safe_bool": True,
|
||||
"unsafe_list": [1, 2, 3],
|
||||
"safe_float": 3.14,
|
||||
}
|
||||
|
||||
context = builder.build(passthrough_data=passthrough_data)
|
||||
|
||||
assert context["safe_string"] == "value"
|
||||
assert context["safe_bool"] is True
|
||||
assert context["safe_float"] == 3.14
|
||||
assert "unsafe_object" not in context
|
||||
assert "unsafe_list" not in context
|
||||
|
||||
def test_passthrough_namespace_allows_none_values(self):
|
||||
from application.templates.namespaces import PassthroughNamespace
|
||||
|
||||
builder = PassthroughNamespace()
|
||||
passthrough_data = {"nullable_field": None}
|
||||
|
||||
context = builder.build(passthrough_data=passthrough_data)
|
||||
|
||||
assert context["nullable_field"] is None
|
||||
|
||||
def test_passthrough_namespace_name(self):
|
||||
from application.templates.namespaces import PassthroughNamespace
|
||||
|
||||
builder = PassthroughNamespace()
|
||||
assert builder.namespace_name == "passthrough"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSourceNamespace:
|
||||
|
||||
def test_source_namespace_build_with_docs(self):
|
||||
from application.templates.namespaces import SourceNamespace
|
||||
|
||||
builder = SourceNamespace()
|
||||
docs = [
|
||||
{"text": "Doc 1", "filename": "file1.txt"},
|
||||
{"text": "Doc 2", "filename": "file2.txt"},
|
||||
]
|
||||
docs_together = "Doc 1 content\n\nDoc 2 content"
|
||||
|
||||
context = builder.build(docs=docs, docs_together=docs_together)
|
||||
|
||||
assert context["documents"] == docs
|
||||
assert context["count"] == 2
|
||||
assert context["content"] == docs_together
|
||||
assert context["summaries"] == docs_together
|
||||
|
||||
def test_source_namespace_build_empty(self):
|
||||
from application.templates.namespaces import SourceNamespace
|
||||
|
||||
builder = SourceNamespace()
|
||||
context = builder.build(docs=None, docs_together=None)
|
||||
|
||||
assert context == {}
|
||||
|
||||
def test_source_namespace_build_docs_only(self):
|
||||
from application.templates.namespaces import SourceNamespace
|
||||
|
||||
builder = SourceNamespace()
|
||||
docs = [{"text": "Doc 1"}]
|
||||
|
||||
context = builder.build(docs=docs)
|
||||
|
||||
assert context["documents"] == docs
|
||||
assert context["count"] == 1
|
||||
assert "content" not in context
|
||||
|
||||
def test_source_namespace_build_docs_together_only(self):
|
||||
from application.templates.namespaces import SourceNamespace
|
||||
|
||||
builder = SourceNamespace()
|
||||
docs_together = "Content here"
|
||||
|
||||
context = builder.build(docs_together=docs_together)
|
||||
|
||||
assert context["content"] == docs_together
|
||||
assert context["summaries"] == docs_together
|
||||
assert "documents" not in context
|
||||
|
||||
def test_source_namespace_name(self):
|
||||
from application.templates.namespaces import SourceNamespace
|
||||
|
||||
builder = SourceNamespace()
|
||||
assert builder.namespace_name == "source"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestToolsNamespace:
|
||||
|
||||
def test_tools_namespace_build_with_memory_data(self):
|
||||
from application.templates.namespaces import ToolsNamespace
|
||||
|
||||
builder = ToolsNamespace()
|
||||
tools_data = {
|
||||
"memory": {"root": "Files:\n- /notes.txt\n- /tasks.txt", "available": True}
|
||||
}
|
||||
|
||||
context = builder.build(tools_data=tools_data)
|
||||
|
||||
assert context["memory"]["root"] == "Files:\n- /notes.txt\n- /tasks.txt"
|
||||
assert context["memory"]["available"] is True
|
||||
|
||||
def test_tools_namespace_build_empty(self):
|
||||
from application.templates.namespaces import ToolsNamespace
|
||||
|
||||
builder = ToolsNamespace()
|
||||
context = builder.build(tools_data=None)
|
||||
|
||||
assert context == {}
|
||||
|
||||
def test_tools_namespace_build_multiple_tools(self):
|
||||
from application.templates.namespaces import ToolsNamespace
|
||||
|
||||
builder = ToolsNamespace()
|
||||
tools_data = {
|
||||
"memory": {"root": "content", "available": True},
|
||||
"search": {"results": ["result1", "result2"]},
|
||||
"api": {"status": "success"},
|
||||
}
|
||||
|
||||
context = builder.build(tools_data=tools_data)
|
||||
|
||||
assert "memory" in context
|
||||
assert "search" in context
|
||||
assert "api" in context
|
||||
assert context["memory"]["root"] == "content"
|
||||
assert context["search"]["results"] == ["result1", "result2"]
|
||||
assert context["api"]["status"] == "success"
|
||||
|
||||
def test_tools_namespace_filters_unsafe_values(self):
|
||||
from application.templates.namespaces import ToolsNamespace
|
||||
|
||||
builder = ToolsNamespace()
|
||||
|
||||
class UnsafeObject:
|
||||
pass
|
||||
|
||||
tools_data = {"safe_tool": {"result": "success"}, "unsafe_tool": UnsafeObject()}
|
||||
|
||||
context = builder.build(tools_data=tools_data)
|
||||
|
||||
assert "safe_tool" in context
|
||||
assert "unsafe_tool" not in context
|
||||
|
||||
def test_tools_namespace_name(self):
|
||||
from application.templates.namespaces import ToolsNamespace
|
||||
|
||||
builder = ToolsNamespace()
|
||||
assert builder.namespace_name == "tools"
|
||||
|
||||
def test_tools_namespace_with_empty_dict(self):
|
||||
from application.templates.namespaces import ToolsNamespace
|
||||
|
||||
builder = ToolsNamespace()
|
||||
context = builder.build(tools_data={})
|
||||
|
||||
assert context == {}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestNamespaceManagerWithTools:
|
||||
|
||||
def test_namespace_manager_includes_tools_in_context(self):
|
||||
from application.templates.namespaces import NamespaceManager
|
||||
|
||||
manager = NamespaceManager()
|
||||
tools_data = {"memory": {"root": "content", "available": True}}
|
||||
|
||||
context = manager.build_context(tools_data=tools_data)
|
||||
|
||||
assert "tools" in context
|
||||
assert context["tools"]["memory"]["root"] == "content"
|
||||
|
||||
def test_namespace_manager_build_context_all_namespaces(self):
|
||||
from application.templates.namespaces import NamespaceManager
|
||||
|
||||
manager = NamespaceManager()
|
||||
context = manager.build_context(
|
||||
request_id="req_123",
|
||||
user_id="user_456",
|
||||
passthrough_data={"key": "value"},
|
||||
docs_together="Document content",
|
||||
tools_data={"memory": {"root": "notes"}},
|
||||
)
|
||||
|
||||
assert "system" in context
|
||||
assert "passthrough" in context
|
||||
assert "source" in context
|
||||
assert "tools" in context
|
||||
assert context["tools"]["memory"]["root"] == "notes"
|
||||
|
||||
def test_namespace_manager_build_context_partial_data(self):
|
||||
from application.templates.namespaces import NamespaceManager
|
||||
|
||||
manager = NamespaceManager()
|
||||
context = manager.build_context(request_id="req_123")
|
||||
|
||||
assert "system" in context
|
||||
assert context["system"]["request_id"] == "req_123"
|
||||
|
||||
|
||||
class TestPromptRendererArtifactParent:
|
||||
"""render_prompt forwards artifact_parent so artifacts.artifact(id) resolves for classic agents."""
|
||||
|
||||
def test_render_resolves_artifact_by_id_with_conversation_parent(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
row = {"id": "art-1", "current_version": 1, "kind": "document", "title": "Report"}
|
||||
version = {"mime_type": "text/html", "filename": "report.html", "size": 12}
|
||||
with patch(
|
||||
"application.storage.db.repositories.artifacts.ArtifactsRepository"
|
||||
) as repo_cls, patch("application.storage.db.session.db_readonly"):
|
||||
repo = repo_cls.return_value
|
||||
repo.get_artifact_in_parent.return_value = row
|
||||
repo.get_version.return_value = version
|
||||
|
||||
rendered = PromptRenderer().render_prompt(
|
||||
"file={{ artifacts.artifact('art-1').filename }}",
|
||||
artifact_parent={"conversation_id": "conv-1"},
|
||||
)
|
||||
|
||||
assert rendered == "file=report.html"
|
||||
_, kwargs = repo.get_artifact_in_parent.call_args
|
||||
assert kwargs.get("conversation_id") == "conv-1"
|
||||
|
||||
def test_render_without_parent_yields_empty_lookup(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
# No artifact_parent (e.g. a first-turn conversation_id is None) must not error;
|
||||
# the lookup returns an empty mapping and the attribute renders blank.
|
||||
rendered = PromptRenderer().render_prompt(
|
||||
"[{{ artifacts.artifact('art-1').filename }}]"
|
||||
)
|
||||
|
||||
assert rendered == "[]"
|
||||
|
||||
def test_namespace_manager_get_builder(self):
|
||||
from application.templates.namespaces import NamespaceManager, SystemNamespace
|
||||
|
||||
manager = NamespaceManager()
|
||||
builder = manager.get_builder("system")
|
||||
|
||||
assert isinstance(builder, SystemNamespace)
|
||||
|
||||
def test_namespace_manager_get_builder_nonexistent(self):
|
||||
from application.templates.namespaces import NamespaceManager
|
||||
|
||||
manager = NamespaceManager()
|
||||
builder = manager.get_builder("nonexistent")
|
||||
|
||||
assert builder is None
|
||||
|
||||
def test_namespace_manager_handles_builder_exceptions(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
from application.templates.namespaces import NamespaceManager
|
||||
|
||||
manager = NamespaceManager()
|
||||
|
||||
with patch.object(
|
||||
manager._builders["system"],
|
||||
"build",
|
||||
side_effect=Exception("Builder error"),
|
||||
):
|
||||
context = manager.build_context()
|
||||
# Namespace should be present but empty when builder fails
|
||||
|
||||
assert "system" in context
|
||||
assert context["system"] == {}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPromptRenderer:
|
||||
|
||||
def test_render_prompt_with_template_syntax(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Hello {{ system.user_id }}, today is {{ system.date }}"
|
||||
|
||||
result = renderer.render_prompt(prompt, user_id="user_123")
|
||||
|
||||
assert "user_123" in result
|
||||
assert "202" in result
|
||||
|
||||
def test_render_prompt_with_passthrough_data(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Company: {{ passthrough.company }}\nUser: {{ passthrough.user_name }}"
|
||||
passthrough_data = {"company": "Acme", "user_name": "John"}
|
||||
|
||||
result = renderer.render_prompt(prompt, passthrough_data=passthrough_data)
|
||||
|
||||
assert "Company: Acme" in result
|
||||
assert "User: John" in result
|
||||
|
||||
def test_render_prompt_with_source_docs(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Use this information:\n{{ source.content }}"
|
||||
docs_together = "Important document content"
|
||||
|
||||
result = renderer.render_prompt(prompt, docs_together=docs_together)
|
||||
|
||||
assert "Use this information:" in result
|
||||
assert "Important document content" in result
|
||||
|
||||
def test_render_prompt_empty_content(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
result = renderer.render_prompt("")
|
||||
|
||||
assert result == ""
|
||||
|
||||
def test_render_prompt_legacy_format_with_summaries(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Context: {summaries}\nQuestion: What is this?"
|
||||
docs_together = "This is the document content"
|
||||
|
||||
result = renderer.render_prompt(prompt, docs_together=docs_together)
|
||||
|
||||
assert "Context: This is the document content" in result
|
||||
|
||||
def test_render_prompt_legacy_format_without_docs(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Context: {summaries}\nQuestion: What is this?"
|
||||
|
||||
result = renderer.render_prompt(prompt)
|
||||
|
||||
# The placeholder must never leak to the model when no docs exist.
|
||||
assert "{summaries}" not in result
|
||||
assert "Question: What is this?" in result
|
||||
|
||||
def test_render_prompt_combined_namespace_variables(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "User: {{ passthrough.user }}, Date: {{ system.date }}, Docs: {{ source.content }}"
|
||||
passthrough_data = {"user": "Alice"}
|
||||
docs_together = "Doc content"
|
||||
|
||||
result = renderer.render_prompt(
|
||||
prompt,
|
||||
passthrough_data=passthrough_data,
|
||||
docs_together=docs_together,
|
||||
)
|
||||
|
||||
assert "User: Alice" in result
|
||||
assert "Date: 202" in result
|
||||
assert "Doc content" in result
|
||||
|
||||
def test_render_prompt_with_tools_data(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Memory contents:\n{{ tools.memory.root }}\n\nStatus: {{ tools.memory.available }}"
|
||||
tools_data = {
|
||||
"memory": {"root": "Files:\n- /notes.txt\n- /tasks.txt", "available": True}
|
||||
}
|
||||
|
||||
result = renderer.render_prompt(prompt, tools_data=tools_data)
|
||||
|
||||
assert "Memory contents:" in result
|
||||
assert "Files:" in result
|
||||
assert "/notes.txt" in result
|
||||
assert "/tasks.txt" in result
|
||||
assert "Status: True" in result
|
||||
|
||||
def test_render_prompt_with_all_namespaces(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = """
|
||||
System: {{ system.date }}
|
||||
User: {{ passthrough.user }}
|
||||
Docs: {{ source.content }}
|
||||
Memory: {{ tools.memory.root }}
|
||||
"""
|
||||
passthrough_data = {"user": "Alice"}
|
||||
docs_together = "Important docs"
|
||||
tools_data = {"memory": {"root": "Notes content", "available": True}}
|
||||
|
||||
result = renderer.render_prompt(
|
||||
prompt,
|
||||
passthrough_data=passthrough_data,
|
||||
docs_together=docs_together,
|
||||
tools_data=tools_data,
|
||||
)
|
||||
|
||||
assert "202" in result
|
||||
assert "Alice" in result
|
||||
assert "Important docs" in result
|
||||
assert "Notes content" in result
|
||||
|
||||
def test_render_prompt_undefined_variable_returns_empty_string(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Hello {{ undefined_var }}"
|
||||
|
||||
result = renderer.render_prompt(prompt)
|
||||
assert result == "Hello "
|
||||
|
||||
def test_render_prompt_with_undefined_variable_in_template(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Hello {{ undefined_name }}"
|
||||
|
||||
result = renderer.render_prompt(prompt)
|
||||
assert result == "Hello "
|
||||
|
||||
def test_validate_template_valid(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
assert renderer.validate_template("Valid {{ variable }}") is True
|
||||
|
||||
def test_validate_template_invalid(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
assert renderer.validate_template("Invalid {{ variable") is False
|
||||
|
||||
def test_extract_variables(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
template = "{{ var1 }} and {{ var2 }}"
|
||||
|
||||
result = renderer.extract_variables(template)
|
||||
|
||||
assert isinstance(result, set)
|
||||
|
||||
def test_uses_template_syntax_detection(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
|
||||
assert renderer._uses_template_syntax("Text with {{ var }}") is True
|
||||
assert renderer._uses_template_syntax("Text with {var}") is False
|
||||
assert renderer._uses_template_syntax("Plain text") is False
|
||||
|
||||
def test_apply_legacy_substitutions(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Use {summaries} to answer"
|
||||
docs_together = "Important info"
|
||||
|
||||
result = renderer._apply_legacy_substitutions(prompt, docs_together)
|
||||
|
||||
assert "Use Important info to answer" in result
|
||||
|
||||
def test_apply_legacy_substitutions_without_docs(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "Use {summaries} to answer"
|
||||
|
||||
result = renderer._apply_legacy_substitutions(prompt, None)
|
||||
|
||||
assert result == "Use to answer"
|
||||
assert "{summaries}" not in result
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPromptRendererIntegration:
|
||||
|
||||
def test_render_prompt_real_world_scenario(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = "You are helping {{ passthrough.company }}.\n\nUser: {{ passthrough.user_name }}\n\nRequest ID: {{ system.request_id }}\n\nDate: {{ system.date }}\n\nReference Documents:\n\n{{ source.content }}\n\nPlease answer the question professionally."
|
||||
|
||||
passthrough_data = {"company": "Tech Corp", "user_name": "Alice"}
|
||||
docs_together = "Document 1: Technical specs\nDocument 2: Requirements"
|
||||
|
||||
result = renderer.render_prompt(
|
||||
prompt,
|
||||
request_id="req_123",
|
||||
user_id="user_456",
|
||||
passthrough_data=passthrough_data,
|
||||
docs_together=docs_together,
|
||||
)
|
||||
|
||||
assert "Tech Corp" in result
|
||||
assert "Alice" in result
|
||||
assert "req_123" in result
|
||||
assert "Technical specs" in result
|
||||
assert "professionally" in result
|
||||
|
||||
def test_render_prompt_multiple_doc_references(self):
|
||||
from application.api.answer.services.prompt_renderer import PromptRenderer
|
||||
|
||||
renderer = PromptRenderer()
|
||||
prompt = """Documents: {{ source.content }} \n\nAlso summaries: {{ source.summaries }}"""
|
||||
docs_together = "Content here"
|
||||
|
||||
result = renderer.render_prompt(prompt, docs_together=docs_together)
|
||||
|
||||
assert result.count("Content here") == 2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.skip(
|
||||
reason="Uses removed MongoDB.get_client() + user_tools mongo collection; "
|
||||
"needs rewrite against UserToolsRepository / pg_conn."
|
||||
)
|
||||
class TestStreamProcessorPromptRendering:
|
||||
|
||||
def test_stream_processor_pre_fetch_docs_none_doc_mode(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {"question": "Test question", "isNoneDoc": True}
|
||||
processor = StreamProcessor(request_data, None)
|
||||
|
||||
docs_together, docs_list = processor.pre_fetch_docs("Test question")
|
||||
|
||||
assert docs_together is None
|
||||
assert docs_list is None
|
||||
|
||||
def test_pre_fetch_tools_disabled_globally(self, monkeypatch):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "ENABLE_TOOL_PREFETCH", False)
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user1"})
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_pre_fetch_tools_disabled_per_request(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {"question": "test", "disable_tool_prefetch": True}
|
||||
processor = StreamProcessor(request_data, {"sub": "user1"})
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_pre_fetch_tools_skips_tool_with_no_actions(self, mock_mongo_db):
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
|
||||
db = mock_mongo_db[settings.MONGO_DB_NAME]
|
||||
db["user_tools"].insert_one(
|
||||
{
|
||||
"name": "memory",
|
||||
"user": "user1",
|
||||
"status": True,
|
||||
"config": {},
|
||||
}
|
||||
)
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user1"})
|
||||
|
||||
with patch(
|
||||
"application.agents.tools.tool_manager.ToolManager"
|
||||
) as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
# Tool has no actions
|
||||
mock_tool.get_actions_metadata.return_value = []
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
# Should return None when tool has no actions
|
||||
assert result is None
|
||||
|
||||
def test_pre_fetch_tools_enabled_by_default(self, mock_mongo_db, monkeypatch):
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
|
||||
db = mock_mongo_db[settings.MONGO_DB_NAME]
|
||||
db["user_tools"].insert_one(
|
||||
{
|
||||
"name": "memory",
|
||||
"user": "user1",
|
||||
"status": True,
|
||||
"config": {},
|
||||
}
|
||||
)
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user1"})
|
||||
|
||||
with patch(
|
||||
"application.agents.tools.tool_manager.ToolManager"
|
||||
) as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance returned by load_tool
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
# Mock get_actions_metadata on the tool instance
|
||||
mock_tool.get_actions_metadata.return_value = [
|
||||
{"name": "memory_ls", "description": "List files", "parameters": {"properties": {}}}
|
||||
]
|
||||
mock_tool.execute_action.return_value = "Directory: /\n- file.txt"
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
assert result is not None
|
||||
assert "memory" in result
|
||||
assert "memory_ls" in result["memory"]
|
||||
|
||||
def test_pre_fetch_tools_no_tools_configured(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user1"})
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_pre_fetch_tools_memory_returns_error(self, mock_mongo_db):
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
|
||||
db = mock_mongo_db[settings.MONGO_DB_NAME]
|
||||
db["user_tools"].insert_one(
|
||||
{
|
||||
"name": "memory",
|
||||
"user": "user1",
|
||||
"status": True,
|
||||
"config": {},
|
||||
}
|
||||
)
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user1"})
|
||||
|
||||
with patch(
|
||||
"application.agents.tools.tool_manager.ToolManager"
|
||||
) as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
mock_tool.get_actions_metadata.return_value = [
|
||||
{"name": "memory_ls", "description": "List files", "parameters": {"properties": {}}}
|
||||
]
|
||||
# Simulate execution error
|
||||
mock_tool.execute_action.side_effect = Exception("Tool error")
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
# Should return None when all actions fail
|
||||
assert result is None
|
||||
|
||||
def test_pre_fetch_tools_memory_returns_empty(self, mock_mongo_db):
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.core.settings import settings
|
||||
|
||||
db = mock_mongo_db[settings.MONGO_DB_NAME]
|
||||
db["user_tools"].insert_one(
|
||||
{
|
||||
"name": "memory",
|
||||
"user": "user1",
|
||||
"status": True,
|
||||
"config": {},
|
||||
}
|
||||
)
|
||||
|
||||
request_data = {"question": "test"}
|
||||
processor = StreamProcessor(request_data, {"sub": "user1"})
|
||||
|
||||
with patch(
|
||||
"application.agents.tools.tool_manager.ToolManager"
|
||||
) as mock_manager_class:
|
||||
mock_manager = MagicMock()
|
||||
mock_manager_class.return_value = mock_manager
|
||||
|
||||
# Mock the tool instance
|
||||
mock_tool = MagicMock()
|
||||
mock_manager.load_tool.return_value = mock_tool
|
||||
|
||||
mock_tool.get_actions_metadata.return_value = [
|
||||
{"name": "memory_ls", "description": "List files", "parameters": {"properties": {}}}
|
||||
]
|
||||
# Return empty string
|
||||
mock_tool.execute_action.return_value = ""
|
||||
|
||||
result = processor.pre_fetch_tools()
|
||||
|
||||
# Empty result should still be included
|
||||
assert result is not None
|
||||
assert "memory" in result
|
||||
@@ -0,0 +1,353 @@
|
||||
"""Tests for application/api/answer/services/stream_processor.py.
|
||||
|
||||
The previous suite was tightly coupled to Mongo (mock_mongo_db fixture,
|
||||
bson.ObjectId, bson.DBRef, find_one, etc.) which no longer exist after the
|
||||
Mongo -> Postgres cutover. Rewriting these ~18 tests against the new
|
||||
repositories (AgentsRepository / PromptsRepository / ConversationsRepository)
|
||||
requires meaningful setup that is best done alongside the migration of the
|
||||
StreamProcessor internals themselves.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# A static 24-hex-char string that is a valid ObjectId hex representation.
|
||||
_STATIC_OID = "507f1f77bcf86cd799439011"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGetPromptFunction:
|
||||
pass
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestStreamProcessorInitialization:
|
||||
pass
|
||||
|
||||
def test_initializes_with_decoded_token(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
conv_id = _STATIC_OID
|
||||
request_data = {
|
||||
"question": "What is Python?",
|
||||
"conversation_id": conv_id,
|
||||
}
|
||||
decoded_token = {"sub": "user_123", "email": "test@example.com"}
|
||||
|
||||
processor = StreamProcessor(request_data, decoded_token)
|
||||
|
||||
assert processor.data == request_data
|
||||
assert processor.decoded_token == decoded_token
|
||||
assert processor.initial_user_id == "user_123"
|
||||
assert processor.conversation_id == request_data["conversation_id"]
|
||||
|
||||
def test_initializes_without_token(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {"question": "Test question"}
|
||||
|
||||
processor = StreamProcessor(request_data, None)
|
||||
|
||||
assert processor.decoded_token is None
|
||||
assert processor.initial_user_id is None
|
||||
assert processor.data == request_data
|
||||
|
||||
def test_initializes_default_attributes(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
processor = StreamProcessor({"question": "Test"}, {"sub": "user_123"})
|
||||
|
||||
assert processor.source == {}
|
||||
assert processor.all_sources == []
|
||||
assert processor.attachments == []
|
||||
assert processor.history == []
|
||||
assert processor.agent_config == {}
|
||||
assert processor.retriever_config == {}
|
||||
assert processor.is_shared_usage is False
|
||||
assert processor.shared_token is None
|
||||
|
||||
def test_extracts_conversation_id_from_request(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
conv_id = _STATIC_OID
|
||||
request_data = {"question": "Test", "conversation_id": conv_id}
|
||||
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
|
||||
assert processor.conversation_id == conv_id
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestStreamProcessorHistoryLoading:
|
||||
pass
|
||||
|
||||
def test_uses_request_history_when_no_conversation_id(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {
|
||||
"question": "What is Python?",
|
||||
"history": [{"prompt": "Hello", "response": "Hi there!"}],
|
||||
}
|
||||
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
|
||||
assert processor.conversation_id is None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestStreamProcessorAgentConfiguration:
|
||||
pass
|
||||
|
||||
def test_uses_default_config_without_api_key(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {"question": "Test"}
|
||||
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
processor._configure_agent()
|
||||
|
||||
assert isinstance(processor.agent_config, dict)
|
||||
assert processor.agent_id is None
|
||||
|
||||
def test_embedded_workflow_without_saved_id(self):
|
||||
"""A preview run with no saved workflow id carries no ``workflow_id``."""
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {
|
||||
"question": "Test",
|
||||
"workflow": {"nodes": [], "edges": []},
|
||||
}
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
processor._configure_agent()
|
||||
|
||||
assert processor.agent_config["agent_type"] == "workflow"
|
||||
assert processor.agent_config["workflow"] == {"nodes": [], "edges": []}
|
||||
assert "workflow_id" not in processor.agent_config
|
||||
|
||||
def test_embedded_workflow_with_saved_id_persists_run(self):
|
||||
"""A saved workflow id alongside the embedded graph is captured so the
|
||||
run can persist a ``workflow_runs`` row for artifact listing."""
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {
|
||||
"question": "Test",
|
||||
"workflow": {"nodes": [], "edges": []},
|
||||
"workflow_id": "11111111-1111-1111-1111-111111111111",
|
||||
}
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
processor._configure_agent()
|
||||
|
||||
assert processor.agent_config["agent_type"] == "workflow"
|
||||
assert (
|
||||
processor.agent_config["workflow_id"]
|
||||
== "11111111-1111-1111-1111-111111111111"
|
||||
)
|
||||
assert processor.agent_config["workflow_owner"] == "user_123"
|
||||
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestStreamProcessorDocPrefetch:
|
||||
pass
|
||||
|
||||
def test_prefetch_skipped_when_no_active_docs(self):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
processor = StreamProcessor(
|
||||
{"question": "Hi there"},
|
||||
{"sub": "user_123"},
|
||||
)
|
||||
processor.initialize()
|
||||
processor.create_retriever = MagicMock()
|
||||
|
||||
docs_together, docs = processor.pre_fetch_docs("Hi there")
|
||||
|
||||
processor.create_retriever.assert_not_called()
|
||||
assert docs_together is None
|
||||
assert docs is None
|
||||
|
||||
def test_prefetch_skipped_when_active_docs_is_default(self):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
processor = StreamProcessor(
|
||||
{"question": "Hi", "active_docs": "default"},
|
||||
{"sub": "user_123"},
|
||||
)
|
||||
processor.initialize()
|
||||
processor.create_retriever = MagicMock()
|
||||
|
||||
docs_together, docs = processor.pre_fetch_docs("Hi")
|
||||
|
||||
processor.create_retriever.assert_not_called()
|
||||
assert docs_together is None
|
||||
assert docs is None
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestStreamProcessorAttachments:
|
||||
pass
|
||||
|
||||
def test_handles_empty_attachments(self):
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
request_data = {"question": "Simple question"}
|
||||
|
||||
processor = StreamProcessor(request_data, {"sub": "user_123"})
|
||||
|
||||
assert processor.attachments == []
|
||||
assert (
|
||||
"attachments" not in processor.data
|
||||
or processor.data.get("attachments") is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestToolPreFetch:
|
||||
"""Tests for tool pre-fetching with saved parameter values."""
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestBuildContinuationFromMessages:
|
||||
"""Stateless tool continuation rebuilt from the resent messages array.
|
||||
|
||||
OpenAI-compatible clients (opencode, etc.) resend the full conversation but
|
||||
carry no conversation_id, so the agent + pending tool calls are reconstructed
|
||||
directly from the request instead of from server-side ``pending_tool_state``.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _make_processor():
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
|
||||
processor = StreamProcessor({"question": ""}, {"sub": "user_123"})
|
||||
fake_agent = MagicMock()
|
||||
fake_agent.tool_executor.get_tools.return_value = {"search": object()}
|
||||
# build_agent touches the DB / config; stub it for this unit test.
|
||||
processor.build_agent = MagicMock(return_value=fake_agent)
|
||||
return processor, fake_agent
|
||||
|
||||
def _messages_with_tool_call(self, arguments):
|
||||
return [
|
||||
{"role": "system", "content": "You are a bot"},
|
||||
{"role": "user", "content": "search the docs"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": arguments},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "found it"},
|
||||
]
|
||||
|
||||
def test_rebuilds_pending_tool_calls_and_prior_messages(self):
|
||||
processor, fake_agent = self._make_processor()
|
||||
messages = self._messages_with_tool_call('{"q": "docs"}')
|
||||
tool_actions = [{"call_id": "call_1", "result": "found it"}]
|
||||
|
||||
(
|
||||
agent,
|
||||
prior_messages,
|
||||
tools_dict,
|
||||
pending_tool_calls,
|
||||
returned_actions,
|
||||
reasoning,
|
||||
) = processor.build_continuation_from_messages(messages, tool_actions)
|
||||
|
||||
assert agent is fake_agent
|
||||
assert tools_dict == fake_agent.tool_executor.get_tools.return_value
|
||||
assert returned_actions is tool_actions
|
||||
assert reasoning == ""
|
||||
# prior_messages stop before the assistant-with-tool_calls message.
|
||||
assert prior_messages == messages[:2]
|
||||
assert len(pending_tool_calls) == 1
|
||||
call = pending_tool_calls[0]
|
||||
assert call["call_id"] == "call_1"
|
||||
assert call["name"] == "search"
|
||||
assert call["tool_name"] == "search"
|
||||
assert call["action_name"] == "search"
|
||||
assert call["arguments"] == {"q": "docs"}
|
||||
processor.build_agent.assert_called_once_with("")
|
||||
|
||||
def test_invalid_json_arguments_default_to_empty_dict(self):
|
||||
processor, _ = self._make_processor()
|
||||
messages = self._messages_with_tool_call("not-json")
|
||||
|
||||
_, _, _, pending_tool_calls, _, _ = (
|
||||
processor.build_continuation_from_messages(messages, [])
|
||||
)
|
||||
|
||||
assert pending_tool_calls[0]["arguments"] == {}
|
||||
|
||||
def test_dict_arguments_passed_through(self):
|
||||
processor, _ = self._make_processor()
|
||||
messages = self._messages_with_tool_call({"already": "parsed"})
|
||||
|
||||
_, _, _, pending_tool_calls, _, _ = (
|
||||
processor.build_continuation_from_messages(messages, [])
|
||||
)
|
||||
|
||||
assert pending_tool_calls[0]["arguments"] == {"already": "parsed"}
|
||||
|
||||
def test_uses_last_assistant_tool_call_message(self):
|
||||
processor, _ = self._make_processor()
|
||||
messages = [
|
||||
{"role": "user", "content": "first"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "old",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "old", "content": "r1"},
|
||||
{"role": "user", "content": "second"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "new",
|
||||
"type": "function",
|
||||
"function": {"name": "search", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "new", "content": "r2"},
|
||||
]
|
||||
|
||||
_, prior_messages, _, pending_tool_calls, _, _ = (
|
||||
processor.build_continuation_from_messages(messages, [])
|
||||
)
|
||||
|
||||
assert pending_tool_calls[0]["call_id"] == "new"
|
||||
# Everything up to (not including) the last assistant tool_calls message.
|
||||
assert prior_messages == messages[:4]
|
||||
|
||||
def test_raises_when_no_assistant_tool_calls(self):
|
||||
processor, _ = self._make_processor()
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "no tools here"},
|
||||
]
|
||||
|
||||
with pytest.raises(ValueError, match="No assistant message with tool_calls"):
|
||||
processor.build_continuation_from_messages(messages, [])
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,238 @@
|
||||
"""Token-usage attribution tests for the always-on inline-persist model.
|
||||
|
||||
Persistence is owned by the per-call decorator in ``application.usage``.
|
||||
``finalize_message`` no longer writes ``token_usage`` rows. These tests
|
||||
exercise the decorator path through ``stream_token_usage`` /
|
||||
``gen_token_usage``:
|
||||
|
||||
1. Every LLM call writes one row, regardless of whether the route saves
|
||||
the conversation.
|
||||
2. ``_token_usage_source`` on the LLM instance flows to the row's
|
||||
``source`` column for cost-attribution dashboards.
|
||||
3. ``_request_id`` on the LLM instance flows to the row's ``request_id``
|
||||
column so ``count_in_range`` can DISTINCT-collapse multi-call agent
|
||||
runs into a single request.
|
||||
4. Calls with no attribution (no ``user_id`` and no ``user_api_key``)
|
||||
warn and skip — the repository would otherwise raise on the
|
||||
``token_usage_attribution_chk`` constraint.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_db_session_for(modules, conn):
|
||||
"""Reroute every named module's ``db_session`` to ``conn``."""
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
patches = [patch(f"{m}.db_session", _yield) for m in modules]
|
||||
for p in patches:
|
||||
p.start()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for p in patches:
|
||||
p.stop()
|
||||
|
||||
|
||||
def _seed_user(conn) -> str:
|
||||
user_id = str(uuid.uuid4())
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO users (user_id) VALUES (:u) "
|
||||
"ON CONFLICT (user_id) DO NOTHING"
|
||||
),
|
||||
{"u": user_id},
|
||||
)
|
||||
return user_id
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDecoratorAlwaysPersists:
|
||||
"""Per-call inline persistence — no opt-in flag."""
|
||||
|
||||
def test_primary_stream_writes_agent_stream_row(self, pg_conn):
|
||||
from application.usage import stream_token_usage
|
||||
|
||||
user_id = _seed_user(pg_conn)
|
||||
|
||||
class _PrimaryLLM:
|
||||
decoded_token = {"sub": user_id}
|
||||
user_api_key = None
|
||||
agent_id = None
|
||||
|
||||
def __init__(self):
|
||||
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
||||
|
||||
@stream_token_usage
|
||||
def _raw(self, model, messages, stream, tools, **kwargs):
|
||||
yield "chunk-a"
|
||||
yield "chunk-b"
|
||||
|
||||
llm = _PrimaryLLM()
|
||||
with _patch_db_session_for(("application.usage",), pg_conn):
|
||||
for _ in llm._raw(
|
||||
"m", [{"role": "user", "content": "hi"}], True, None,
|
||||
):
|
||||
pass
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT prompt_tokens, generated_tokens, source, request_id "
|
||||
"FROM token_usage WHERE user_id = :u"
|
||||
),
|
||||
{"u": user_id},
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
assert row[2] == "agent_stream"
|
||||
assert row[3] is None # No request_id stamped on this LLM.
|
||||
assert row[0] > 0
|
||||
assert row[1] > 0
|
||||
|
||||
def test_side_channel_source_flows_to_row(self, pg_conn):
|
||||
"""``_token_usage_source`` overrides the default ``agent_stream``."""
|
||||
from application.usage import stream_token_usage
|
||||
|
||||
user_id = _seed_user(pg_conn)
|
||||
|
||||
class _RagLLM:
|
||||
decoded_token = {"sub": user_id}
|
||||
user_api_key = None
|
||||
agent_id = None
|
||||
_token_usage_source = "rag_condense"
|
||||
|
||||
def __init__(self):
|
||||
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
||||
|
||||
@stream_token_usage
|
||||
def _raw(self, model, messages, stream, tools, **kwargs):
|
||||
yield "chunk"
|
||||
|
||||
llm = _RagLLM()
|
||||
with _patch_db_session_for(("application.usage",), pg_conn):
|
||||
for _ in llm._raw("m", [{"role": "user", "content": "q"}], True, None):
|
||||
pass
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT source FROM token_usage WHERE user_id = :u"
|
||||
),
|
||||
{"u": user_id},
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
assert row[0] == "rag_condense"
|
||||
|
||||
def test_request_id_propagates_to_row(self, pg_conn):
|
||||
"""``_request_id`` on the LLM (stamped by the route) lands in
|
||||
``token_usage.request_id`` so ``count_in_range`` can DISTINCT it.
|
||||
"""
|
||||
from application.usage import stream_token_usage
|
||||
|
||||
user_id = _seed_user(pg_conn)
|
||||
request_id = f"req-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
class _PrimaryLLM:
|
||||
decoded_token = {"sub": user_id}
|
||||
user_api_key = None
|
||||
agent_id = None
|
||||
|
||||
def __init__(self):
|
||||
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
||||
self._request_id = request_id
|
||||
|
||||
@stream_token_usage
|
||||
def _raw(self, model, messages, stream, tools, **kwargs):
|
||||
yield "chunk"
|
||||
|
||||
llm = _PrimaryLLM()
|
||||
with _patch_db_session_for(("application.usage",), pg_conn):
|
||||
# Call twice — the route invokes the LLM once per tool round.
|
||||
for _ in llm._raw("m", [{"role": "user", "content": "q"}], True, None):
|
||||
pass
|
||||
for _ in llm._raw("m", [{"role": "user", "content": "q2"}], True, None):
|
||||
pass
|
||||
|
||||
rows = pg_conn.execute(
|
||||
text(
|
||||
"SELECT request_id FROM token_usage WHERE user_id = :u"
|
||||
),
|
||||
{"u": user_id},
|
||||
).fetchall()
|
||||
assert len(rows) == 2
|
||||
assert all(r[0] == request_id for r in rows)
|
||||
|
||||
def test_zero_count_call_is_skipped(self, pg_conn):
|
||||
from application.usage import gen_token_usage
|
||||
|
||||
user_id = _seed_user(pg_conn)
|
||||
|
||||
class _EmptyLLM:
|
||||
decoded_token = {"sub": user_id}
|
||||
user_api_key = None
|
||||
agent_id = None
|
||||
|
||||
def __init__(self):
|
||||
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
||||
|
||||
@gen_token_usage
|
||||
def _raw(self, model, messages, stream, tools, **kwargs):
|
||||
return None # empty result → 0 generated tokens, 0 prompt tokens
|
||||
|
||||
llm = _EmptyLLM()
|
||||
with _patch_db_session_for(("application.usage",), pg_conn):
|
||||
llm._raw("m", [], False, None)
|
||||
|
||||
n = pg_conn.execute(
|
||||
text("SELECT count(*) FROM token_usage WHERE user_id = :u"),
|
||||
{"u": user_id},
|
||||
).scalar()
|
||||
assert n == 0
|
||||
|
||||
def test_no_attribution_warns_and_skips(self, pg_conn, caplog):
|
||||
"""No user_id and no api_key → log a warning, don't insert.
|
||||
|
||||
The repository would otherwise raise on the attribution CHECK
|
||||
constraint; the decorator skips before that to keep the stream
|
||||
running.
|
||||
"""
|
||||
from application.usage import stream_token_usage
|
||||
|
||||
class _OrphanLLM:
|
||||
decoded_token = None
|
||||
user_api_key = None
|
||||
agent_id = None
|
||||
|
||||
def __init__(self):
|
||||
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
||||
|
||||
@stream_token_usage
|
||||
def _raw(self, model, messages, stream, tools, **kwargs):
|
||||
yield "chunk"
|
||||
|
||||
llm = _OrphanLLM()
|
||||
with _patch_db_session_for(
|
||||
("application.usage",), pg_conn,
|
||||
), caplog.at_level(logging.WARNING, logger="application.usage"):
|
||||
for _ in llm._raw("m", [{"role": "user", "content": "q"}], True, None):
|
||||
pass
|
||||
|
||||
n = pg_conn.execute(text("SELECT count(*) FROM token_usage")).scalar()
|
||||
# New attribution rows specifically for this orphan path: nothing
|
||||
# should land. The fixture pins state, so an existing baseline is
|
||||
# 0 by default.
|
||||
assert n == 0
|
||||
assert any(
|
||||
"no user_id/api_key" in r.message
|
||||
for r in caplog.records
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,113 @@
|
||||
"""Unit tests for application/api/answer/services/conversation_service.py.
|
||||
|
||||
Additional coverage beyond tests/api/answer/services/test_conversation_service.py:
|
||||
- save_conversation: index-based update, metadata persistence, agent key tracking
|
||||
- update_compression_metadata
|
||||
- append_compression_message
|
||||
- get_compression_metadata
|
||||
- Edge cases: None token, empty summary, shared_with access
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestConversationServiceGetExtended:
|
||||
pass
|
||||
|
||||
def test_handles_exception_gracefully(self, mock_mongo_db):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
service = ConversationService()
|
||||
# Pass an invalid ObjectId
|
||||
result = service.get_conversation("not-an-objectid", "user_123")
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSaveConversationExtended:
|
||||
pass
|
||||
|
||||
def test_raises_for_none_token(self, mock_mongo_db):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
service = ConversationService()
|
||||
with pytest.raises(ValueError, match="Invalid or missing authentication"):
|
||||
service.save_conversation(
|
||||
conversation_id=None,
|
||||
question="Q",
|
||||
response="A",
|
||||
thought="",
|
||||
sources=[],
|
||||
tool_calls=[],
|
||||
llm=Mock(),
|
||||
model_id="m",
|
||||
decoded_token=None,
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestUpdateCompressionMetadata:
|
||||
pass
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestAppendCompressionMessage:
|
||||
pass
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGetCompressionMetadata:
|
||||
pass
|
||||
|
||||
def test_returns_none_for_missing_conversation(self, mock_mongo_db):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
service = ConversationService()
|
||||
result = service.get_compression_metadata("507f1f77bcf86cd799439011")
|
||||
assert result is None
|
||||
|
||||
def test_handles_invalid_id(self, mock_mongo_db):
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
service = ConversationService()
|
||||
result = service.get_compression_metadata("invalid-id")
|
||||
assert result is None
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Coverage gap tests (lines 233-237, 258, 261)
|
||||
# =====================================================================
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestConversationServiceGaps:
|
||||
pass
|
||||
|
||||
def test_append_compression_message_empty_summary_skips(self, mock_mongo_db):
|
||||
"""Cover: empty summary does not insert."""
|
||||
from application.api.answer.services.conversation_service import (
|
||||
ConversationService,
|
||||
)
|
||||
|
||||
service = ConversationService()
|
||||
service.conversations_collection = MagicMock()
|
||||
|
||||
service.append_compression_message(
|
||||
"507f1f77bcf86cd799439011", {"compressed_summary": ""}
|
||||
)
|
||||
service.conversations_collection.update_one.assert_not_called()
|
||||
@@ -0,0 +1,190 @@
|
||||
"""Integration tests for the end-to-end snapshot+tail handoff.
|
||||
|
||||
Exercises the publisher → journal → reconnect-reader round-trip without
|
||||
mocking the journal layer, so a regression in any of:
|
||||
- complete_stream's _emit closure
|
||||
- record_event's commit-per-call contract
|
||||
- the shared snapshot read (``event_replay.read_snapshot_lines``)
|
||||
- the reconnect route's ownership SQL (``async_sse._user_owns_message``)
|
||||
- message_events repo SQL
|
||||
would surface here as a failed integration assertion.
|
||||
|
||||
The live tail itself (pub/sub) and the full async HTTP route are covered by
|
||||
``scripts/e2e_async_sse.py`` against a real Redis + uvicorn. These tests pin
|
||||
the DB-backed halves against a transactional ``pg_conn`` fixture.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid as _uuid
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text as sql_text
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_journal_session(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.streaming.message_journal.db_session", _yield
|
||||
), patch(
|
||||
"application.streaming.event_replay.db_readonly", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _seed_message(conn, user_id: str | None = None):
|
||||
user_id = user_id or f"u-{_uuid.uuid4().hex[:6]}"
|
||||
conv_id = _uuid.uuid4()
|
||||
msg_id = _uuid.uuid4()
|
||||
conn.execute(sql_text("INSERT INTO users (user_id) VALUES (:u)"), {"u": user_id})
|
||||
conn.execute(
|
||||
sql_text(
|
||||
"INSERT INTO conversations (id, user_id, name) VALUES (:id, :u, 't')"
|
||||
),
|
||||
{"id": conv_id, "u": user_id},
|
||||
)
|
||||
conn.execute(
|
||||
sql_text(
|
||||
"INSERT INTO conversation_messages (id, conversation_id, user_id, position) "
|
||||
"VALUES (:id, :c, :u, 0)"
|
||||
),
|
||||
{"id": msg_id, "c": conv_id, "u": user_id},
|
||||
)
|
||||
return user_id, str(msg_id)
|
||||
|
||||
|
||||
def _emitted_ids(lines: list[str]) -> list[int]:
|
||||
"""Extract the ``id:`` sequence numbers from formatted SSE frames."""
|
||||
return sorted(
|
||||
int(line.split("\n", 1)[0].split(": ", 1)[1])
|
||||
for line in lines
|
||||
if line.startswith("id: ")
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
class TestSnapshotPlusTailRoundTrip:
|
||||
def test_record_event_then_snapshot_returns_what_was_journaled(
|
||||
self, pg_conn,
|
||||
):
|
||||
"""End-to-end of the journal half: ``record_event`` writes through a
|
||||
real ``MessageEventsRepository``, ``read_snapshot_lines`` reads the
|
||||
snapshot back via the same repo and formats it for the wire — the
|
||||
exact primitive the async reader replays through.
|
||||
"""
|
||||
from application.streaming.event_replay import read_snapshot_lines
|
||||
from application.streaming.message_journal import record_event
|
||||
|
||||
_, message_id = _seed_message(pg_conn)
|
||||
|
||||
with _patch_journal_session(pg_conn):
|
||||
record_event(message_id, 0, "answer", {"type": "answer", "answer": "A"})
|
||||
record_event(message_id, 1, "answer", {"type": "answer", "answer": "B"})
|
||||
record_event(message_id, 2, "end", {"type": "end"})
|
||||
|
||||
lines, max_seq, terminal = read_snapshot_lines(message_id, None)
|
||||
|
||||
assert len(lines) == 3
|
||||
assert "id: 0" in lines[0] and '"answer": "A"' in lines[0]
|
||||
assert "id: 1" in lines[1] and '"answer": "B"' in lines[1]
|
||||
assert "id: 2" in lines[2] and '"type": "end"' in lines[2]
|
||||
assert max_seq == 2
|
||||
# A terminal event in the snapshot tells the reader to close.
|
||||
assert terminal is True
|
||||
|
||||
def test_snapshot_resumes_past_last_event_id(self, pg_conn):
|
||||
from application.streaming.event_replay import read_snapshot_lines
|
||||
from application.streaming.message_journal import record_event
|
||||
|
||||
_, message_id = _seed_message(pg_conn)
|
||||
|
||||
with _patch_journal_session(pg_conn):
|
||||
for seq in range(5):
|
||||
record_event(
|
||||
message_id, seq, "answer", {"type": "answer", "answer": str(seq)}
|
||||
)
|
||||
|
||||
# Client says it has seen up through seq=2; expect 3 + 4.
|
||||
lines, max_seq, terminal = read_snapshot_lines(message_id, 2)
|
||||
|
||||
assert _emitted_ids(lines) == [3, 4]
|
||||
assert max_seq == 4
|
||||
assert terminal is False
|
||||
|
||||
def test_ownership_sql_accepts_owner_and_rejects_others(self, pg_conn):
|
||||
"""The async route's ownership gate runs real SQL against
|
||||
``conversation_messages`` — the owner passes, everyone else 404s.
|
||||
"""
|
||||
from application.api import async_sse
|
||||
|
||||
user_id, message_id = _seed_message(pg_conn)
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield pg_conn
|
||||
|
||||
with patch("application.api.async_sse.db_readonly", _yield):
|
||||
assert async_sse._user_owns_message(message_id, user_id) is True
|
||||
assert async_sse._user_owns_message(message_id, "different-user") is False
|
||||
# A well-formed but unknown id is also not owned.
|
||||
assert async_sse._user_owns_message(str(_uuid.uuid4()), user_id) is False
|
||||
|
||||
def test_snapshot_read_is_user_scoped(self, pg_conn):
|
||||
"""``read_snapshot_lines(..., user_id=)`` re-asserts ownership at the
|
||||
data layer: the owner gets the journal rows, a non-owner gets none.
|
||||
"""
|
||||
from application.streaming.event_replay import read_snapshot_lines
|
||||
from application.streaming.message_journal import record_event
|
||||
|
||||
user_id, message_id = _seed_message(pg_conn)
|
||||
|
||||
with _patch_journal_session(pg_conn):
|
||||
record_event(message_id, 0, "answer", {"type": "answer", "answer": "x"})
|
||||
record_event(message_id, 1, "end", {"type": "end"})
|
||||
|
||||
owner_lines, _, owner_terminal = read_snapshot_lines(
|
||||
message_id, None, user_id
|
||||
)
|
||||
other_lines, _, other_terminal = read_snapshot_lines(
|
||||
message_id, None, "different-user"
|
||||
)
|
||||
# Unscoped (user_id=None) still returns everything.
|
||||
unscoped_lines, _, _ = read_snapshot_lines(message_id, None)
|
||||
|
||||
assert _emitted_ids(owner_lines) == [0, 1] and owner_terminal is True
|
||||
assert other_lines == [] and other_terminal is False
|
||||
assert _emitted_ids(unscoped_lines) == [0, 1]
|
||||
|
||||
def test_watchdog_is_user_scoped(self, pg_conn):
|
||||
"""``_check_producer_liveness(..., user_id)`` only sees the caller's
|
||||
own row; a non-owner reads as missing (terminal), not as the row's
|
||||
real status.
|
||||
"""
|
||||
from application.streaming.event_replay import _check_producer_liveness
|
||||
|
||||
# Seed a row and flip it to a terminal status the watchdog reports.
|
||||
user_id, message_id = _seed_message(pg_conn)
|
||||
pg_conn.execute(
|
||||
sql_text(
|
||||
"UPDATE conversation_messages SET status='complete' WHERE id = :id"
|
||||
),
|
||||
{"id": message_id},
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield pg_conn
|
||||
|
||||
with patch("application.streaming.event_replay.db_readonly", _yield):
|
||||
owner = _check_producer_liveness(message_id, user_id, 90.0)
|
||||
other = _check_producer_liveness(message_id, "different-user", 90.0)
|
||||
|
||||
# Owner sees the real terminal state; non-owner sees "missing".
|
||||
assert owner == {"type": "end"}
|
||||
assert other is not None and other.get("code") == "message_missing"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,177 @@
|
||||
"""Tests for per-source search exposure partitioning (E1 / D11)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.answer.services.stream_processor import StreamProcessor
|
||||
from application.storage.db.source_config import RetrievalConfig
|
||||
|
||||
|
||||
def _processor() -> StreamProcessor:
|
||||
"""A bare StreamProcessor with only the fields the exposure helpers touch.
|
||||
|
||||
Built via ``__new__`` to skip the DB-heavy ``__init__``; the exposure
|
||||
methods read ``all_sources`` / ``source`` / ``agent_config`` only.
|
||||
"""
|
||||
sp = StreamProcessor.__new__(StreamProcessor)
|
||||
sp.all_sources = []
|
||||
sp.source = {}
|
||||
sp.agent_config = {}
|
||||
sp.retriever_config = {"retriever_name": "classic", "chunks": 2, "doc_token_limit": 50000}
|
||||
sp.data = {}
|
||||
return sp
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestExposurePartition:
|
||||
def test_no_config_all_default_to_prefetch(self):
|
||||
sp = _processor()
|
||||
sp.all_sources = [{"id": "a", "retrieval": None}, {"id": "b", "retrieval": None}]
|
||||
prefetch, agentic = sp._exposure_partition()
|
||||
assert [e["id"] for e in prefetch] == ["a", "b"]
|
||||
assert agentic == []
|
||||
|
||||
def test_mixed_partition(self):
|
||||
sp = _processor()
|
||||
sp.all_sources = [
|
||||
{"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
|
||||
{"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
|
||||
{"id": "c", "retrieval": RetrievalConfig()}, # default prefetch
|
||||
]
|
||||
prefetch, agentic = sp._exposure_partition()
|
||||
assert sorted(e["id"] for e in prefetch) == ["a", "c"]
|
||||
assert [e["id"] for e in agentic] == ["b"]
|
||||
|
||||
def test_exposure_of_dict_and_model_and_missing(self):
|
||||
sp = _processor()
|
||||
assert sp._exposure_of(RetrievalConfig(exposure="agentic_tool")) == "agentic_tool"
|
||||
assert sp._exposure_of({"exposure": "agentic_tool"}) == "agentic_tool"
|
||||
assert sp._exposure_of(None) == "prefetch"
|
||||
assert sp._exposure_of({}) == "prefetch"
|
||||
|
||||
def test_source_for_docs(self):
|
||||
sp = _processor()
|
||||
assert sp._source_for_docs(["a", "b"]) == {"active_docs": ["a", "b"]}
|
||||
assert sp._source_for_docs([]) == {}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestBuildAgentExposure:
|
||||
def _agentic_processor(self):
|
||||
sp = _processor()
|
||||
sp.agent_config = {"agent_type": "agentic"}
|
||||
sp.initialize = MagicMock()
|
||||
sp.pre_fetch_docs = MagicMock(return_value=("docs_together", ["doc"]))
|
||||
sp.pre_fetch_tools = MagicMock(return_value={"tool": {}})
|
||||
sp.create_agent = MagicMock(return_value="AGENT")
|
||||
return sp
|
||||
|
||||
def test_all_prefetch_is_today_no_partition(self):
|
||||
# No agentic_tool source → today's behavior: no pre-fetch, no
|
||||
# agentic_sources passed (tool exposes all sources).
|
||||
sp = self._agentic_processor()
|
||||
sp.all_sources = [
|
||||
{"id": "a", "retrieval": RetrievalConfig()},
|
||||
{"id": "b", "retrieval": None},
|
||||
]
|
||||
result = sp.build_agent("q")
|
||||
assert result == "AGENT"
|
||||
sp.pre_fetch_docs.assert_not_called()
|
||||
_, kwargs = sp.create_agent.call_args
|
||||
assert "agentic_sources" not in kwargs
|
||||
|
||||
def test_mixed_prefetches_and_scopes_tool(self):
|
||||
sp = self._agentic_processor()
|
||||
sp.all_sources = [
|
||||
{"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
|
||||
{"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
|
||||
]
|
||||
sp.build_agent("q")
|
||||
# Pre-fetch ran scoped to the prefetch subset.
|
||||
sp.pre_fetch_docs.assert_called_once()
|
||||
_, pf_kwargs = sp.pre_fetch_docs.call_args
|
||||
assert pf_kwargs.get("exposure") == "prefetch"
|
||||
# create_agent received only the agentic_tool subset as the tool sources.
|
||||
_, kwargs = sp.create_agent.call_args
|
||||
assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"]
|
||||
|
||||
def _classic_processor(self):
|
||||
sp = _processor()
|
||||
sp.agent_config = {"agent_type": "classic"}
|
||||
sp.initialize = MagicMock()
|
||||
sp.pre_fetch_docs = MagicMock(return_value=("docs_together", ["doc"]))
|
||||
sp.pre_fetch_tools = MagicMock(return_value=None)
|
||||
sp.create_agent = MagicMock(return_value="AGENT")
|
||||
return sp
|
||||
|
||||
def test_classic_default_prefetches_all_no_partition(self):
|
||||
# Default classic (all prefetch / no config): byte-identical to today —
|
||||
# unscoped pre-fetch and no agentic_sources, so no search tool is added.
|
||||
sp = self._classic_processor()
|
||||
sp.all_sources = [
|
||||
{"id": "a", "retrieval": RetrievalConfig()},
|
||||
{"id": "b", "retrieval": None},
|
||||
]
|
||||
result = sp.build_agent("q")
|
||||
assert result == "AGENT"
|
||||
sp.pre_fetch_docs.assert_called_once_with("q")
|
||||
_, kwargs = sp.create_agent.call_args
|
||||
assert "agentic_sources" not in kwargs
|
||||
|
||||
def test_classic_no_per_source_detail_is_today(self):
|
||||
# Single-source / no-config requests carry no per-source detail; classic
|
||||
# must fall back to the unscoped pre-fetch (no exposure scoping).
|
||||
sp = self._classic_processor()
|
||||
sp.all_sources = []
|
||||
sp.source = {"active_docs": ["a"]}
|
||||
sp.build_agent("q")
|
||||
sp.pre_fetch_docs.assert_called_once_with("q")
|
||||
_, kwargs = sp.create_agent.call_args
|
||||
assert "agentic_sources" not in kwargs
|
||||
|
||||
def test_classic_mixed_prefetches_and_scopes_tool(self):
|
||||
# Classic with an agentic_tool source now pre-fetches only the prefetch
|
||||
# subset and exposes the agentic_tool subset via the search tool.
|
||||
sp = self._classic_processor()
|
||||
sp.all_sources = [
|
||||
{"id": "a", "retrieval": RetrievalConfig(exposure="prefetch")},
|
||||
{"id": "b", "retrieval": RetrievalConfig(exposure="agentic_tool")},
|
||||
]
|
||||
sp.build_agent("q")
|
||||
sp.pre_fetch_docs.assert_called_once()
|
||||
_, pf_kwargs = sp.pre_fetch_docs.call_args
|
||||
assert pf_kwargs.get("exposure") == "prefetch"
|
||||
_, kwargs = sp.create_agent.call_args
|
||||
assert [e["id"] for e in kwargs["agentic_sources"]] == ["b"]
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestNonAgentSourceConfig:
|
||||
"""The non-agent chat path must also load per-source config so exposure
|
||||
(and other per-source overrides) are honored, not just the agent path."""
|
||||
|
||||
def test_active_docs_loads_per_source_retrieval_config(self):
|
||||
sp = StreamProcessor.__new__(StreamProcessor)
|
||||
sp._agent_data = None
|
||||
sp.data = {"active_docs": "s1"}
|
||||
sp.initial_user_id = "user1"
|
||||
sp.source = {}
|
||||
sp.all_sources = []
|
||||
fake_source = {
|
||||
"config": {"retrieval": {"exposure": "agentic_tool", "chunks": 7}}
|
||||
}
|
||||
with patch(
|
||||
"application.api.answer.services.stream_processor.db_readonly"
|
||||
), patch(
|
||||
"application.api.answer.services.stream_processor.SourcesRepository"
|
||||
) as repo:
|
||||
repo.return_value.get.return_value = fake_source
|
||||
sp._configure_source()
|
||||
assert sp.source == {"active_docs": "s1"}
|
||||
assert len(sp.all_sources) == 1
|
||||
assert sp.all_sources[0]["id"] == "s1"
|
||||
assert sp.all_sources[0]["retrieval"].exposure == "agentic_tool"
|
||||
assert sp.all_sources[0]["retrieval"].chunks == 7
|
||||
Reference in New Issue
Block a user