|
|
| from __future__ import annotations
|
|
|
| import json
|
| import os
|
| import unittest
|
|
|
| os.environ.setdefault("BORDERLESS_INFERENCE_MODE", "hub")
|
| os.environ.setdefault("BORDERLESS_PRELOAD_MODEL", "0")
|
|
|
| from langchain_core.messages import AIMessage, ToolMessage
|
|
|
| from ui.agent.minicpm.messages import (
|
| append_tool_instructions,
|
| normalize_messages,
|
| )
|
|
|
|
|
| class NormalizeMessagesTests(unittest.TestCase):
|
| def test_dict_messages(self) -> None:
|
| messages = [
|
| {"role": "system", "content": "You are helpful."},
|
| {"role": "user", "content": "Hello"},
|
| ]
|
| normalized = normalize_messages(messages)
|
| self.assertEqual(normalized[0]["role"], "system")
|
| self.assertEqual(normalized[1]["content"], "Hello")
|
|
|
| def test_tool_message_becomes_user_context(self) -> None:
|
| messages = [
|
| ToolMessage(content='{"results": []}', tool_call_id="call-1"),
|
| ]
|
| normalized = normalize_messages(messages)
|
| self.assertEqual(normalized[0]["role"], "user")
|
| self.assertIn("Tool result", normalized[0]["content"])
|
|
|
| def test_ai_message(self) -> None:
|
| messages = [AIMessage(content="Done.")]
|
| normalized = normalize_messages(messages)
|
| self.assertEqual(normalized[0]["role"], "assistant")
|
| self.assertEqual(normalized[0]["content"], "Done.")
|
|
|
|
|
| class ToolInstructionTests(unittest.TestCase):
|
| def test_appends_tool_block_to_system(self) -> None:
|
| tools = [
|
| {
|
| "type": "function",
|
| "function": {
|
| "name": "search_immigration_info",
|
| "description": "Search immigration sources",
|
| "parameters": {
|
| "type": "object",
|
| "properties": {"query": {"type": "string"}},
|
| },
|
| },
|
| }
|
| ]
|
| messages = [
|
| {"role": "system", "content": "Base prompt"},
|
| {"role": "user", "content": "Research Canada"},
|
| ]
|
| updated = append_tool_instructions(messages, tools)
|
| self.assertIn("search_immigration_info", updated[0]["content"])
|
| self.assertIn(json.dumps(tools[0]["function"]["parameters"]), updated[0]["content"])
|
|
|
|
|
| if __name__ == "__main__":
|
| unittest.main()
|
|
|
