feat: add sessions and i/o on traces (#233)

Author: marcklingen

langfuse/api/client.py

@@ -31,7 +31,7 @@ def __init__(
         username: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
         password: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
         timeout: typing.Optional[float] = 60,
-        httpx_client: typing.Optional[httpx.Client] = None
+        httpx_client: typing.Optional[httpx.Client] = None,
     ):
         self._client_wrapper = SyncClientWrapper(
             base_url=base_url,
@@ -68,7 +68,7 @@ def __init__(
         username: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
         password: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
         timeout: typing.Optional[float] = 60,
-        httpx_client: typing.Optional[httpx.AsyncClient] = None
+        httpx_client: typing.Optional[httpx.AsyncClient] = None,
     ):
         self._client_wrapper = AsyncClientWrapper(
             base_url=base_url,

langfuse/api/resources/commons/types/trace.py

@@ -22,9 +22,7 @@ class Trace(pydantic.BaseModel):
     version: typing.Optional[str] = None
     user_id: typing.Optional[str] = pydantic.Field(alias="userId", default=None)
     metadata: typing.Optional[typing.Any] = None
-    public: typing.Optional[bool] = pydantic.Field(
-        default=None, description="Public traces are accessible via url without login"
-    )
+    public: typing.Optional[bool] = pydantic.Field(default=None, description="Public traces are accessible via url without login")
 
     def json(self, **kwargs: typing.Any) -> str:
         kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}

langfuse/api/resources/sessions/__init__.py

@@ -1,2 +1 @@
 # This file was auto-generated by Fern from our API Definition.
-

langfuse/openai.py

@@ -74,11 +74,12 @@ 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, **kwargs):
+    def __init__(self, name=None, metadata=None, trace_id=None, session_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.kwargs = kwargs
 
     def get_langfuse_args(self):
@@ -104,12 +105,19 @@ def _get_langfuse_data_from_kwargs(resource: OpenAiDefinition, langfuse: Langfus
     if name is not None and not isinstance(name, str):
         raise TypeError("name must be a string")
 
-    trace_id = kwargs.get("trace_id", "OpenAI-generation")
+    trace_id = kwargs.get("trace_id", None)
     if trace_id is not None and not isinstance(trace_id, str):
         raise TypeError("trace_id must be a string")
+    
+    session_id = kwargs.get("session_id", None)
+    if session_id is not None and not isinstance(session_id, str):
+        raise TypeError("session_id must be a string")
 
     if trace_id:
-        langfuse.trace(CreateTrace(id=trace_id))
+        langfuse.trace(CreateTrace(id=trace_id, sessionId=session_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(CreateTrace(sessionId=session_id)).id
 
     metadata = kwargs.get("metadata", {})
 

tests/test_core_sdk.py

@@ -402,7 +402,7 @@ def test_create_trace_and_generation():
     trace_name = create_uuid()
     generationId = create_uuid()
 
-    trace = langfuse.trace(CreateTrace(name=trace_name))
+    trace = langfuse.trace(CreateTrace(name=trace_name, input={"key": "value"}, sessionId="test"))
     trace.generation(CreateGeneration(id=generationId, name="generation"))
 
     langfuse.flush()
@@ -417,8 +417,9 @@ def test_create_trace_and_generation():
 
     generation = getTrace.observations[0]
     assert generation.name == "generation"
-    assert generation.trace_id == trace.id
+    assert generation.trace_id == getTrace.id
     assert generation.start_time is not None
+    assert getTrace.input == {"key": "value"}
 
 
 def test_create_generation_and_trace():