支持输入dict创建json输出

baidubce · Jan 22, 2025 · d435512 · d435512
1 parent 4289098
commit d435512
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/python/core/component.py b/python/core/component.py
@@ -559,6 +559,8 @@ def create_output(cls, type, text, role="tool", name="", visible_scope="all", ra
                 key_list = ["detail", "steps"]
             elif type == "function_call":
                 key_list = ["thought", "name", "arguments"]
+            elif type == "json":
+                key_list = ["data"]
             else:
                 raise ValueError("Unknown type: {}".format(type))
             assert all(key in text for key in key_list), "all keys:{} must be included in the text field".format(key_list)

diff --git a/python/tests/test_base_component.py b/python/tests/test_base_component.py
@@ -33,6 +33,7 @@ def test_valid_output_with_dict(self):
         output9 = self.component.create_output(type="plan", text={"detail": "hello", "steps":[{"name": "1", "arguments": {"query": "a", "chat_history": "world"}}]})
         output10 = self.component.create_output(type="function_call", text={"thought": "hello", "name": "AppBuilder", "arguments": {"query": "a", "chat_history": "world"}})
         output11 = self.component.create_output(type="references", text={"type": "engine", "doc_id": "1", "content": "hello, world", "title": "Have a nice day", "source": "bing", "extra": {"key": "value"}})
+        output12 = self.component.create_output(type="json", text={"key": "value"})
         self.assertIsInstance(output1, ComponentOutput)
         self.assertIsInstance(output2, ComponentOutput)
         self.assertIsInstance(output3, ComponentOutput)
@@ -44,6 +45,7 @@ def test_valid_output_with_dict(self):
         self.assertIsInstance(output9, ComponentOutput)
         self.assertIsInstance(output10, ComponentOutput)
         self.assertIsInstance(output11, ComponentOutput)
+        self.assertIsInstance(output12, ComponentOutput)
         self.assertEqual(output11.content[0].text.extra["key"], "value")
 
     def test_valid_output_type_with_same_key(self):