@@ -464,6 +464,154 @@ def test_nonstreaming_chat_completion(
464464 assert span ["data" ]["gen_ai.usage.total_tokens" ] == 30
465465
466466
467+ @pytest .mark .skipif (
468+ OPENAI_VERSION <= (1 , 1 , 0 ),
469+ reason = "OpenAI versions <=1.1.0 do not support the tools parameter." ,
470+ )
471+ @pytest .mark .parametrize ("span_streaming" , [True , False ])
472+ @pytest .mark .parametrize ("stream_gen_ai_spans" , [True , False ])
473+ @pytest .mark .parametrize (
474+ "data_collection,include_prompts,expected_present,expected_absent" ,
475+ [
476+ pytest .param (
477+ {"gen_ai" : {"inputs" : True }},
478+ True ,
479+ {
480+ SPANDATA .GEN_AI_SYSTEM_INSTRUCTIONS : json .dumps (
481+ [{"type" : "text" , "content" : "You are a helpful assistant." }]
482+ ),
483+ SPANDATA .GEN_AI_REQUEST_MESSAGES : safe_serialize (
484+ [{"role" : "user" , "content" : "hello" }]
485+ ),
486+ SPANDATA .GEN_AI_REQUEST_AVAILABLE_TOOLS : safe_serialize (EXAMPLE_TOOLS ),
487+ },
488+ [],
489+ id = "inputs-enabled" ,
490+ ),
491+ pytest .param (
492+ {"gen_ai" : {"inputs" : False }},
493+ True ,
494+ {},
495+ [
496+ SPANDATA .GEN_AI_SYSTEM_INSTRUCTIONS ,
497+ SPANDATA .GEN_AI_REQUEST_MESSAGES ,
498+ SPANDATA .GEN_AI_REQUEST_AVAILABLE_TOOLS ,
499+ ],
500+ id = "inputs-disabled" ,
501+ ),
502+ pytest .param (
503+ {},
504+ True ,
505+ {
506+ SPANDATA .GEN_AI_SYSTEM_INSTRUCTIONS : json .dumps (
507+ [{"type" : "text" , "content" : "You are a helpful assistant." }]
508+ ),
509+ SPANDATA .GEN_AI_REQUEST_MESSAGES : safe_serialize (
510+ [{"role" : "user" , "content" : "hello" }]
511+ ),
512+ SPANDATA .GEN_AI_REQUEST_AVAILABLE_TOOLS : safe_serialize (EXAMPLE_TOOLS ),
513+ },
514+ [],
515+ id = "gen-ai-omitted-defaults-to-enabled" ,
516+ ),
517+ pytest .param (
518+ {"gen_ai" : {"inputs" : True }},
519+ False ,
520+ {},
521+ [
522+ SPANDATA .GEN_AI_SYSTEM_INSTRUCTIONS ,
523+ SPANDATA .GEN_AI_REQUEST_MESSAGES ,
524+ SPANDATA .GEN_AI_REQUEST_AVAILABLE_TOOLS ,
525+ ],
526+ id = "include-prompts-disabled-overrides-inputs-enabled" ,
527+ ),
528+ ],
529+ )
530+ def test_completions_api_data_collection (
531+ sentry_init ,
532+ capture_events ,
533+ capture_items ,
534+ data_collection ,
535+ include_prompts ,
536+ expected_present ,
537+ expected_absent ,
538+ nonstreaming_chat_completions_model_response ,
539+ stream_gen_ai_spans ,
540+ span_streaming ,
541+ ):
542+ sentry_init (
543+ integrations = [OpenAIIntegration (include_prompts = include_prompts )],
544+ disabled_integrations = [StdlibIntegration ],
545+ traces_sample_rate = 1.0 ,
546+ _experiments = {"data_collection" : data_collection },
547+ stream_gen_ai_spans = stream_gen_ai_spans ,
548+ trace_lifecycle = "stream" if span_streaming else "static" ,
549+ )
550+
551+ client = OpenAI (api_key = "z" )
552+ client .chat .completions ._post = mock .Mock (
553+ return_value = nonstreaming_chat_completions_model_response (
554+ response_id = "chat-id" ,
555+ response_model = "gpt-3.5-turbo" ,
556+ message_content = "the model response" ,
557+ created = 10000000 ,
558+ usage = CompletionUsage (
559+ prompt_tokens = 20 ,
560+ completion_tokens = 10 ,
561+ total_tokens = 30 ,
562+ ),
563+ )
564+ )
565+
566+ create_kwargs = {
567+ "model" : "some-model" ,
568+ "messages" : [
569+ {"role" : "system" , "content" : "You are a helpful assistant." },
570+ {"role" : "user" , "content" : "hello" },
571+ ],
572+ "max_tokens" : 100 ,
573+ "presence_penalty" : 0.1 ,
574+ "frequency_penalty" : 0.2 ,
575+ "temperature" : 0.7 ,
576+ "top_p" : 0.9 ,
577+ "tools" : EXAMPLE_TOOLS ,
578+ }
579+
580+ if span_streaming or stream_gen_ai_spans :
581+ items = capture_items ("span" )
582+
583+ with start_transaction (name = "openai tx" ):
584+ client .chat .completions .create (** create_kwargs )
585+
586+ sentry_sdk .flush ()
587+ (span ,) = (item .payload for item in items )
588+ span_data = span ["attributes" ]
589+ else :
590+ events = capture_events ()
591+
592+ with start_transaction (name = "openai tx" ):
593+ client .chat .completions .create (** create_kwargs )
594+
595+ (transaction ,) = events
596+ (span ,) = transaction ["spans" ]
597+ span_data = span ["data" ]
598+
599+ assert span_data [SPANDATA .GEN_AI_OPERATION_NAME ] == "chat"
600+ assert span_data [SPANDATA .GEN_AI_REQUEST_MODEL ] == "some-model"
601+ assert span_data [SPANDATA .GEN_AI_REQUEST_MAX_TOKENS ] == 100
602+ assert span_data [SPANDATA .GEN_AI_REQUEST_PRESENCE_PENALTY ] == 0.1
603+ assert span_data [SPANDATA .GEN_AI_REQUEST_FREQUENCY_PENALTY ] == 0.2
604+ assert span_data [SPANDATA .GEN_AI_REQUEST_TEMPERATURE ] == 0.7
605+ assert span_data [SPANDATA .GEN_AI_REQUEST_TOP_P ] == 0.9
606+ assert span_data [SPANDATA .GEN_AI_SYSTEM ] == "openai"
607+
608+ for key , value in expected_present .items ():
609+ assert span_data [key ] == value
610+
611+ for key in expected_absent :
612+ assert key not in span_data
613+
614+
467615@pytest .mark .parametrize ("span_streaming" , [True , False ])
468616@pytest .mark .parametrize ("stream_gen_ai_spans" , [True , False ])
469617@pytest .mark .asyncio
0 commit comments