Skip to content

Commit

Permalink
add uninstrument in tests + remove any none values from span attributes
Browse files Browse the repository at this point in the history
  • Loading branch information
alizenhom committed Oct 22, 2024
1 parent 8fc4336 commit 9e273f6
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,6 @@ def traced_method(wrapped, instance, args, kwargs):
name=span_name, kind=SpanKind.CLIENT, attributes=span_attributes
)
if span.is_recording():
_set_input_attributes(span, span_attributes)
set_event_prompt(span, json.dumps(llm_prompts))

try:
Expand Down Expand Up @@ -84,12 +83,6 @@ def traced_method(wrapped, instance, args, kwargs):
return traced_method


@silently_fail
def _set_input_attributes(span, attributes):
for field, value in attributes.items():
set_span_attribute(span, field, value)


@silently_fail
def _set_response_attributes(span, result):
set_span_attribute(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,7 @@ def get_llm_request_attributes(
operation_name=GenAIAttributes.GenAiOperationNameValues.CHAT.value,
):

return {
attributes = {
GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name,
GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.OPENAI.value,
GenAIAttributes.GEN_AI_REQUEST_MODEL: kwargs.get(
Expand All @@ -160,3 +160,6 @@ def get_llm_request_attributes(
"frequency_penalty"
),
}

# filter out None values
return {k: v for k, v in attributes.items() if v is not None}
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,11 @@ def instrument():
OpenAIInstrumentor().instrument()


@pytest.fixture(scope="session", autouse=True)
def uninstrument():
OpenAIInstrumentor().uninstrument()


def scrub_response_headers(response):
"""
This scrubs sensitive response headers. Note they are case-sensitive!
Expand Down

0 comments on commit 9e273f6

Please sign in to comment.