From 47cd6b5dca3c4e283aefb4d3299425cc2bc80b1c Mon Sep 17 00:00:00 2001
From: Parker Stafford <52351508+Parker-Stafford@users.noreply.github.com>
Date: Thu, 14 Nov 2024 08:32:18 -0800
Subject: [PATCH] fix(js): set spans from instrumentation active span on
 context in OpenAI instrumentor (#1121)

---
 js/.changeset/dirty-candles-brake.md                       | 5 +++++
 .../src/instrumentation.ts                                 | 7 +++----
 2 files changed, 8 insertions(+), 4 deletions(-)
 create mode 100644 js/.changeset/dirty-candles-brake.md

diff --git a/js/.changeset/dirty-candles-brake.md b/js/.changeset/dirty-candles-brake.md
new file mode 100644
index 000000000..31258521d
--- /dev/null
+++ b/js/.changeset/dirty-candles-brake.md
@@ -0,0 +1,5 @@
+---
+"@arizeai/openinference-instrumentation-openai": patch
+---
+
+fix: propgate context to spans created as a result of work done within openai calls
diff --git a/js/packages/openinference-instrumentation-openai/src/instrumentation.ts b/js/packages/openinference-instrumentation-openai/src/instrumentation.ts
index 7bdfd4219..16d029ca8 100644
--- a/js/packages/openinference-instrumentation-openai/src/instrumentation.ts
+++ b/js/packages/openinference-instrumentation-openai/src/instrumentation.ts
@@ -175,13 +175,12 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               },
             },
           );
-
           const execContext = getExecContext(span);
           const execPromise = safeExecuteInTheMiddle<
             ReturnType<ChatCompletionCreateType>
           >(
             () => {
-              return context.with(execContext, () => {
+              return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
               });
             },
@@ -263,7 +262,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             ReturnType<CompletionsCreateType>
           >(
             () => {
-              return context.with(execContext, () => {
+              return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
               });
             },
@@ -335,7 +334,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
             ReturnType<EmbeddingsCreateType>
           >(
             () => {
-              return context.with(execContext, () => {
+              return context.with(trace.setSpan(execContext, span), () => {
                 return original.apply(this, args);
               });
             },