Repository: deepset-ai/haystack

Issue:
Add TTFT (time to first token) to Langfuse traces
**Is your feature request related to a problem? Please describe.**
We are developing several chatbot-like applications that require streaming the response from LLM. There are a couple of metrics to look at, and one of which is the TTFT (time to first token) indicating how long the user needs to wait before seeing something in the output dialog box. However, due to [the way that the tracing spans are handled in the pipeline](https://github.com/deepset-ai/haystack/blob/main/haystack/core/pipeline/pipeline.py#L43-L85), the run invocation inside the component does not have direct access to the span, so we are not able to log this information to the tracer.

**Describe the solution you'd like**
The simplest solution would be adding visibility to tracing span from component run() method. This could be a context variable that the methods inside of the component have access to, but I am not very confident about the exact approach here.

**Describe alternatives you've considered**
The only temporary solution right now is to directly manipulate low-level tracing sdks inside the streaming callback function, and make a special callback function such that it uploads the timestamp upon receiving the first SSE.

**Additional context**
Add any other context or screenshots about the feature request here.

Discussion:
@vblagoje: Note to self:

TTFT in Langfuse is automatically calculated when `completion_start_time` (the timestamp of the first token) is provided to generation span i.e. just call `update` on the generation span with this key/value i.e `completion_start_time=datetime.now()`

This could be done by attaching custom stream callback on chat generator (most likely from our LangfuseTracer), consuming first token in the callback and calling generation span update - perhaps directly in that callback.

To be investigated how we can eventually do this in async calls as well. @vblagoje: @LastRemote and @julian-risch 

This was, in fact, not so [hard to do](https://github.com/deepset-ai/haystack/compare/main...add_ttft). Forget the recommendation above. We need to simply timestamp the first chunk received from the LLM. And we should do this across all LLM chat generators. When streaming we don't have data for prompt and completion tokens available and what's really interesting is that even if we set "prompt_tokens" and "completion_tokens" to 0 (see branch above) Langfuse somehow correctly counts them. I'll check with them how this is actually done. Here is the [trace](https://cloud.langfuse.com/project/clvdk4bcb000099hzwhts3g66/traces/affa4b43-9f8c-4f14-ae13-a6e90a72bfb3?observation=79ad1391-7eba-42a7-a0de-d79ad3159155) depicted below.

 

![Image](https://github.com/user-attachments/assets/e4cde5f5-5f6f-41ec-a8ba-ef65ad040690)
@LastRemote: > @LastRemote and @julian-risch
> 
> This was, in fact, not so [hard to do](https://github.com/deepset-ai/haystack/compare/main...add_ttft). Forget the recommendation above. We need to simply timestamp the first chunk received from the LLM. And we should do this across all LLM chat generators. When streaming we don't have data for prompt and completion tokens available and what's really interesting is that even if we set "prompt_tokens" and "completion_tokens" to 0 (see branch above) Langfuse somehow correctly counts them. I'll check with them how this is actually done.
> 
> ![Image](https://private-user-images.githubusercontent.com/458335/371060284-e4cde5f5-5f6f-41ec-a8ba-ef65ad040690.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjczNDE4NzgsIm5iZiI6MTcyNzM0MTU3OCwicGF0aCI6Ii80NTgzMzUvMzcxMDYwMjg0LWU0Y2RlNWY1LTVmNmYtNDFlYy1hOGJhLWVmNjVhZDA0MDY5MC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQwOTI2JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MDkyNlQwOTA2MThaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT1lNDVkOWI4MTdkMTMxNmUxOGU2YTIwNGFmMGRkODBiMWNhODIwZjM5YTkwOGY2NDNkM2U3ZTQyNzgxM2U2NmZjJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.fFp15I9iNxbnl7zl8FO43j-4Grkf1_IPCGb8Xz-1bVs)

Interesting, I have never thought about this approach, but I guess it should work. So basically we store `completion_first_chunk` as a part of the usage meta so it can be accessed inside of `haystack.component.output`?

By the way, for OpenAI and the latest versions of Azure OpenAI models, if you set `stream_options` [accordingly](https://platform.openai.com/docs/api-reference/chat/create#chat-create-stream_options), the last streaming chunk will include the actual usage data. Additionally Langfuse will automatically count the usage tokens if the model is from OpenAI or Claude (see screenshot below).

<img width="1397" alt="Screenshot 2024-09-26 at 17 14 49" src="https://github.com/user-attachments/assets/a1ab799c-98ce-47df-9554-b44a591396ef">
@vblagoje: Aha nice @LastRemote that's why I collected those meta chunks hoping that one day this will work. The change in langfuse tracer was minimal as well. One LOC change in `langfuse/tracer.py` at ~151, we need to:

```
span._span.update(usage=meta.get("usage") or None,
                                  model=meta.get("model"),
                                  completion_start_time=meta.get("usage", {}).get("completion_start_time"))
```

I'll speak to @julian-risch about scheduling this change in the near future but if you wish - feel free to create a PR that updates our chat generators with this change and we'll take it from there - I can review the PR and we can try out various chat generators together. @LastRemote: @vblagoje Okay sure, I will make a PR.@vblagoje: @LastRemote before we open a bunch of PRs or group together all the changes into one PR for all chat generators - let's do one trial PR and set the standard for other chat generators. @LastRemote: @vblagoje Sorry for the delayed response. I took some days off last week. Here we go: https://github.com/deepset-ai/haystack/pull/8444

Note that I only made minimal changes since I am not exactly sure how to enable `include_usage` in OpenAI SDK. I have a customized OpenAI implementation based on httpx that works, but that would require a complete refactor.@LastRemote: By the way, I also attempted to support Anthropic (including Bedrock Anthropic) models. It seems there's a mismatch in the usage data format between the Langfuse API and Anthropic [when updating the Langfuse span](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/langfuse/src/haystack_integrations/tracing/langfuse/tracer.py#L151), which causes the operation to fail. I believe it would be better for Langfuse to provide direct support for the raw Anthropic format.@vblagoje: Hey @LastRemote I'll check it out and get back to you. The most recent Anthropic haystack [release](https://pypi.org/project/anthropic-haystack/) should work our of the box. I don't think Bedrock Anthropic works with Langfuse atm.@vblagoje: Hey @LastRemote how do we stand here? What else do we need to do to consider this issue done? @LastRemote: > Hey @LastRemote how do we stand here? What else do we need to do to consider this issue done?

Hi @vblagoje , thanks for the reminder. We should close this. 

An extra note for future readers: we only implemented TTFT tracking for `OpenAIChatGenerator` as all generators in the near future would be deprecated. See #8444 for more information.

Solution patch:
diff --git a/haystack/components/generators/chat/openai.py b/haystack/components/generators/chat/openai.py
index 8f54b7ad36..6013880e74 100644
--- a/haystack/components/generators/chat/openai.py
+++ b/haystack/components/generators/chat/openai.py
@@ -5,6 +5,7 @@
 import copy
 import json
 import os
+from datetime import datetime
 from typing import Any, Callable, Dict, List, Optional, Union
 
 from openai import OpenAI, Stream
@@ -63,7 +64,7 @@ class OpenAIChatGenerator:
     ```
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-positional-arguments
         self,
         api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"),
         model: str = "gpt-4o-mini",
@@ -222,11 +223,15 @@ def run(
                 raise ValueError("Cannot stream multiple responses, please set n=1.")
             chunks: List[StreamingChunk] = []
             chunk = None
+            _first_token = True
 
             # pylint: disable=not-an-iterable
             for chunk in chat_completion:
                 if chunk.choices and streaming_callback:
                     chunk_delta: StreamingChunk = self._build_chunk(chunk)
+                    if _first_token:
+                        _first_token = False
+                        chunk_delta.meta["completion_start_time"] = datetime.now().isoformat()
                     chunks.append(chunk_delta)
                     streaming_callback(chunk_delta)  # invoke callback with the chunk_delta
             completions = [self._connect_chunks(chunk, chunks)]
@@ -280,7 +285,12 @@ def _connect_chunks(self, chunk: Any, chunks: List[StreamingChunk]) -> ChatMessa
                         payload["function"]["arguments"] += delta.arguments or ""
             complete_response = ChatMessage.from_assistant(json.dumps(payloads))
         else:
-            complete_response = ChatMessage.from_assistant("".join([chunk.content for chunk in chunks]))
+            total_content = ""
+            total_meta = {}
+            for streaming_chunk in chunks:
+                total_content += streaming_chunk.content
+                total_meta.update(streaming_chunk.meta)
+            complete_response = ChatMessage.from_assistant(total_content, meta=total_meta)
         complete_response.meta.update(
             {
                 "model": chunk.model,
diff --git a/releasenotes/notes/openai-ttft-42b1ad551b542930.yaml b/releasenotes/notes/openai-ttft-42b1ad551b542930.yaml
new file mode 100644
index 0000000000..bc39b96f7f
--- /dev/null
+++ b/releasenotes/notes/openai-ttft-42b1ad551b542930.yaml
@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Add TTFT (Time-to-First-Token) support for OpenAI generators. This
+    captures the time taken to generate the first token from the model and
+    can be used to analyze the latency of the application.


Test patch:
diff --git a/test/components/generators/chat/test_openai.py b/test/components/generators/chat/test_openai.py
index c3332aae97..9d5122d069 100644
--- a/test/components/generators/chat/test_openai.py
+++ b/test/components/generators/chat/test_openai.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import logging
 import os
+from unittest.mock import patch
 
 import pytest
 from openai import OpenAIError
@@ -219,7 +220,8 @@ def streaming_callback(chunk: StreamingChunk) -> None:
         assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
         assert "Hello" in response["replies"][0].content  # see mock_chat_completion_chunk
 
-    def test_run_with_streaming_callback_in_run_method(self, chat_messages, mock_chat_completion_chunk):
+    @patch("haystack.components.generators.chat.openai.datetime")
+    def test_run_with_streaming_callback_in_run_method(self, mock_datetime, chat_messages, mock_chat_completion_chunk):
         streaming_callback_called = False
 
         def streaming_callback(chunk: StreamingChunk) -> None:
@@ -240,6 +242,13 @@ def streaming_callback(chunk: StreamingChunk) -> None:
         assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
         assert "Hello" in response["replies"][0].content  # see mock_chat_completion_chunk
 
+        assert hasattr(response["replies"][0], "meta")
+        assert isinstance(response["replies"][0].meta, dict)
+        assert (
+            response["replies"][0].meta["completion_start_time"]
+            == mock_datetime.now.return_value.isoformat.return_value
+        )
+
     def test_check_abnormal_completions(self, caplog):
         caplog.set_level(logging.INFO)
         component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))