Skip to content

Commit 27694a0

Browse files
hassiebphassiebp
committed
add environment to stateful clients
1 parent 3c522ba commit 27694a0

File tree

7 files changed

+68
-35
lines changed

7 files changed

+68
-35
lines changed

‎langfuse/client.py‎

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1508,7 +1508,12 @@ def trace(
15081508
self._log_memory_usage()
15091509

15101510
return StatefulTraceClient(
1511-
self.client, new_id, StateType.TRACE, new_id, self.task_manager
1511+
self.client,
1512+
new_id,
1513+
StateType.TRACE,
1514+
new_id,
1515+
self.task_manager,
1516+
self.environment,
15121517
)
15131518

15141519
def _log_memory_usage(self):
@@ -1792,6 +1797,7 @@ def span(
17921797
StateType.OBSERVATION,
17931798
new_trace_id,
17941799
self.task_manager,
1800+
self.environment,
17951801
)
17961802

17971803
def event(
@@ -1892,6 +1898,7 @@ def event(
18921898
StateType.OBSERVATION,
18931899
new_trace_id,
18941900
self.task_manager,
1901+
self.environment,
18951902
)
18961903

18971904
def generation(
@@ -2039,6 +2046,7 @@ def generation(
20392046
StateType.OBSERVATION,
20402047
new_trace_id,
20412048
self.task_manager,
2049+
self.environment,
20422050
)
20432051

20442052
def _generate_trace(self, trace_id: str, name: str):
@@ -2172,7 +2180,6 @@ def __init__(
21722180
self.log.warning(
21732181
f'Invalid environment specified "{environment}" that does not match validation pattern ("{ENVIRONMENT_PATTERN}"). Setting will be ignored.'
21742182
)
2175-
self.environment = None
21762183

21772184
def _add_state_to_event(self, body: dict):
21782185
if self.state_type == StateType.OBSERVATION:
@@ -2304,7 +2311,8 @@ def generation(
23042311
generation_id,
23052312
StateType.OBSERVATION,
23062313
self.trace_id,
2307-
task_manager=self.task_manager,
2314+
self.task_manager,
2315+
self.environment,
23082316
)
23092317

23102318
def span(
@@ -2394,7 +2402,8 @@ def span(
23942402
span_id,
23952403
StateType.OBSERVATION,
23962404
self.trace_id,
2397-
task_manager=self.task_manager,
2405+
self.task_manager,
2406+
self.environment,
23982407
)
23992408

24002409
@overload
@@ -2627,9 +2636,10 @@ def __init__(
26272636
state_type: StateType,
26282637
trace_id: str,
26292638
task_manager: TaskManager,
2639+
environment: Optional[str] = None,
26302640
):
26312641
"""Initialize the StatefulGenerationClient."""
2632-
super().__init__(client, id, state_type, trace_id, task_manager)
2642+
super().__init__(client, id, state_type, trace_id, task_manager, environment)
26332643

26342644
# WHEN CHANGING THIS METHOD, UPDATE END() FUNCTION ACCORDINGLY
26352645
def update(
@@ -2741,7 +2751,8 @@ def update(
27412751
self.id,
27422752
StateType.OBSERVATION,
27432753
self.trace_id,
2744-
task_manager=self.task_manager,
2754+
self.task_manager,
2755+
self.environment,
27452756
)
27462757

27472758
def end(
@@ -2846,9 +2857,10 @@ def __init__(
28462857
state_type: StateType,
28472858
trace_id: str,
28482859
task_manager: TaskManager,
2860+
environment: Optional[str] = None,
28492861
):
28502862
"""Initialize the StatefulSpanClient."""
2851-
super().__init__(client, id, state_type, trace_id, task_manager)
2863+
super().__init__(client, id, state_type, trace_id, task_manager, environment)
28522864

28532865
# WHEN CHANGING THIS METHOD, UPDATE END() FUNCTION ACCORDINGLY
28542866
def update(
@@ -2932,7 +2944,8 @@ def update(
29322944
self.id,
29332945
StateType.OBSERVATION,
29342946
self.trace_id,
2935-
task_manager=self.task_manager,
2947+
self.task_manager,
2948+
self.environment,
29362949
)
29372950

29382951
def end(
@@ -3005,7 +3018,8 @@ def end(
30053018
self.id,
30063019
StateType.OBSERVATION,
30073020
self.trace_id,
3008-
task_manager=self.task_manager,
3021+
self.task_manager,
3022+
self.environment,
30093023
)
30103024

30113025
def get_langchain_handler(self, update_parent: bool = False):
@@ -3044,9 +3058,10 @@ def __init__(
30443058
state_type: StateType,
30453059
trace_id: str,
30463060
task_manager: TaskManager,
3061+
environment: Optional[str] = None,
30473062
):
30483063
"""Initialize the StatefulTraceClient."""
3049-
super().__init__(client, id, state_type, trace_id, task_manager)
3064+
super().__init__(client, id, state_type, trace_id, task_manager, environment)
30503065
self.task_manager = task_manager
30513066

30523067
def update(
@@ -3137,7 +3152,8 @@ def update(
31373152
self.id,
31383153
StateType.TRACE,
31393154
self.trace_id,
3140-
task_manager=self.task_manager,
3155+
self.task_manager,
3156+
self.environment,
31413157
)
31423158

31433159
def get_langchain_handler(self, update_parent: bool = False):

‎langfuse/decorators/langfuse_decorator.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,7 @@ def _prepare_call(
344344
task_manager=self.client_instance.task_manager,
345345
client=self.client_instance.client,
346346
state_type=StateType.OBSERVATION,
347+
environment=self.client_instance.environment,
347348
)
348349
self._set_root_trace_id(provided_parent_trace_id)
349350

@@ -354,6 +355,7 @@ def _prepare_call(
354355
task_manager=self.client_instance.task_manager,
355356
client=self.client_instance.client,
356357
state_type=StateType.TRACE,
358+
environment=self.client_instance.environment,
357359
)
358360
self._set_root_trace_id(provided_parent_trace_id)
359361

@@ -490,6 +492,7 @@ def _handle_call_result(
490492
task_manager=self.client_instance.task_manager,
491493
client=self.client_instance.client,
492494
state_type=StateType.TRACE,
495+
environment=self.client_instance.environment,
493496
)
494497
trace_client.update(**observation_params)
495498

‎langfuse/llama_index/_event_handler.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,31 @@
1-
from typing import Optional, Any, Union, Mapping
1+
from typing import Any, Mapping, Optional, Union
2+
from uuid import uuid4 as create_uuid
23

34
from langfuse.client import (
45
Langfuse,
56
StatefulGenerationClient,
67
StateType,
78
)
89
from langfuse.utils import _get_timestamp
10+
911
from ._context import InstrumentorContext
10-
from uuid import uuid4 as create_uuid
1112

1213
try:
1314
from llama_index.core.base.llms.types import (
1415
ChatResponse,
1516
CompletionResponse,
1617
)
18+
from llama_index.core.instrumentation.event_handlers import BaseEventHandler
1719
from llama_index.core.instrumentation.events import BaseEvent
1820
from llama_index.core.instrumentation.events.embedding import (
19-
EmbeddingStartEvent,
2021
EmbeddingEndEvent,
22+
EmbeddingStartEvent,
2123
)
22-
from llama_index.core.instrumentation.event_handlers import BaseEventHandler
2324
from llama_index.core.instrumentation.events.llm import (
24-
LLMCompletionEndEvent,
25-
LLMCompletionStartEvent,
2625
LLMChatEndEvent,
2726
LLMChatStartEvent,
27+
LLMCompletionEndEvent,
28+
LLMCompletionStartEvent,
2829
)
2930
from llama_index.core.utilities.token_counting import TokenCounter
3031

@@ -148,6 +149,7 @@ def _get_generation_client(self, id: str) -> StatefulGenerationClient:
148149
trace_id=trace_id,
149150
task_manager=self._langfuse.task_manager,
150151
state_type=StateType.OBSERVATION,
152+
environment=self._langfuse.environment,
151153
)
152154

153155

‎langfuse/llama_index/_instrumentor.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,7 @@ def _get_trace_client(self, trace_id: str) -> StatefulTraceClient:
288288
trace_id=trace_id,
289289
task_manager=self._langfuse.task_manager,
290290
state_type=StateType.TRACE,
291+
environment=self._langfuse.environment,
291292
)
292293

293294
@property

‎langfuse/llama_index/_span_handler.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,32 @@
11
import inspect
2-
from typing import Optional, Any, Tuple, Generator, AsyncGenerator
32
import uuid
3+
from logging import getLogger
4+
from typing import Any, AsyncGenerator, Generator, Optional, Tuple
5+
6+
from pydantic import BaseModel
47

58
from langfuse.client import (
69
Langfuse,
7-
StatefulSpanClient,
10+
StatefulClient,
811
StatefulGenerationClient,
12+
StatefulSpanClient,
913
StateType,
10-
StatefulClient,
1114
)
1215

13-
from logging import getLogger
14-
from pydantic import BaseModel
15-
1616
from ._context import InstrumentorContext
1717

1818
logger = getLogger(__name__)
1919

2020
try:
21-
from llama_index.core.instrumentation.span_handlers import BaseSpanHandler
22-
from llama_index.core.instrumentation.span import BaseSpan
23-
from llama_index.core.base.embeddings.base import BaseEmbedding
2421
from llama_index.core.base.base_query_engine import BaseQueryEngine
25-
from llama_index.core.llms import LLM, ChatResponse
22+
from llama_index.core.base.embeddings.base import BaseEmbedding
2623
from llama_index.core.base.response.schema import (
27-
StreamingResponse,
2824
AsyncStreamingResponse,
25+
StreamingResponse,
2926
)
27+
from llama_index.core.instrumentation.span import BaseSpan
28+
from llama_index.core.instrumentation.span_handlers import BaseSpanHandler
29+
from llama_index.core.llms import LLM, ChatResponse
3030
from llama_index.core.workflow import Context
3131

3232
except ImportError:
@@ -215,6 +215,7 @@ def _get_generation_client(self, id: str) -> StatefulGenerationClient:
215215
trace_id=trace_id,
216216
task_manager=self._langfuse_client.task_manager,
217217
state_type=StateType.OBSERVATION,
218+
environment=self._langfuse_client.environment,
218219
)
219220

220221
def _get_span_client(self, id: str) -> StatefulSpanClient:
@@ -231,6 +232,7 @@ def _get_span_client(self, id: str) -> StatefulSpanClient:
231232
trace_id=trace_id,
232233
task_manager=self._langfuse_client.task_manager,
233234
state_type=StateType.OBSERVATION,
235+
environment=self._langfuse_client.environment,
234236
)
235237

236238
def _parse_generation_input(

‎langfuse/llama_index/llama_index.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,34 @@
1+
import logging
12
from collections import defaultdict
23
from contextvars import ContextVar
3-
from typing import Any, Dict, List, Optional, Union, Tuple, Callable, Generator
4+
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
45
from uuid import uuid4
5-
import logging
6+
67
import httpx
8+
from pydantic import BaseModel
79

810
from langfuse.client import (
11+
StatefulGenerationClient,
912
StatefulSpanClient,
1013
StatefulTraceClient,
11-
StatefulGenerationClient,
1214
StateType,
1315
)
16+
from langfuse.types import TraceMetadata
17+
from langfuse.utils.base_callback_handler import LangfuseBaseCallbackHandler
1418
from langfuse.utils.error_logging import (
1519
auto_decorate_methods_with,
1620
catch_and_log_errors,
1721
)
18-
from langfuse.types import TraceMetadata
19-
from langfuse.utils.base_callback_handler import LangfuseBaseCallbackHandler
22+
2023
from .utils import CallbackEvent, ParsedLLMEndPayload
21-
from pydantic import BaseModel
2224

2325
try:
2426
from llama_index.core.callbacks.base_handler import (
2527
BaseCallbackHandler as LlamaIndexBaseCallbackHandler,
2628
)
2729
from llama_index.core.callbacks.schema import (
28-
CBEventType,
2930
BASE_TRACE_EVENT,
31+
CBEventType,
3032
EventPayload,
3133
)
3234
from llama_index.core.utilities.token_counting import TokenCounter
@@ -165,6 +167,7 @@ def set_root(
165167
StateType.TRACE,
166168
root.trace_id,
167169
root.task_manager,
170+
root.environment,
168171
)
169172

170173
self._task_manager = root.task_manager

‎tests/test_core_sdk.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1538,7 +1538,13 @@ def test_environment_from_constructor():
15381538
api_wrapper = LangfuseAPI()
15391539

15401540
trace = langfuse.trace(name="test_environment")
1541+
sleep(0.1)
1542+
trace.update(name="updated_name")
1543+
15411544
generation = trace.generation(name="test_gen")
1545+
sleep(0.1)
1546+
generation.update(name="test_gen_1")
1547+
15421548
score_id = create_uuid()
15431549
langfuse.score(id=score_id, trace_id=trace.id, name="test_score", value=1)
15441550

0 commit comments

Comments
 (0)