Skip to content

Commit ff73608

Browse files
michaelsafyanmichaelsafyan
committed
Reformat with ruff.
1 parent 0b9c5bb commit ff73608

File tree

8 files changed

+221
-167
lines changed

8 files changed

+221
-167
lines changed

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/allowlist_util.py‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,15 @@
1313
# limitations under the License.
1414

1515
import os
16-
from typing import Optional, Set, Callable, List, Union
16+
from typing import Callable, List, Optional, Set, Union
1717

1818
ALLOWED = True
1919
DENIED = False
2020

21+
2122
def _parse_env_list(s: str) -> Set[str]:
2223
result = set()
23-
for entry in s.split(','):
24+
for entry in s.split(","):
2425
stripped_entry = entry.strip()
2526
if not stripped_entry:
2627
continue
@@ -29,17 +30,19 @@ def _parse_env_list(s: str) -> Set[str]:
2930

3031

3132
class AllowList:
32-
3333
def __init__(
3434
self,
3535
includes: Optional[Union[Set[str], List[str]]] = None,
3636
excludes: Optional[Union[Set[str], List[str]]] = None,
37-
if_none_match: Optional[Callable[str, bool]] = None):
37+
if_none_match: Optional[Callable[str, bool]] = None,
38+
):
3839
self._includes = set(includes or [])
3940
self._excludes = set(excludes or [])
40-
self._include_all = '*' in self._includes
41-
self._exclude_all = '*' in self._excludes
42-
assert (not self._include_all) or (not self._exclude_all), "Can't have '*' in both includes and excludes."
41+
self._include_all = "*" in self._includes
42+
self._exclude_all = "*" in self._excludes
43+
assert (not self._include_all) or (
44+
not self._exclude_all
45+
), "Can't have '*' in both includes and excludes."
4346

4447
def allowed(self, x: str):
4548
if self._exclude_all:
@@ -50,12 +53,10 @@ def allowed(self, x: str):
5053

5154
@staticmethod
5255
def from_env(
53-
includes_env_var: str,
54-
excludes_env_var: Optional[str] = None):
55-
includes = _parse_env_list(os.getenv(includes_env_var) or '')
56-
excludes = set()
57-
if excludes_env_var:
58-
excludes = _parse_env_list(os.getenv(excludes_env_var) or '')
59-
return AllowList(
60-
includes=includes,
61-
excludes=excludes)
56+
includes_env_var: str, excludes_env_var: Optional[str] = None
57+
):
58+
includes = _parse_env_list(os.getenv(includes_env_var) or "")
59+
excludes = set()
60+
if excludes_env_var:
61+
excludes = _parse_env_list(os.getenv(excludes_env_var) or "")
62+
return AllowList(includes=includes, excludes=excludes)

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/dict_util.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,9 @@ def _flatten_compound_value(
8383
)
8484
if func_output is None:
8585
return {}
86-
elif _is_primitive(func_output) or _is_homogenous_primitive_list(func_output):
86+
elif _is_primitive(func_output) or _is_homogenous_primitive_list(
87+
func_output
88+
):
8789
return {key: func_output}
8890
else:
8991
value = func_output

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py‎

