From 29c4b6011aff3e4d52caf2de816640dde3998a6e Mon Sep 17 00:00:00 2001 From: bryan-agicap <145330610+bryan-agicap@users.noreply.github.com> Date: Mon, 12 Feb 2024 22:13:19 +0100 Subject: [PATCH 1/3] Add Support for `user_id` argument for OpenAI Chat completions, generations, spans and events (#289) Co-authored-by: Max Deichmann --- langfuse/openai.py | 12 +++++++++--- tests/test_core_sdk.py | 1 - tests/test_openai.py | 19 +++++++++++++++++++ 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/langfuse/openai.py b/langfuse/openai.py index 62e977ec2..c1c49d48f 100644 --- a/langfuse/openai.py +++ b/langfuse/openai.py @@ -83,13 +83,14 @@ def __init__(self, module: str, object: str, method: str, type: str, sync: bool) class OpenAiArgsExtractor: def __init__( - self, name=None, metadata=None, trace_id=None, session_id=None, **kwargs + self, name=None, metadata=None, trace_id=None, session_id=None, user_id=None, **kwargs ): self.args = {} self.args["name"] = name self.args["metadata"] = metadata self.args["trace_id"] = trace_id self.args["session_id"] = session_id + self.args["user_id"] = user_id self.kwargs = kwargs def get_langfuse_args(self): @@ -125,11 +126,15 @@ def _get_langfuse_data_from_kwargs( if session_id is not None and not isinstance(session_id, str): raise TypeError("session_id must be a string") + user_id = kwargs.get("user_id", None) + if user_id is not None and not isinstance(user_id, str): + raise TypeError("user_id must be a string") + if trace_id: - langfuse.trace(id=trace_id, session_id=session_id) + langfuse.trace(id=trace_id, session_id=session_id, user_id=user_id) elif session_id: # If a session_id is provided but no trace_id, we should create a trace using the SDK and then use its trace_id - trace_id = langfuse.trace(session_id=session_id).id + trace_id = langfuse.trace(session_id=session_id, user_id=user_id).id metadata = kwargs.get("metadata", {}) @@ -165,6 +170,7 @@ def _get_langfuse_data_from_kwargs( "name": name, "metadata": metadata, "trace_id": trace_id, + "user_id": user_id, "start_time": start_time, "input": prompt, "model_parameters": modelParameters, diff --git a/tests/test_core_sdk.py b/tests/test_core_sdk.py index f36c2358d..ffc0c14f1 100644 --- a/tests/test_core_sdk.py +++ b/tests/test_core_sdk.py @@ -149,7 +149,6 @@ def test_create_generation(): timestamp = _get_timestamp() generation_id = create_uuid() - langfuse.generation( id=generation_id, name="query-generation", diff --git a/tests/test_openai.py b/tests/test_openai.py index b84cedfdc..e65fe8926 100644 --- a/tests/test_openai.py +++ b/tests/test_openai.py @@ -223,6 +223,25 @@ def test_openai_chat_completion_fail(): openai.api_key = os.environ["OPENAI_API_KEY"] +def test_openai_chat_completion_with_user_id(): + api = get_api() + user_id = create_uuid() + completion = chat_func( + name="user-creation", + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "1 + 1 = "}], + temperature=0, + metadata={"someKey": "someResponse"}, + user_id=user_id, + ) + + assert len(completion.choices) != 0 + traces = api.trace.list(limit=10, page=1, user_id=user_id) + + assert len(traces.data) == 1 + assert traces.data[0].user_id == user_id + + def test_openai_chat_completion_without_extra_param(): completion = chat_func( model="gpt-3.5-turbo", From 1bd6144272d294b246d43faa7ce2c024e06344e5 Mon Sep 17 00:00:00 2001 From: Max Deichmann Date: Mon, 12 Feb 2024 22:22:55 +0100 Subject: [PATCH 2/3] fix test --- tests/test_openai.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_openai.py b/tests/test_openai.py index e65fe8926..b4f88b83f 100644 --- a/tests/test_openai.py +++ b/tests/test_openai.py @@ -226,6 +226,7 @@ def test_openai_chat_completion_fail(): def test_openai_chat_completion_with_user_id(): api = get_api() user_id = create_uuid() + trace_id = create_uuid() completion = chat_func( name="user-creation", model="gpt-3.5-turbo", @@ -233,13 +234,13 @@ def test_openai_chat_completion_with_user_id(): temperature=0, metadata={"someKey": "someResponse"}, user_id=user_id, + trace_id=trace_id, ) assert len(completion.choices) != 0 - traces = api.trace.list(limit=10, page=1, user_id=user_id) + traces = api.trace.get(trace_id) - assert len(traces.data) == 1 - assert traces.data[0].user_id == user_id + assert traces.user_id == user_id def test_openai_chat_completion_without_extra_param(): From dae1309670e20a58492f17665e4b8a08311db11b Mon Sep 17 00:00:00 2001 From: Max Deichmann Date: Mon, 12 Feb 2024 22:32:35 +0100 Subject: [PATCH 3/3] push --- tests/test_openai.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_openai.py b/tests/test_openai.py index b4f88b83f..09c54b3f5 100644 --- a/tests/test_openai.py +++ b/tests/test_openai.py @@ -237,6 +237,8 @@ def test_openai_chat_completion_with_user_id(): trace_id=trace_id, ) + openai.flush_langfuse() + assert len(completion.choices) != 0 traces = api.trace.get(trace_id)