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 (
@@ -589,7 +590,6 @@ def create_score(
                 value,
                 type,
                 step_id,
-                generation_id,
                 dataset_experiment_item_id,
                 comment,
                 tags,
@@ -711,8 +711,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 +740,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,
@@ -1808,7 +1818,6 @@ async def create_score(
                 value,
                 type,
                 step_id,
-                generation_id,
                 dataset_experiment_item_id,
                 comment,
                 tags,
@@ -1849,7 +1858,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/api/gql.py

@@ -395,7 +395,6 @@
                 id
                 projectId
                 stepId
-                generationId
                 datasetExperimentItemId
                 type
                 updatedAt
@@ -422,7 +421,6 @@
     $type: ScoreType!,
     $value: Float!,
     $stepId: String,
-    $generationId: String,
     $datasetExperimentItemId: String,
     $comment: String,
     $tags: [String!],
@@ -433,7 +431,6 @@
         type: $type,
         value: $value,
         stepId: $stepId,
-        generationId: $generationId,
         datasetExperimentItemId: $datasetExperimentItemId,
         comment: $comment,
         tags: $tags,
@@ -443,7 +440,6 @@
         type,
         value,
         stepId,
-        generationId,
         datasetExperimentItemId,
         comment,
         tags,
@@ -467,7 +463,6 @@
         type,
         value,
         stepId,
-        generationId,
         datasetExperimentItemId,
         comment
     }

literalai/api/score_helpers.py

@@ -42,7 +42,6 @@ def create_score_helper(
     value: float,
     type: ScoreType,
     step_id: Optional[str] = None,
-    generation_id: Optional[str] = None,
     dataset_experiment_item_id: Optional[str] = None,
     comment: Optional[str] = None,
     tags: Optional[List[str]] = None,
@@ -52,7 +51,6 @@ def create_score_helper(
         "type": type,
         "value": value,
         "stepId": step_id,
-        "generationId": generation_id,
         "datasetExperimentItemId": dataset_experiment_item_id,
         "comment": comment,
         "tags": tags,
@@ -73,7 +71,6 @@ def create_scores_fields_builder(scores: List[ScoreDict]):
         $type_{id}: ScoreType!
         $value_{id}: Float!
         $stepId_{id}: String
-        $generationId_{id}: String
         $datasetExperimentItemId_{id}: String
         $scorer_{id}: String
         $comment_{id}: String
@@ -91,7 +88,6 @@ def create_scores_args_builder(scores: List[ScoreDict]):
                 type: $type_{id}
                 value: $value_{id}
                 stepId: $stepId_{id}
-                generationId: $generationId_{id}
                 datasetExperimentItemId: $datasetExperimentItemId_{id}
                 scorer: $scorer_{id}
                 comment: $comment_{id}

literalai/my_types.py

@@ -246,7 +246,6 @@ class ScoreDict(TypedDict, total=False):
     type: ScoreType
     value: float
     stepId: Optional[str]
-    generationId: Optional[str]
     datasetExperimentItemId: Optional[str]
     comment: Optional[str]
     tags: Optional[List[str]]
@@ -269,7 +268,6 @@ class Score(Utils):
     type: ScoreType
     value: float
     step_id: Optional[str]
-    generation_id: Optional[str]
     dataset_experiment_item_id: Optional[str]
     comment: Optional[str]
     tags: Optional[List[str]]
@@ -282,7 +280,6 @@ def to_dict(self):
             "type": self.type,
             "value": self.value,
             "stepId": self.step_id,
-            "generationId": self.generation_id,
             "datasetExperimentItemId": self.dataset_experiment_item_id,
             "comment": self.comment,
             "tags": self.tags,
@@ -295,7 +292,6 @@ def from_dict(cls, score_dict: ScoreDict) -> "Score":
         type = score_dict.get("type", "HUMAN")
         value = score_dict.get("value", 0.0)
         step_id = score_dict.get("stepId", "")
-        generation_id = score_dict.get("generationId", "")
         dataset_experiment_item_id = score_dict.get("datasetExperimentItemId", "")
         comment = score_dict.get("comment", "")
         tags = score_dict.get("tags", [])
@@ -306,7 +302,6 @@ def from_dict(cls, score_dict: ScoreDict) -> "Score":
             type=type,
             value=value,
             step_id=step_id,
-            generation_id=generation_id,
             dataset_experiment_item_id=dataset_experiment_item_id,
             comment=comment,
             tags=tags,

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