fix: bad type conversion (Chainlit#1693)

* fix: bad type conversion
Progressive-Insurance · Jan 13, 2025 · 4f67278 · 4f67278
1 parent 534b73a
commit 4f67278
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/backend/chainlit/data/literalai.py b/backend/chainlit/data/literalai.py
@@ -39,8 +39,6 @@
 class LiteralToChainlitConverter:
     @classmethod
     def steptype_to_steptype(cls, step_type: Optional[StepType]) -> TrueStepType:
-        if step_type in ["user_message", "assistant_message", "system_message"]:
-            return "undefined"
         return cast(TrueStepType, step_type or "undefined")
 
     @classmethod

diff --git a/backend/tests/data/test_literalai.py b/backend/tests/data/test_literalai.py
@@ -777,14 +777,16 @@ async def test_update_step(
 
 def test_steptype_to_steptype():
     assert (
-        LiteralToChainlitConverter.steptype_to_steptype("user_message") == "undefined"
+        LiteralToChainlitConverter.steptype_to_steptype("user_message")
+        == "user_message"
     )
     assert (
         LiteralToChainlitConverter.steptype_to_steptype("assistant_message")
-        == "undefined"
+        == "assistant_message"
     )
     assert (
-        LiteralToChainlitConverter.steptype_to_steptype("system_message") == "undefined"
+        LiteralToChainlitConverter.steptype_to_steptype("system_message")
+        == "system_message"
     )
     assert LiteralToChainlitConverter.steptype_to_steptype("tool") == "tool"
     assert LiteralToChainlitConverter.steptype_to_steptype(None) == "undefined"
@@ -990,7 +992,7 @@ def test_step_to_step():
     assert isinstance(chainlit_step, Step)
     assert chainlit_step.id == "test_step_id"
     assert chainlit_step.thread_id == "test_thread_id"
-    assert chainlit_step.type == "undefined"
+    assert chainlit_step.type == "user_message"
     assert chainlit_step.name == "Test Step"
     assert chainlit_step.input == "test input"
     assert chainlit_step.output == "test output"