wehub-resource-sync
fed8b2eed7
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
1289 lines
49 KiB
Python
1289 lines
49 KiB
Python
"""Unit tests for application/api/answer/routes/base.py — BaseAnswerResource.
|
|
|
|
Additional coverage beyond tests/api/answer/routes/test_base.py:
|
|
- _prepare_tool_calls_for_logging: truncation, non-dict items
|
|
- complete_stream: tool_calls, thoughts, structured output, metadata,
|
|
isNoneDoc, GeneratorExit handling, compression metadata
|
|
- process_response_stream: structured answer, incomplete stream
|
|
- error_stream_generate: format
|
|
- check_usage: string boolean parsing ("True" strings)
|
|
"""
|
|
|
|
import json
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestPrepareToolCallsForLogging:
|
|
pass
|
|
|
|
def test_empty_list(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
assert resource._prepare_tool_calls_for_logging([]) == []
|
|
|
|
def test_none_returns_empty(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
assert resource._prepare_tool_calls_for_logging(None) == []
|
|
|
|
def test_truncates_long_result(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
tool_calls = [{"result": "x" * 20000}]
|
|
prepared = resource._prepare_tool_calls_for_logging(tool_calls, max_chars=100)
|
|
assert len(prepared[0]["result"]) == 100
|
|
|
|
def test_truncates_result_full(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
tool_calls = [{"result_full": "y" * 20000}]
|
|
prepared = resource._prepare_tool_calls_for_logging(tool_calls, max_chars=50)
|
|
assert len(prepared[0]["result_full"]) == 50
|
|
|
|
def test_non_dict_items_wrapped(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
tool_calls = ["string_item", 42]
|
|
prepared = resource._prepare_tool_calls_for_logging(tool_calls)
|
|
assert prepared[0] == {"result": "string_item"}
|
|
assert prepared[1] == {"result": "42"}
|
|
|
|
def test_preserves_short_results(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
tool_calls = [{"tool_name": "search", "result": "short text"}]
|
|
prepared = resource._prepare_tool_calls_for_logging(tool_calls)
|
|
assert prepared[0]["result"] == "short text"
|
|
assert prepared[0]["tool_name"] == "search"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamToolCalls:
|
|
pass
|
|
|
|
def test_streams_tool_calls(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": "Using tool..."},
|
|
{"tool_calls": [{"name": "search", "result": "found"}]},
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Search for X",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
tool_chunks = [s for s in stream if '"type": "tool_calls"' in s]
|
|
assert len(tool_chunks) == 1
|
|
|
|
def test_streams_thought_events(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(
|
|
[
|
|
{"thought": "Let me think..."},
|
|
{"answer": "Here is the answer"},
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
thought_chunks = [s for s in stream if '"type": "thought"' in s]
|
|
assert len(thought_chunks) == 1
|
|
assert "Let me think" in thought_chunks[0]
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamStructuredOutput:
|
|
pass
|
|
|
|
def test_streams_structured_answer(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": '{"key": "value"}',
|
|
"structured": True,
|
|
"schema": {"type": "object"},
|
|
},
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
structured_chunks = [
|
|
s for s in stream if '"type": "structured_answer"' in s
|
|
]
|
|
assert len(structured_chunks) == 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamMetadata:
|
|
pass
|
|
|
|
def test_metadata_collected(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(
|
|
[
|
|
{"metadata": {"search_query": "test"}},
|
|
{"answer": "result"},
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
# Should not crash, metadata handled silently
|
|
answer_chunks = [s for s in stream if '"type": "answer"' in s]
|
|
assert len(answer_chunks) == 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamIsNoneDoc:
|
|
pass
|
|
|
|
def test_isNoneDoc_sets_source_to_none(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": "answer"},
|
|
{"sources": [{"text": "doc", "source": "real_source"}]},
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
isNoneDoc=True,
|
|
should_persist=False,
|
|
)
|
|
)
|
|
# Verify stream completes without error
|
|
end_chunks = [s for s in stream if '"type": "end"' in s]
|
|
assert len(end_chunks) == 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamErrorType:
|
|
pass
|
|
|
|
def test_error_type_event_sanitized(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(
|
|
[
|
|
{"type": "error", "error": "API key invalid: sk-xxx"},
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
error_chunks = [s for s in stream if '"type": "error"' in s]
|
|
assert len(error_chunks) == 1
|
|
|
|
def test_non_error_type_event_passed_through(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(
|
|
[
|
|
{"type": "custom_event", "data": "value"},
|
|
]
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
custom_chunks = [s for s in stream if '"type": "custom_event"' in s]
|
|
assert len(custom_chunks) == 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestProcessResponseStreamExtended:
|
|
pass
|
|
|
|
def test_handles_structured_answer(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
stream = [
|
|
f'data: {json.dumps({"type": "structured_answer", "answer": "{}", "structured": True, "schema": None})}\n\n',
|
|
f'data: {json.dumps({"type": "id", "id": str(uuid.uuid4())})}\n\n',
|
|
f'data: {json.dumps({"type": "end"})}\n\n',
|
|
]
|
|
result = resource.process_response_stream(iter(stream))
|
|
assert result["answer"] == "{}"
|
|
assert result.get("extra", {}).get("structured") is True
|
|
|
|
def test_handles_tool_calls_event(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
stream = [
|
|
f'data: {json.dumps({"type": "answer", "answer": "result"})}\n\n',
|
|
f'data: {json.dumps({"type": "tool_calls", "tool_calls": [{"name": "t1"}]})}\n\n',
|
|
f'data: {json.dumps({"type": "id", "id": "conv1"})}\n\n',
|
|
f'data: {json.dumps({"type": "end"})}\n\n',
|
|
]
|
|
result = resource.process_response_stream(iter(stream))
|
|
assert result["tool_calls"] == [{"name": "t1"}]
|
|
|
|
def test_incomplete_stream(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
stream = [
|
|
f'data: {json.dumps({"type": "answer", "answer": "partial"})}\n\n',
|
|
]
|
|
result = resource.process_response_stream(iter(stream))
|
|
assert result["error"] == "Stream ended unexpectedly"
|
|
|
|
def test_handles_thought_event(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
stream = [
|
|
f'data: {json.dumps({"type": "thought", "thought": "thinking..."})}\n\n',
|
|
f'data: {json.dumps({"type": "end"})}\n\n',
|
|
]
|
|
result = resource.process_response_stream(iter(stream))
|
|
assert result["thought"] == "thinking..."
|
|
|
|
def test_handles_id_prefixed_chunks(self, mock_mongo_db, flask_app):
|
|
"""``complete_stream`` emits ``id: <seq>\\n`` before each
|
|
``data:`` line so reconnects can resume. The non-streaming
|
|
``/api/answer`` consumer must skip the ``id:`` header (and the
|
|
informational ``message_id`` event) without breaking JSON
|
|
decoding.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
stream = [
|
|
'id: 0\n'
|
|
f'data: {json.dumps({"type": "message_id", "message_id": "abc"})}\n\n',
|
|
'id: 1\n'
|
|
f'data: {json.dumps({"type": "answer", "answer": "Hello "})}\n\n',
|
|
'id: 2\n'
|
|
f'data: {json.dumps({"type": "answer", "answer": "world"})}\n\n',
|
|
'id: 3\n'
|
|
f'data: {json.dumps({"type": "id", "id": "conv-1"})}\n\n',
|
|
'id: 4\n'
|
|
f'data: {json.dumps({"type": "end"})}\n\n',
|
|
]
|
|
result = resource.process_response_stream(iter(stream))
|
|
assert result["error"] is None
|
|
assert result["answer"] == "Hello world"
|
|
assert result["conversation_id"] == "conv-1"
|
|
|
|
def test_skips_keepalive_comment_lines(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
stream = [
|
|
': keepalive\n\n',
|
|
'id: 0\n'
|
|
f'data: {json.dumps({"type": "answer", "answer": "ok"})}\n\n',
|
|
'id: 1\n'
|
|
f'data: {json.dumps({"type": "end"})}\n\n',
|
|
]
|
|
result = resource.process_response_stream(iter(stream))
|
|
assert result["answer"] == "ok"
|
|
assert result["error"] is None
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCheckUsageStringBooleans:
|
|
pass
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamCompressionMetadata:
|
|
"""Cover lines 307-319 (compression metadata persistence in complete_stream)."""
|
|
|
|
def test_compression_metadata_persisted(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": "compressed answer"},
|
|
]
|
|
)
|
|
mock_agent.compression_metadata = {"ratio": 2.5}
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_conversation.return_value = "conv123"
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv123",
|
|
"message_id": "msg123",
|
|
"request_id": "req123",
|
|
}
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
)
|
|
|
|
# Verify compression metadata was persisted
|
|
resource.conversation_service.update_compression_metadata.assert_called_once_with(
|
|
"conv123", {"ratio": 2.5}
|
|
)
|
|
resource.conversation_service.append_compression_message.assert_called_once()
|
|
assert mock_agent.compression_saved is True
|
|
end_chunks = [s for s in stream if '"type": "end"' in s]
|
|
assert len(end_chunks) == 1
|
|
|
|
def test_compression_metadata_error_handled(self, mock_mongo_db, flask_app):
|
|
"""Cover lines 318-322: compression metadata persistence 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([{"answer": "answer"}])
|
|
mock_agent.compression_metadata = {"ratio": 2.5}
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_conversation.return_value = "conv123"
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv123",
|
|
"message_id": "msg123",
|
|
"request_id": "req123",
|
|
}
|
|
resource.conversation_service.update_compression_metadata.side_effect = (
|
|
Exception("db error")
|
|
)
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
)
|
|
|
|
# Stream should still complete despite compression error
|
|
end_chunks = [s for s in stream if '"type": "end"' in s]
|
|
assert len(end_chunks) == 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamLogTruncation:
|
|
"""Cover line 354: log data truncation for long values."""
|
|
|
|
def test_long_response_truncated_in_log(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
long_answer = "x" * 20000
|
|
mock_agent.gen.return_value = iter([{"answer": long_answer}])
|
|
mock_agent.tool_calls = []
|
|
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=False,
|
|
)
|
|
)
|
|
|
|
end_chunks = [s for s in stream if '"type": "end"' in s]
|
|
assert len(end_chunks) == 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamGeneratorExit:
|
|
"""Cover lines 360-416 (GeneratorExit handling in complete_stream)."""
|
|
|
|
def test_generator_exit_saves_partial_response(self, mock_mongo_db, flask_app):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
|
|
def gen_with_answers():
|
|
yield {"answer": "partial"}
|
|
yield {"answer": " answer"}
|
|
# Simulating a long stream that gets interrupted
|
|
yield {"answer": " more"}
|
|
|
|
mock_agent.gen.return_value = gen_with_answers()
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_conversation.return_value = "conv1"
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": "msg1",
|
|
"request_id": "req1",
|
|
}
|
|
|
|
gen = resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
|
|
# Drain the early ``message_id`` event before reading the
|
|
# ``partial`` chunk that this test is asserting on.
|
|
next(gen)
|
|
chunk = next(gen)
|
|
assert "partial" in chunk
|
|
gen.close() # This triggers GeneratorExit
|
|
|
|
def test_generator_exit_with_compression_metadata(self, mock_mongo_db, flask_app):
|
|
"""Cover lines 393-411: GeneratorExit with compression metadata."""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
|
|
def gen_answers():
|
|
yield {"answer": "partial answer"}
|
|
|
|
mock_agent.gen.return_value = gen_answers()
|
|
mock_agent.compression_metadata = {"ratio": 3.0}
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_conversation.return_value = "conv1"
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": "msg1",
|
|
"request_id": "req1",
|
|
}
|
|
|
|
gen = resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
isNoneDoc=True,
|
|
)
|
|
|
|
# Skip past the early ``message_id`` event.
|
|
next(gen)
|
|
next(gen)
|
|
gen.close()
|
|
|
|
def test_generator_exit_save_error_handled(self, mock_mongo_db, flask_app):
|
|
"""Cover lines 412-415: exception during partial save."""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
|
|
def gen_answers():
|
|
yield {"answer": "partial"}
|
|
|
|
mock_agent.gen.return_value = gen_answers()
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_conversation.side_effect = Exception(
|
|
"save error"
|
|
)
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": "msg1",
|
|
"request_id": "req1",
|
|
}
|
|
|
|
gen = resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
|
|
# Skip past the early ``message_id`` event.
|
|
next(gen)
|
|
next(gen)
|
|
gen.close() # Should not crash even with save error
|
|
|
|
def test_generator_exit_before_any_response_journals_error_not_end(
|
|
self, mock_mongo_db, flask_app,
|
|
):
|
|
"""A client disconnect right after the early ``message_id`` frame
|
|
leaves ``response_full`` empty, so finalize never runs. The abort
|
|
handler must journal ``error`` (not ``end``) and flip the row to
|
|
``failed`` — otherwise a reconnecting client sees a terminal
|
|
``end`` for a row whose DB status is still non-terminal and the
|
|
UI parks on a blank successful answer.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
|
|
# An agent that's primed but never yields anything before the
|
|
# client disconnects — keeps ``response_full`` empty.
|
|
def gen_never_yields():
|
|
if False:
|
|
yield # pragma: no cover
|
|
return
|
|
|
|
mock_agent.gen.return_value = gen_never_yields()
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": "msg1",
|
|
"request_id": "req1",
|
|
}
|
|
|
|
journaled: list[tuple] = []
|
|
|
|
def _capture_record(message_id, sequence_no, event_type, payload):
|
|
journaled.append((message_id, sequence_no, event_type, payload))
|
|
return True
|
|
|
|
with patch(
|
|
"application.api.answer.routes.base.record_event",
|
|
side_effect=_capture_record,
|
|
):
|
|
gen = resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
# Drain the early ``message_id`` event, then close before
|
|
# the agent yields anything.
|
|
next(gen)
|
|
gen.close()
|
|
|
|
# The early message_id frame got journaled (seq 0); the abort
|
|
# handler must follow with an ``error`` event (NOT ``end``).
|
|
terminal_events = [
|
|
(et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
|
|
]
|
|
assert len(terminal_events) == 1, (
|
|
f"expected exactly one terminal journal write, got {terminal_events}"
|
|
)
|
|
assert terminal_events[0][0] == "error", (
|
|
f"expected ``error`` terminal but got ``end``: {terminal_events}"
|
|
)
|
|
payload = terminal_events[0][1]
|
|
assert payload.get("type") == "error"
|
|
assert payload.get("code") == "client_disconnect"
|
|
|
|
# And the DB row should have been flipped to ``failed`` via
|
|
# finalize_message. The mocked service records the call.
|
|
finalize_calls = (
|
|
resource.conversation_service.finalize_message.call_args_list
|
|
)
|
|
assert len(finalize_calls) == 1
|
|
assert finalize_calls[0].kwargs.get("status") == "failed"
|
|
|
|
def test_generator_exit_after_response_still_journals_end(
|
|
self, mock_mongo_db, flask_app,
|
|
):
|
|
"""Regression guard: a disconnect AFTER partial response was
|
|
produced and ``finalize_message`` succeeded must still journal
|
|
``end`` (the row matches ``complete``). Only the empty-response
|
|
branch flips to ``error``.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
from application.storage.db.repositories.conversations import (
|
|
MessageUpdateOutcome,
|
|
)
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
|
|
def gen_answers():
|
|
yield {"answer": "partial"}
|
|
|
|
mock_agent.gen.return_value = gen_answers()
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.finalize_message.return_value = (
|
|
MessageUpdateOutcome.UPDATED
|
|
)
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": "msg1",
|
|
"request_id": "req1",
|
|
}
|
|
|
|
journaled: list[tuple] = []
|
|
|
|
def _capture_record(message_id, sequence_no, event_type, payload):
|
|
journaled.append((message_id, sequence_no, event_type, payload))
|
|
return True
|
|
|
|
with patch(
|
|
"application.api.answer.routes.base.record_event",
|
|
side_effect=_capture_record,
|
|
):
|
|
gen = resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
next(gen) # message_id frame
|
|
next(gen) # answer frame (consumes ``partial``)
|
|
gen.close()
|
|
|
|
terminal_events = [
|
|
(et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
|
|
]
|
|
assert terminal_events and terminal_events[0][0] == "end", (
|
|
f"finalize succeeded but abort journaled {terminal_events}"
|
|
)
|
|
|
|
def test_generator_exit_after_normal_finalize_already_complete_journals_end(
|
|
self, mock_mongo_db, flask_app,
|
|
):
|
|
"""Regression for the race where the normal-path finalize wins
|
|
against a client disconnect.
|
|
|
|
Trace: agent finishes, ``complete_stream`` runs the normal-path
|
|
``finalize_message`` at base.py:632 and flips the row to
|
|
``complete``. The client TCP-resets before the ``end`` frame can
|
|
be journaled. The GeneratorExit handler calls ``finalize_message``
|
|
again — and the repository, gated by ``only_if_non_terminal``,
|
|
reports ``ALREADY_COMPLETE`` because the row is already at the
|
|
target state. The abort handler must journal ``end``, not
|
|
``error``: the DB says ``complete``, the reconnecting client
|
|
must see the same.
|
|
|
|
Without the fix the abort handler treated ``ALREADY_COMPLETE``
|
|
as a failure and journaled ``error`` — a reconnect would then
|
|
replay a successful completion as a failed answer.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
from application.storage.db.repositories.conversations import (
|
|
MessageUpdateOutcome,
|
|
)
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
|
|
def gen_answers():
|
|
yield {"answer": "partial"}
|
|
|
|
mock_agent.gen.return_value = gen_answers()
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
|
|
resource.conversation_service = MagicMock()
|
|
# Normal-path call returns UPDATED; abort-handler call sees
|
|
# the row already terminal and returns ALREADY_COMPLETE.
|
|
resource.conversation_service.finalize_message.side_effect = [
|
|
MessageUpdateOutcome.UPDATED,
|
|
MessageUpdateOutcome.ALREADY_COMPLETE,
|
|
]
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": "msg1",
|
|
"request_id": "req1",
|
|
}
|
|
|
|
journaled: list[tuple] = []
|
|
|
|
def _capture_record(message_id, sequence_no, event_type, payload):
|
|
journaled.append((message_id, sequence_no, event_type, payload))
|
|
return True
|
|
|
|
with patch(
|
|
"application.api.answer.routes.base.record_event",
|
|
side_effect=_capture_record,
|
|
):
|
|
gen = resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
next(gen) # message_id frame
|
|
next(gen) # answer frame
|
|
# Pull the ``id`` frame so the normal-path finalize at
|
|
# line 632 actually runs (it sits between the answer
|
|
# frame yield and the id frame yield).
|
|
next(gen) # id frame — runs normal-path finalize first
|
|
gen.close() # GeneratorExit at the id frame yield
|
|
|
|
terminal_events = [
|
|
(et, pl) for (_, _, et, pl) in journaled if et in ("end", "error")
|
|
]
|
|
assert len(terminal_events) == 1, (
|
|
f"expected exactly one terminal journal write, got "
|
|
f"{terminal_events}"
|
|
)
|
|
assert terminal_events[0][0] == "end", (
|
|
f"row was already ``complete`` but abort journaled "
|
|
f"{terminal_events[0]} — reconnect would surface a "
|
|
f"successful answer as failed"
|
|
)
|
|
|
|
# Both finalize_message calls were made: the normal-path
|
|
# one and the abort-handler one. Asserting on side_effect
|
|
# consumption ensures the test really exercised both
|
|
# branches.
|
|
assert (
|
|
resource.conversation_service.finalize_message.call_count == 2
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def _patch_db_session(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,
|
|
), patch(
|
|
# ``record_event`` opens its own short-lived ``db_session`` for
|
|
# cross-connection visibility. In tests we route it back to the
|
|
# same ``pg_conn`` so the journal write can see the message row
|
|
# the conversation_service just wrote in this transaction.
|
|
"application.streaming.message_journal.db_session",
|
|
_yield,
|
|
), patch(
|
|
# ``complete_stream`` reads ``latest_sequence_no`` via
|
|
# ``db_readonly`` to seed continuation runs. Same patch reason
|
|
# as the journal — keep the read on the same pg_conn so it sees
|
|
# uncommitted writes from this transaction.
|
|
"application.api.answer.routes.base.db_readonly",
|
|
_yield,
|
|
):
|
|
yield
|
|
|
|
|
|
def _extract_sse_data(chunk: str) -> str:
|
|
"""Pull the ``data:`` payload from an SSE record, ignoring any
|
|
``id:`` header introduced by the journal wiring.
|
|
"""
|
|
for line in chunk.split("\n"):
|
|
if line.startswith("data:"):
|
|
return line[len("data:") :].lstrip()
|
|
return ""
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCompleteStreamWalAcceptance:
|
|
"""Acceptance for the WAL pre-persist behaviour: when the LLM raises
|
|
immediately, the user question is still queryable from PG with
|
|
status='failed' and a meaningful error in metadata."""
|
|
|
|
def test_failed_llm_persists_question_with_failed_status(
|
|
self, pg_conn, flask_app,
|
|
):
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
from application.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.side_effect = RuntimeError("LLM upstream failed")
|
|
|
|
with _patch_db_session(pg_conn):
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="why does the WAL matter?",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u-acceptance"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
)
|
|
error_chunks = [s for s in stream if '"type": "error"' in s]
|
|
assert len(error_chunks) == 1
|
|
|
|
from sqlalchemy import text as sql_text
|
|
convs = pg_conn.execute(
|
|
sql_text("SELECT id FROM conversations WHERE user_id = :u"),
|
|
{"u": "u-acceptance"},
|
|
).fetchall()
|
|
assert len(convs) == 1
|
|
conv_id = str(convs[0][0])
|
|
msgs = ConversationsRepository(pg_conn).get_messages(conv_id)
|
|
assert len(msgs) == 1
|
|
assert msgs[0]["prompt"] == "why does the WAL matter?"
|
|
assert msgs[0]["status"] == "failed"
|
|
assert "RuntimeError" in msgs[0]["metadata"]["error"]
|
|
assert "LLM upstream failed" in msgs[0]["metadata"]["error"]
|
|
|
|
def test_tool_approval_event_only_fires_when_state_saved(
|
|
self, pg_conn, flask_app,
|
|
):
|
|
"""A `tool.approval.required` notification with no resumable
|
|
``pending_tool_state`` row would deep-link the user to a 404.
|
|
Gate the publish on save_state actually committing.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
# Drive the agent into the paused branch with a single
|
|
# ``tool_calls_pending`` event.
|
|
mock_agent.gen.return_value = iter(
|
|
[
|
|
{
|
|
"type": "tool_calls_pending",
|
|
"data": {"pending_tool_calls": [{"call_id": "c1"}]},
|
|
}
|
|
]
|
|
)
|
|
mock_agent._pending_continuation = {
|
|
"messages": [],
|
|
"tools_dict": {},
|
|
"pending_tool_calls": [{"call_id": "c1"}],
|
|
}
|
|
mock_agent.tool_calls = []
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
|
|
published: list = []
|
|
|
|
def _capture(*args, **kwargs):
|
|
published.append(args)
|
|
|
|
with _patch_db_session(pg_conn), patch(
|
|
"application.api.answer.routes.base.publish_user_event",
|
|
side_effect=_capture,
|
|
), patch(
|
|
"application.api.answer.services.continuation_service."
|
|
"ContinuationService.save_state",
|
|
side_effect=RuntimeError("PG outage"),
|
|
):
|
|
list(
|
|
resource.complete_stream(
|
|
question="run my tool",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u-tap"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
)
|
|
|
|
# No tool.approval.required publish when save_state failed.
|
|
event_types = [a[1] for a in published if len(a) >= 2]
|
|
assert "tool.approval.required" not in event_types
|
|
|
|
def test_continuation_seeds_sequence_no_from_journal_high_water_mark(
|
|
self, pg_conn, flask_app,
|
|
):
|
|
"""A resumed (tool-actions) stream must continue numbering past
|
|
the original run's max ``sequence_no``. Otherwise the second
|
|
invocation collides on the duplicate-PK and silently drops
|
|
every journal write past the resume point.
|
|
|
|
We pre-seed the journal with synthetic rows simulating an
|
|
original run, then invoke ``complete_stream`` with
|
|
``_continuation`` set and assert seq numbering picks up past
|
|
the high-water mark.
|
|
"""
|
|
import uuid as _uuid
|
|
|
|
from sqlalchemy import text as sql_text
|
|
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
from application.storage.db.repositories.message_events import (
|
|
MessageEventsRepository,
|
|
)
|
|
|
|
# Pre-seed: insert a parent conversation + message + a few
|
|
# journal rows so latest_sequence_no returns 7.
|
|
user_id = "u-resume"
|
|
conv_id = _uuid.uuid4()
|
|
message_id = _uuid.uuid4()
|
|
pg_conn.execute(
|
|
sql_text("INSERT INTO users (user_id) VALUES (:u)"),
|
|
{"u": user_id},
|
|
)
|
|
pg_conn.execute(
|
|
sql_text(
|
|
"INSERT INTO conversations (id, user_id, name) "
|
|
"VALUES (:id, :u, 'pre-seed')"
|
|
),
|
|
{"id": conv_id, "u": user_id},
|
|
)
|
|
pg_conn.execute(
|
|
sql_text(
|
|
"INSERT INTO conversation_messages "
|
|
"(id, conversation_id, user_id, position, prompt) "
|
|
"VALUES (:id, :c, :u, 0, 'q')"
|
|
),
|
|
{"id": message_id, "c": conv_id, "u": user_id},
|
|
)
|
|
repo = MessageEventsRepository(pg_conn)
|
|
for seq in range(8): # rows 0..7
|
|
repo.record(str(message_id), seq, "answer", {"chunk": str(seq)})
|
|
original_max = repo.latest_sequence_no(str(message_id))
|
|
assert original_max == 7
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
cont_agent = MagicMock()
|
|
cont_agent.gen_continuation.return_value = iter(
|
|
[{"answer": "d"}, {"answer": "e"}]
|
|
)
|
|
cont_agent.tool_calls = []
|
|
cont_agent.compression_metadata = None
|
|
cont_agent.compression_saved = False
|
|
cont_agent.tool_executor = None
|
|
|
|
with _patch_db_session(pg_conn):
|
|
list(
|
|
resource.complete_stream(
|
|
question="",
|
|
agent=cont_agent,
|
|
conversation_id=str(conv_id),
|
|
user_api_key=None,
|
|
decoded_token={"sub": user_id},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
_continuation={
|
|
"messages": [],
|
|
"tools_dict": {},
|
|
"pending_tool_calls": [],
|
|
"tool_actions": [],
|
|
"reserved_message_id": str(message_id),
|
|
"request_id": "req-resume",
|
|
},
|
|
)
|
|
)
|
|
|
|
new_max = MessageEventsRepository(pg_conn).latest_sequence_no(
|
|
str(message_id)
|
|
)
|
|
# Continuation extended the journal — the new high-water mark
|
|
# is strictly greater than the seeded ``original_max=7``,
|
|
# confirming the allocator picked up past the resume point.
|
|
assert new_max is not None and new_max > original_max
|
|
|
|
def test_request_id_consistent_across_sse_event_and_wal_row(
|
|
self, pg_conn, flask_app,
|
|
):
|
|
"""The early ``message_id`` SSE event reports the same
|
|
``request_id`` that ``save_user_question`` writes onto the WAL
|
|
row, so client-side correlation, ``token_usage`` joins, and
|
|
``count_in_range``'s DISTINCT all line up.
|
|
"""
|
|
from application.api.answer.routes.base import BaseAnswerResource
|
|
from application.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
|
|
with flask_app.app_context():
|
|
resource = BaseAnswerResource()
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = iter([{"answer": "ok"}])
|
|
mock_agent.tool_calls = []
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
|
|
with _patch_db_session(pg_conn):
|
|
stream = list(
|
|
resource.complete_stream(
|
|
question="hello",
|
|
agent=mock_agent,
|
|
conversation_id=None,
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u-request-id"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
)
|
|
|
|
sse_events = [
|
|
json.loads(_extract_sse_data(s))
|
|
for s in stream
|
|
if "data:" in s
|
|
]
|
|
early_events = [e for e in sse_events if e.get("type") == "message_id"]
|
|
assert len(early_events) == 1
|
|
sse_request_id = early_events[0]["request_id"]
|
|
assert sse_request_id
|
|
|
|
from sqlalchemy import text as sql_text
|
|
convs = pg_conn.execute(
|
|
sql_text("SELECT id FROM conversations WHERE user_id = :u"),
|
|
{"u": "u-request-id"},
|
|
).fetchall()
|
|
assert len(convs) == 1
|
|
msgs = ConversationsRepository(pg_conn).get_messages(str(convs[0][0]))
|
|
assert len(msgs) == 1
|
|
assert msgs[0]["request_id"] == sse_request_id
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestStreamingHeartbeatSeed:
|
|
"""Regression guard: the reserved row must carry a fresh
|
|
``last_heartbeat_at`` from generation start so the reconciler doesn't fall
|
|
back to ``timestamp`` (creation time) on slow LLM cold-starts or while a
|
|
reasoning model streams only ``thought`` chunks. The heartbeat is seeded
|
|
once before the first chunk and re-stamped when the row flips to
|
|
``streaming``; the ``pending → streaming`` status transition itself still
|
|
fires exactly once on the first ``answer``/``sources`` chunk.
|
|
"""
|
|
|
|
def test_heartbeat_seeded_at_generation_start_and_on_first_chunk(
|
|
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": "a"}, {"answer": "b"}]
|
|
)
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
mock_agent.tool_executor = None
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_conversation.return_value = "conv1"
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": "msg1",
|
|
"request_id": "req1",
|
|
}
|
|
|
|
list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
)
|
|
|
|
# update_message_status flips the row to ``streaming`` exactly
|
|
# once (idempotent via streaming_marked) on the first answer chunk.
|
|
status_calls = [
|
|
c for c in resource.conversation_service
|
|
.update_message_status.call_args_list
|
|
if c.args[1] == "streaming"
|
|
]
|
|
assert len(status_calls) == 1
|
|
assert status_calls[0].args[0] == "msg1"
|
|
|
|
# Heartbeat is stamped twice for this 2-answer-chunk stream: once at
|
|
# the generation-start seed (before the first chunk, while still
|
|
# ``pending``) and once when the first answer chunk flips the row to
|
|
# ``streaming``. The second answer chunk does NOT re-stamp — the
|
|
# per-chunk ``_heartbeat_streaming`` pump is throttled by
|
|
# STREAM_HEARTBEAT_INTERVAL and the two chunks fall inside one
|
|
# interval under real ``time.monotonic``.
|
|
hb_calls = (
|
|
resource.conversation_service.heartbeat_message.call_args_list
|
|
)
|
|
assert len(hb_calls) == 2
|
|
assert all(c.args[0] == "msg1" for c in hb_calls)
|
|
|
|
def test_heartbeat_seed_skipped_without_reserved_message_id(
|
|
self, mock_mongo_db, flask_app,
|
|
):
|
|
"""No DB-backed message row → no heartbeat call (and no 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([{"answer": "a"}])
|
|
mock_agent.compression_metadata = None
|
|
mock_agent.compression_saved = False
|
|
mock_agent.tool_calls = []
|
|
mock_agent.tool_executor = None
|
|
|
|
resource.conversation_service = MagicMock()
|
|
resource.conversation_service.save_conversation.return_value = "conv1"
|
|
# save_user_question returns no message_id → reservation absent
|
|
resource.conversation_service.save_user_question.return_value = {
|
|
"conversation_id": "conv1",
|
|
"message_id": None,
|
|
"request_id": "req1",
|
|
}
|
|
|
|
list(
|
|
resource.complete_stream(
|
|
question="Q",
|
|
agent=mock_agent,
|
|
conversation_id="conv1",
|
|
user_api_key=None,
|
|
decoded_token={"sub": "u"},
|
|
should_persist=True,
|
|
model_id="gpt-4",
|
|
)
|
|
)
|
|
|
|
resource.conversation_service.heartbeat_message.assert_not_called()
|