feat: support user_id argument for OpenAI Chat completions (#354)

Co-authored-by: bryan-agicap <145330610+bryan-agicap@users.noreply.github.com>

Author: maxdeichmann

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,

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",

tests/test_openai.py

@@ -223,6 +223,28 @@ 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()
+    trace_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,
+        trace_id=trace_id,
+    )
+
+    openai.flush_langfuse()
+
+    assert len(completion.choices) != 0
+    traces = api.trace.get(trace_id)
+
+    assert traces.user_id == user_id
+
+
 def test_openai_chat_completion_without_extra_param():
     completion = chat_func(
         model="gpt-3.5-turbo",