Skip to content
This repository was archived by the owner on Aug 5, 2025. It is now read-only.

Commit 89e966c

Browse files
clementsirieixclementsirieix
authored
Merge pull request #107 from Chainlit/clement/eng-1746-update-sdks-to-send-the-root-run
feat: update sdks to send the root run
2 parents 81ae21d + 6c64445 commit 89e966c

File tree

9 files changed

+151
-14
lines changed

9 files changed

+151
-14
lines changed

‎literalai/api/__init__.py‎

Lines changed: 32 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,18 @@
1010
Literal,
1111
Optional,
1212
TypeVar,
13-
Union, cast,
13+
Union,
14+
cast,
1415
)
1516

1617
from typing_extensions import deprecated
1718

1819
from literalai.context import active_steps_var, active_thread_var
1920
from literalai.evaluation.dataset import Dataset, DatasetType
20-
from literalai.evaluation.dataset_experiment import DatasetExperiment, DatasetExperimentItem
21+
from literalai.evaluation.dataset_experiment import (
22+
DatasetExperiment,
23+
DatasetExperimentItem,
24+
)
2125
from literalai.observability.filter import (
2226
generations_filters,
2327
generations_order_by,
@@ -51,7 +55,10 @@
5155
get_dataset_item_helper,
5256
update_dataset_helper,
5357
)
54-
from literalai.api.generation_helpers import create_generation_helper, get_generations_helper
58+
from literalai.api.generation_helpers import (
59+
create_generation_helper,
60+
get_generations_helper,
61+
)
5562
from literalai.api.prompt_helpers import (
5663
create_prompt_helper,
5764
create_prompt_lineage_helper,
@@ -101,8 +108,20 @@
101108
Environment,
102109
PaginatedResponse,
103110
)
104-
from literalai.observability.generation import GenerationMessage, CompletionGeneration, ChatGeneration
105-
from literalai.observability.step import Step, StepDict, StepType, ScoreType, ScoreDict, Score, Attachment
111+
from literalai.observability.generation import (
112+
GenerationMessage,
113+
CompletionGeneration,
114+
ChatGeneration,
115+
)
116+
from literalai.observability.step import (
117+
Step,
118+
StepDict,
119+
StepType,
120+
ScoreType,
121+
ScoreDict,
122+
Score,
123+
Attachment,
124+
)
106125

107126
logger = logging.getLogger(__name__)
108127

