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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (18) hide show
  1. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/CHANGELOG.md +15 -1
  2. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/PKG-INFO +2 -3
  3. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/example/basic_rag.py +1 -0
  4. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/example/chat.py +11 -1
  5. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/pyproject.toml +12 -9
  6. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/src/haystack_integrations/components/connectors/langfuse/langfuse_connector.py +14 -2
  7. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/src/haystack_integrations/tracing/langfuse/tracer.py +75 -44
  8. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/tests/test_langfuse_span.py +2 -0
  9. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/tests/test_tracer.py +71 -34
  10. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/tests/test_tracing.py +31 -14
  11. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/.gitignore +0 -0
  12. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/LICENSE.txt +0 -0
  13. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/README.md +0 -0
  14. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/example/requirements.txt +0 -0
  15. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/pydoc/config.yml +0 -0
  16. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/src/haystack_integrations/components/connectors/langfuse/__init__.py +0 -0
  17. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/src/haystack_integrations/tracing/langfuse/__init__.py +0 -0
  18. {langfuse_haystack-0.4.0 → langfuse_haystack-0.6.0}/tests/__init__.py +0 -0
@@ -1,6 +1,20 @@
1
1
  # Changelog
2
2
 
3
- ## [unreleased]
3
+ ## [integrations/langfuse-v0.5.0] - 2024-10-01
4
+
5
+ ### ⚙️ Miscellaneous Tasks
6
+
7
+ - Update ruff linting scripts and settings (#1105)
8
+
9
+ ### Langfuse
10
+
11
+ - Add invocation_context to identify traces (#1089)
12
+
13
+ ## [integrations/langfuse-v0.4.0] - 2024-09-17
14
+
15
+ ### 🚀 Features
16
+
17
+ - Langfuse - support generation span for more LLMs (#1087)
4
18
 
5
19
  ### 🚜 Refactor
6
20
 
@@ -1,12 +1,11 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: langfuse-haystack
3
- Version: 0.4.0
3
+ Version: 0.6.0
4
4
  Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/langfuse#readme
5
5
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
6
6
  Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/langfuse
7
7
  Author-email: deepset GmbH <info@deepset.ai>
8
- License-Expression: Apache-2.0
9
- License-File: LICENSE.txt
8
+ License: Apache-2.0
10
9
  Classifier: Development Status :: 4 - Beta
11
10
  Classifier: Programming Language :: Python
12
11
  Classifier: Programming Language :: Python :: 3.8
@@ -10,6 +10,7 @@ from haystack.components.embedders import SentenceTransformersDocumentEmbedder,
10
10
  from haystack.components.generators import OpenAIGenerator
11
11
  from haystack.components.retrievers import InMemoryEmbeddingRetriever
12
12
  from haystack.document_stores.in_memory import InMemoryDocumentStore
13
+
13
14
  from haystack_integrations.components.connectors.langfuse import LangfuseConnector
14
15
 
15
16
 
@@ -49,6 +49,16 @@ if __name__ == "__main__":
49
49
  ChatMessage.from_user("Tell me about {{location}}"),
50
50
  ]
51
51
 
52
- response = pipe.run(data={"prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages}})
52
+ response = pipe.run(
53
+ data={
54
+ "prompt_builder": {
55
+ "template_variables": {"location": "Berlin"},
56
+ "template": messages,
57
+ },
58
+ "tracer": {
59
+ "invocation_context": {"some_key": "some_value"},
60
+ },
61
+ }
62
+ )
53
63
  print(response["llm"]["replies"][0])
54
64
  print(response["tracer"]["trace_url"])
@@ -42,6 +42,7 @@ root = "../.."
42
42
  git_describe_command = 'git describe --tags --match="integrations/langfuse-v[0-9]*"'
43
43
 
44
44
  [tool.hatch.envs.default]
45
+ installer = "uv"
45
46
  dependencies = [
46
47
  "coverage[toml]>=6.5",
47
48
  "pytest",
@@ -64,13 +65,14 @@ docs = ["pydoc-markdown pydoc/config.yml"]
64
65
  python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
65
66
 
66
67
  [tool.hatch.envs.lint]
68
+ installer = "uv"
67
69
  detached = true
68
- dependencies = ["black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
70
+ dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
69
71
 
70
72
  [tool.hatch.envs.lint.scripts]
71
73
  typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
72
- style = ["ruff check {args:. --exclude tests/}", "black --check --diff {args:.}"]
73
- fmt = ["black {args:.}", "ruff --fix {args:. --exclude tests/}", "style"]
74
+ style = ["ruff check {args:}", "black --check --diff {args:.}"]
75
+ fmt = ["black {args:.}", "ruff check --fix {args:}", "style"]
74
76
  all = ["style", "typing"]
75
77
 
76
78
  [tool.hatch.metadata]
@@ -84,7 +86,10 @@ skip-string-normalization = true
84
86
  [tool.ruff]
85
87
  target-version = "py38"
86
88
  line-length = 120
87
- lint.select = [
89
+ exclude = ["example", "tests"]
90
+
91
+ [tool.ruff.lint]
92
+ select = [
88
93
  "A",
89
94
  "ARG",
90
95
  "B",
@@ -110,8 +115,7 @@ lint.select = [
110
115
  "W",
111
116
  "YTT",
112
117
  ]
113
-
114
- lint.ignore = [
118
+ ignore = [
115
119
  # Allow non-abstract empty methods in abstract base classes
116
120
  "B027",
117
121
  # Ignore checks for possible passwords
@@ -127,14 +131,13 @@ lint.ignore = [
127
131
  # Asserts
128
132
  "S101",
129
133
  ]
130
- lint.unfixable = [
134
+ unfixable = [
131
135
  # Don't touch unused imports
132
136
  "F401",
133
137
  ]
134
- extend-exclude = ["tests", "example"]
135
138
 
136
139
  [tool.ruff.lint.isort]
137
- known-first-party = ["src"]
140
+ known-first-party = ["haystack_integrations"]
138
141
 
139
142
  [tool.ruff.lint.flake8-tidy-imports]
140
143
  ban-relative-imports = "parents"
@@ -1,8 +1,12 @@
1
- from haystack import component, tracing
1
+ from typing import Any, Dict, Optional
2
+
3
+ from haystack import component, logging, tracing
2
4
 
3
5
  from haystack_integrations.tracing.langfuse import LangfuseTracer
4
6
  from langfuse import Langfuse
5
7
 
8
+ logger = logging.getLogger(__name__)
9
+
6
10
 
7
11
  @component
8
12
  class LangfuseConnector:
@@ -105,12 +109,20 @@ class LangfuseConnector:
105
109
  tracing.enable_tracing(self.tracer)
106
110
 
107
111
  @component.output_types(name=str, trace_url=str)
108
- def run(self):
112
+ def run(self, invocation_context: Optional[Dict[str, Any]] = None):
109
113
  """
110
114
  Runs the LangfuseConnector component.
111
115
 
116
+ :param invocation_context: A dictionary with additional context for the invocation. This parameter
117
+ is useful when users want to mark this particular invocation with additional information, e.g.
118
+ a run id from their own execution framework, user id, etc. These key-value pairs are then visible
119
+ in the Langfuse traces.
112
120
  :returns: A dictionary with the following keys:
113
121
  - `name`: The name of the tracing component.
114
122
  - `trace_url`: The URL to the tracing data.
115
123
  """
124
+ logger.debug(
125
+ "Langfuse tracer invoked with the following context: '{invocation_context}'",
126
+ invocation_context=invocation_context,
127
+ )
116
128
  return {"name": self.name, "trace_url": self.tracer.get_trace_url()}
@@ -1,7 +1,10 @@
1
1
  import contextlib
2
2
  import os
3
- from typing import Any, Dict, Iterator, Optional, Union
3
+ from contextvars import ContextVar
4
+ from datetime import datetime
5
+ from typing import Any, Dict, Iterator, List, Optional, Union
4
6
 
7
+ from haystack import logging
5
8
  from haystack.components.generators.openai_utils import _convert_message_to_openai_format
6
9
  from haystack.dataclasses import ChatMessage
7
10
  from haystack.tracing import Span, Tracer, tracer
@@ -9,6 +12,8 @@ from haystack.tracing import utils as tracing_utils
9
12
 
10
13
  import langfuse
11
14
 
15
+ logger = logging.getLogger(__name__)
16
+
12
17
  HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR = "HAYSTACK_LANGFUSE_ENFORCE_FLUSH"
13
18
  _SUPPORTED_GENERATORS = [
14
19
  "AzureOpenAIGenerator",
@@ -28,6 +33,17 @@ _SUPPORTED_CHAT_GENERATORS = [
28
33
  ]
29
34
  _ALL_SUPPORTED_GENERATORS = _SUPPORTED_GENERATORS + _SUPPORTED_CHAT_GENERATORS
30
35
 
36
+ # These are the keys used by Haystack for traces and span.
37
+ # We keep them here to avoid making typos when using them.
38
+ _PIPELINE_RUN_KEY = "haystack.pipeline.run"
39
+ _COMPONENT_NAME_KEY = "haystack.component.name"
40
+ _COMPONENT_TYPE_KEY = "haystack.component.type"
41
+ _COMPONENT_OUTPUT_KEY = "haystack.component.output"
42
+
43
+ # Context var used to keep track of tracing related info.
44
+ # This mainly useful for parents spans.
45
+ tracing_context_var: ContextVar[Dict[Any, Any]] = ContextVar("tracing_context", default={})
46
+
31
47
 
32
48
  class LangfuseSpan(Span):
33
49
  """
@@ -82,7 +98,7 @@ class LangfuseSpan(Span):
82
98
 
83
99
  self._data[key] = value
84
100
 
85
- def raw_span(self) -> Any:
101
+ def raw_span(self) -> "Union[langfuse.client.StatefulSpanClient, langfuse.client.StatefulTraceClient]":
86
102
  """
87
103
  Return the underlying span instance.
88
104
 
@@ -111,75 +127,90 @@ class LangfuseTracer(Tracer):
111
127
  and only accessible to the Langfuse account owner.
112
128
  """
113
129
  self._tracer = tracer
114
- self._context: list[LangfuseSpan] = []
130
+ self._context: List[LangfuseSpan] = []
115
131
  self._name = name
116
132
  self._public = public
117
133
  self.enforce_flush = os.getenv(HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR, "true").lower() == "true"
118
134
 
119
135
  @contextlib.contextmanager
120
- def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> Iterator[Span]:
121
- """
122
- Start and manage a new trace span.
123
- :param operation_name: The name of the operation.
124
- :param tags: A dictionary of tags to attach to the span.
125
- :return: A context manager yielding the span.
126
- """
136
+ def trace(
137
+ self, operation_name: str, tags: Optional[Dict[str, Any]] = None, parent_span: Optional[Span] = None
138
+ ) -> Iterator[Span]:
127
139
  tags = tags or {}
128
- span_name = tags.get("haystack.component.name", operation_name)
129
-
130
- if tags.get("haystack.component.type") in _ALL_SUPPORTED_GENERATORS:
131
- span = LangfuseSpan(self.current_span().raw_span().generation(name=span_name))
140
+ span_name = tags.get(_COMPONENT_NAME_KEY, operation_name)
141
+
142
+ # Create new span depending whether there's a parent span or not
143
+ if not parent_span:
144
+ if operation_name != _PIPELINE_RUN_KEY:
145
+ logger.warning(
146
+ "Creating a new trace without a parent span is not recommended for operation '{operation_name}'.",
147
+ operation_name=operation_name,
148
+ )
149
+ # Create a new trace if no parent span is provided
150
+ span = LangfuseSpan(
151
+ self._tracer.trace(
152
+ name=self._name,
153
+ public=self._public,
154
+ id=tracing_context_var.get().get("trace_id"),
155
+ user_id=tracing_context_var.get().get("user_id"),
156
+ session_id=tracing_context_var.get().get("session_id"),
157
+ tags=tracing_context_var.get().get("tags"),
158
+ version=tracing_context_var.get().get("version"),
159
+ )
160
+ )
161
+ elif tags.get(_COMPONENT_TYPE_KEY) in _ALL_SUPPORTED_GENERATORS:
162
+ span = LangfuseSpan(parent_span.raw_span().generation(name=span_name))
132
163
  else:
133
- span = LangfuseSpan(self.current_span().raw_span().span(name=span_name))
164
+ span = LangfuseSpan(parent_span.raw_span().span(name=span_name))
134
165
 
135
166
  self._context.append(span)
136
167
  span.set_tags(tags)
137
168
 
138
169
  yield span
139
170
 
140
- if tags.get("haystack.component.type") in _SUPPORTED_GENERATORS:
141
- meta = span._data.get("haystack.component.output", {}).get("meta")
171
+ # Update span metadata based on component type
172
+ if tags.get(_COMPONENT_TYPE_KEY) in _SUPPORTED_GENERATORS:
173
+ # Haystack returns one meta dict for each message, but the 'usage' value
174
+ # is always the same, let's just pick the first item
175
+ meta = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("meta")
142
176
  if meta:
143
- # Haystack returns one meta dict for each message, but the 'usage' value
144
- # is always the same, let's just pick the first item
145
177
  m = meta[0]
146
178
  span._span.update(usage=m.get("usage") or None, model=m.get("model"))
147
- elif tags.get("haystack.component.type") in _SUPPORTED_CHAT_GENERATORS:
148
- replies = span._data.get("haystack.component.output", {}).get("replies")
179
+ elif tags.get(_COMPONENT_TYPE_KEY) in _SUPPORTED_CHAT_GENERATORS:
180
+ replies = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("replies")
149
181
  if replies:
150
182
  meta = replies[0].meta
151
- span._span.update(usage=meta.get("usage") or None, model=meta.get("model"))
152
-
153
- pipeline_input = tags.get("haystack.pipeline.input_data", None)
154
- if pipeline_input:
155
- span._span.update(input=tags["haystack.pipeline.input_data"])
156
- pipeline_output = tags.get("haystack.pipeline.output_data", None)
157
- if pipeline_output:
158
- span._span.update(output=tags["haystack.pipeline.output_data"])
159
-
160
- span.raw_span().end()
183
+ completion_start_time = meta.get("completion_start_time")
184
+ if completion_start_time:
185
+ try:
186
+ completion_start_time = datetime.fromisoformat(completion_start_time)
187
+ except ValueError:
188
+ logger.error(f"Failed to parse completion_start_time: {completion_start_time}")
189
+ completion_start_time = None
190
+ span._span.update(
191
+ usage=meta.get("usage") or None,
192
+ model=meta.get("model"),
193
+ completion_start_time=completion_start_time,
194
+ )
195
+
196
+ raw_span = span.raw_span()
197
+ if isinstance(raw_span, langfuse.client.StatefulSpanClient):
198
+ raw_span.end()
161
199
  self._context.pop()
162
200
 
163
- if len(self._context) == 1:
164
- # The root span has to be a trace, which need to be removed from the context after the pipeline run
165
- self._context.pop()
166
-
167
- if self.enforce_flush:
168
- self.flush()
201
+ if self.enforce_flush:
202
+ self.flush()
169
203
 
170
204
  def flush(self):
171
205
  self._tracer.flush()
172
206
 
173
- def current_span(self) -> Span:
207
+ def current_span(self) -> Optional[Span]:
174
208
  """
175
- Return the currently active span.
209
+ Return the current active span.
176
210
 
177
- :return: The currently active span.
211
+ :return: The current span if available, else None.
178
212
  """
179
- if not self._context:
180
- # The root span has to be a trace
181
- self._context.append(LangfuseSpan(self._tracer.trace(name=self._name, public=self._public)))
182
- return self._context[-1]
213
+ return self._context[-1] if self._context else None
183
214
 
184
215
  def get_trace_url(self) -> str:
185
216
  """
@@ -3,7 +3,9 @@ import os
3
3
  os.environ["HAYSTACK_CONTENT_TRACING_ENABLED"] = "true"
4
4
 
5
5
  from unittest.mock import Mock
6
+
6
7
  from haystack.dataclasses import ChatMessage
8
+
7
9
  from haystack_integrations.tracing.langfuse.tracer import LangfuseSpan
8
10
 
9
11
 
@@ -1,9 +1,43 @@
1
- import os
2
- from unittest.mock import Mock, MagicMock, patch
1
+ import datetime
2
+ from unittest.mock import MagicMock, Mock, patch
3
3
 
4
+ from haystack.dataclasses import ChatMessage
4
5
  from haystack_integrations.tracing.langfuse.tracer import LangfuseTracer
5
6
 
6
7
 
8
+ class MockSpan:
9
+ def __init__(self):
10
+ self._data = {}
11
+ self._span = self
12
+ self.operation_name = "operation_name"
13
+
14
+ def raw_span(self):
15
+ return self
16
+
17
+ def span(self, name=None):
18
+ # assert correct operation name passed to the span
19
+ assert name == "operation_name"
20
+ return self
21
+
22
+ def update(self, **kwargs):
23
+ self._data.update(kwargs)
24
+
25
+ def generation(self, name=None):
26
+ return self
27
+
28
+ def end(self):
29
+ pass
30
+
31
+
32
+ class MockTracer:
33
+
34
+ def trace(self, name, **kwargs):
35
+ return MockSpan()
36
+
37
+ def flush(self):
38
+ pass
39
+
40
+
7
41
  class TestLangfuseTracer:
8
42
 
9
43
  # LangfuseTracer can be initialized with a Langfuse instance, a name and a boolean value for public.
@@ -35,7 +69,7 @@ class TestLangfuseTracer:
35
69
  tracer = LangfuseTracer(tracer=mock_tracer, name="Haystack", public=False)
36
70
 
37
71
  with tracer.trace("operation_name", tags={"tag1": "value1", "tag2": "value2"}) as span:
38
- assert len(tracer._context) == 2, "The trace span should have been added to the the root context span"
72
+ assert len(tracer._context) == 1, "The trace span should have been added to the the root context span"
39
73
  assert span.raw_span().operation_name == "operation_name"
40
74
  assert span.raw_span().metadata == {"tag1": "value1", "tag2": "value2"}
41
75
 
@@ -45,37 +79,6 @@ class TestLangfuseTracer:
45
79
 
46
80
  # check that update method is called on the span instance with the provided key value pairs
47
81
  def test_update_span_with_pipeline_input_output_data(self):
48
- class MockTracer:
49
-
50
- def trace(self, name, **kwargs):
51
- return MockSpan()
52
-
53
- def flush(self):
54
- pass
55
-
56
- class MockSpan:
57
- def __init__(self):
58
- self._data = {}
59
- self._span = self
60
- self.operation_name = "operation_name"
61
-
62
- def raw_span(self):
63
- return self
64
-
65
- def span(self, name=None):
66
- # assert correct operation name passed to the span
67
- assert name == "operation_name"
68
- return self
69
-
70
- def update(self, **kwargs):
71
- self._data.update(kwargs)
72
-
73
- def generation(self, name=None):
74
- return self
75
-
76
- def end(self):
77
- pass
78
-
79
82
  tracer = LangfuseTracer(tracer=MockTracer(), name="Haystack", public=False)
80
83
  with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.input_data": "hello"}) as span:
81
84
  assert span.raw_span()._data["metadata"] == {"haystack.pipeline.input_data": "hello"}
@@ -83,6 +86,40 @@ class TestLangfuseTracer:
83
86
  with tracer.trace(operation_name="operation_name", tags={"haystack.pipeline.output_data": "bye"}) as span:
84
87
  assert span.raw_span()._data["metadata"] == {"haystack.pipeline.output_data": "bye"}
85
88
 
89
+ def test_trace_generation(self):
90
+ tracer = LangfuseTracer(tracer=MockTracer(), name="Haystack", public=False)
91
+ tags = {
92
+ "haystack.component.type": "OpenAIChatGenerator",
93
+ "haystack.component.output": {
94
+ "replies": [
95
+ ChatMessage.from_assistant(
96
+ "", meta={"completion_start_time": "2021-07-27T16:02:08.012345", "model": "test_model"}
97
+ )
98
+ ]
99
+ },
100
+ }
101
+ with tracer.trace(operation_name="operation_name", tags=tags) as span:
102
+ ...
103
+ assert span.raw_span()._data["usage"] is None
104
+ assert span.raw_span()._data["model"] == "test_model"
105
+ assert span.raw_span()._data["completion_start_time"] == datetime.datetime(2021, 7, 27, 16, 2, 8, 12345)
106
+
107
+ def test_trace_generation_invalid_start_time(self):
108
+ tracer = LangfuseTracer(tracer=MockTracer(), name="Haystack", public=False)
109
+ tags = {
110
+ "haystack.component.type": "OpenAIChatGenerator",
111
+ "haystack.component.output": {
112
+ "replies": [
113
+ ChatMessage.from_assistant("", meta={"completion_start_time": "foobar", "model": "test_model"}),
114
+ ]
115
+ },
116
+ }
117
+ with tracer.trace(operation_name="operation_name", tags=tags) as span:
118
+ ...
119
+ assert span.raw_span()._data["usage"] is None
120
+ assert span.raw_span()._data["model"] == "test_model"
121
+ assert span.raw_span()._data["completion_start_time"] is None
122
+
86
123
  def test_update_span_gets_flushed_by_default(self):
87
124
  tracer_mock = Mock()
88
125
 
@@ -1,14 +1,17 @@
1
1
  import os
2
- import pytest
2
+ import random
3
+ import time
3
4
  from urllib.parse import urlparse
5
+
6
+ import pytest
4
7
  import requests
5
- from requests.auth import HTTPBasicAuth
6
8
  from haystack import Pipeline
7
9
  from haystack.components.builders import ChatPromptBuilder
8
- from haystack.dataclasses import ChatMessage
9
- from haystack_integrations.components.connectors.langfuse import LangfuseConnector
10
10
  from haystack.components.generators.chat import OpenAIChatGenerator
11
+ from haystack.dataclasses import ChatMessage
12
+ from requests.auth import HTTPBasicAuth
11
13
 
14
+ from haystack_integrations.components.connectors.langfuse import LangfuseConnector
12
15
  from haystack_integrations.components.generators.anthropic import AnthropicChatGenerator
13
16
  from haystack_integrations.components.generators.cohere import CohereChatGenerator
14
17
 
@@ -40,23 +43,37 @@ def test_tracing_integration(llm_class, env_var, expected_trace):
40
43
  ChatMessage.from_user("Tell me about {{location}}"),
41
44
  ]
42
45
 
43
- response = pipe.run(data={"prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages}})
46
+ response = pipe.run(
47
+ data={
48
+ "prompt_builder": {"template_variables": {"location": "Berlin"}, "template": messages},
49
+ "tracer": {"invocation_context": {"user_id": "user_42"}},
50
+ }
51
+ )
44
52
  assert "Berlin" in response["llm"]["replies"][0].content
45
53
  assert response["tracer"]["trace_url"]
46
54
 
47
- url = "https://cloud.langfuse.com/api/public/traces/"
48
55
  trace_url = response["tracer"]["trace_url"]
49
56
  uuid = os.path.basename(urlparse(trace_url).path)
57
+ url = f"https://cloud.langfuse.com/api/public/traces/{uuid}"
50
58
 
51
- try:
52
- response = requests.get(
53
- url + uuid, auth=HTTPBasicAuth(os.environ["LANGFUSE_PUBLIC_KEY"], os.environ["LANGFUSE_SECRET_KEY"])
59
+ # Poll the Langfuse API a bit as the trace might not be ready right away
60
+ attempts = 5
61
+ delay = 1
62
+ while attempts >= 0:
63
+ res = requests.get(
64
+ url, auth=HTTPBasicAuth(os.environ["LANGFUSE_PUBLIC_KEY"], os.environ["LANGFUSE_SECRET_KEY"])
54
65
  )
55
- assert response.status_code == 200, f"Failed to retrieve data from Langfuse API: {response.status_code}"
66
+ if attempts > 0 and res.status_code != 200:
67
+ attempts -= 1
68
+ time.sleep(delay)
69
+ delay *= 2
70
+ continue
71
+ assert res.status_code == 200, f"Failed to retrieve data from Langfuse API: {res.status_code}"
56
72
 
57
73
  # check if the trace contains the expected LLM name
58
- assert expected_trace in str(response.content)
74
+ assert expected_trace in str(res.content)
59
75
  # check if the trace contains the expected generation span
60
- assert "GENERATION" in str(response.content)
61
- except requests.exceptions.RequestException as e:
62
- pytest.fail(f"Failed to retrieve data from Langfuse API: {e}")
76
+ assert "GENERATION" in str(res.content)
77
+ # check if the trace contains the expected user_id
78
+ assert "user_42" in str(res.content)
79
+ break