Comments

Author: stellasia

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -85,6 +85,8 @@ def get_messages(
             if i["role"] == "system":
                 system_instruction = i["content"]
             else:
+                if i["role"] not in ("user", "assistant"):
+                    raise ValueError(f"Unknown role: {i['role']}")
                 messages.append(
                     self.anthropic.types.MessageParam(
                         role=i["role"],

src/neo4j_graphrag/llm/base.py

@@ -159,6 +159,7 @@ def _invoke_with_tools(
     ) -> ToolCallResponse:
         raise NotImplementedError("This LLM provider does not support tool calling.")
 
+    @async_rate_limit_handler
     async def ainvoke_with_tools(
         self,
         input: str,

src/neo4j_graphrag/llm/cohere_llm.py

@@ -81,12 +81,14 @@ def get_messages(
         for i in input:
             if i["role"] == "system":
                 messages.append(self.cohere.SystemChatMessageV2(content=i["content"]))
-            if i["role"] == "user":
+            elif i["role"] == "user":
                 messages.append(self.cohere.UserChatMessageV2(content=i["content"]))
-            if i["role"] == "assistant":
+            elif i["role"] == "assistant":
                 messages.append(
                     self.cohere.AssistantChatMessageV2(content=i["content"])
                 )
+            else:
+                raise ValueError(f"Unknown role: {i['role']}")
         return messages
 
     def _invoke(

src/neo4j_graphrag/llm/mistralai_llm.py

@@ -86,6 +86,7 @@ def get_messages(
             if m["role"] == "assistant":
                 messages.append(AssistantMessage(content=m["content"]))
                 continue
+            raise ValueError(f"Unknown role: {m['role']}")
         return messages
 
     def _invoke(

src/neo4j_graphrag/llm/vertexai_llm.py

@@ -97,25 +97,27 @@ def get_messages(
         messages = []
         system_instruction = self.system_instruction
         for message in input:
-            if message.get("role") == "system":
+            role = message.get("role")
+            if role == "system":
                 system_instruction = message.get("content")
                 continue
-            if message.get("role") == "user":
+            if role == "user":
                 messages.append(
                     Content(
                         role="user",
                         parts=[Part.from_text(message.get("content", ""))],
                     )
                 )
                 continue
-            if message.get("role") == "assistant":
+            if role == "assistant":
                 messages.append(
                     Content(
                         role="model",
                         parts=[Part.from_text(message.get("content", ""))],
                     )
                 )
                 continue
+            raise ValueError(f"Unknown role: {role}")
         return system_instruction, messages
 
     def _invoke(

tests/unit/llm/test_anthropic_llm.py

@@ -83,6 +83,15 @@ def test_anthropic_llm_get_messages_without_system_instructions() -> None:
         assert actual["content"] == expected["content"]
 
 
+def test_anthropic_llm_get_messages_unknown_role() -> None:
+    llm = AnthropicLLM(api_key="my key", model_name="claude")
+    message_history = [
+        LLMMessage(**{"role": "unknown role", "content": "Usually around 6am."}),  # type: ignore[typeddict-item]
+    ]
+    with pytest.raises(ValueError, match="Unknown role"):
+        llm.get_messages(message_history)
+
+
 def test_anthropic_invoke_happy_path(mock_anthropic: Mock) -> None:
     mock_anthropic.Anthropic.return_value.messages.create.return_value = MagicMock(
         content=[MagicMock(text="generated text")]

tests/unit/llm/test_vertexai_llm.py

@@ -16,7 +16,6 @@
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
-from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm.types import ToolCallResponse
 from neo4j_graphrag.llm.vertexai_llm import VertexAILLM
 from neo4j_graphrag.tool import Tool
@@ -89,16 +88,14 @@ def test_vertexai_get_messages(GenerativeModelMock: MagicMock) -> None:
 
 @patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel")
 def test_vertexai_get_messages_validation_error(GenerativeModelMock: MagicMock) -> None:
-    system_instruction = "You are a helpful assistant."
     model_name = "gemini-1.5-flash-001"
-    question = "hi!"
     message_history = [
         LLMMessage(**{"role": "model", "content": "hello!"}),  # type: ignore[typeddict-item]
     ]
 
-    llm = VertexAILLM(model_name=model_name, system_instruction=system_instruction)
-    with pytest.raises(LLMGenerationError, match="Input validation failed"):
-        llm.invoke(question, message_history)
+    llm = VertexAILLM(model_name=model_name)
+    with pytest.raises(ValueError, match="Unknown role"):
+        llm.get_messages(message_history)
 
 
 @pytest.mark.asyncio