@@ -678,8 +697,7 @@ def upload_file(
678697
fields: Dict = request_dict.get("fields", {})
679698
object_key: Optional[str] = fields.get("key")
680699
upload_type: Literal["raw", "multipart"] = cast(
681-
Literal["raw", "multipart"],
682-
request_dict.get("uploadType", "multipart")
700+
Literal["raw", "multipart"], request_dict.get("uploadType", "multipart")
683701
)
684702
signed_url: Optional[str] = json_res.get("signedUrl")
685703

@@ -843,6 +861,7 @@ def create_step(
843861
parent_id: Optional[str] = None,
844862
name: Optional[str] = None,
845863
tags: Optional[List[str]] = None,
864+
root_run_id: Optional[str] = None,
846865
):
847866
"""
848867
Creates a new step with the specified parameters.
@@ -858,6 +877,7 @@ def create_step(
858877
parent_id (Optional[str]): The ID of the parent step, if any.
859878
name (Optional[str]): The name of the step.
860879
tags (Optional[List[str]]): Tags associated with the step.
880+
root_run_id (Optional[str]): The ID of the root run, if any.
861881
862882
Returns:
863883
The result of the GraphQL helper function for creating a step.
@@ -874,6 +894,7 @@ def create_step(
874894
parent_id=parent_id,
875895
name=name,
876896
tags=tags,
897+
root_run_id=root_run_id,
877898
)
878899
)
879900

@@ -1912,8 +1933,7 @@ async def upload_file(
19121933
fields: Dict = request_dict.get("fields", {})
19131934
object_key: Optional[str] = fields.get("key")
19141935
upload_type: Literal["raw", "multipart"] = cast(
1915-
Literal["raw", "multipart"],
1916-
request_dict.get("uploadType", "multipart")
1936+
Literal["raw", "multipart"], request_dict.get("uploadType", "multipart")
19171937
)
19181938
signed_url: Optional[str] = json_res.get("signedUrl")
19191939

@@ -2069,6 +2089,7 @@ async def create_step(
20692089
parent_id: Optional[str] = None,
20702090
name: Optional[str] = None,
20712091
tags: Optional[List[str]] = None,
2092+
root_run_id: Optional[str] = None,
20722093
):
20732094
"""
20742095
Asynchronously creates a new step with the specified parameters.
@@ -2084,6 +2105,7 @@ async def create_step(
20842105
parent_id (Optional[str]): The ID of the parent step, if any.
20852106
name (Optional[str]): The name of the step.
20862107
tags (Optional[List[str]]): Tags associated with the step.
2108+
root_run_id (Optional[str]): The ID of the root run, if any.
20872109
20882110
Returns:
20892111
The result of the GraphQL helper function for creating a step.
@@ -2100,6 +2122,7 @@ async def create_step(
21002122
parent_id=parent_id,
21012123
name=name,
21022124
tags=tags,
2125+
root_run_id=root_run_id,
21032126
)
21042127
)
21052128

‎literalai/api/gql.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
STEP_FIELDS = """
77
id
88
threadId
9+
rootRunId
910
parentId
1011
startTime
1112
endTime
@@ -559,6 +560,7 @@
559560
"""
560561
mutation CreateStep(
561562
$threadId: String,
563+
$rootRunId: String,
562564
$type: StepType,
563565
$startTime: DateTime,
564566
$endTime: DateTime,
@@ -572,6 +574,7 @@
572574
) {
573575
createStep(
574576
threadId: $threadId,
577+
rootRunId: $rootRunId,
575578
type: $type,
576579
startTime: $startTime,
577580
endTime: $endTime,
@@ -878,12 +881,14 @@
878881
$input: Json!
879882
$expectedOutput: Json
880883
$metadata: Json
884+
$generationId: String
881885
) {
882886
createDatasetItem(
883887
datasetId: $datasetId
884888
input: $input
885889
expectedOutput: $expectedOutput
886890
metadata: $metadata
891+
generationId: $generationId
887892
) {
888893
id
889894
createdAt
@@ -892,6 +897,7 @@
892897
input
893898
expectedOutput
894899
intermediarySteps
900+
stepId
895901
}
896902
}
897903
"""
@@ -906,6 +912,7 @@
906912
input
907913
expectedOutput
908914
intermediarySteps
915+
stepId
909916
}
910917
}
911918
"""
@@ -920,6 +927,7 @@
920927
input
921928
expectedOutput
922929
intermediarySteps
930+
stepId
923931
}
924932
}
925933
"""
@@ -942,6 +950,7 @@
942950
input
943951
expectedOutput
944952
intermediarySteps
953+
stepId
945954
}
946955
}
947956
"""
@@ -964,6 +973,7 @@
964973
input
965974
expectedOutput
966975
intermediarySteps
976+
stepId
967977
}
968978
}
969979
"""
@@ -1080,6 +1090,7 @@ def steps_query_variables_builder(steps):
10801090
for id in range(len(steps)):
10811091
generated += f"""$id_{id}: String!
10821092
$threadId_{id}: String
1093+
$rootRunId_{id}: String
10831094
$type_{id}: StepType
10841095
$startTime_{id}: DateTime
10851096
$endTime_{id}: DateTime
@@ -1104,6 +1115,7 @@ def steps_ingest_steps_builder(steps):
11041115
step{id}: ingestStep(
11051116
id: $id_{id}
11061117
threadId: $threadId_{id}
1118+
rootRunId: $rootRunId_{id}
11071119
startTime: $startTime_{id}
11081120
endTime: $endTime_{id}
11091121
type: $type_{id}

‎literalai/api/step_helpers.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ def create_step_helper(
1818
parent_id: Optional[str] = None,
1919
name: Optional[str] = None,
2020
tags: Optional[List[str]] = None,
21+
root_run_id: Optional[str] = None,
2122
):
2223
variables = {
2324
"threadId": thread_id,
@@ -30,6 +31,7 @@ def create_step_helper(
3031
"parentId": parent_id,
3132
"name": name,
3233
"tags": tags,
34+
"root_run_id": root_run_id,
3335
}
3436

3537
def process_response(response):

‎literalai/client.py‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from literalai.api import AsyncLiteralAPI, LiteralAPI
55
from literalai.callback.langchain_callback import get_langchain_callback
6-
from literalai.context import active_steps_var, active_thread_var
6+
from literalai.context import active_steps_var, active_thread_var, active_root_run_var
77
from literalai.environment import EnvContextManager, env_decorator
88
from literalai.event_processor import EventProcessor
99
from literalai.evaluation.experiment_item_run import (
@@ -20,7 +20,8 @@
2020
Step,
2121
StepContextManager,
2222
TrueStepType,
23-
step_decorator, Attachment,
23+
step_decorator,
24+
Attachment,
2425
)
2526
from literalai.observability.thread import ThreadContextManager, thread_decorator
2627

@@ -153,6 +154,7 @@ def step(
153154
id: Optional[str] = None,
154155
parent_id: Optional[str] = None,
155156
thread_id: Optional[str] = None,
157+
root_run_id: Optional[str] = None,
156158
**kwargs,
157159
):
158160
"""
@@ -167,6 +169,7 @@ def step(
167169
id (Optional[str]): The id of the step to create.
168170
parent_id (Optional[str]): The id of the parent step.
169171
thread_id (Optional[str]): The id of the parent thread.
172+
root_run_id (Optional[str]): The id of the root run.
170173
171174
Returns:
172175
The wrapper for the step's context.
@@ -180,6 +183,7 @@ def step(
180183
id=id,
181184
parent_id=parent_id,
182185
thread_id=thread_id,
186+
root_run_id=root_run_id,
183187
**kwargs,
184188
)
185189
else:
@@ -190,6 +194,7 @@ def step(
190194
id=id,
191195
parent_id=parent_id,
192196
thread_id=thread_id,
197+
root_run_id=root_run_id,
193198
**kwargs,
194199
)
195200

@@ -201,6 +206,7 @@ def run(
201206
id: Optional[str] = None,
202207
parent_id: Optional[str] = None,
203208
thread_id: Optional[str] = None,
209+
root_run_id: Optional[str] = None,
204210
):
205211
"""
206212
Creates a run where all the subsequents steps will be logged. Works as a decorator or a ContextManager.
@@ -211,6 +217,7 @@ def run(
211217
id (Optional[str]): The id of the step to create.
212218
parent_id (Optional[str]): The id of the parent step.
213219
thread_id (Optional[str]): The id of the parent thread.
220+
root_run_id (Optional[str]): The id of the root run.
214221
215222
Returns:
216223
The wrapper for the step's context.
@@ -222,6 +229,7 @@ def run(
222229
id=id,
223230
parent_id=parent_id,
224231
thread_id=thread_id,
232+
root_run_id=root_run_id,
225233
)
226234

227235
def message(
@@ -235,6 +243,7 @@ def message(
235243
attachments: List[Attachment] = [],
236244
tags: Optional[List[str]] = None,
237245
metadata: Dict = {},
246+
root_run_id: Optional[str] = None,
238247
):
239248
"""
240249
Creates a conversational message step and sends it to Literal AI.
@@ -251,6 +260,7 @@ def message(
251260
attachments (List[Attachment]): A list of attachments to append to the message.
252261
tags (Optional[List[str]]): A list of tags to add to the message.
253262
metadata (Dict): Metadata to add to the message, in key-value pairs.
263+
root_run_id (Optional[str]): The id of the root run.
254264
255265
Returns:
256266
Message: the created message.
@@ -266,6 +276,7 @@ def message(
266276
tags=tags,
267277
metadata=metadata,
268278
processor=self.event_processor,
279+
root_run_id=root_run_id,
269280
)
270281
step.end()
271282

@@ -335,6 +346,7 @@ def start_step(
335346
id: Optional[str] = None,
336347
parent_id: Optional[str] = None,
337348
thread_id: Optional[str] = None,
349+
root_run_id: Optional[str] = None,
338350
**kwargs,
339351
):
340352
"""
@@ -348,6 +360,7 @@ def start_step(
348360
id (Optional[str]): The id of the step to create.
349361
parent_id (Optional[str]): The id of the parent step.
350362
thread_id (Optional[str]): The id of the parent thread.
363+
root_run_id (Optional[str]): The id of the root run.
351364
352365
Returns:
353366
Step: the created step.
@@ -359,6 +372,7 @@ def start_step(
359372
parent_id=parent_id,
360373
thread_id=thread_id,
361374
processor=self.event_processor,
375+
root_run_id=root_run_id,
362376
**kwargs,
363377
)
364378
step.start()
@@ -380,12 +394,19 @@ def get_current_thread(self):
380394
"""
381395
return active_thread_var.get()
382396

397+
def get_current_root_run(self):
398+
"""
399+
Gets the current root run from the context.
400+
"""
401+
return active_root_run_var.get()
402+
383403
def reset_context(self):
384404
"""
385405
Resets the context, forgetting active steps & setting current thread to None.
386406
"""
387407
active_steps_var.set([])
388408
active_thread_var.set(None)
409+
active_root_run_var.set(None)
389410

390411
def flush_and_stop(self):
391412
"""

‎literalai/context.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
active_steps_var = ContextVar[List["Step"]]("active_steps", default=[])
99
active_thread_var = ContextVar[Optional["Thread"]]("active_thread", default=None)
10+
active_root_run_var = ContextVar[Optional["Step"]]("active_root_run_var", default=None)
1011

1112
active_experiment_item_run_id_var = ContextVar[Optional[str]](
1213
"active_experiment_item_run", default=None

0 commit comments

Comments
 (0)