Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Langfuse: add invocation_context to identify traces #1089

Merged
merged 6 commits into from
Oct 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion integrations/langfuse/example/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,16 @@
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"])
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -105,12 +109,20 @@ def __init__(self, name: str, public: bool = False):
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()}
9 changes: 8 additions & 1 deletion integrations/langfuse/tests/test_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,12 @@ 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"]

Expand All @@ -65,5 +70,7 @@ def test_tracing_integration(llm_class, env_var, expected_trace):
assert expected_trace in str(response.content)
# check if the trace contains the expected generation span
assert "GENERATION" in str(response.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}")