Skip to content

Commit 40e4f15

Browse files
committed
ruff formatting
1 parent b007681 commit 40e4f15

File tree

2 files changed

+34
-23
lines changed

2 files changed

+34
-23
lines changed

instrumentation-genai/opentelemetry-instrumentation-langchain/src/opentelemetry/instrumentation/langchain/span_manager.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,10 @@
1717
from typing import Dict, List, Optional
1818
from uuid import UUID
1919

20-
from langchain_core.callbacks import BaseCallbackHandler # type: ignore
21-
from langchain_core.messages import BaseMessage # type: ignore
22-
from langchain_core.outputs import LLMResult # type: ignore
23-
2420
from opentelemetry.context import Context, get_current
2521
from opentelemetry.semconv._incubating.attributes import (
2622
gen_ai_attributes as GenAI,
2723
)
28-
2924
from opentelemetry.trace import Span, SpanKind, Tracer, set_span_in_context
3025

3126

@@ -36,6 +31,7 @@ class _SpanState:
3631
start_time: float = field(default_factory=time.time)
3732
children: List[UUID] = field(default_factory=list)
3833

34+
3935
class SpanManager:
4036
def __init__(
4137
self,
@@ -102,4 +98,4 @@ def end_span(self, run_id: UUID) -> None:
10298
state.span.end()
10399

104100
def get_span(self, run_id: UUID) -> Span:
105-
return self.spans[run_id].span
101+
return self.spans[run_id].span

instrumentation-genai/opentelemetry-instrumentation-langchain/tests/test_span_manager.py

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,13 @@
22
import uuid
33

44
import pytest
5-
from opentelemetry.trace import SpanKind, get_tracer
6-
from opentelemetry.trace.span import Span
75

86
from opentelemetry.instrumentation.langchain.span_manager import (
97
SpanManager,
108
_SpanState,
119
)
10+
from opentelemetry.trace import SpanKind, get_tracer
11+
from opentelemetry.trace.span import Span
1212

1313

1414
class TestSpanManager:
@@ -28,8 +28,9 @@ def handler(self, tracer):
2828
(uuid.uuid4(), True), # Parent in spans
2929
],
3030
)
31-
32-
def test_create_span(self, handler, tracer, parent_run_id, parent_in_spans):
31+
def test_create_span(
32+
self, handler, tracer, parent_run_id, parent_in_spans
33+
):
3334
# Arrange
3435
run_id = uuid.uuid4()
3536
span_name = "test_span"
@@ -40,17 +41,25 @@ def test_create_span(self, handler, tracer, parent_run_id, parent_in_spans):
4041
# Setup parent if needed
4142
if parent_run_id is not None and parent_in_spans:
4243
parent_mock_span = unittest.mock.Mock(spec=Span)
43-
handler.spans[parent_run_id] = _SpanState(span=parent_mock_span, context=None)
44+
handler.spans[parent_run_id] = _SpanState(
45+
span=parent_mock_span, context=None
46+
)
4447

45-
with unittest.mock.patch.object(
46-
tracer, "start_span", return_value=mock_span
47-
) as mock_start_span, unittest.mock.patch(
48-
"opentelemetry.instrumentation.langchain.span_manager.set_span_in_context"
49-
) as mock_set_span_in_context, unittest.mock.patch(
50-
"opentelemetry.instrumentation.langchain.span_manager.get_current"
51-
) as mock_get_current:
48+
with (
49+
unittest.mock.patch.object(
50+
tracer, "start_span", return_value=mock_span
51+
) as mock_start_span,
52+
unittest.mock.patch(
53+
"opentelemetry.instrumentation.langchain.span_manager.set_span_in_context"
54+
) as mock_set_span_in_context,
55+
unittest.mock.patch(
56+
"opentelemetry.instrumentation.langchain.span_manager.get_current"
57+
) as mock_get_current,
58+
):
5259
# Act
53-
result = handler.create_span(run_id, parent_run_id, span_name, kind)
60+
result = handler.create_span(
61+
run_id, parent_run_id, span_name, kind
62+
)
5463

5564
# Assert
5665
assert result == mock_span
@@ -59,11 +68,17 @@ def test_create_span(self, handler, tracer, parent_run_id, parent_in_spans):
5968

6069
# Verify parent-child relationship
6170
if parent_run_id is not None and parent_in_spans:
62-
mock_set_span_in_context.assert_called_once_with(handler.spans[parent_run_id].span)
71+
mock_set_span_in_context.assert_called_once_with(
72+
handler.spans[parent_run_id].span
73+
)
6374
mock_start_span.assert_called_once_with(
64-
name=span_name, kind=kind, context=mock_set_span_in_context.return_value
75+
name=span_name,
76+
kind=kind,
77+
context=mock_set_span_in_context.return_value,
6578
)
6679
assert run_id in handler.spans[parent_run_id].children
6780
else:
68-
mock_start_span.assert_called_once_with(name=span_name, kind=kind)
69-
mock_set_span_in_context.assert_not_called()
81+
mock_start_span.assert_called_once_with(
82+
name=span_name, kind=kind
83+
)
84+
mock_set_span_in_context.assert_not_called()

0 commit comments

Comments
 (0)