langfuse-haystack 0.5.0__tar.gz → 0.6.0__tar.gz

{langfuse_haystack-0.5.0 → langfuse_haystack-0.6.0}/CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## [integrations/langfuse-v0.5.0] - 2024-10-01
+
+### ⚙️ Miscellaneous Tasks
+
+- Update ruff linting scripts and settings (#1105)
+
+### Langfuse
+
+- Add invocation_context to identify traces (#1089)
+
 ## [integrations/langfuse-v0.4.0] - 2024-09-17
 
 ### 🚀 Features

{langfuse_haystack-0.5.0 → langfuse_haystack-0.6.0}/PKG-INFO

@@ -1,12 +1,11 @@
 Metadata-Version: 2.3
 Name: langfuse-haystack
-Version: 0.5.0
+Version: 0.6.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
 Author-email: deepset GmbH <info@deepset.ai>
-License-Expression: Apache-2.0
-License-File: LICENSE.txt
+License: Apache-2.0
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3.8

{langfuse_haystack-0.5.0 → langfuse_haystack-0.6.0}/pyproject.toml

@@ -42,6 +42,7 @@ root = "../.."
 git_describe_command = 'git describe --tags --match="integrations/langfuse-v[0-9]*"'
 
 [tool.hatch.envs.default]
+installer = "uv"
 dependencies = [
   "coverage[toml]>=6.5",
   "pytest",
@@ -64,8 +65,9 @@ docs = ["pydoc-markdown pydoc/config.yml"]
 python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
 
 [tool.hatch.envs.lint]
+installer = "uv"
 detached = true
-dependencies = ["black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
+dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
 
 [tool.hatch.envs.lint.scripts]
 typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"

{langfuse_haystack-0.5.0 → langfuse_haystack-0.6.0}/src/haystack_integrations/tracing/langfuse/tracer.py

@@ -1,7 +1,10 @@
 import contextlib
 import os
-from typing import Any, Dict, Iterator, Optional, Union
+from contextvars import ContextVar
+from datetime import datetime
+from typing import Any, Dict, Iterator, List, Optional, Union
 
+from haystack import logging
 from haystack.components.generators.openai_utils import _convert_message_to_openai_format
 from haystack.dataclasses import ChatMessage
 from haystack.tracing import Span, Tracer, tracer
@@ -9,6 +12,8 @@ from haystack.tracing import utils as tracing_utils
 
 import langfuse
 
+logger = logging.getLogger(__name__)
+
 HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR = "HAYSTACK_LANGFUSE_ENFORCE_FLUSH"
 _SUPPORTED_GENERATORS = [
     "AzureOpenAIGenerator",
@@ -28,6 +33,17 @@ _SUPPORTED_CHAT_GENERATORS = [
 ]
 _ALL_SUPPORTED_GENERATORS = _SUPPORTED_GENERATORS + _SUPPORTED_CHAT_GENERATORS
 
+# These are the keys used by Haystack for traces and span.
+# We keep them here to avoid making typos when using them.
+_PIPELINE_RUN_KEY = "haystack.pipeline.run"
+_COMPONENT_NAME_KEY = "haystack.component.name"
+_COMPONENT_TYPE_KEY = "haystack.component.type"
+_COMPONENT_OUTPUT_KEY = "haystack.component.output"
+
+# Context var used to keep track of tracing related info.
+# This mainly useful for parents spans.
+tracing_context_var: ContextVar[Dict[Any, Any]] = ContextVar("tracing_context", default={})
+
 
 class LangfuseSpan(Span):
     """
@@ -82,7 +98,7 @@ class LangfuseSpan(Span):
 
         self._data[key] = value
 
-    def raw_span(self) -> Any:
+    def raw_span(self) -> "Union[langfuse.client.StatefulSpanClient, langfuse.client.StatefulTraceClient]":
         """
         Return the underlying span instance.
 
@@ -111,75 +127,90 @@ class LangfuseTracer(Tracer):
         and only accessible to the Langfuse account owner.
         """
         self._tracer = tracer
-        self._context: list[LangfuseSpan] = []
+        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.
-        """
+    def trace(
+        self, operation_name: str, tags: Optional[Dict[str, Any]] = None, parent_span: Optional[Span] = None
+    ) -> Iterator[Span]:
         tags = tags or {}
-        span_name = tags.get("haystack.component.name", operation_name)
-
-        if tags.get("haystack.component.type") in _ALL_SUPPORTED_GENERATORS:
-            span = LangfuseSpan(self.current_span().raw_span().generation(name=span_name))
+        span_name = tags.get(_COMPONENT_NAME_KEY, operation_name)
+
+        # Create new span depending whether there's a parent span or not
+        if not parent_span:
+            if operation_name != _PIPELINE_RUN_KEY:
+                logger.warning(
+                    "Creating a new trace without a parent span is not recommended for operation '{operation_name}'.",
+                    operation_name=operation_name,
+                )
+            # Create a new trace if no parent span is provided
+            span = LangfuseSpan(
+                self._tracer.trace(
+                    name=self._name,
+                    public=self._public,
+                    id=tracing_context_var.get().get("trace_id"),
+                    user_id=tracing_context_var.get().get("user_id"),
+                    session_id=tracing_context_var.get().get("session_id"),
+                    tags=tracing_context_var.get().get("tags"),
+                    version=tracing_context_var.get().get("version"),
+                )
+            )
+        elif tags.get(_COMPONENT_TYPE_KEY) in _ALL_SUPPORTED_GENERATORS:
+            span = LangfuseSpan(parent_span.raw_span().generation(name=span_name))
         else:
-            span = LangfuseSpan(self.current_span().raw_span().span(name=span_name))
+            span = LangfuseSpan(parent_span.raw_span().span(name=span_name))
 
         self._context.append(span)
         span.set_tags(tags)
 
         yield span
 
-        if tags.get("haystack.component.type") in _SUPPORTED_GENERATORS:
-            meta = span._data.get("haystack.component.output", {}).get("meta")
+        # Update span metadata based on component type
+        if tags.get(_COMPONENT_TYPE_KEY) in _SUPPORTED_GENERATORS:
+            # Haystack returns one meta dict for each message, but the 'usage' value
+            # is always the same, let's just pick the first item
+            meta = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("meta")
             if meta:
-                # 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") or None, model=m.get("model"))
-        elif tags.get("haystack.component.type") in _SUPPORTED_CHAT_GENERATORS:
-            replies = span._data.get("haystack.component.output", {}).get("replies")
+        elif tags.get(_COMPONENT_TYPE_KEY) in _SUPPORTED_CHAT_GENERATORS:
+            replies = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("replies")
             if replies:
                 meta = replies[0].meta
-                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:
-            span._span.update(input=tags["haystack.pipeline.input_data"])
-        pipeline_output = tags.get("haystack.pipeline.output_data", None)
-        if pipeline_output:
-            span._span.update(output=tags["haystack.pipeline.output_data"])
-
-        span.raw_span().end()
+                completion_start_time = meta.get("completion_start_time")
+                if completion_start_time:
+                    try:
+                        completion_start_time = datetime.fromisoformat(completion_start_time)
+                    except ValueError:
+                        logger.error(f"Failed to parse completion_start_time: {completion_start_time}")
+                        completion_start_time = None
+                span._span.update(
+                    usage=meta.get("usage") or None,
+                    model=meta.get("model"),
+                    completion_start_time=completion_start_time,
+                )
+
+        raw_span = span.raw_span()
+        if isinstance(raw_span, langfuse.client.StatefulSpanClient):
+            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()
+        if self.enforce_flush:
+            self.flush()
 
     def flush(self):
         self._tracer.flush()
 
-    def current_span(self) -> Span:
+    def current_span(self) -> Optional[Span]:
         """
-        Return the currently active span.
+        Return the current active span.
 
-        :return: The currently active span.
+        :return: The current span if available, else None.
         """
-        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]
+        return self._context[-1] if self._context else None
 
     def get_trace_url(self) -> str:
         """

{langfuse_haystack-0.5.0 → langfuse_haystack-0.6.0}/tests/test_tracer.py

@@ -1,9 +1,43 @@
-import os
+import datetime
 from unittest.mock import MagicMock, Mock, patch
 
+from haystack.dataclasses import ChatMessage
 from haystack_integrations.tracing.langfuse.tracer import LangfuseTracer
 
 
+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
+
+
+class MockTracer:
+
+    def trace(self, name, **kwargs):
+        return MockSpan()
+
+    def flush(self):
+        pass
+
+
 class TestLangfuseTracer:
 
     #  LangfuseTracer can be initialized with a Langfuse instance, a name and a boolean value for public.
@@ -35,7 +69,7 @@ class TestLangfuseTracer:
             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 len(tracer._context) == 1, "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"}
 
@@ -45,37 +79,6 @@ class TestLangfuseTracer:
 
     # 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"}
@@ -83,6 +86,40 @@ class TestLangfuseTracer:
         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_trace_generation(self):
+        tracer = LangfuseTracer(tracer=MockTracer(), name="Haystack", public=False)
+        tags = {
+            "haystack.component.type": "OpenAIChatGenerator",
+            "haystack.component.output": {
+                "replies": [
+                    ChatMessage.from_assistant(
+                        "", meta={"completion_start_time": "2021-07-27T16:02:08.012345", "model": "test_model"}
+                    )
+                ]
+            },
+        }
+        with tracer.trace(operation_name="operation_name", tags=tags) as span:
+            ...
+        assert span.raw_span()._data["usage"] is None
+        assert span.raw_span()._data["model"] == "test_model"
+        assert span.raw_span()._data["completion_start_time"] == datetime.datetime(2021, 7, 27, 16, 2, 8, 12345)
+
+    def test_trace_generation_invalid_start_time(self):
+        tracer = LangfuseTracer(tracer=MockTracer(), name="Haystack", public=False)
+        tags = {
+            "haystack.component.type": "OpenAIChatGenerator",
+            "haystack.component.output": {
+                "replies": [
+                    ChatMessage.from_assistant("", meta={"completion_start_time": "foobar", "model": "test_model"}),
+                ]
+            },
+        }
+        with tracer.trace(operation_name="operation_name", tags=tags) as span:
+            ...
+        assert span.raw_span()._data["usage"] is None
+        assert span.raw_span()._data["model"] == "test_model"
+        assert span.raw_span()._data["completion_start_time"] is None
+
     def test_update_span_gets_flushed_by_default(self):
         tracer_mock = Mock()
 

{langfuse_haystack-0.5.0 → langfuse_haystack-0.6.0}/tests/test_tracing.py

@@ -52,25 +52,28 @@ def test_tracing_integration(llm_class, env_var, expected_trace):
     assert "Berlin" in response["llm"]["replies"][0].content
     assert response["tracer"]["trace_url"]
 
-    # add a random delay between 1 and 3 seconds to make sure the trace is flushed
-    # and that the trace is available in Langfuse when we fetch it below
-    time.sleep(random.uniform(1, 3))
-
-    url = "https://cloud.langfuse.com/api/public/traces/"
     trace_url = response["tracer"]["trace_url"]
     uuid = os.path.basename(urlparse(trace_url).path)
+    url = f"https://cloud.langfuse.com/api/public/traces/{uuid}"
 
-    try:
-        response = requests.get(
-            url + uuid, auth=HTTPBasicAuth(os.environ["LANGFUSE_PUBLIC_KEY"], os.environ["LANGFUSE_SECRET_KEY"])
+    # Poll the Langfuse API a bit as the trace might not be ready right away
+    attempts = 5
+    delay = 1
+    while attempts >= 0:
+        res = requests.get(
+            url, auth=HTTPBasicAuth(os.environ["LANGFUSE_PUBLIC_KEY"], os.environ["LANGFUSE_SECRET_KEY"])
         )
-        assert response.status_code == 200, f"Failed to retrieve data from Langfuse API: {response.status_code}"
+        if attempts > 0 and res.status_code != 200:
+            attempts -= 1
+            time.sleep(delay)
+            delay *= 2
+            continue
+        assert res.status_code == 200, f"Failed to retrieve data from Langfuse API: {res.status_code}"
 
         # check if the trace contains the expected LLM name
-        assert expected_trace in str(response.content)
+        assert expected_trace in str(res.content)
         # check if the trace contains the expected generation span
-        assert "GENERATION" in str(response.content)
+        assert "GENERATION" in str(res.content)
         # check if the trace contains the expected user_id
-        assert "user_42" in str(response.content)
-    except requests.exceptions.RequestException as e:
-        pytest.fail(f"Failed to retrieve data from Langfuse API: {e}")
+        assert "user_42" in str(res.content)
+        break