infer step id and thread id when creating an attachment

literalai/api/__init__.py

@@ -14,6 +14,7 @@
 
 from typing_extensions import deprecated
 
+from literalai.context import active_steps_var, active_thread_var
 from literalai.dataset import DatasetType
 from literalai.dataset_experiment import DatasetExperiment, DatasetExperimentItem
 from literalai.filter import (
@@ -711,8 +712,8 @@ def upload_file(
 
     def create_attachment(
         self,
-        thread_id: str,
-        step_id: str,
+        thread_id: Optional[str] = None,
+        step_id: Optional[str] = None,
         id: Optional[str] = None,
         metadata: Optional[Dict] = None,
         mime: Optional[str] = None,
@@ -740,6 +741,16 @@ def create_attachment(
         Returns:
             Attachment: The created or updated attachment object.
         """
+        if not thread_id:
+            if active_thread := active_thread_var.get(None):
+                thread_id = active_thread.id
+
+        if not step_id:
+            if active_steps := active_steps_var.get([]):
+                step_id = active_steps[-1].id
+            else:
+                raise Exception("No step_id provided and no active step found.")
+
         (
             query,
             description,
@@ -1849,7 +1860,7 @@ async def delete_score(self, id: str):
     async def upload_file(
         self,
         content: Union[bytes, str],
-        thread_id: str,
+        thread_id: Optional[str] = None,
         mime: Optional[str] = "application/octet-stream",
     ) -> Dict:
         """

literalai/api/attachment_helpers.py

@@ -7,8 +7,8 @@
 
 
 def create_attachment_helper(
-    thread_id: str,
     step_id: str,
+    thread_id: Optional[str] = None,
     id: Optional[str] = None,
     metadata: Optional[Dict] = None,
     mime: Optional[str] = None,

literalai/step.py

@@ -60,7 +60,7 @@ class StepDict(TypedDict, total=False):
 
 
 class Step(Utils):
-    id: Optional[str] = None
+    id: str
     name: Optional[str] = ""
     type: Optional[StepType] = None
     metadata: Optional[Dict] = None
@@ -169,7 +169,7 @@ def from_dict(cls, step_dict: StepDict) -> "Step":
 
         step = cls(name=name, type=step_type, thread_id=thread_id)
 
-        step.id = step_dict.get("id")
+        step.id = step_dict.get("id") or ""
         step.input = step_dict.get("input", None)
         step.error = step_dict.get("error", None)
         step.output = step_dict.get("output", None)
@@ -234,7 +234,7 @@ async def __aenter__(self):
             id=self.id,
             parent_id=self.parent_id,
             thread_id=self.thread_id,
-            **self.kwargs
+            **self.kwargs,
         )
         return self.step
 
@@ -251,7 +251,7 @@ def __enter__(self) -> Step:
             id=self.id,
             parent_id=self.parent_id,
             thread_id=self.thread_id,
-            **self.kwargs
+            **self.kwargs,
         )
         return self.step
 
@@ -271,7 +271,7 @@ def step_decorator(
     parent_id: Optional[str] = None,
     thread_id: Optional[str] = None,
     ctx_manager: Optional[StepContextManager] = None,
-    **decorator_kwargs
+    **decorator_kwargs,
 ):
     if not name:
         name = func.__name__
@@ -283,7 +283,7 @@ def step_decorator(
             id=id,
             parent_id=parent_id,
             thread_id=thread_id,
-            **decorator_kwargs
+            **decorator_kwargs,
         )
     else:
         ctx_manager.step_name = name