mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 22:14:46 +00:00
chore: api and tool tests
This commit is contained in:
Alex committed
1 parent
9f3d9ab860
commit
fe185e5b8d
18 files changed
+3891
No files matched your search
@@ -0,0 +1,235 @@
|
||||
"""Tests for application/api/answer/routes/answer.py"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from bson import ObjectId
|
||||
|
||||
|
||||
@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 = str(ObjectId())
|
||||
processor.agent_config = {}
|
||||
processor.agent_id = str(ObjectId())
|
||||
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(ObjectId())
|
||||
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=(conv_id, "Hello", [], [], "", 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=(None, None, None, None, None, "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(ObjectId())
|
||||
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=(
|
||||
conv_id,
|
||||
'{"key": "val"}',
|
||||
[],
|
||||
[],
|
||||
"",
|
||||
None,
|
||||
{"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(ObjectId())
|
||||
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=(
|
||||
conv_id,
|
||||
"answer text",
|
||||
[{"title": "src"}],
|
||||
[{"tool": "t"}],
|
||||
"thinking...",
|
||||
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,195 @@
|
||||
"""Tests for application/api/answer/routes/stream.py"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from bson import ObjectId
|
||||
|
||||
|
||||
@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 = str(ObjectId())
|
||||
processor.agent_config = {}
|
||||
processor.agent_id = str(ObjectId())
|
||||
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(ObjectId())
|
||||
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"]
|
||||
Whitespace-only changes.
@@ -0,0 +1,303 @@
|
||||
"""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 "function_call" in messages[3]["content"][0]
|
||||
assert messages[4]["role"] == "tool"
|
||||
assert "function_response" in messages[4]["content"][0]
|
||||
|
||||
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,
|
||||
)
|
||||
tool_msg = messages[3]["content"][0]
|
||||
call_id = tool_msg["function_call"]["call_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,447 @@
|
||||
"""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,
|
||||
)
|
||||
|
||||
mock_threshold_checker.should_compress.assert_called_once_with(
|
||||
sample_conversation, "gpt-4", 1000
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
mock_get_provider.assert_called_with("gpt-3.5-turbo")
|
||||
|
||||
@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
|
||||
)
|
||||
|
||||
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,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 == []
|
||||
Reference in new issue
Block a user