Lines changed: 53 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,9 @@ def _add_request_options_to_span(
185185
},
186186
)
187187
for key, value in attributes.items():
188-
if key.startswith(CUSTOM_LLM_REQUEST_PREFIX) and not allow_list.allowed(key):
188+
if key.startswith(
189+
CUSTOM_LLM_REQUEST_PREFIX
190+
) and not allow_list.allowed(key):
189191
# The allowlist is used to control inclusion of the dynamic keys.
190192
continue
191193
span.set_attribute(key, value)
@@ -223,7 +225,9 @@ def __init__(
223225
self._content_recording_enabled = is_content_recording_enabled()
224226
self._response_index = 0
225227
self._candidate_index = 0
226-
self._generate_content_config_key_allowlist = generate_content_config_key_allowlist or AllowList()
228+
self._generate_content_config_key_allowlist = (
229+
generate_content_config_key_allowlist or AllowList()
230+
)
227231

228232
def start_span_as_current_span(
229233
self, model_name, function_name, end_on_exit=True
@@ -246,7 +250,9 @@ def process_request(
246250
config: Optional[GenerateContentConfigOrDict],
247251
):
248252
span = trace.get_current_span()
249-
_add_request_options_to_span(span, config, self._generate_content_config_key_allowlist)
253+
_add_request_options_to_span(
254+
span, config, self._generate_content_config_key_allowlist
255+
)
250256
self._maybe_log_system_instruction(config=config)
251257
self._maybe_log_user_prompt(contents)
252258

@@ -507,7 +513,9 @@ def _record_duration_metric(self):
507513

508514

509515
def _create_instrumented_generate_content(
510-
snapshot: _MethodsSnapshot, otel_wrapper: OTelWrapper, generate_content_config_key_allowlist: Optional[AllowList] = None
516+
snapshot: _MethodsSnapshot,
517+
otel_wrapper: OTelWrapper,
518+
generate_content_config_key_allowlist: Optional[AllowList] = None,
511519
):
512520
wrapped_func = snapshot.generate_content
513521

@@ -521,7 +529,10 @@ def instrumented_generate_content(
521529
**kwargs: Any,
522530
) -> GenerateContentResponse:
523531
helper = _GenerateContentInstrumentationHelper(
524-
self, otel_wrapper, model, generate_content_config_key_allowlist=generate_content_config_key_allowlist,
532+
self,
533+
otel_wrapper,
534+
model,
535+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
525536
)
526537
with helper.start_span_as_current_span(
527538
model, "google.genai.Models.generate_content"
@@ -547,7 +558,9 @@ def instrumented_generate_content(
547558

548559

549560
def _create_instrumented_generate_content_stream(
550-
snapshot: _MethodsSnapshot, otel_wrapper: OTelWrapper, generate_content_config_key_allowlist: Optional[AllowList] = None
561+
snapshot: _MethodsSnapshot,
562+
otel_wrapper: OTelWrapper,
563+
generate_content_config_key_allowlist: Optional[AllowList] = None,
551564
):
552565
wrapped_func = snapshot.generate_content_stream
553566

@@ -561,7 +574,10 @@ def instrumented_generate_content_stream(
561574
**kwargs: Any,
562575
) -> Iterator[GenerateContentResponse]:
563576
helper = _GenerateContentInstrumentationHelper(
564-
self, otel_wrapper, model, generate_content_config_key_allowlist=generate_content_config_key_allowlist
577+
self,
578+
otel_wrapper,
579+
model,
580+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
565581
)
566582
with helper.start_span_as_current_span(
567583
model, "google.genai.Models.generate_content_stream"
@@ -587,7 +603,9 @@ def instrumented_generate_content_stream(
587603

588604

589605
def _create_instrumented_async_generate_content(
590-
snapshot: _MethodsSnapshot, otel_wrapper: OTelWrapper, generate_content_config_key_allowlist: Optional[AllowList] = None
606+
snapshot: _MethodsSnapshot,
607+
otel_wrapper: OTelWrapper,
608+
generate_content_config_key_allowlist: Optional[AllowList] = None,
591609
):
592610
wrapped_func = snapshot.async_generate_content
593611

@@ -601,7 +619,10 @@ async def instrumented_generate_content(
601619
**kwargs: Any,
602620
) -> GenerateContentResponse:
603621
helper = _GenerateContentInstrumentationHelper(
604-
self, otel_wrapper, model, generate_content_config_key_allowlist=generate_content_config_key_allowlist,
622+
self,
623+
otel_wrapper,
624+
model,
625+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
605626
)
606627
with helper.start_span_as_current_span(
607628
model, "google.genai.AsyncModels.generate_content"
@@ -628,7 +649,9 @@ async def instrumented_generate_content(
628649

629650
# Disabling type checking because this is not yet implemented and tested fully.
630651
def _create_instrumented_async_generate_content_stream( # type: ignore
631-
snapshot: _MethodsSnapshot, otel_wrapper: OTelWrapper, generate_content_config_key_allowlist: Optional[AllowList] = None
652+
snapshot: _MethodsSnapshot,
653+
otel_wrapper: OTelWrapper,
654+
generate_content_config_key_allowlist: Optional[AllowList] = None,
632655
):
633656
wrapped_func = snapshot.async_generate_content_stream
634657

@@ -642,7 +665,10 @@ async def instrumented_generate_content_stream(
642665
**kwargs: Any,
643666
) -> Awaitable[AsyncIterator[GenerateContentResponse]]: # type: ignore
644667
helper = _GenerateContentInstrumentationHelper(
645-
self, otel_wrapper, model, generate_content_config_key_allowlist=generate_content_config_key_allowlist
668+
self,
669+
otel_wrapper,
670+
model,
671+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
646672
)
647673
with helper.start_span_as_current_span(
648674
model,
@@ -688,20 +714,27 @@ def uninstrument_generate_content(snapshot: object):
688714

689715
def instrument_generate_content(
690716
otel_wrapper: OTelWrapper,
691-
generate_content_config_key_allowlist: Optional[AllowList]=None) -> object:
717+
generate_content_config_key_allowlist: Optional[AllowList] = None,
718+
) -> object:
692719
snapshot = _MethodsSnapshot()
693720
Models.generate_content = _create_instrumented_generate_content(
694-
snapshot, otel_wrapper, generate_content_config_key_allowlist=generate_content_config_key_allowlist,
721+
snapshot,
722+
otel_wrapper,
723+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
695724
)
696-
Models.generate_content_stream = (
697-
_create_instrumented_generate_content_stream(snapshot, otel_wrapper, generate_content_config_key_allowlist=generate_content_config_key_allowlist)
725+
Models.generate_content_stream = _create_instrumented_generate_content_stream(
726+
snapshot,
727+
otel_wrapper,
728+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
698729
)
699730
AsyncModels.generate_content = _create_instrumented_async_generate_content(
700-
snapshot, otel_wrapper, generate_content_config_key_allowlist=generate_content_config_key_allowlist
731+
snapshot,
732+
otel_wrapper,
733+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
701734
)
702-
AsyncModels.generate_content_stream = (
703-
_create_instrumented_async_generate_content_stream(
704-
snapshot, otel_wrapper, generate_content_config_key_allowlist=generate_content_config_key_allowlist,
705-
)
735+
AsyncModels.generate_content_stream = _create_instrumented_async_generate_content_stream(
736+
snapshot,
737+
otel_wrapper,
738+
generate_content_config_key_allowlist=generate_content_config_key_allowlist,
706739
)
707740
return snapshot

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/instrumentor.py‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,17 @@
2828

2929

3030
class GoogleGenAiSdkInstrumentor(BaseInstrumentor):
31-
def __init__(self, generate_content_config_key_allowlist: Optional[AllowList] = None):
31+
def __init__(
32+
self, generate_content_config_key_allowlist: Optional[AllowList] = None
33+
):
3234
self._generate_content_snapshot = None
33-
self._generate_content_config_key_allowlist = (generate_content_config_key_allowlist or AllowList.from_env(
34-
'OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_INCLUDES',
35-
excludes_env_var='OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_EXCLUDES'))
35+
self._generate_content_config_key_allowlist = (
36+
generate_content_config_key_allowlist
37+
or AllowList.from_env(
38+
"OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_INCLUDES",
39+
excludes_env_var="OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_EXCLUDES",
40+
)
41+
)
3642

3743
# Inherited, abstract function from 'BaseInstrumentor'. Even though 'self' is
3844
# not used in the definition, a method is required per the API contract.
@@ -54,7 +60,8 @@ def _instrument(self, **kwargs: Any):
5460
)
5561
self._generate_content_snapshot = instrument_generate_content(
5662
otel_wrapper,
57-
generate_content_config_key_allowlist=self._generate_content_config_key_allowlist)
63+
generate_content_config_key_allowlist=self._generate_content_config_key_allowlist,
64+
)
5865

5966
def _uninstrument(self, **kwargs: Any):
6067
uninstrument_generate_content(self._generate_content_snapshot)

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common/base.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,9 @@ def setUp(self):
3636
self._instrumentor_args = {}
3737

3838
def _lazy_init(self):
39-
self._instrumentation_context = InstrumentationContext(**self._instrumentor_args)
39+
self._instrumentation_context = InstrumentationContext(
40+
**self._instrumentor_args
41+
)
4042
self._instrumentation_context.install()
4143

4244
def set_instrumentor_constructor_kwarg(self, key, value):

‎instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/test_config_span_attributes.py‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@
1515
import os
1616
from unittest import mock
1717

18-
from opentelemetry.instrumentation.google_genai.allowlist_util import AllowList
1918
from google.genai.types import GenerateContentConfig
19+
20+
from opentelemetry.instrumentation.google_genai.allowlist_util import AllowList
21+
2022
from .base import TestCase
2123

2224

@@ -96,7 +98,9 @@ def test_option_reflected_to_span_attribute_top_p(self):
9698
span = self.generate_and_get_span(config={"top_p": 10})
9799
self.assertEqual(span.attributes["gen_ai.request.top_p"], 10)
98100

99-
@mock.patch.dict(os.environ, {"OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_INCLUDES": "*"})
101+
@mock.patch.dict(
102+
os.environ, {"OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_INCLUDES": "*"}
103+
)
100104
def test_option_not_reflected_to_span_attribute_system_instruction(self):
101105
span = self.generate_and_get_span(
102106
config={"system_instruction": "Yadda yadda yadda"}
@@ -110,7 +114,9 @@ def test_option_not_reflected_to_span_attribute_system_instruction(self):
110114
if isinstance(value, str):
111115
self.assertNotIn("Yadda yadda yadda", value)
112116

113-
@mock.patch.dict(os.environ, {"OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_INCLUDES": "*"})
117+
@mock.patch.dict(
118+
os.environ, {"OTEL_GOOGLE_GENAI_GENERATE_CONTENT_CONFIG_INCLUDES": "*"}
119+
)
114120
def test_option_reflected_to_span_attribute_automatic_func_calling(self):
115121
span = self.generate_and_get_span(
116122
config={
@@ -140,8 +146,8 @@ def test_dynamic_config_options_not_included_without_allow_list(self):
140146

141147
def test_can_supply_allow_list_via_instrumentor_constructor(self):
142148
self.set_instrumentor_constructor_kwarg(
143-
"generate_content_config_key_allowlist",
144-
AllowList(includes=["*"]))
149+
"generate_content_config_key_allowlist", AllowList(includes=["*"])
150+
)
145151
span = self.generate_and_get_span(
146152
config={
147153
"automatic_function_calling": {

0 commit comments

Comments
 (0)