QA for processor changes

Andrei Bratu · Andrei Bratu · commit 776f824978ce · 2024-12-20T15:09:48.000+02:00
diff --git a/src/otel/constants.ts b/src/otel/constants.ts
@@ -4,6 +4,6 @@ export const HUMANLOOP_FILE_KEY = "humanloop.file";
 export const HUMANLOOP_LOG_KEY = "humanloop.log";
 export const HUMANLOOP_FILE_TYPE_KEY = "humanloop.file_type";
 export const HUMANLOOP_PATH_KEY = "humanloop.file.path";
-export const HUMANLOOP_WRAPPED_FUNCTION_NAME = "humanloop.file.function_name";
+export const HUMANLOOP_META_FUNCTION_NAME = "humanloop.meta.function_name";
 export const HUMANLOOP_PARENT_SPAN_CTX_KEY = "humanloop.context.parentSpanId";
 export const HUMANLOOP_TRACE_FLOW_CTX_KEY = "humanloop.context.traceFlow";
diff --git a/src/otel/processor.ts b/src/otel/processor.ts
@@ -14,7 +14,7 @@ import {
     HUMANLOOP_FILE_KEY,
     HUMANLOOP_FILE_TYPE_KEY,
     HUMANLOOP_LOG_KEY,
-    HUMANLOOP_WRAPPED_FUNCTION_NAME,
+    HUMANLOOP_META_FUNCTION_NAME,
 } from "./constants";
 import {
     NestedDict,
@@ -44,7 +44,7 @@ export class HumanloopSpanProcessor implements SpanProcessor {
 
     async forceFlush(): Promise<void> {}
 
-    onStart(span: Span, parentContext: Context): void {
+    onStart(span: Span, _: Context): void {
         // Handle stream case: when Prompt instrumented function calls a provider with streaming: true
         // The instrumentor span will end only when the ChunksResponse is consumed, which can happen
         // after the span created by the Prompt utility finishes. To handle this, we register all instrumentor
@@ -61,22 +61,23 @@ export class HumanloopSpanProcessor implements SpanProcessor {
     async shutdown(): Promise<void> {}
 
     /**
-     * Handles spans at the end of their lifecycle.
-     * Enriches Humanloop spans and send both HL and
+     * Handles spans at the end of their lifecycle. Enriches Humanloop spans and send both HL and
      * non-HL spans to the exporter.
      */
     onEnd(span: ReadableSpan): void {
         if (isHumanloopSpan(span)) {
             new Promise<void>((resolve) => {
-                while (true) {
+                const checkChildrenSpans = () => {
                     const childrenSpans = this.children.get(span.spanContext().spanId);
                     if (
-                        (childrenSpans || [])?.every((childSpan) => childSpan.complete)
+                        (childrenSpans || []).every((childSpan) => childSpan.complete)
                     ) {
-                        break;
+                        resolve();
+                    } else {
+                        setTimeout(checkChildrenSpans, 100);
                     }
-                }
-                resolve();
+                };
+                checkChildrenSpans();
             }).then((_) => {
                 // All children/ instrumentor spans have arrived, we can process the
                 // Humanloop parent span owning them
@@ -126,13 +127,21 @@ export class HumanloopSpanProcessor implements SpanProcessor {
                         : childSpan,
                 ),
             );
-        }
 
-        this.spanExporter.export([span], (result: ExportResult) => {
-            if (result.code !== ExportResultCode.SUCCESS) {
-                console.error("Failed to export span:", result.error);
-            }
-        });
+            // Export the instrumentor span
+            this.spanExporter.export([span], (result: ExportResult) => {
+                if (result.code !== ExportResultCode.SUCCESS) {
+                    console.error("Failed to export span:", result.error);
+                }
+            });
+        } else {
+            // Unknown span, export as it is
+            this.spanExporter.export([span], (result: ExportResult) => {
+                if (result.code !== ExportResultCode.SUCCESS) {
+                    console.error("Failed to export span:", result.error);
+                }
+            });
+        }
     }
 
     /**
@@ -190,7 +199,7 @@ export class HumanloopSpanProcessor implements SpanProcessor {
             const prompt = (hlFile.prompt || {}) as unknown as PromptKernelRequest;
             if (!("model" in prompt) || !prompt.model) {
                 const functionName =
-                    promptSpan.attributes[HUMANLOOP_WRAPPED_FUNCTION_NAME];
+                    promptSpan.attributes[HUMANLOOP_META_FUNCTION_NAME];
                 throw Error(
                     `Error in ${functionName}: the LLM provider and model could not be inferred. Call one of the supported providers in your prompt function definition or define them in the promptKernel argument of the prompt() function wrapper.`,
                 );
diff --git a/src/utilities/flow.ts b/src/utilities/flow.ts
@@ -5,10 +5,10 @@ import { FlowKernelRequest } from "../api/types/FlowKernelRequest";
 import {
     HUMANLOOP_FILE_TYPE_KEY,
     HUMANLOOP_LOG_KEY,
+    HUMANLOOP_META_FUNCTION_NAME,
     HUMANLOOP_PARENT_SPAN_CTX_KEY,
     HUMANLOOP_PATH_KEY,
     HUMANLOOP_TRACE_FLOW_CTX_KEY,
-    HUMANLOOP_WRAPPED_FUNCTION_NAME,
     NestedDict,
     generateSpanId,
     jsonifyIfNotString,
@@ -67,7 +67,7 @@ export function flowUtilityFactory<I, M, O>(
             // Add span attributes
             span = span.setAttribute(HUMANLOOP_PATH_KEY, path || func.name);
             span = span.setAttribute(HUMANLOOP_FILE_TYPE_KEY, "flow");
-            span = span.setAttribute(HUMANLOOP_WRAPPED_FUNCTION_NAME, func.name);
+            span = span.setAttribute(HUMANLOOP_META_FUNCTION_NAME, func.name);
 
             if (version) {
                 writeToOpenTelemetrySpan(
diff --git a/src/utilities/prompt.ts b/src/utilities/prompt.ts
@@ -6,10 +6,10 @@ import { Humanloop } from "../index";
 import {
     HUMANLOOP_FILE_TYPE_KEY,
     HUMANLOOP_LOG_KEY,
+    HUMANLOOP_META_FUNCTION_NAME,
     HUMANLOOP_PARENT_SPAN_CTX_KEY,
     HUMANLOOP_PATH_KEY,
     HUMANLOOP_TRACE_FLOW_CTX_KEY,
-    HUMANLOOP_WRAPPED_FUNCTION_NAME,
     NestedDict,
     generateSpanId,
     jsonifyIfNotString,
@@ -87,7 +87,7 @@ export function promptUtilityFactory<I, M, O>(
             // Add span attributes
             span = span.setAttribute(HUMANLOOP_PATH_KEY, path || func.name);
             span = span.setAttribute(HUMANLOOP_FILE_TYPE_KEY, "prompt");
-            span = span.setAttribute(HUMANLOOP_WRAPPED_FUNCTION_NAME, func.name);
+            span = span.setAttribute(HUMANLOOP_META_FUNCTION_NAME, func.name);
 
             if (version) {
                 writeToOpenTelemetrySpan(
diff --git a/src/utilities/tool.ts b/src/utilities/tool.ts
@@ -12,10 +12,10 @@ import {
     HUMANLOOP_FILE_KEY,
     HUMANLOOP_FILE_TYPE_KEY,
     HUMANLOOP_LOG_KEY,
+    HUMANLOOP_META_FUNCTION_NAME,
     HUMANLOOP_PARENT_SPAN_CTX_KEY,
     HUMANLOOP_PATH_KEY,
     HUMANLOOP_TRACE_FLOW_CTX_KEY,
-    HUMANLOOP_WRAPPED_FUNCTION_NAME,
 } from "../otel/constants";
 import { ToolCallableType } from "./types";
 
@@ -78,7 +78,7 @@ export function toolUtilityFactory<I, O>(
             // Add span attributes
             span.setAttribute(HUMANLOOP_PATH_KEY, path || func.name);
             span.setAttribute(HUMANLOOP_FILE_TYPE_KEY, "tool");
-            span = span.setAttribute(HUMANLOOP_WRAPPED_FUNCTION_NAME, func.name);
+            span = span.setAttribute(HUMANLOOP_META_FUNCTION_NAME, func.name);
 
             // @ts-ignore
             // Execute the wrapped function in the appropriate context
@@ -143,8 +143,9 @@ function validateArgumentsAgainstSchema(toolKernel: ToolKernelRequest, inputs?:
         if (inputs === undefined) {
             return;
         }
+        console.log("BAI", parameters);
         throw new Error(
-            `Tool function ${toolKernel.function?.name} received inputs when the JSON schema defines none.`,
+            `Tool function ${toolKernel.function?.name} received inputs when the JSON schema defines none`,
         );
     }
 
diff --git a/tests/utilities/flow.test.ts b/tests/utilities/flow.test.ts
@@ -165,6 +165,6 @@ describe("flow decorator", () => {
 
             expect(createToolLogResponse.mock.calls).toHaveLength(1);
         },
-        10 * 1000,
+        20 * 1000,
     );
 });
diff --git a/tests/utilities/prompt.test.ts b/tests/utilities/prompt.test.ts
@@ -128,6 +128,11 @@ describe("prompt decorator", () => {
                 { provider, model },
                 callLLMMessages() as Humanloop.ChatMessage[],
             );
+
+            // The HLProcessor contains waits for Instrumentor spans to enrich
+            // the parent HL Spans in an async manner
+            await new Promise((resolve) => setTimeout(resolve, 1000));
+
             const spans = exporter.getFinishedSpans();
             expect(spans.length).toBe(2);
             expect(isHumanloopSpan(spans[0])).toBeFalsy();
@@ -148,6 +153,10 @@ describe("prompt decorator", () => {
                 callLLMMessages() as Humanloop.ChatMessage[],
             );
 
+            // The HLProcessor contains waits for Instrumentor spans to enrich
+            // the parent HL Spans in an async manner
+            await new Promise((resolve) => setTimeout(resolve, 1000));
+
             const spans = exporter.getFinishedSpans();
             expect(spans.length).toBe(2);
 
@@ -181,6 +190,10 @@ describe("prompt decorator", () => {
                 callLLMMessages() as Humanloop.ChatMessage[],
             );
 
+            // The HLProcessor contains waits for Instrumentor spans to enrich
+            // the parent HL Spans in an async manner
+            await new Promise((resolve) => setTimeout(resolve, 1000));
+
             expect(exporter.getFinishedSpans().length).toBe(2);
 
             const promptKernel = readFromOpenTelemetrySpan(
diff --git a/tests/utilities/tool.test.ts b/tests/utilities/tool.test.ts
@@ -1,9 +1,11 @@
 import { Validator } from "jsonschema";
+import { add } from "lodash";
 
 import {
     HUMANLOOP_FILE_KEY,
     HUMANLOOP_FILE_TYPE_KEY,
     HUMANLOOP_LOG_KEY,
+    HUMANLOOP_META_FUNCTION_NAME,
     readFromOpenTelemetrySpan,
 } from "../../src/otel";
 import { toolUtilityFactory } from "../../src/utilities/tool";
@@ -57,9 +59,14 @@ describe("tool decorator", () => {
                     description: "Perform arithmetic operations on two numbers.",
                     strict: true,
                     parameters: {
-                        operation: "string",
-                        num1: "number",
-                        num2: "number",
+                        type: "object",
+                        required: ["operation", "num1", "num2"],
+                        additionalProperties: false,
+                        properties: {
+                            operation: { type: "string" },
+                            num1: { type: "number" },
+                            num2: { type: "number" },
+                        },
                     },
                 },
             },
@@ -83,9 +90,14 @@ describe("tool decorator", () => {
                     description: "Perform arithmetic operations on two numbers.",
                     strict: true,
                     parameters: {
-                        operation: "string",
-                        num1: "number",
-                        num2: "number",
+                        type: "object",
+                        required: ["operation", "num1", "num2"],
+                        additionalProperties: false,
+                        properties: {
+                            operation: { type: "string" },
+                            num1: { type: "number" },
+                            num2: { type: "number" },
+                        },
                     },
                 },
             },
@@ -103,12 +115,18 @@ describe("tool decorator", () => {
             name: "calculator",
             description: "Perform arithmetic operations on two numbers.",
             parameters: {
-                operation: "string",
-                num1: "number",
-                num2: "number",
+                type: "object",
+                required: ["operation", "num1", "num2"],
+                additionalProperties: false,
+                properties: {
+                    operation: { type: "string" },
+                    num1: { type: "number" },
+                    num2: { type: "number" },
+                },
             },
             strict: true,
         });
+        expect(span.attributes[HUMANLOOP_META_FUNCTION_NAME] === "calculator");
         // @ts-ignore
         new Validator().validate(log, calculatorDecorated.jsonSchema);
     });
@@ -144,8 +162,13 @@ describe("tool decorator", () => {
                     description: "Add two numbers.",
                     strict: true,
                     parameters: {
-                        a: "number",
-                        b: "number",
+                        type: "object",
+                        required: ["a", "b"],
+                        properties: {
+                            a: { type: "number" },
+                            b: { type: "number" },
+                        },
+                        additionalProperties: false,
                     },
                 },
             },
