Skip to content

Commit 702d277

Browse files
Steffen911Steffen911
committed
chore: typing
1 parent f15d538 commit 702d277

3 files changed

Lines changed: 31 additions & 22 deletions

File tree

‎langfuse/_client/client.py‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
Any,
1818
Callable,
1919
Dict,
20+
Generator,
2021
List,
2122
Literal,
2223
Optional,
@@ -198,6 +199,7 @@ class Langfuse:
198199
_resources: Optional[LangfuseResourceManager] = None
199200
_mask: Optional[MaskFunction] = None
200201
_otel_tracer: otel_trace_api.Tracer
202+
_host: str
201203

202204
def __init__(
203205
self,
@@ -220,8 +222,10 @@ def __init__(
220222
additional_headers: Optional[Dict[str, str]] = None,
221223
tracer_provider: Optional[TracerProvider] = None,
222224
):
223-
self._host = host or cast(
224-
str, os.environ.get(LANGFUSE_HOST, "https://cloud.langfuse.com")
225+
self._host = (
226+
host
227+
if host is not None
228+
else os.environ.get(LANGFUSE_HOST, "https://cloud.langfuse.com")
225229
)
226230
self._environment = environment or cast(
227231
str, os.environ.get(LANGFUSE_TRACING_ENVIRONMENT)
@@ -3466,7 +3470,9 @@ def clear_prompt_cache(self) -> None:
34663470
self._resources.prompt_cache.clear()
34673471

34683472
@_agnosticcontextmanager
3469-
def session(self, id: str, *, as_baggage: bool = False) -> _AgnosticContextManager:
3473+
def session(
3474+
self, id: str, *, as_baggage: bool = False
3475+
) -> Generator[None, None, None]:
34703476
"""Create a session context manager that propagates session_id to all child spans.
34713477
34723478
Args:
@@ -3518,7 +3524,7 @@ def session(self, id: str, *, as_baggage: bool = False) -> _AgnosticContextManag
35183524
otel_context_api.detach(baggage_token)
35193525

35203526
@_agnosticcontextmanager
3521-
def user(self, id: str, *, as_baggage: bool = False) -> _AgnosticContextManager:
3527+
def user(self, id: str, *, as_baggage: bool = False) -> Generator[None, None, None]:
35223528
"""Create a user context manager that propagates user_id to all child spans.
35233529
35243530
Args:
@@ -3571,8 +3577,8 @@ def user(self, id: str, *, as_baggage: bool = False) -> _AgnosticContextManager:
35713577

35723578
@_agnosticcontextmanager
35733579
def metadata(
3574-
self, *, as_baggage: bool = False, **kwargs
3575-
) -> _AgnosticContextManager:
3580+
self, *, as_baggage: bool = False, **kwargs: Any
3581+
) -> Generator[None, None, None]:
35763582
"""Create a metadata context manager that propagates metadata to all child spans.
35773583
35783584
Args:
@@ -3618,8 +3624,8 @@ def metadata(
36183624
# Set baggage if requested
36193625
baggage_tokens = []
36203626
if as_baggage:
3621-
current_baggage = otel_baggage_api.get_all()
3622-
new_baggage = current_baggage
3627+
# Start with None context and chain baggage settings
3628+
new_baggage = None
36233629

36243630
# Add each metadata key-value pair to baggage
36253631
for key, value in kwargs.items():
@@ -3634,7 +3640,7 @@ def metadata(
36343640
)
36353641

36363642
# Attach the new baggage context
3637-
if new_baggage != current_baggage:
3643+
if new_baggage is not None:
36383644
baggage_token = otel_context_api.attach(new_baggage)
36393645
baggage_tokens.append(baggage_token)
36403646

‎langfuse/_client/span.py‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
TYPE_CHECKING,
2222
Any,
2323
Dict,
24+
Generator,
2425
List,
2526
Literal,
2627
Optional,
@@ -1134,7 +1135,7 @@ def start_as_current_observation( # type: ignore[misc]
11341135
@_agnosticcontextmanager
11351136
def session(
11361137
self, id: str, *, as_baggage: bool = False
1137-
) -> "_AgnosticContextManager":
1138+
) -> Generator[None, None, None]:
11381139
"""Create a session context manager that propagates session_id to all child spans.
11391140
11401141
Args:
@@ -1185,7 +1186,7 @@ def session(
11851186
otel_context_api.detach(baggage_token)
11861187

11871188
@_agnosticcontextmanager
1188-
def user(self, id: str, *, as_baggage: bool = False) -> "_AgnosticContextManager":
1189+
def user(self, id: str, *, as_baggage: bool = False) -> Generator[None, None, None]:
11891190
"""Create a user context manager that propagates user_id to all child spans.
11901191
11911192
Args:
@@ -1237,8 +1238,8 @@ def user(self, id: str, *, as_baggage: bool = False) -> "_AgnosticContextManager
12371238

12381239
@_agnosticcontextmanager
12391240
def metadata(
1240-
self, *, as_baggage: bool = False, **kwargs
1241-
) -> "_AgnosticContextManager":
1241+
self, *, as_baggage: bool = False, **kwargs: Any
1242+
) -> Generator[None, None, None]:
12421243
"""Create a metadata context manager that propagates metadata to all child spans.
12431244
12441245
Args:
@@ -1283,8 +1284,8 @@ def metadata(
12831284
# Set baggage if requested
12841285
baggage_tokens = []
12851286
if as_baggage:
1286-
current_baggage = otel_baggage_api.get_all()
1287-
new_baggage = current_baggage
1287+
# Start with None context and chain baggage settings
1288+
new_baggage = None
12881289

12891290
# Add each metadata key-value pair to baggage
12901291
for key, value in kwargs.items():
@@ -1299,7 +1300,7 @@ def metadata(
12991300
)
13001301

13011302
# Attach the new baggage context
1302-
if new_baggage != current_baggage:
1303+
if new_baggage is not None:
13031304
baggage_token = otel_context_api.attach(new_baggage)
13041305
baggage_tokens.append(baggage_token)
13051306

‎langfuse/_client/span_processor.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,9 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
151151
baggage_entries = baggage.get_all(context=current_context)
152152
for key, value in baggage_entries.items():
153153
# Check if this baggage entry is already present as a span attribute
154-
if not hasattr(span.attributes, key) or span.attributes.get(key) != value:
154+
if not hasattr(span.attributes, key) or (
155+
span.attributes is not None and span.attributes.get(key) != value
156+
):
155157
propagated_attributes[key] = value
156158
langfuse_logger.debug(
157159
f"Propagated baggage key '{key}' = '{value}' to span '{span.name}'"
@@ -167,9 +169,9 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
167169
attr_key = ctx_key.replace("langfuse.ctx.", "")
168170

169171
# Only propagate if not already set on span
170-
if (
171-
not hasattr(span.attributes, attr_key)
172-
or span.attributes.get(attr_key) != value
172+
if not hasattr(span.attributes, attr_key) or (
173+
span.attributes is not None
174+
and span.attributes.get(attr_key) != value
173175
):
174176
propagated_attributes[attr_key] = value
175177
langfuse_logger.debug(
@@ -203,7 +205,7 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
203205
# Only propagate if not already set or different
204206
existing_metadata = (
205207
span.attributes.get("metadata")
206-
if hasattr(span, "attributes")
208+
if hasattr(span, "attributes") and span.attributes is not None
207209
else None
208210
)
209211
if existing_metadata != metadata_json:
@@ -222,7 +224,7 @@ def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None
222224

223225
# Set all propagated attributes on the span
224226
for key, value in propagated_attributes.items():
225-
span.set_attribute(key, value)
227+
span.set_attribute(key, value) # type: ignore[arg-type]
226228

227229
return super().on_start(span, parent_context)
228230

0 commit comments

Comments
 (0)