From 2bf094b950b426b80734e32601593fc3a665886c Mon Sep 17 00:00:00 2001
From: TomTom101 <thobra@gmail.com>
Date: Thu, 30 May 2024 10:43:33 +0200
Subject: [PATCH] test(langchain): Rewrite to unittest, make it meaningful

---
 tests/test_langchain.py | 113 +++++++++++++++++++++++++---------------
 1 file changed, 71 insertions(+), 42 deletions(-)

diff --git a/tests/test_langchain.py b/tests/test_langchain.py
index e204c56a2..0f3ccdc86 100644
--- a/tests/test_langchain.py
+++ b/tests/test_langchain.py
@@ -1,8 +1,11 @@
-import pytest
+import unittest
+
 from langchain.prompts import ChatPromptTemplate
 from langchain_core.language_models import FakeStreamingListLLM
 
-from pipecat.frames.frames import (StopTaskFrame, TranscriptionFrame,
+from pipecat.frames.frames import (LLMFullResponseEndFrame,
+                                   LLMFullResponseStartFrame, StopTaskFrame,
+                                   TextFrame, TranscriptionFrame,
                                    UserStartedSpeakingFrame,
                                    UserStoppedSpeakingFrame)
 from pipecat.pipeline.pipeline import Pipeline
@@ -10,48 +13,74 @@ from pipecat.pipeline.runner import PipelineRunner
 from pipecat.pipeline.task import PipelineTask
 from pipecat.processors.aggregators.llm_response import (
     LLMAssistantResponseAggregator, LLMUserResponseAggregator)
-from pipecat.processors.logger import FrameLogger
+from pipecat.processors.frame_processor import FrameProcessor
 from pipecat.services.langchain import LangchainProcessor
 
 
-@pytest.fixture
-def fake_llm():
-    responses = ["Hello dear human"]
-    return FakeStreamingListLLM(responses=responses)
+class TestLangchain(unittest.IsolatedAsyncioTestCase):
+
+    class MockProcessor(FrameProcessor):
+        def __init__(self, name):
+            self.name = name
+            self.token: list[str] = []
+            # Start collecting tokens when we see the start frame
+            self.start_collecting = False
+
+        def __str__(self):
+            return self.name
+
+        async def process_frame(self, frame, direction):
+            if isinstance(frame, LLMFullResponseStartFrame):
+                self.start_collecting = True
+            elif isinstance(frame, TextFrame) and self.start_collecting:
+                self.token.append(frame.text)
+            elif isinstance(frame, LLMFullResponseEndFrame):
+                self.start_collecting = False
+
+            await self.push_frame(frame, direction)
+
+    def setUp(self):
+        self.expected_response = "Hello dear human"
+        self.fake_llm = FakeStreamingListLLM(responses=[self.expected_response])
+        self.mock_proc = self.MockProcessor("token_collector")
+
+    async def test_langchain(self):
+
+        messages = [("system", "Say hello to {name}"), ("human", "{input}")]
+        prompt = ChatPromptTemplate.from_messages(messages).partial(name="Thomas")
+        chain = prompt | self.fake_llm
+        proc = LangchainProcessor(chain=chain)
+
+        tma_in = LLMUserResponseAggregator(messages)
+        tma_out = LLMAssistantResponseAggregator(messages)
+
+        pipeline = Pipeline(
+            [
+                tma_in,
+                proc,
+                self.mock_proc,
+                tma_out,
+            ]
+        )
+
+        task = PipelineTask(pipeline)
+        await task.queue_frames(
+            [
+                UserStartedSpeakingFrame(),
+                TranscriptionFrame(text="Hi World", user_id="user", timestamp="now"),
+                UserStoppedSpeakingFrame(),
+                StopTaskFrame(),
+            ]
+        )
+
+        runner = PipelineRunner()
+        await runner.run(task)
+        self.assertEqual("".join(self.mock_proc.token), self.expected_response)
+        # TODO: Address this issue
+        # This next one would fail with:
+        # AssertionError: ' H e l l o   d e a r   h u m a n' != 'Hello dear human'
+        # self.assertEqual(tma_out.messages[-1]["content"], self.expected_response)
 
 
-@pytest.mark.asyncio
-async def test_langchain(fake_llm: FakeStreamingListLLM):
-    fl_in = FrameLogger("Inner")
-    fl_out = FrameLogger("Outer")
-
-    messages = [("system", "Say hello to {name}"), ("human", "{input}")]
-    prompt = ChatPromptTemplate.from_messages(messages).partial(name="Thomas")
-    chain = prompt | fake_llm
-    proc = LangchainProcessor(chain=chain)
-
-    tma_in = LLMUserResponseAggregator(messages)
-    tma_out = LLMAssistantResponseAggregator(messages)
-
-    pipeline = Pipeline(
-        [
-            fl_in,
-            tma_in,
-            proc,
-            tma_out,
-            fl_out,
-        ]
-    )
-
-    task = PipelineTask(pipeline)
-    await task.queue_frames(
-        [
-            UserStartedSpeakingFrame(),
-            TranscriptionFrame(text="Hi World", user_id="user", timestamp="now"),
-            UserStoppedSpeakingFrame(),
-            StopTaskFrame(),
-        ]
-    )
-
-    runner = PipelineRunner()
-    await runner.run(task)
+if __name__ == "__main__":
+    unittest.main()