|
|
| from __future__ import annotations
|
|
|
| import json
|
| import os
|
| import unittest
|
| from unittest.mock import patch
|
|
|
| os.environ.setdefault("BORDERLESS_INFERENCE_MODE", "hub")
|
| os.environ.setdefault("BORDERLESS_PRELOAD_MODEL", "0")
|
|
|
| from langchain_core.messages import AIMessage, ToolMessage
|
|
|
| from ui.agent.graph.nodes.helpers import (
|
| extract_assistant_text,
|
| extract_thinking_text,
|
| research_tool_calls,
|
| )
|
| from ui.agent.synthesis import synthesize_finding_from_tool_messages
|
|
|
|
|
| class ExtractAssistantTextTests(unittest.TestCase):
|
| def test_plain_string_content(self) -> None:
|
| message = AIMessage(content="## Eligibility\nYou may qualify.")
|
| self.assertEqual(extract_assistant_text(message), "## Eligibility\nYou may qualify.")
|
|
|
| def test_list_content_blocks(self) -> None:
|
| message = AIMessage(
|
| content=[
|
| {"type": "text", "text": "Line one"},
|
| {"type": "text", "text": "Line two"},
|
| ]
|
| )
|
| self.assertEqual(extract_assistant_text(message), "Line one\nLine two")
|
|
|
| def test_reasoning_only_message(self) -> None:
|
| message = AIMessage(content="", additional_kwargs={"reasoning_content": "Final report"})
|
| self.assertEqual(extract_assistant_text(message), "Final report")
|
|
|
| def test_strips_think_wrappers(self) -> None:
|
| open_tag = "<" + "think" + ">"
|
| close_tag = "</" + "think" + ">"
|
| message = AIMessage(
|
| content=(
|
| f"{open_tag}internal reasoning{close_tag}\n\n"
|
| "## Findings\nCanada Express Entry."
|
| )
|
| )
|
| extracted = extract_assistant_text(message)
|
| self.assertIn("## Findings", extracted)
|
| self.assertNotIn("internal reasoning", extracted)
|
|
|
| def test_prefers_content_over_reasoning(self) -> None:
|
| message = AIMessage(
|
| content="Visible answer",
|
| additional_kwargs={"reasoning_content": "Hidden reasoning"},
|
| )
|
| self.assertEqual(extract_assistant_text(message), "Visible answer")
|
| self.assertEqual(extract_thinking_text(message), "Hidden reasoning")
|
|
|
| def test_extract_thinking_text_from_think_tags(self) -> None:
|
| open_tag = "<" + "think" + ">"
|
| close_tag = "</" + "think" + ">"
|
| message = AIMessage(
|
| content=(
|
| f"{open_tag}Compare Canada and Germany pathways.{close_tag}\n\n"
|
| "## Recommended Countries\n- Canada"
|
| )
|
| )
|
| self.assertIn("Canada", extract_assistant_text(message))
|
| self.assertEqual(
|
| extract_thinking_text(message),
|
| "Compare Canada and Germany pathways.",
|
| )
|
|
|
|
|
| class ResearchToolCallsTests(unittest.TestCase):
|
| def test_parses_tool_calls_from_reasoning(self) -> None:
|
| reasoning = 'search_immigration_info{"query": "Canada Express Entry"}'
|
| message = AIMessage(content="", additional_kwargs={"reasoning_content": reasoning})
|
| calls = research_tool_calls(message)
|
| self.assertEqual(len(calls), 1)
|
| self.assertEqual(calls[0][0], "search_immigration_info")
|
| self.assertEqual(calls[0][1]["query"], "Canada Express Entry")
|
|
|
|
|
| class SynthesizeFindingTests(unittest.TestCase):
|
| def test_builds_fallback_from_search_results(self) -> None:
|
| todo = {"id": 1, "country": "Canada", "methods": "Express Entry"}
|
| payload = {
|
| "query": "Canada Express Entry",
|
| "num_results": 1,
|
| "results": [
|
| {
|
| "title": "Express Entry",
|
| "url": "https://www.canada.ca/en/immigration/express-entry.html",
|
| "highlights": ["Skilled workers may apply through Express Entry."],
|
| }
|
| ],
|
| }
|
| tool_messages = [ToolMessage(content=json.dumps(payload), tool_call_id="abc")]
|
| summary = synthesize_finding_from_tool_messages(todo, tool_messages)
|
| self.assertIn("Canada", summary)
|
| self.assertIn("Express Entry", summary)
|
| self.assertIn("canada.ca", summary)
|
| self.assertNotEqual(summary, "")
|
|
|
| def test_returns_empty_without_tool_evidence(self) -> None:
|
| todo = {"id": 1, "country": "Canada", "methods": "Express Entry"}
|
| self.assertEqual(
|
| synthesize_finding_from_tool_messages(todo, []),
|
| "",
|
| )
|
|
|
|
|
| class ResearcherFallbackTests(unittest.TestCase):
|
| @patch("ui.agent.graph.nodes.researcher.build_llm")
|
| @patch("ui.agent.graph.nodes.researcher.get_stream_writer")
|
| @patch("ui.agent.graph.nodes.researcher.execute_tool_call")
|
| def test_uses_tool_fallback_when_llm_summary_empty(
|
| self,
|
| mock_execute_tool_call: unittest.mock.Mock,
|
| mock_get_stream_writer: unittest.mock.Mock,
|
| mock_build_llm: unittest.mock.Mock,
|
| ) -> None:
|
| from ui.agent.graph.nodes.researcher import researcher_node
|
|
|
| mock_get_stream_writer.return_value = lambda _event: None
|
| mock_execute_tool_call.return_value = json.dumps(
|
| {
|
| "query": "Germany Blue Card",
|
| "num_results": 1,
|
| "results": [
|
| {
|
| "title": "EU Blue Card",
|
| "url": "https://www.make-it-in-germany.com/en/visa-residence/eu-blue-card",
|
| "highlights": ["The EU Blue Card is for qualified professionals."],
|
| }
|
| ],
|
| }
|
| )
|
|
|
| search_response = AIMessage(
|
| content="",
|
| tool_calls=[
|
| {
|
| "name": "search_immigration_info",
|
| "args": {"query": "Germany Blue Card"},
|
| "id": "call-1",
|
| }
|
| ],
|
| )
|
| empty_summary = AIMessage(content="")
|
|
|
| mock_llm = unittest.mock.Mock()
|
| mock_bound = unittest.mock.Mock()
|
| mock_bound.invoke.side_effect = [search_response, empty_summary]
|
| mock_llm.bind_tools.return_value = mock_bound
|
| mock_build_llm.return_value = mock_llm
|
|
|
| task = {
|
| "todo": {"id": 2, "country": "Germany", "methods": "EU Blue Card"},
|
| "profile_summary": "Software engineer from India.",
|
| }
|
| config = {"configurable": {"hf_token": "test-token"}}
|
|
|
| result = researcher_node(task, config)
|
|
|
| summary = result["findings"][0]["summary"]
|
| self.assertNotEqual(summary, "No findings could be produced for this to-do.")
|
| self.assertIn("Germany", summary)
|
|
|
|
|
| if __name__ == "__main__":
|
| unittest.main()
|
|
|
