Skip to content

Commit 5ea4d7c

Browse files
zbirenbaumzbirenbaum
committed
lint check
1 parent 3e15f51 commit 5ea4d7c

File tree

1 file changed

+80
-64
lines changed
  • packages/nvidia_nat_opentelemetry/src/nat/plugins/opentelemetry

1 file changed

+80
-64
lines changed

‎packages/nvidia_nat_opentelemetry/src/nat/plugins/opentelemetry/register.py‎

Lines changed: 80 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515

1616
import logging
1717
import os
18-
from collections.abc import AsyncGenerator
1918

2019
from pydantic import Field
2120

@@ -35,12 +34,15 @@ class LangfuseTelemetryExporter(BatchConfigMixin, TelemetryExporterBaseConfig, n
3534
"""A telemetry exporter to transmit traces to externally hosted langfuse service."""
3635

3736
endpoint: str = Field(description="The langfuse OTEL endpoint (/api/public/otel/v1/traces)")
38-
public_key: SerializableSecretStr = Field(description="The Langfuse public key",
39-
default_factory=lambda: SerializableSecretStr(""))
40-
secret_key: SerializableSecretStr = Field(description="The Langfuse secret key",
41-
default_factory=lambda: SerializableSecretStr(""))
42-
resource_attributes: dict[str, str] = Field(default_factory=dict,
43-
description="The resource attributes to add to the span")
37+
public_key: SerializableSecretStr = Field(
38+
description="The Langfuse public key", default_factory=lambda: SerializableSecretStr("")
39+
)
40+
secret_key: SerializableSecretStr = Field(
41+
description="The Langfuse secret key", default_factory=lambda: SerializableSecretStr("")
42+
)
43+
resource_attributes: dict[str, str] = Field(
44+
default_factory=dict, description="The resource attributes to add to the span"
45+
)
4446

4547

4648
@register_telemetry_exporter(config_type=LangfuseTelemetryExporter)
@@ -59,13 +61,15 @@ async def langfuse_telemetry_exporter(config: LangfuseTelemetryExporter, builder
5961
auth_header = base64.b64encode(credentials).decode("utf-8")
6062
headers = {"Authorization": f"Basic {auth_header}"}
6163

62-
yield OTLPSpanAdapterExporter(endpoint=config.endpoint,
63-
headers=headers,
64-
batch_size=config.batch_size,
65-
flush_interval=config.flush_interval,
66-
max_queue_size=config.max_queue_size,
67-
drop_on_overflow=config.drop_on_overflow,
68-
shutdown_timeout=config.shutdown_timeout)
64+
yield OTLPSpanAdapterExporter(
65+
endpoint=config.endpoint,
66+
headers=headers,
67+
batch_size=config.batch_size,
68+
flush_interval=config.flush_interval,
69+
max_queue_size=config.max_queue_size,
70+
drop_on_overflow=config.drop_on_overflow,
71+
shutdown_timeout=config.shutdown_timeout,
72+
)
6973

7074

7175
class LangsmithTelemetryExporter(BatchConfigMixin, CollectorConfigMixin, TelemetryExporterBaseConfig, name="langsmith"):
@@ -75,13 +79,15 @@ class LangsmithTelemetryExporter(BatchConfigMixin, CollectorConfigMixin, Telemet
7579
description="The langsmith OTEL endpoint",
7680
default="https://api.smith.langchain.com/otel/v1/traces",
7781
)
78-
api_key: SerializableSecretStr = Field(description="The Langsmith API key",
79-
default_factory=lambda: SerializableSecretStr(""))
80-
workspace_id: str = Field(default="",
81-
description="The Langsmith workspace ID. "
82-
"Falls back to LANGSMITH_WORKSPACE_ID env var if not set.")
83-
resource_attributes: dict[str, str] = Field(default_factory=dict,
84-
description="The resource attributes to add to the span")
82+
api_key: SerializableSecretStr = Field(
83+
description="The Langsmith API key", default_factory=lambda: SerializableSecretStr("")
84+
)
85+
workspace_id: str = Field(
86+
default="", description="The Langsmith workspace ID. Falls back to LANGSMITH_WORKSPACE_ID env var if not set."
87+
)
88+
resource_attributes: dict[str, str] = Field(
89+
default_factory=dict, description="The resource attributes to add to the span"
90+
)
8591

8692

8793
@register_telemetry_exporter(config_type=LangsmithTelemetryExporter)
@@ -95,27 +101,30 @@ async def langsmith_telemetry_exporter(config: LangsmithTelemetryExporter, build
95101
raise ValueError("API key is required for langsmith")
96102

97103
headers = {"x-api-key": api_key, "Langsmith-Project": config.project}
98-
workspace_id = config.workspace_id or os.environ.get("LANGSMITH_WORKSPACE_ID") or os.environ.get(
99-
"LANGCHAIN_WORKSPACE_ID")
104+
workspace_id = (
105+
config.workspace_id or os.environ.get("LANGSMITH_WORKSPACE_ID") or os.environ.get("LANGCHAIN_WORKSPACE_ID")
106+
)
100107
if workspace_id:
101108
headers["X-Tenant-Id"] = workspace_id
102-
yield OTLPSpanAdapterExporter(endpoint=config.endpoint,
103-
headers=headers,
104-
batch_size=config.batch_size,
105-
flush_interval=config.flush_interval,
106-
max_queue_size=config.max_queue_size,
107-
drop_on_overflow=config.drop_on_overflow,
108-
shutdown_timeout=config.shutdown_timeout)
109-
110-
111-
class OtelCollectorTelemetryExporter(BatchConfigMixin,
112-
CollectorConfigMixin,
113-
TelemetryExporterBaseConfig,
114-
name="otelcollector"):
109+
yield OTLPSpanAdapterExporter(
110+
endpoint=config.endpoint,
111+
headers=headers,
112+
batch_size=config.batch_size,
113+
flush_interval=config.flush_interval,
114+
max_queue_size=config.max_queue_size,
115+
drop_on_overflow=config.drop_on_overflow,
116+
shutdown_timeout=config.shutdown_timeout,
117+
)
118+
119+
120+
class OtelCollectorTelemetryExporter(
121+
BatchConfigMixin, CollectorConfigMixin, TelemetryExporterBaseConfig, name="otelcollector"
122+
):
115123
"""A telemetry exporter to transmit traces to externally hosted otel collector service."""
116124

117-
resource_attributes: dict[str, str] = Field(default_factory=dict,
118-
description="The resource attributes to add to the span")
125+
resource_attributes: dict[str, str] = Field(
126+
default_factory=dict, description="The resource attributes to add to the span"
127+
)
119128

120129

121130
@register_telemetry_exporter(config_type=OtelCollectorTelemetryExporter)
@@ -136,22 +145,26 @@ async def otel_telemetry_exporter(config: OtelCollectorTelemetryExporter, builde
136145
# Merge defaults with config, giving precedence to config
137146
merged_resource_attributes = {**default_resource_attributes, **config.resource_attributes}
138147

139-
yield OTLPSpanAdapterExporter(endpoint=config.endpoint,
140-
resource_attributes=merged_resource_attributes,
141-
batch_size=config.batch_size,
142-
flush_interval=config.flush_interval,
143-
max_queue_size=config.max_queue_size,
144-
drop_on_overflow=config.drop_on_overflow,
145-
shutdown_timeout=config.shutdown_timeout)
148+
yield OTLPSpanAdapterExporter(
149+
endpoint=config.endpoint,
150+
resource_attributes=merged_resource_attributes,
151+
batch_size=config.batch_size,
152+
flush_interval=config.flush_interval,
153+
max_queue_size=config.max_queue_size,
154+
drop_on_overflow=config.drop_on_overflow,
155+
shutdown_timeout=config.shutdown_timeout,
156+
)
146157

147158

148159
class PatronusTelemetryExporter(BatchConfigMixin, CollectorConfigMixin, TelemetryExporterBaseConfig, name="patronus"):
149160
"""A telemetry exporter to transmit traces to Patronus service."""
150161

151-
api_key: SerializableSecretStr = Field(description="The Patronus API key",
152-
default_factory=lambda: SerializableSecretStr(""))
153-
resource_attributes: dict[str, str] = Field(default_factory=dict,
154-
description="The resource attributes to add to the span")
162+
api_key: SerializableSecretStr = Field(
163+
description="The Patronus API key", default_factory=lambda: SerializableSecretStr("")
164+
)
165+
resource_attributes: dict[str, str] = Field(
166+
default_factory=dict, description="The resource attributes to add to the span"
167+
)
155168

156169

157170
@register_telemetry_exporter(config_type=PatronusTelemetryExporter)
@@ -168,21 +181,25 @@ async def patronus_telemetry_exporter(config: PatronusTelemetryExporter, builder
168181
"x-api-key": api_key,
169182
"pat-project-name": config.project,
170183
}
171-
yield OTLPSpanAdapterExporter(endpoint=config.endpoint,
172-
headers=headers,
173-
batch_size=config.batch_size,
174-
flush_interval=config.flush_interval,
175-
max_queue_size=config.max_queue_size,
176-
drop_on_overflow=config.drop_on_overflow,
177-
shutdown_timeout=config.shutdown_timeout,
178-
protocol="grpc")
184+
yield OTLPSpanAdapterExporter(
185+
endpoint=config.endpoint,
186+
headers=headers,
187+
batch_size=config.batch_size,
188+
flush_interval=config.flush_interval,
189+
max_queue_size=config.max_queue_size,
190+
drop_on_overflow=config.drop_on_overflow,
191+
shutdown_timeout=config.shutdown_timeout,
192+
protocol="grpc",
193+
)
179194

180195

181196
class GalileoTelemetryExporter(BatchConfigMixin, CollectorConfigMixin, TelemetryExporterBaseConfig, name="galileo"):
182197
"""A telemetry exporter to transmit traces to externally hosted galileo service."""
183198

184-
endpoint: str = Field(description="The galileo endpoint to export telemetry traces.",
185-
default="https://app.galileo.ai/api/galileo/otel/traces")
199+
endpoint: str = Field(
200+
description="The galileo endpoint to export telemetry traces.",
201+
default="https://app.galileo.ai/api/galileo/otel/traces",
202+
)
186203
logstream: str = Field(description="The logstream name to group the telemetry traces.")
187204
api_key: SerializableSecretStr = Field(description="The api key to authenticate with the galileo service.")
188205

@@ -217,16 +234,15 @@ class WeaveOtelTelemetryExporter(BatchConfigMixin, TelemetryExporterBaseConfig,
217234
description="The W&B Weave OTel endpoint",
218235
default="https://trace.wandb.ai/otel/v1/traces",
219236
)
220-
api_key: SerializableSecretStr = Field(description="The W&B API key",
221-
default_factory=lambda: SerializableSecretStr(""))
237+
api_key: SerializableSecretStr = Field(
238+
description="The W&B API key", default_factory=lambda: SerializableSecretStr("")
239+
)
222240
project: str = Field(description="The W&B project name.")
223241
entity: str = Field(description="The W&B username or team name.")
224242

225243

226244
@register_telemetry_exporter(config_type=WeaveOtelTelemetryExporter)
227-
async def weave_otel_telemetry_exporter(
228-
config: WeaveOtelTelemetryExporter, builder: Builder,
229-
) -> AsyncGenerator["OTLPSpanAdapterExporter", None]:
245+
async def weave_otel_telemetry_exporter(config: WeaveOtelTelemetryExporter, builder: Builder):
230246
"""Create a Weave OTel telemetry exporter."""
231247

232248
from nat.plugins.opentelemetry import OTLPSpanAdapterExporter

0 commit comments

Comments
 (0)