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

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

@@ -1,6 +1,20 @@
 # Changelog
 
-## [unreleased]
+## [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 - support generation span for more LLMs (#1087)
 
 ### 🚜 Refactor
 

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

@@ -1,12 +1,11 @@
 Metadata-Version: 2.3
 Name: langfuse-haystack
-Version: 0.4.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.4.0 → langfuse_haystack-0.6.0}/example/basic_rag.py

@@ -10,6 +10,7 @@ from haystack.components.embedders import SentenceTransformersDocumentEmbedder,
 from haystack.components.generators import OpenAIGenerator
 from haystack.components.retrievers import InMemoryEmbeddingRetriever
 from haystack.document_stores.in_memory import InMemoryDocumentStore
+
 from haystack_integrations.components.connectors.langfuse import LangfuseConnector
 
 

{langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/example/chat.py

@@ -49,6 +49,16 @@ if __name__ == "__main__":
         ChatMessage.from_user("Tell me about {{location}}"),
     ]
 
-    response = pipe.run(data={"prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages}})
+    response = pipe.run(
+        data={
+            "prompt_builder": {
+                "template_variables": {"location": "Berlin"},
+                "template": messages,
+            },
+            "tracer": {
+                "invocation_context": {"some_key": "some_value"},
+            },
+        }
+    )
     print(response["llm"]["replies"][0])
     print(response["tracer"]["trace_url"])

{langfuse_haystack-0.4.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,13 +65,14 @@ 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}"
-style = ["ruff check {args:. --exclude tests/}", "black --check --diff {args:.}"]
-fmt = ["black {args:.}", "ruff --fix {args:. --exclude tests/}", "style"]
+style = ["ruff check {args:}", "black --check --diff {args:.}"]
+fmt = ["black {args:.}", "ruff check --fix {args:}", "style"]
 all = ["style", "typing"]
 
 [tool.hatch.metadata]
@@ -84,7 +86,10 @@ skip-string-normalization = true
 [tool.ruff]
 target-version = "py38"
 line-length = 120
-lint.select = [
+exclude = ["example", "tests"]
+
+[tool.ruff.lint]
+select = [
   "A",
   "ARG",
   "B",
@@ -110,8 +115,7 @@ lint.select = [
   "W",
   "YTT",
 ]
-
-lint.ignore = [
+ignore = [
   # Allow non-abstract empty methods in abstract base classes
   "B027",
   # Ignore checks for possible passwords
@@ -127,14 +131,13 @@ lint.ignore = [
   # Asserts
   "S101",
 ]
-lint.unfixable = [
+unfixable = [
   # Don't touch unused imports
   "F401",
 ]
-extend-exclude = ["tests", "example"]
 
 [tool.ruff.lint.isort]
-known-first-party = ["src"]
+known-first-party = ["haystack_integrations"]
 
 [tool.ruff.lint.flake8-tidy-imports]
 ban-relative-imports = "parents"

{langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/src/haystack_integrations/components/connectors/langfuse/langfuse_connector.py

@@ -1,8 +1,12 @@
-from haystack import component, tracing
+from typing import Any, Dict, Optional
+
+from haystack import component, logging, tracing
 
 from haystack_integrations.tracing.langfuse import LangfuseTracer
 from langfuse import Langfuse
 
+logger = logging.getLogger(__name__)
+
 
 @component
 class LangfuseConnector:
@@ -105,12 +109,20 @@ class LangfuseConnector:
         tracing.enable_tracing(self.tracer)
 
     @component.output_types(name=str, trace_url=str)
-    def run(self):
+    def run(self, invocation_context: Optional[Dict[str, Any]] = None):
         """
         Runs the LangfuseConnector component.
 
+        :param invocation_context: A dictionary with additional context for the invocation. This parameter
+            is useful when users want to mark this particular invocation with additional information, e.g.
+            a run id from their own execution framework, user id, etc. These key-value pairs are then visible
+            in the Langfuse traces.
         :returns: A dictionary with the following keys:
             - `name`: The name of the tracing component.
             - `trace_url`: The URL to the tracing data.
         """
+        logger.debug(
+            "Langfuse tracer invoked with the following context: '{invocation_context}'",
+            invocation_context=invocation_context,
+        )
         return {"name": self.name, "trace_url": self.tracer.get_trace_url()}

{langfuse_haystack-0.4.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.4.0 → langfuse_haystack-0.6.0}/tests/test_langfuse_span.py

@@ -3,7 +3,9 @@ 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
 
 

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

@@ -1,9 +1,43 @@
-import os
-from unittest.mock import Mock, MagicMock, patch
+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.4.0 → langfuse_haystack-0.6.0}/tests/test_tracing.py

@@ -1,14 +1,17 @@
 import os
-import pytest
+import random
+import time
 from urllib.parse import urlparse
+
+import pytest
 import requests
-from requests.auth import HTTPBasicAuth
 from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
-from haystack.dataclasses import ChatMessage
-from haystack_integrations.components.connectors.langfuse import LangfuseConnector
 from haystack.components.generators.chat import OpenAIChatGenerator
+from haystack.dataclasses import ChatMessage
+from requests.auth import HTTPBasicAuth
 
+from haystack_integrations.components.connectors.langfuse import LangfuseConnector
 from haystack_integrations.components.generators.anthropic import AnthropicChatGenerator
 from haystack_integrations.components.generators.cohere import CohereChatGenerator
 
@@ -40,23 +43,37 @@ def test_tracing_integration(llm_class, env_var, expected_trace):
         ChatMessage.from_user("Tell me about {{location}}"),
     ]
 
-    response = pipe.run(data={"prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages}})
+    response = pipe.run(
+        data={
+            "prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages},
+            "tracer": {"invocation_context": {"user_id": "user_42"}},
+        }
+    )
     assert "Berlin" in response["llm"]["replies"][0].content
     assert response["tracer"]["trace_url"]
 
-    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)
-    except requests.exceptions.RequestException as e:
-        pytest.fail(f"Failed to retrieve data from Langfuse API: {e}")
+        assert "GENERATION" in str(res.content)
+        # check if the trace contains the expected user_id
+        assert "user_42" in str(res.content)
+        break