convert to inputs and outputs

Arize-ai · Jan 11, 2024 · eb6c4ee · eb6c4ee
1 parent e63c78d
commit eb6c4ee
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 24 deletions.
diff --git a/js/packages/openinference-instrumentation-openai/src/instrumentation.ts b/js/packages/openinference-instrumentation-openai/src/instrumentation.ts
@@ -170,11 +170,9 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
               [SemanticConventions.OPENINFERENCE_SPAN_KIND]:
                 OpenInferenceSpanKind.LLM,
               [SemanticConventions.LLM_MODEL_NAME]: body.model,
-              [SemanticConventions.INPUT_VALUE]: JSON.stringify(body),
-              [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON,
               [SemanticConventions.LLM_INVOCATION_PARAMETERS]:
                 JSON.stringify(invocationParameters),
-              ...getLLMPromptAttributes(body),
+              ...getCompletionInputValueAndMimeType(body),
             },
           });
           const execContext = trace.setSpan(context.active(), span);
@@ -206,7 +204,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<typeof openai> {
                 [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.JSON,
                 // Override the model from the value sent by the server
                 [SemanticConventions.LLM_MODEL_NAME]: result.model,
-                ...getCompletionLLMMessageAttributes(result),
+                ...getCompletionOutputValueAndMimeType(result),
                 ...getUsageAttributes(result),
               });
               span.setStatus({ code: SpanStatusCode.OK });
@@ -348,16 +346,22 @@ function getLLMInputMessagesAttributes(
 }
 
 /**
- * Converts the body of a completions request to LLM input messages
+ * Converts the body of a completions request to input attributes
  */
-function getLLMPromptAttributes(body: CompletionCreateParamsBase): Attributes {
+function getCompletionInputValueAndMimeType(body: CompletionCreateParamsBase): Attributes {
   if (typeof body.prompt === "string") {
     return {
-      [SemanticConventions.LLM_PROMPTS]: [body.prompt],
+      [SemanticConventions.INPUT_VALUE]: body.prompt,
+      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT,
     };
   } else if (isPromptStringArray(body.prompt)) {
+    const prompt = body.prompt[0]; // Only single prompts are currently supported
+    if (prompt === undefined) {
+      return {};
+    }
     return {
-      [SemanticConventions.LLM_PROMPTS]: body.prompt,
+      [SemanticConventions.INPUT_VALUE]: prompt,
+      [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT,
     };
   }
   // Other cases in which the prompt is a token or array of tokens are currently unsupported
@@ -405,20 +409,17 @@ function getChatCompletionLLMOutputMessagesAttributes(
 }
 
 /**
- * Converts the completion result to LLM output attributes
+ * Converts the completion result to output attributes
  */
-function getCompletionLLMMessageAttributes(completion: Completion): Attributes {
+function getCompletionOutputValueAndMimeType(completion: Completion): Attributes {
   // Right now support just the first choice
   const choice = completion.choices[0];
   if (!choice) {
     return {};
   }
-  const indexPrefix = `${SemanticConventions.LLM_OUTPUT_MESSAGES}.0`;
   return {
-    [`${indexPrefix}.${SemanticConventions.MESSAGE_CONTENT}`]: String(
-      choice.text,
-    ),
-    [`${indexPrefix}.${SemanticConventions.MESSAGE_ROLE}`]: "assistant",
+    [SemanticConventions.OUTPUT_VALUE]: String(choice.text),
+    [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.TEXT,
   };
 }
 

diff --git a/js/packages/openinference-instrumentation-openai/test/openai.test.ts b/js/packages/openinference-instrumentation-openai/test/openai.test.ts
@@ -135,21 +135,16 @@ describe("OpenAIInstrumentation", () => {
     expect(span.name).toBe("OpenAI Completions");
     expect(span.attributes).toMatchInlineSnapshot(`
       {
-        "input.mime_type": "application/json",
-        "input.value": "{"prompt":"Say this is a test","model":"gpt-3.5-turbo-instruct"}",
+        "input.mime_type": "text/plain",
+        "input.value": "Say this is a test",
         "llm.invocation_parameters": "{"model":"gpt-3.5-turbo-instruct"}",
         "llm.model_name": "gpt-3.5-turbo-instruct",
-        "llm.output_messages.0.message.content": "This is a test",
-        "llm.output_messages.0.message.role": "assistant",
-        "llm.prompts": [
-          "Say this is a test",
-        ],
         "llm.token_count.completion": 5,
         "llm.token_count.prompt": 12,
         "llm.token_count.total": 17,
         "openinference.span.kind": "llm",
-        "output.mime_type": "application/json",
-        "output.value": "{"id":"cmpl-8fZu1H3VijJUWev9asnxaYyQvJTC9","object":"text_completion","created":1704920149,"model":"gpt-3.5-turbo-instruct","choices":[{"text":"This is a test","index":0,"logprobs":null,"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}",
+        "output.mime_type": "text/plain",
+        "output.value": "This is a test",
       }
     `);
   });