fix(haystack): ensure compatibility with haystack 2.9 (#1205)

Co-authored-by: Alexander Song <axiomofjoy@gmail.com>
Arize-ai · Jan 17, 2025 · 6ee2ebf · 6ee2ebf
1 parent f5ff3a0
commit 6ee2ebf
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 37 deletions.
diff --git a/python/instrumentation/openinference-instrumentation-haystack/pyproject.toml b/python/instrumentation/openinference-instrumentation-haystack/pyproject.toml
@@ -8,7 +8,7 @@ dynamic = ["version"]
 description = "OpenInference Haystack Instrumentation"
 readme = "README.md"
 license = "Apache-2.0"
-requires-python = ">=3.8, <3.13"
+requires-python = ">=3.9, <3.13"
 authors = [
     { name = "OpenInference Authors", email = "oss@arize.com" },
 ]
@@ -37,14 +37,14 @@ dependencies = [
 
 [project.optional-dependencies]
 instruments = [
-    "haystack-ai >= 2.0.0",
+    "haystack-ai >= 2.9.0",
 ]
 test = [
-    "haystack-ai==2.3.1",
+    "haystack-ai==2.9.0",
     "cohere-haystack",
     "opentelemetry-sdk",
+    "vcrpy>=7.0.0",
     "pytest-recording",
-    "respx",
 ]
 
 [project.urls]

diff --git a/...inference-instrumentation-haystack/src/openinference/instrumentation/haystack/__init__.py b/...inference-instrumentation-haystack/src/openinference/instrumentation/haystack/__init__.py
@@ -13,7 +13,7 @@
 
 logger = logging.getLogger(__name__)
 
-_instruments = ("haystack-ai >= 2.0.0",)
+_instruments = ("haystack-ai >= 2.9.0",)
 
 
 class HaystackInstrumentor(BaseInstrumentor):  # type: ignore[misc]

diff --git a/...nference-instrumentation-haystack/src/openinference/instrumentation/haystack/_wrappers.py b/...nference-instrumentation-haystack/src/openinference/instrumentation/haystack/_wrappers.py
@@ -1,4 +1,3 @@
-import json
 from abc import ABC
 from enum import Enum, auto
 from inspect import BoundArguments, Parameter, signature
@@ -380,7 +379,7 @@ def _get_llm_input_message_attributes(arguments: Mapping[str, Any]) -> Iterator[
         map(lambda x: isinstance(x, ChatMessage), messages)
     ):
         for message_index, message in enumerate(messages):
-            if (content := message.content) is not None:
+            if (content := message.text) is not None:
                 yield f"{LLM_INPUT_MESSAGES}.{message_index}.{MESSAGE_CONTENT}", content
             if (role := message.role) is not None:
                 yield f"{LLM_INPUT_MESSAGES}.{message_index}.{MESSAGE_ROLE}", role
@@ -408,25 +407,20 @@ def _get_llm_output_message_attributes(response: Mapping[str, Any]) -> Iterator[
             ):
                 continue
             if finish_reason == "tool_calls":
-                try:
-                    tool_calls = json.loads(reply.content)
-                except json.JSONDecodeError:
-                    continue
+                tool_calls = reply.tool_calls
                 for tool_call_index, tool_call in enumerate(tool_calls):
-                    if (function_call := tool_call.get("function")) is None:
-                        continue
-                    if (tool_call_arguments_json := function_call.get("arguments")) is not None:
+                    if (tool_call_arguments := tool_call.arguments) is not None:
                         yield (
                             f"{LLM_OUTPUT_MESSAGES}.{reply_index}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
-                            tool_call_arguments_json,
+                            safe_json_dumps(tool_call_arguments),
                         )
-                    if (tool_name := function_call.get("name")) is not None:
+                    if (tool_name := tool_call.tool_name) is not None:
                         yield (
                             f"{LLM_OUTPUT_MESSAGES}.{reply_index}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
                             tool_name,
                         )
             else:
-                yield f"{LLM_OUTPUT_MESSAGES}.{reply_index}.{MESSAGE_CONTENT}", reply.content
+                yield f"{LLM_OUTPUT_MESSAGES}.{reply_index}.{MESSAGE_CONTENT}", reply.text
             yield f"{LLM_OUTPUT_MESSAGES}.{reply_index}.{MESSAGE_ROLE}", reply.role.value
         elif isinstance(reply, str):
             yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", reply