306 lines
11 KiB
Python
306 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from services.conversation_history import ConversationRecorder
|
|
|
|
|
|
class ConversationRecorderTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_background_messages_keep_order_and_ignore_streaming_deltas(self):
|
|
recorder = ConversationRecorder("conv_test")
|
|
recorder.record_transport_message = AsyncMock()
|
|
events = [
|
|
{"type": "transcript", "content": "user"},
|
|
{"type": "assistant-text-end", "content": "assistant"},
|
|
]
|
|
for event in events:
|
|
recorder.record_transport_message_later(event)
|
|
recorder.record_transport_message_later({"type": "assistant-text-delta"})
|
|
await recorder.flush_pending_messages()
|
|
self.assertEqual(
|
|
[call.args[0] for call in recorder.record_transport_message.await_args_list],
|
|
events,
|
|
)
|
|
self.assertIsNone(recorder._message_writer)
|
|
|
|
async def test_slow_database_does_not_block_audio(self):
|
|
from pipecat.frames.frames import EndFrame, OutputAudioRawFrame, OutputTransportMessageUrgentFrame
|
|
from pipecat.pipeline.pipeline import Pipeline
|
|
from pipecat.pipeline.worker import PipelineWorker
|
|
from pipecat.processors.frame_processor import FrameProcessor
|
|
from pipecat.workers.runner import WorkerRunner
|
|
from services.pipecat.processors import ConversationHistoryProcessor
|
|
|
|
entered, release, audio_received = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
|
recorder = ConversationRecorder("conv_test")
|
|
|
|
async def slow_write(message):
|
|
entered.set()
|
|
await release.wait()
|
|
|
|
recorder.record_transport_message = slow_write
|
|
|
|
class AudioSink(FrameProcessor):
|
|
async def process_frame(self, frame, direction):
|
|
await super().process_frame(frame, direction)
|
|
if isinstance(frame, OutputAudioRawFrame):
|
|
audio_received.set()
|
|
await self.push_frame(frame, direction)
|
|
|
|
worker = PipelineWorker(
|
|
Pipeline([ConversationHistoryProcessor(recorder), AudioSink()]), enable_rtvi=False,
|
|
)
|
|
runner = WorkerRunner(handle_sigint=False)
|
|
await runner.add_workers(worker)
|
|
task = asyncio.create_task(runner.run())
|
|
try:
|
|
await worker.queue_frame(OutputTransportMessageUrgentFrame(message={"type": "transcript"}))
|
|
await asyncio.wait_for(entered.wait(), 3)
|
|
await worker.queue_frame(OutputAudioRawFrame(audio=b"\0" * 320, sample_rate=16000, num_channels=1))
|
|
await asyncio.wait_for(audio_received.wait(), 1)
|
|
self.assertFalse(release.is_set())
|
|
finally:
|
|
release.set()
|
|
await recorder.flush_pending_messages()
|
|
await worker.queue_frame(EndFrame())
|
|
await asyncio.wait_for(task, 3)
|
|
|
|
async def test_history_buffer_is_bounded_and_slow_writer_is_cancelled(self):
|
|
started = asyncio.Event()
|
|
cancelled = asyncio.Event()
|
|
with patch("services.conversation_history.MAX_PENDING_HISTORY_MESSAGES", 2):
|
|
recorder = ConversationRecorder("conv_test")
|
|
|
|
async def blocked_write(message):
|
|
started.set()
|
|
try:
|
|
await asyncio.Event().wait()
|
|
finally:
|
|
cancelled.set()
|
|
|
|
recorder.record_transport_message = blocked_write
|
|
recorder.record_transport_message_later({"type": "transcript"})
|
|
await started.wait()
|
|
for _ in range(4):
|
|
recorder.record_transport_message_later({"type": "transcript"})
|
|
self.assertEqual(recorder._pending_messages.qsize(), 2)
|
|
with patch("services.conversation_history.HISTORY_FLUSH_TIMEOUT_SECONDS", 0.01):
|
|
await recorder.flush_pending_messages()
|
|
self.assertTrue(cancelled.is_set())
|
|
self.assertTrue(recorder._pending_messages.empty())
|
|
recorder.record_transport_message_later({"type": "transcript"})
|
|
self.assertIsNone(recorder._message_writer)
|
|
|
|
async def test_completed_conversation_queues_enabled_analysis(self):
|
|
conversation = SimpleNamespace(
|
|
status="active",
|
|
ended_at=None,
|
|
message_count=0,
|
|
analysis_status="none",
|
|
analysis_data={},
|
|
analysis_error="",
|
|
)
|
|
|
|
class FakeSession:
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
async def get(self, _model, _session_id):
|
|
return conversation
|
|
|
|
async def commit(self):
|
|
return None
|
|
|
|
plan = {
|
|
"enabled": True,
|
|
"model_resource_id": "model_001",
|
|
"fields": [{"name": "customer_name", "type": "string"}],
|
|
}
|
|
recorder = ConversationRecorder("conv_test", plan)
|
|
recorder._sequence = 1
|
|
with patch(
|
|
"services.conversation_history.SessionLocal",
|
|
return_value=FakeSession(),
|
|
):
|
|
await recorder._finish(status="completed")
|
|
|
|
self.assertEqual(conversation.analysis_status, "pending")
|
|
self.assertEqual(conversation.analysis_data, {"plan": plan})
|
|
|
|
async def test_finish_waits_for_database_cleanup_when_cancelled(self):
|
|
recorder = ConversationRecorder("conv_test")
|
|
started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
completed = asyncio.Event()
|
|
|
|
async def finish_database_write(*, status: str):
|
|
self.assertEqual(status, "failed")
|
|
started.set()
|
|
await release.wait()
|
|
completed.set()
|
|
|
|
recorder._finish = finish_database_write
|
|
task = asyncio.create_task(recorder.finish(status="failed"))
|
|
await started.wait()
|
|
task.cancel()
|
|
await asyncio.sleep(0)
|
|
|
|
self.assertFalse(task.done())
|
|
release.set()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
self.assertTrue(completed.is_set())
|
|
|
|
async def test_fixed_reply_transcript_keeps_workflow_metadata(self):
|
|
recorder = ConversationRecorder("conv_test")
|
|
recorder._append = AsyncMock()
|
|
|
|
await recorder.record_transport_message(
|
|
{
|
|
"type": "transcript",
|
|
"role": "assistant",
|
|
"content": "请稍等，我正在处理。",
|
|
"timestamp": "2026-07-14T10:00:00+08:00",
|
|
"source": "workflow-fixed-reply",
|
|
"nodeId": "agent_service",
|
|
}
|
|
)
|
|
|
|
recorder._append.assert_awaited_once_with(
|
|
"assistant",
|
|
"请稍等，我正在处理。",
|
|
"2026-07-14T10:00:00+08:00",
|
|
{
|
|
"source": "workflow-fixed-reply",
|
|
"node_id": "agent_service",
|
|
},
|
|
)
|
|
|
|
async def test_workflow_trace_is_persisted_without_counting_as_message(self):
|
|
conversation = SimpleNamespace(
|
|
extra={"workflow": {"revision": "sha256:test"}},
|
|
message_count=3,
|
|
)
|
|
|
|
class FakeSession:
|
|
committed = False
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
async def get(self, _model, _session_id):
|
|
return conversation
|
|
|
|
async def commit(self):
|
|
self.committed = True
|
|
|
|
session = FakeSession()
|
|
recorder = ConversationRecorder("conv_test")
|
|
event = {
|
|
"type": "workflow-event",
|
|
"eventId": "wfe_test",
|
|
"event": "node_entered",
|
|
"timestamp": "2026-08-01T10:00:00+08:00",
|
|
"workflowRevision": "sha256:test",
|
|
"transitionId": 0,
|
|
"nodeId": "start",
|
|
}
|
|
|
|
with patch(
|
|
"services.conversation_history.SessionLocal",
|
|
return_value=session,
|
|
):
|
|
await recorder.record_transport_message(event)
|
|
await recorder.record_transport_message(event)
|
|
|
|
self.assertTrue(session.committed)
|
|
self.assertEqual(conversation.message_count, 3)
|
|
self.assertEqual(len(conversation.extra["workflowTrace"]), 1)
|
|
saved = conversation.extra["workflowTrace"][0]
|
|
self.assertEqual(saved["sessionId"], "conv_test")
|
|
self.assertEqual(saved["sequence"], 1)
|
|
|
|
async def test_internal_message_stage_transport_trace_is_not_persisted_twice(self):
|
|
recorder = ConversationRecorder("conv_test")
|
|
event = {
|
|
"type": "workflow-event",
|
|
"eventId": "wfe_internal",
|
|
"event": "client_tool_completed",
|
|
"functionName": "show_message",
|
|
"internalStage": True,
|
|
}
|
|
|
|
with patch("services.conversation_history.SessionLocal") as session_local:
|
|
await recorder.record_transport_message(event)
|
|
|
|
session_local.assert_not_called()
|
|
self.assertEqual(recorder._trace_sequence, 0)
|
|
|
|
async def test_image_is_persisted_as_message_with_artifact(self):
|
|
conversation = SimpleNamespace(message_count=0)
|
|
|
|
class FakeSession:
|
|
def __init__(self):
|
|
self.added = []
|
|
self.committed = False
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
def add(self, value):
|
|
self.added.append(value)
|
|
|
|
async def flush(self):
|
|
return None
|
|
|
|
async def get(self, _model, _session_id):
|
|
return conversation
|
|
|
|
async def commit(self):
|
|
self.committed = True
|
|
|
|
session = FakeSession()
|
|
recorder = ConversationRecorder("conv_test")
|
|
with (
|
|
patch("services.conversation_history.SessionLocal", return_value=session),
|
|
patch("services.conversation_history.put_object") as put_object,
|
|
):
|
|
await recorder._record_image(
|
|
b"jpeg-data",
|
|
input_id="input_photo",
|
|
timestamp="2026-08-05T10:00:00+08:00",
|
|
mime_type="image/jpeg",
|
|
content="帮我看看",
|
|
source="uploaded_asset",
|
|
)
|
|
|
|
self.assertTrue(session.committed)
|
|
self.assertEqual(len(session.added), 2)
|
|
message, artifact = session.added
|
|
self.assertEqual(message.content_type, "image")
|
|
self.assertEqual(message.role, "user")
|
|
self.assertEqual(message.content, "帮我看看")
|
|
self.assertEqual(message.extra["input_id"], "input_photo")
|
|
self.assertEqual(message.extra["source"], "uploaded_asset")
|
|
self.assertEqual(artifact.message_id, message.id)
|
|
self.assertEqual(artifact.kind, "image")
|
|
self.assertEqual(artifact.size_bytes, len(b"jpeg-data"))
|
|
self.assertEqual(conversation.message_count, 1)
|
|
put_object.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|