Merge pull request #40 from lakeraai/fix-tests

Fix multi-message behavior
lakeraai · Jan 13, 2025 · 13a4f56 · 13a4f56
2 parents d308ac4 + cbf81a2
commit 13a4f56
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/lakera_chainguard/lakera_chainguard.py b/lakera_chainguard/lakera_chainguard.py
@@ -171,26 +171,25 @@ def _convert_to_lakera_guard_input(
             if isinstance(prompt, PromptValue):
                 prompt = prompt.to_messages()
             if isinstance(prompt, List):
-                formatted_input = [
-                    {"role": "system", "content": ""},
-                    {"role": "user", "content": ""},
-                    {"role": "assistant", "content": ""},
-                ]
-                # For system, human, assistant, we put the last message of each
-                # type in the guard input
+                user_message = ""
+                formatted_input = []
                 for message in prompt:
                     if not isinstance(
                         message, (HumanMessage, SystemMessage, AIMessage)
                     ) or not isinstance(message.content, str):
                         raise TypeError("Input type not supported by Lakera Guard.")
+
+                    role = "assistant"
                     if isinstance(message, SystemMessage):
-                        formatted_input[0]["content"] = message.content
+                        role = "system"
                     elif isinstance(message, HumanMessage):
-                        formatted_input[1]["content"] = message.content
-                    else:  # must be AIMessage
-                        formatted_input[2]["content"] = message.content
+                        user_message = message.content
+                        role = "user"
+
+                    formatted_input.append({"role": role, "content": message.content})
+
                 if self.endpoint != "prompt_injection":
-                    return formatted_input[1]["content"]
+                    return user_message
                 return formatted_input
             else:
                 return str(prompt)

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lakera-chainguard"
-version = "0.1.4"
+version = "0.1.5"
 description = "Guard your LangChain applications against prompt injection with Lakera Guard"
 homepage = "https://lakeraai.github.io/chainguard/"
 repository = "https://github.com/lakeraai/chainguard"