File size: 2,293 Bytes
b1198f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import sys
import unittest
from pathlib import Path
from langchain_core.messages import AIMessage, HumanMessage
ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from src.agents.agents import DietarySpecialist, OutputMerger
class StubLLM:
def __init__(self, content):
self.content = content
self.calls = []
async def ainvoke(self, messages, config=None):
self.calls.append(messages)
return AIMessage(content=self.content)
async def astream(self, messages, config=None):
self.calls.append(messages)
yield AIMessage(content=self.content)
class TestChapter6(unittest.IsolatedAsyncioTestCase):
async def test_output_merger_uses_latest_specialist_output(self):
llm = StubLLM("merged response")
merger = OutputMerger()
merger.llm = llm
state = {
"messages": [
HumanMessage(content="What should I do next?"),
AIMessage(content="Older specialist response"),
AIMessage(content="Older user-facing answer"),
],
"clinician_outputs": [
"Diagnosis says continue monitoring.",
"Treatment says adjust medication dose.",
],
}
result = await merger.run(state)
self.assertEqual(result["messages"][-1].content, "merged response")
self.assertEqual(result["metrics"][0]["agent"], "OutputMerger")
merged_prompt = "\n".join(
message.content for message in llm.calls[-1]
)
self.assertIn("Treatment says adjust medication dose.", merged_prompt)
self.assertNotIn("Older specialist response", merged_prompt)
async def test_dietary_specialist_returns_metrics(self):
llm = StubLLM("dietary guidance")
specialist = DietarySpecialist()
specialist.llm = llm
state = {
"messages": [HumanMessage(content="Suggest a low-carb meal plan for today.")]
}
result = await specialist.run(state)
self.assertEqual(result["messages"][-1].content, "dietary guidance")
self.assertEqual(result["metrics"][0]["agent"], "DietarySpecialist")
if __name__ == "__main__":
unittest.main()
|