langfuse-haystack 0.0.4__tar.gz → 0.1.0__tar.gz

{langfuse_haystack-0.0.4 → langfuse_haystack-0.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: langfuse-haystack
-Version: 0.0.4
+Version: 0.1.0
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/langfuse#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
 Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/langfuse
@@ -17,7 +17,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.8
-Requires-Dist: haystack-ai
+Requires-Dist: haystack-ai>=2.1.0
 Requires-Dist: langfuse
 Description-Content-Type: text/markdown
 

{langfuse_haystack-0.0.4 → langfuse_haystack-0.1.0}/pydoc/config.yml

@@ -15,12 +15,12 @@ processors:
   - type: smart
   - type: crossref
 renderer:
-  type: haystack_pydoc_tools.renderers.ReadmePreviewRenderer
+  type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer
   excerpt: Langfuse integration for Haystack
   category_slug: integrations-api
   title: langfuse
   slug: integrations-langfuse
-  order: 135
+  order: 136
   markdown:
     descriptive_class_title: false
     classdef_code_block: false

{langfuse_haystack-0.0.4 → langfuse_haystack-0.1.0}/pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dependencies = [
-  "haystack-ai",
+  "haystack-ai>=2.1.0",
   "langfuse"
 ]
 

{langfuse_haystack-0.0.4 → langfuse_haystack-0.1.0}/src/haystack_integrations/components/connectors/langfuse/langfuse_connector.py

@@ -20,6 +20,31 @@ class LangfuseConnector:
     In addition, you need to set the `HAYSTACK_CONTENT_TRACING_ENABLED` environment variable to `true` in order to
     enable Haystack tracing in your pipeline.
 
+    Lastly, you may disable flushing the data after each component by setting the `HAYSTACK_LANGFUSE_ENFORCE_FLUSH`
+    environent variable to `false`. By default, the data is flushed after each component and blocks the thread until
+    the data is sent to Langfuse. **Caution**: Disabling this feature may result in data loss if the program crashes
+    before the data is sent to Langfuse. Make sure you will call langfuse.flush() explicitly before the program exits.
+    E.g. by using tracer.actual_tracer.flush():
+
+    ```python
+    from haystack.tracing import tracer
+
+    try:
+        # your code here
+    finally:
+        tracer.actual_tracer.flush()
+    ```
+    or in FastAPI by defining a shutdown event handler:
+    ```python
+    from haystack.tracing import tracer
+
+    # ...
+
+    @app.on_event("shutdown")
+    async def shutdown_event():
+        tracer.actual_tracer.flush()
+    ```
+
     Here is an example of how to use it:
 
     ```python

{langfuse_haystack-0.0.4 → langfuse_haystack-0.1.0}/src/haystack_integrations/tracing/langfuse/tracer.py

@@ -1,4 +1,5 @@
 import contextlib
+import os
 from typing import Any, Dict, Iterator, Optional, Union
 
 from haystack.dataclasses import ChatMessage
@@ -7,31 +8,54 @@ from haystack.tracing import utils as tracing_utils
 
 import langfuse
 
+HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR = "HAYSTACK_LANGFUSE_ENFORCE_FLUSH"
+
 
 class LangfuseSpan(Span):
+    """
+    Internal class representing a bridge between the Haystack span tracing API and Langfuse.
+    """
+
     def __init__(self, span: "Union[langfuse.client.StatefulSpanClient, langfuse.client.StatefulTraceClient]") -> None:
+        """
+        Initialize a LangfuseSpan instance.
+
+        :param span: The span instance managed by Langfuse.
+        """
         self._span = span
         # locally cache tags
         self._data: Dict[str, Any] = {}
 
     def set_tag(self, key: str, value: Any) -> None:
+        """
+        Set a generic tag for this span.
+
+        :param key: The tag key.
+        :param value: The tag value.
+        """
         coerced_value = tracing_utils.coerce_tag_value(value)
         self._span.update(metadata={key: coerced_value})
         self._data[key] = value
 
     def set_content_tag(self, key: str, value: Any) -> None:
+        """
+        Set a content-specific tag for this span.
+
+        :param key: The content tag key.
+        :param value: The content tag value.
+        """
         if not tracer.is_content_tracing_enabled:
             return
         if key.endswith(".input"):
             if "messages" in value:
-                messages = [self.to_openai_format(m) for m in value["messages"]]
+                messages = [m.to_openai_format() for m in value["messages"]]
                 self._span.update(input=messages)
             else:
                 self._span.update(input=value)
         elif key.endswith(".output"):
             if "replies" in value:
                 if all(isinstance(r, ChatMessage) for r in value["replies"]):
-                    replies = [self.to_openai_format(m) for m in value["replies"]]
+                    replies = [m.to_openai_format() for m in value["replies"]]
                 else:
                     replies = value["replies"]
                 self._span.update(output=replies)
@@ -41,31 +65,47 @@ class LangfuseSpan(Span):
         self._data[key] = value
 
     def raw_span(self) -> Any:
+        """
+        Return the underlying span instance.
+
+        :return: The Langfuse span instance.
+        """
         return self._span
 
     def get_correlation_data_for_logs(self) -> Dict[str, Any]:
         return {}
 
-    def to_openai_format(self, m: ChatMessage) -> Dict[str, Any]:
-        """
-        Remove after haystack 2.0.1 has been released and use the `to_openai_format` method from the ChatMessage class
-        """
-        msg = {"role": m.role.value, "content": m.content}
-        if m.name:
-            msg["name"] = m.name
-
-        return msg
-
 
 class LangfuseTracer(Tracer):
+    """
+    Internal class representing a bridge between the Haystack tracer and Langfuse.
+    """
+
     def __init__(self, tracer: "langfuse.Langfuse", name: str = "Haystack", public: bool = False) -> None:
+        """
+        Initialize a LangfuseTracer instance.
+
+        :param tracer: The Langfuse tracer instance.
+        :param name: The name of the pipeline or component. This name will be used to identify the tracing run on the
+            Langfuse dashboard.
+        :param public: Whether the tracing data should be public or private. If set to `True`, the tracing data will
+        be publicly accessible to anyone with the tracing URL. If set to `False`, the tracing data will be private
+        and only accessible to the Langfuse account owner.
+        """
         self._tracer = tracer
         self._context: list[LangfuseSpan] = []
         self._name = name
         self._public = public
+        self.enforce_flush = os.getenv(HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR, "true").lower() == "true"
 
     @contextlib.contextmanager
     def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> Iterator[Span]:
+        """
+        Start and manage a new trace span.
+        :param operation_name: The name of the operation.
+        :param tags: A dictionary of tags to attach to the span.
+        :return: A context manager yielding the span.
+        """
         tags = tags or {}
         span_name = tags.get("haystack.component.name", operation_name)
 
@@ -85,12 +125,12 @@ class LangfuseTracer(Tracer):
                 # Haystack returns one meta dict for each message, but the 'usage' value
                 # is always the same, let's just pick the first item
                 m = meta[0]
-                span._span.update(usage=m.get("usage"), model=m.get("model"))
+                span._span.update(usage=m.get("usage") or None, model=m.get("model"))
         elif tags.get("haystack.component.type") == "OpenAIChatGenerator":
             replies = span._data.get("haystack.component.output", {}).get("replies")
             if replies:
                 meta = replies[0].meta
-                span._span.update(usage=meta.get("usage"), model=meta.get("model"))
+                span._span.update(usage=meta.get("usage") or None, model=meta.get("model"))
 
         pipeline_input = tags.get("haystack.pipeline.input_data", None)
         if pipeline_input:
@@ -101,13 +141,31 @@ class LangfuseTracer(Tracer):
 
         span.raw_span().end()
         self._context.pop()
+
+        if len(self._context) == 1:
+            # The root span has to be a trace, which need to be removed from the context after the pipeline run
+            self._context.pop()
+
+            if self.enforce_flush:
+                self.flush()
+
+    def flush(self):
         self._tracer.flush()
 
     def current_span(self) -> Span:
+        """
+        Return the currently active span.
+
+        :return: The currently active span.
+        """
         if not self._context:
             # The root span has to be a trace
             self._context.append(LangfuseSpan(self._tracer.trace(name=self._name, public=self._public)))
         return self._context[-1]
 
     def get_trace_url(self) -> str:
+        """
+        Return the URL to the tracing data.
+        :return: The URL to the tracing data.
+        """
         return self._tracer.get_trace_url()

langfuse_haystack-0.1.0/tests/test_langfuse_span.py

@@ -0,0 +1,65 @@
+import os
+
+os.environ["HAYSTACK_CONTENT_TRACING_ENABLED"] = "true"
+
+from unittest.mock import Mock
+from haystack.dataclasses import ChatMessage
+from haystack_integrations.tracing.langfuse.tracer import LangfuseSpan
+
+
+class TestLangfuseSpan:
+
+    #  LangfuseSpan can be initialized with a span object
+    def test_initialized_with_span_object(self):
+        mock_span = Mock()
+        span = LangfuseSpan(mock_span)
+        assert span.raw_span() == mock_span
+
+    #  set_tag method can update metadata of the span object
+    def test_set_tag_updates_metadata(self):
+        mock_span = Mock()
+        span = LangfuseSpan(mock_span)
+
+        span.set_tag("key", "value")
+        mock_span.update.assert_called_once_with(metadata={"key": "value"})
+        assert span._data["key"] == "value"
+
+    #  set_content_tag method can update input and output of the span object
+    def test_set_content_tag_updates_input_and_output(self):
+        mock_span = Mock()
+
+        span = LangfuseSpan(mock_span)
+        span.set_content_tag("input_key", "input_value")
+        assert span._data["input_key"] == "input_value"
+
+        mock_span.reset_mock()
+        span.set_content_tag("output_key", "output_value")
+        assert span._data["output_key"] == "output_value"
+
+    # set_content_tag method can update input and output of the span object with messages/replies
+    def test_set_content_tag_updates_input_and_output_with_messages(self):
+        mock_span = Mock()
+
+        # test message input
+        span = LangfuseSpan(mock_span)
+        span.set_content_tag("key.input", {"messages": [ChatMessage.from_user("message")]})
+        assert mock_span.update.call_count == 1
+        # check we converted ChatMessage to OpenAI format
+        assert mock_span.update.call_args_list[0][1] == {"input": [{"role": "user", "content": "message"}]}
+        assert span._data["key.input"] == {"messages": [ChatMessage.from_user("message")]}
+
+        # test replies ChatMessage list
+        mock_span.reset_mock()
+        span.set_content_tag("key.output", {"replies": [ChatMessage.from_system("reply")]})
+        assert mock_span.update.call_count == 1
+        # check we converted ChatMessage to OpenAI format
+        assert mock_span.update.call_args_list[0][1] == {"output": [{"role": "system", "content": "reply"}]}
+        assert span._data["key.output"] == {"replies": [ChatMessage.from_system("reply")]}
+
+        # test replies string list
+        mock_span.reset_mock()
+        span.set_content_tag("key.output", {"replies": ["reply1", "reply2"]})
+        assert mock_span.update.call_count == 1
+        # check we handle properly string list replies
+        assert mock_span.update.call_args_list[0][1] == {"output": ["reply1", "reply2"]}
+        assert span._data["key.output"] == {"replies": ["reply1", "reply2"]}

langfuse_haystack-0.1.0/tests/test_tracer.py

@@ -0,0 +1,114 @@
+import os
+from unittest.mock import Mock, MagicMock, patch
+
+from haystack_integrations.tracing.langfuse.tracer import LangfuseTracer
+
+
+class TestLangfuseTracer:
+
+    #  LangfuseTracer can be initialized with a Langfuse instance, a name and a boolean value for public.
+    def test_initialization(self):
+        langfuse_instance = Mock()
+        tracer = LangfuseTracer(tracer=langfuse_instance, name="Haystack", public=True)
+        assert tracer._tracer == langfuse_instance
+        assert tracer._context == []
+        assert tracer._name == "Haystack"
+        assert tracer._public
+
+    # check that the trace method is called on the tracer instance with the provided operation name and tags
+    # check that the span is added to the context and removed after the context manager exits
+    def test_create_new_span(self):
+        mock_raw_span = MagicMock()
+        mock_raw_span.operation_name = "operation_name"
+        mock_raw_span.metadata = {"tag1": "value1", "tag2": "value2"}
+
+        with patch("haystack_integrations.tracing.langfuse.tracer.LangfuseSpan") as MockLangfuseSpan:
+            mock_span_instance = MockLangfuseSpan.return_value
+            mock_span_instance.raw_span.return_value = mock_raw_span
+
+            mock_context_manager = MagicMock()
+            mock_context_manager.__enter__.return_value = mock_span_instance
+
+            mock_tracer = MagicMock()
+            mock_tracer.trace.return_value = mock_context_manager
+
+            tracer = LangfuseTracer(tracer=mock_tracer, name="Haystack", public=False)
+
+            with tracer.trace("operation_name", tags={"tag1": "value1", "tag2": "value2"}) as span:
+                assert len(tracer._context) == 2, "The trace span should have been added to the the root context span"
+                assert span.raw_span().operation_name == "operation_name"
+                assert span.raw_span().metadata == {"tag1": "value1", "tag2": "value2"}
+
+            assert (
+                len(tracer._context) == 0
+            ), "The trace span should have been popped, and the root span is closed as well"
+
+    # check that update method is called on the span instance with the provided key value pairs
+    def test_update_span_with_pipeline_input_output_data(self):
+        class MockTracer:
+
+            def trace(self, name, **kwargs):
+                return MockSpan()
+
+            def flush(self):
+                pass
+
+        class MockSpan:
+            def __init__(self):
+                self._data = {}
+                self._span = self
+                self.operation_name = "operation_name"
+
+            def raw_span(self):
+                return self
+
+            def span(self, name=None):
+                # assert correct operation name passed to the span
+                assert name == "operation_name"
+                return self
+
+            def update(self, **kwargs):
+                self._data.update(kwargs)
+
+            def generation(self, name=None):
+                return self
+
+            def end(self):
+                pass
+
+        tracer = LangfuseTracer(tracer=MockTracer(), name="Haystack", public=False)
+        with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.input_data": "hello"}) as span:
+            assert span.raw_span()._data["metadata"] == {"haystack.pipeline.input_data": "hello"}
+
+        with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.output_data": "bye"}) as span:
+            assert span.raw_span()._data["metadata"] == {"haystack.pipeline.output_data": "bye"}
+
+    def test_update_span_gets_flushed_by_default(self):
+        tracer_mock = Mock()
+
+        tracer = LangfuseTracer(tracer=tracer_mock, name="Haystack", public=False)
+        with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.input_data": "hello"}) as span:
+            pass
+
+        tracer_mock.flush.assert_called_once()
+
+    def test_update_span_flush_disable(self, monkeypatch):
+        monkeypatch.setenv("HAYSTACK_LANGFUSE_ENFORCE_FLUSH", "false")
+        tracer_mock = Mock()
+
+        from haystack_integrations.tracing.langfuse.tracer import LangfuseTracer
+
+        tracer = LangfuseTracer(tracer=tracer_mock, name="Haystack", public=False)
+        with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.input_data": "hello"}) as span:
+            pass
+
+        tracer_mock.flush.assert_not_called()
+
+    def test_context_is_empty_after_tracing(self):
+        tracer_mock = Mock()
+
+        tracer = LangfuseTracer(tracer=tracer_mock, name="Haystack", public=False)
+        with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.input_data": "hello"}) as span:
+            pass
+
+        assert tracer._context == []