Fixed Message type

Future-House · Dec 18, 2024 · 31419c8 · 31419c8
1 parent 420203b
commit 31419c8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/ldp/alg/optimizer/ape.py b/ldp/alg/optimizer/ape.py
@@ -194,7 +194,9 @@ def aggregate_trajectory(self, trajectory: Trajectory) -> None:
                     continue
                 # (x: first prompt's user message's content, y: AI response's content)
                 x = next(
-                    cast(str, m.content) for m in result.prompt if m.role == "user"
+                    cast(str, m.content)
+                    for m in result.prompt
+                    if (isinstance(m, Message) and m.role == "user")
                 )
                 y = cast(str, result.messages[0].content)