Fix tests

Author: oskarhane

src/neo4j_graphrag/experimental/pipeline/pipeline.py

@@ -249,8 +249,8 @@ def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph:
                     caption=f"{node.component.__class__.__name__}: {n}({comp_inputs})",
                     size=20,  # Component nodes are larger
                     color="#4C8BF5",  # Blue for component nodes
-                    caption_align=CaptionAlignment.CENTER,
-                    caption_size=12,
+                    caption_alignment=CaptionAlignment.CENTER,
+                    caption_size=3,
                 )
             )
             node_counter += 1
@@ -265,8 +265,8 @@ def get_neo4j_viz_graph(self, hide_unused_outputs: bool = True) -> NeoVizGraph:
                         caption=o,
                         size=10,  # Output nodes are smaller
                         color="#34A853",  # Green for output nodes
-                        caption_align=CaptionAlignment.CENTER,
-                        caption_size=10,
+                        caption_alignment=CaptionAlignment.CENTER,
+                        caption_size=3,
                     )
                 )
                 # Connect component to its output

src/neo4j_graphrag/llm/anthropic_llm.py

@@ -112,8 +112,10 @@ def invoke(
         """
         try:
             if isinstance(message_history, MessageHistory):
-                message_history = message_history.messages
-            messages = self.get_messages(input, message_history)
+                message_history.add_message(LLMMessage(role="user", content=input))
+                messages = message_history
+            else:
+                messages = self.get_messages(input, message_history)
             response = self.client.messages.create(
                 model=self.model_name,
                 system=system_instruction or self.anthropic.NOT_GIVEN,
@@ -148,8 +150,10 @@ async def ainvoke(
         """
         try:
             if isinstance(message_history, MessageHistory):
-                message_history = message_history.messages
-            messages = self.get_messages(input, message_history)
+                message_history.add_message(LLMMessage(role="user", content=input))
+                messages = message_history
+            else:
+                messages = self.get_messages(input, message_history)
             response = await self.async_client.messages.create(
                 model=self.model_name,
                 system=system_instruction or self.anthropic.NOT_GIVEN,