From 23f931e0a39184399127cb6326c5d8b62a38cb2e Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 13 Apr 2026 15:46:46 +0200 Subject: [PATCH 1/3] feat: Add instrumentation for `@huggingface/inference` --- .env.example | 1 + .github/workflows/e2e-canary.yaml | 1 + .github/workflows/integration-tests.yaml | 3 + e2e/README.md | 1 + e2e/config/pr-comment-scenarios.json | 10 + e2e/helpers/scenario-installer.ts | 1 + .../huggingface-v281.log-payloads.json | 525 +++++++++++++++ .../huggingface-v281.span-events.json | 184 ++++++ .../huggingface-v3150.log-payloads.json | 618 ++++++++++++++++++ .../huggingface-v3150.span-events.json | 204 ++++++ .../huggingface-v41315.log-payloads.json | 618 ++++++++++++++++++ .../huggingface-v41315.span-events.json | 204 ++++++ .../huggingface-instrumentation/assertions.ts | 321 +++++++++ .../huggingface-instrumentation/package.json | 18 + .../pnpm-lock.yaml | 65 ++ .../scenario.huggingface-v281.mjs | 9 + .../scenario.huggingface-v281.ts | 9 + .../scenario.huggingface-v3150.mjs | 5 + .../scenario.huggingface-v3150.ts | 5 + .../scenario.impl.mjs | 221 +++++++ .../huggingface-instrumentation/scenario.mjs | 5 + .../scenario.test.ts | 60 ++ .../huggingface-instrumentation/scenario.ts | 5 + e2e/scripts/run-canary-tests-docker.mjs | 1 + .../auto-instrumentations/bundler/plugin.ts | 2 + .../bundler/webpack-loader.ts | 2 + .../configs/huggingface.ts | 245 +++++++ js/src/auto-instrumentations/hook.mts | 2 + js/src/auto-instrumentations/index.ts | 1 + js/src/exports.ts | 1 + .../instrumentation/braintrust-plugin.test.ts | 58 ++ js/src/instrumentation/braintrust-plugin.ts | 14 + .../plugins/huggingface-channels.ts | 57 ++ .../plugins/huggingface-plugin.ts | 513 +++++++++++++++ js/src/instrumentation/registry.test.ts | 1 + js/src/instrumentation/registry.ts | 2 + js/src/vendor-sdk-types/huggingface.ts | 175 +++++ js/src/wrappers/huggingface.test.ts | 398 +++++++++++ js/src/wrappers/huggingface.ts | 315 +++++++++ turbo.json | 24 +- 40 files changed, 4896 insertions(+), 8 deletions(-) create mode 100644 e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.log-payloads.json create mode 100644 e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.span-events.json create mode 100644 e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.log-payloads.json create mode 100644 e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.span-events.json create mode 100644 e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.log-payloads.json create mode 100644 e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.span-events.json create mode 100644 e2e/scenarios/huggingface-instrumentation/assertions.ts create mode 100644 e2e/scenarios/huggingface-instrumentation/package.json create mode 100644 e2e/scenarios/huggingface-instrumentation/pnpm-lock.yaml create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.mjs create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.ts create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.mjs create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.ts create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.impl.mjs create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.mjs create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.test.ts create mode 100644 e2e/scenarios/huggingface-instrumentation/scenario.ts create mode 100644 js/src/auto-instrumentations/configs/huggingface.ts create mode 100644 js/src/instrumentation/plugins/huggingface-channels.ts create mode 100644 js/src/instrumentation/plugins/huggingface-plugin.ts create mode 100644 js/src/vendor-sdk-types/huggingface.ts create mode 100644 js/src/wrappers/huggingface.test.ts create mode 100644 js/src/wrappers/huggingface.ts diff --git a/.env.example b/.env.example index 7ade92540..95f314478 100644 --- a/.env.example +++ b/.env.example @@ -4,3 +4,4 @@ ANTHROPIC_API_KEY= GEMINI_API_KEY= OPENROUTER_API_KEY= MISTRAL_API_KEY= +HUGGINGFACE_API_KEY= diff --git a/.github/workflows/e2e-canary.yaml b/.github/workflows/e2e-canary.yaml index b6156d448..fe900fac6 100644 --- a/.github/workflows/e2e-canary.yaml +++ b/.github/workflows/e2e-canary.yaml @@ -37,6 +37,7 @@ jobs: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} + HUGGINGFACE_API_KEY: ${{ secrets.HUGGINGFACE_API_KEY }} run: pnpm test:e2e:canary - name: Create or update nightly canary issue if: ${{ failure() && github.event_name == 'schedule' }} diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index 120301dd7..9e80e67ec 100644 --- a/.github/workflows/integration-tests.yaml +++ b/.github/workflows/integration-tests.yaml @@ -31,6 +31,7 @@ jobs: GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} + HUGGINGFACE_API_KEY: ${{ secrets.HUGGINGFACE_API_KEY }} steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 @@ -59,6 +60,7 @@ jobs: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} + HUGGINGFACE_API_KEY: ${{ secrets.HUGGINGFACE_API_KEY }} steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 - uses: denoland/setup-deno@ff4860f9d7236f320afa0f82b7e6457384805d05 # v2.0.4 @@ -108,6 +110,7 @@ jobs: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} + HUGGINGFACE_API_KEY: ${{ secrets.HUGGINGFACE_API_KEY }} steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4.3.1 - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 diff --git a/e2e/README.md b/e2e/README.md index 3ec5a95f7..6d01ef2d5 100644 --- a/e2e/README.md +++ b/e2e/README.md @@ -117,6 +117,7 @@ Non-hermetic scenarios require provider credentials in addition to the mock Brai - `GEMINI_API_KEY` or `GOOGLE_API_KEY` - `OPENROUTER_API_KEY` - `MISTRAL_API_KEY` +- `HUGGINGFACE_API_KEY` `claude-agent-sdk-instrumentation` also uses `ANTHROPIC_API_KEY`, because it runs the real Claude Agent SDK against Anthropic in the same style as the existing live Anthropic wrapper coverage. diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index 8a84f7f32..d0140429e 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -33,6 +33,16 @@ { "variantKey": "google-genai-v1460", "label": "v1.46.0" } ] }, + { + "scenarioDirName": "huggingface-instrumentation", + "label": "HuggingFace Instrumentation", + "metadataScenario": "huggingface-instrumentation", + "variants": [ + { "variantKey": "huggingface-v281", "label": "v2.8.1" }, + { "variantKey": "huggingface-v3150", "label": "v3.15.0" }, + { "variantKey": "huggingface-v41315", "label": "v4.13.15" } + ] + }, { "scenarioDirName": "mistral-instrumentation", "label": "Mistral Instrumentation", diff --git a/e2e/helpers/scenario-installer.ts b/e2e/helpers/scenario-installer.ts index bd387f701..940acae28 100644 --- a/e2e/helpers/scenario-installer.ts +++ b/e2e/helpers/scenario-installer.ts @@ -27,6 +27,7 @@ const INSTALL_SECRET_ENV_VARS = [ "OPENAI_API_KEY", "OPENROUTER_API_KEY", "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY", ] as const; const cleanupDirs = new Set(); diff --git a/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.log-payloads.json b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.log-payloads.json new file mode 100644 index 000000000..1424e1775 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.log-payloads.json @@ -0,0 +1,525 @@ +[ + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runTracedScenario", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "scenario": "huggingface-instrumentation", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 0, + "name": "huggingface-instrumentation-root", + "type": "task" + }, + "span_id": "" + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "" + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "chat", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 1, + "name": "huggingface-chat-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": [ + { + "content": "Reply with exactly OK.", + "role": "user" + } + ], + "log_id": "g", + "metadata": { + "endpointUrl": "https://router.huggingface.co", + "max_tokens": 16, + "model": "meta-llama/Llama-3.1-8B-Instruct", + "provider": "huggingface", + "temperature": 0 + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 2, + "name": "huggingface.chat_completion", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_accepted_prediction_tokens": 0, + "completion_reasoning_tokens": 0, + "completion_rejected_prediction_tokens": 0, + "completion_tokens": 3, + "prompt_cached_tokens": 0, + "prompt_tokens": 40, + "tokens": 43 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "created": 0, + "id": "", + "model": "llama3.1-8b", + "object": "chat.completion" + }, + "output": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ], + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "chat-stream", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 3, + "name": "huggingface-chat-stream-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": [ + { + "content": "Reply with exactly OK.", + "role": "user" + } + ], + "log_id": "g", + "metadata": { + "endpointUrl": "https://router.huggingface.co", + "max_tokens": 16, + "model": "meta-llama/Llama-3.1-8B-Instruct", + "provider": "huggingface", + "temperature": 0 + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 4, + "name": "huggingface.chat_completion_stream", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_accepted_prediction_tokens": 0, + "completion_reasoning_tokens": 0, + "completion_rejected_prediction_tokens": 0, + "completion_tokens": 3, + "prompt_cached_tokens": 0, + "prompt_tokens": 40, + "time_to_first_token": "", + "tokens": 43 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "created": 0, + "id": "", + "model": "llama3.1-8b", + "object": "chat.completion.chunk" + }, + "output": { + "choices": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ] + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "text-generation-stream", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 5, + "name": "huggingface-text-generation-stream-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "The capital of France is", + "log_id": "g", + "metadata": { + "endpointUrl": "https://router.huggingface.co/featherless-ai/v1/completions", + "max_tokens": 4, + "model": "meta-llama/Llama-3.1-8B", + "provider": "huggingface" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 6, + "name": "huggingface.text_generation_stream", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 5, + "time_to_first_token": "", + "tokens": 9 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "finish_reason": "length" + }, + "output": { + "finish_reason": "length", + "generated_text": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "feature-extraction", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 7, + "name": "huggingface-feature-extraction-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "Paris France", + "log_id": "g", + "metadata": { + "endpointUrl": "https://router.huggingface.co/hf-inference/models/thenlper/gte-large/pipeline/feature-extraction", + "model": "thenlper/gte-large", + "provider": "huggingface" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 8, + "name": "huggingface.feature_extraction", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": {}, + "output": { + "embedding_length": 1024 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + } +] diff --git a/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.span-events.json b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.span-events.json new file mode 100644 index 000000000..b99a90162 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v281.span-events.json @@ -0,0 +1,184 @@ +[ + { + "has_input": false, + "has_output": false, + "metadata": { + "scenario": "huggingface-instrumentation" + }, + "metric_keys": [], + "name": "huggingface-instrumentation-root", + "root_span_id": "", + "span_id": "", + "span_parents": [], + "type": "task" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat" + }, + "metric_keys": [], + "name": "huggingface-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "endpointUrl": "https://router.huggingface.co", + "model": "llama3.1-8b", + "provider": "huggingface" + }, + "metric_keys": [ + "completion_accepted_prediction_tokens", + "completion_reasoning_tokens", + "completion_rejected_prediction_tokens", + "completion_tokens", + "prompt_cached_tokens", + "prompt_tokens", + "tokens" + ], + "name": "huggingface.chat_completion", + "output": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ], + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat-stream" + }, + "metric_keys": [], + "name": "huggingface-chat-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "endpointUrl": "https://router.huggingface.co", + "model": "llama3.1-8b", + "provider": "huggingface" + }, + "metric_keys": [ + "completion_accepted_prediction_tokens", + "completion_reasoning_tokens", + "completion_rejected_prediction_tokens", + "completion_tokens", + "prompt_cached_tokens", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "huggingface.chat_completion_stream", + "output": null, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + null, + null, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "text-generation-stream" + }, + "metric_keys": [], + "name": "huggingface-text-generation-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "endpointUrl": "https://router.huggingface.co/featherless-ai/v1/completions", + "finish_reason": "length", + "model": "meta-llama/Llama-3.1-8B", + "provider": "huggingface" + }, + "metric_keys": [ + "completion_tokens", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "huggingface.text_generation_stream", + "output": { + "finish_reason": "length", + "generated_text": "" + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "feature-extraction" + }, + "metric_keys": [], + "name": "huggingface-feature-extraction-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "endpointUrl": "https://router.huggingface.co/hf-inference/models/thenlper/gte-large/pipeline/feature-extraction", + "model": "thenlper/gte-large", + "provider": "huggingface" + }, + "metric_keys": [], + "name": "huggingface.feature_extraction", + "output": { + "embedding_length": 1024 + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + } +] diff --git a/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.log-payloads.json b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.log-payloads.json new file mode 100644 index 000000000..1e9d02b65 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.log-payloads.json @@ -0,0 +1,618 @@ +[ + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runTracedScenario", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "scenario": "huggingface-instrumentation", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 0, + "name": "huggingface-instrumentation-root", + "type": "task" + }, + "span_id": "" + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "" + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "chat", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 1, + "name": "huggingface-chat-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": [ + { + "content": "Reply with exactly OK.", + "role": "user" + } + ], + "log_id": "g", + "metadata": { + "max_tokens": 16, + "model": "meta-llama/Llama-3.1-8B-Instruct", + "provider": "featherless-ai", + "temperature": 0 + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 2, + "name": "huggingface.chat_completion", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 40, + "tokens": 41 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "created": 0, + "id": "", + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "object": "chat.completion" + }, + "output": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ], + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "chat-stream", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 3, + "name": "huggingface-chat-stream-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": [ + { + "content": "Reply with exactly OK.", + "role": "user" + } + ], + "log_id": "g", + "metadata": { + "max_tokens": 16, + "model": "meta-llama/Llama-3.1-8B-Instruct", + "provider": "featherless-ai", + "temperature": 0 + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 4, + "name": "huggingface.chat_completion_stream", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 40, + "time_to_first_token": "", + "tokens": 41 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "created": 0, + "id": "", + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "object": "chat.completion.chunk" + }, + "output": { + "choices": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ] + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "text-generation", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 5, + "name": "huggingface-text-generation-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "The capital of France is", + "log_id": "g", + "metadata": { + "model": "meta-llama/Llama-3.1-8B", + "parameters": { + "do_sample": false, + "max_new_tokens": 4, + "return_full_text": false + }, + "provider": "featherless-ai" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 6, + "name": "huggingface.text_generation", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": {}, + "output": { + "generated_text": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "text-generation-stream", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 7, + "name": "huggingface-text-generation-stream-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "The capital of France is", + "log_id": "g", + "metadata": { + "model": "meta-llama/Llama-3.1-8B", + "parameters": { + "do_sample": false, + "max_new_tokens": 4, + "return_full_text": false + }, + "provider": "featherless-ai" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 8, + "name": "huggingface.text_generation_stream", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 5, + "time_to_first_token": "", + "tokens": 9 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "finish_reason": "length" + }, + "output": { + "finish_reason": "length", + "generated_text": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "feature-extraction", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 9, + "name": "huggingface-feature-extraction-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "Paris France", + "log_id": "g", + "metadata": { + "model": "thenlper/gte-large", + "provider": "hf-inference" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 10, + "name": "huggingface.feature_extraction", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": {}, + "output": { + "embedding_length": 1024 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + } +] diff --git a/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.span-events.json b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.span-events.json new file mode 100644 index 000000000..42b07d833 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v3150.span-events.json @@ -0,0 +1,204 @@ +[ + { + "has_input": false, + "has_output": false, + "metadata": { + "scenario": "huggingface-instrumentation" + }, + "metric_keys": [], + "name": "huggingface-instrumentation-root", + "root_span_id": "", + "span_id": "", + "span_parents": [], + "type": "task" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat" + }, + "metric_keys": [], + "name": "huggingface-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "provider": "featherless-ai" + }, + "metric_keys": [ + "completion_tokens", + "prompt_tokens", + "tokens" + ], + "name": "huggingface.chat_completion", + "output": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ], + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat-stream" + }, + "metric_keys": [], + "name": "huggingface-chat-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "provider": "featherless-ai" + }, + "metric_keys": [ + "completion_tokens", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "huggingface.chat_completion_stream", + "output": null, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "text-generation" + }, + "metric_keys": [], + "name": "huggingface-text-generation-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "meta-llama/Llama-3.1-8B", + "provider": "featherless-ai" + }, + "metric_keys": [], + "name": "huggingface.text_generation", + "output": { + "generated_text": "" + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "text-generation-stream" + }, + "metric_keys": [], + "name": "huggingface-text-generation-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "finish_reason": "length", + "model": "meta-llama/Llama-3.1-8B", + "provider": "featherless-ai" + }, + "metric_keys": [ + "completion_tokens", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "huggingface.text_generation_stream", + "output": { + "finish_reason": "length", + "generated_text": "" + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "feature-extraction" + }, + "metric_keys": [], + "name": "huggingface-feature-extraction-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "thenlper/gte-large", + "provider": "hf-inference" + }, + "metric_keys": [], + "name": "huggingface.feature_extraction", + "output": { + "embedding_length": 1024 + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + } +] diff --git a/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.log-payloads.json b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.log-payloads.json new file mode 100644 index 000000000..1e9d02b65 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.log-payloads.json @@ -0,0 +1,618 @@ +[ + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runTracedScenario", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "scenario": "huggingface-instrumentation", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 0, + "name": "huggingface-instrumentation-root", + "type": "task" + }, + "span_id": "" + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "" + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "chat", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 1, + "name": "huggingface-chat-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": [ + { + "content": "Reply with exactly OK.", + "role": "user" + } + ], + "log_id": "g", + "metadata": { + "max_tokens": 16, + "model": "meta-llama/Llama-3.1-8B-Instruct", + "provider": "featherless-ai", + "temperature": 0 + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 2, + "name": "huggingface.chat_completion", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 40, + "tokens": 41 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "created": 0, + "id": "", + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "object": "chat.completion" + }, + "output": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ], + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "chat-stream", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 3, + "name": "huggingface-chat-stream-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": [ + { + "content": "Reply with exactly OK.", + "role": "user" + } + ], + "log_id": "g", + "metadata": { + "max_tokens": 16, + "model": "meta-llama/Llama-3.1-8B-Instruct", + "provider": "featherless-ai", + "temperature": 0 + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 4, + "name": "huggingface.chat_completion_stream", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 40, + "time_to_first_token": "", + "tokens": 41 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "created": 0, + "id": "", + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "object": "chat.completion.chunk" + }, + "output": { + "choices": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ] + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "text-generation", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 5, + "name": "huggingface-text-generation-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "The capital of France is", + "log_id": "g", + "metadata": { + "model": "meta-llama/Llama-3.1-8B", + "parameters": { + "do_sample": false, + "max_new_tokens": 4, + "return_full_text": false + }, + "provider": "featherless-ai" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 6, + "name": "huggingface.text_generation", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": {}, + "output": { + "generated_text": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "text-generation-stream", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 7, + "name": "huggingface-text-generation-stream-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "The capital of France is", + "log_id": "g", + "metadata": { + "model": "meta-llama/Llama-3.1-8B", + "parameters": { + "do_sample": false, + "max_new_tokens": 4, + "return_full_text": false + }, + "provider": "featherless-ai" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 8, + "name": "huggingface.text_generation_stream", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "completion_tokens": 4, + "prompt_tokens": 5, + "time_to_first_token": "", + "tokens": 9 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metadata": { + "finish_reason": "length" + }, + "output": { + "finish_reason": "length", + "generated_text": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": false, + "context": { + "caller_filename": "/e2e/helpers/provider-runtime.mjs", + "caller_functionname": "runOperation", + "caller_lineno": 0 + }, + "created": "", + "id": "", + "log_id": "g", + "metadata": { + "operation": "feature-extraction", + "testRunId": "" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 9, + "name": "huggingface-feature-extraction-operation" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "context": {}, + "created": "", + "id": "", + "input": "Paris France", + "log_id": "g", + "metadata": { + "model": "thenlper/gte-large", + "provider": "hf-inference" + }, + "metrics": { + "start": "" + }, + "project_id": "", + "root_span_id": "", + "span_attributes": { + "exec_counter": 10, + "name": "huggingface.feature_extraction", + "type": "llm" + }, + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": {}, + "output": { + "embedding_length": 1024 + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + }, + { + "_is_merge": true, + "id": "", + "log_id": "g", + "metrics": { + "end": "" + }, + "project_id": "", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ] + } +] diff --git a/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.span-events.json b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.span-events.json new file mode 100644 index 000000000..42b07d833 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/__snapshots__/huggingface-v41315.span-events.json @@ -0,0 +1,204 @@ +[ + { + "has_input": false, + "has_output": false, + "metadata": { + "scenario": "huggingface-instrumentation" + }, + "metric_keys": [], + "name": "huggingface-instrumentation-root", + "root_span_id": "", + "span_id": "", + "span_parents": [], + "type": "task" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat" + }, + "metric_keys": [], + "name": "huggingface-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "provider": "featherless-ai" + }, + "metric_keys": [ + "completion_tokens", + "prompt_tokens", + "tokens" + ], + "name": "huggingface.chat_completion", + "output": [ + { + "content": "", + "finish_reason": "stop", + "index": 0, + "role": "assistant" + } + ], + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "chat-stream" + }, + "metric_keys": [], + "name": "huggingface-chat-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "meta-llama/Meta-Llama-3.1-8B-Instruct", + "provider": "featherless-ai" + }, + "metric_keys": [ + "completion_tokens", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "huggingface.chat_completion_stream", + "output": null, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "text-generation" + }, + "metric_keys": [], + "name": "huggingface-text-generation-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "meta-llama/Llama-3.1-8B", + "provider": "featherless-ai" + }, + "metric_keys": [], + "name": "huggingface.text_generation", + "output": { + "generated_text": "" + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "text-generation-stream" + }, + "metric_keys": [], + "name": "huggingface-text-generation-stream-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "finish_reason": "length", + "model": "meta-llama/Llama-3.1-8B", + "provider": "featherless-ai" + }, + "metric_keys": [ + "completion_tokens", + "prompt_tokens", + "time_to_first_token", + "tokens" + ], + "name": "huggingface.text_generation_stream", + "output": { + "finish_reason": "length", + "generated_text": "" + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "feature-extraction" + }, + "metric_keys": [], + "name": "huggingface-feature-extraction-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "thenlper/gte-large", + "provider": "hf-inference" + }, + "metric_keys": [], + "name": "huggingface.feature_extraction", + "output": { + "embedding_length": 1024 + }, + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + } +] diff --git a/e2e/scenarios/huggingface-instrumentation/assertions.ts b/e2e/scenarios/huggingface-instrumentation/assertions.ts new file mode 100644 index 000000000..209587de7 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/assertions.ts @@ -0,0 +1,321 @@ +import { beforeAll, describe, expect, test } from "vitest"; +import { type Json } from "../../helpers/normalize"; +import type { CapturedLogEvent } from "../../helpers/mock-braintrust-server"; +import { + formatJsonFileSnapshot, + resolveFileSnapshotPath, +} from "../../helpers/file-snapshot"; +import { withScenarioHarness } from "../../helpers/scenario-harness"; +import { + findLatestChildSpan, + findLatestSpan, +} from "../../helpers/trace-selectors"; +import { + payloadRowsForRootSpan, + summarizeWrapperContract, +} from "../../helpers/wrapper-contract"; + +import { ROOT_NAME, SCENARIO_NAME } from "./scenario.impl.mjs"; + +type RunHuggingFaceScenario = (harness: { + runNodeScenarioDir: (options: { + entry: string; + nodeArgs: string[]; + runContext?: { variantKey: string }; + scenarioDir: string; + timeoutMs: number; + }) => Promise; + runScenarioDir: (options: { + entry: string; + runContext?: { variantKey: string }; + scenarioDir: string; + timeoutMs: number; + }) => Promise; +}) => Promise; + +function isRecord(value: Json | undefined): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function summarizeChatOutput(output: Json | undefined): Json { + if (!Array.isArray(output)) { + return null; + } + + return output.map((choice) => { + if (!isRecord(choice as Json)) { + return choice as Json; + } + + const message = isRecord(choice.message as Json) + ? (choice.message as Record) + : undefined; + const content = message?.content; + return { + content: + typeof content === "string" + ? "" + : Array.isArray(content) + ? "" + : (content ?? null), + finish_reason: choice.finish_reason ?? null, + index: choice.index ?? null, + role: message?.role ?? null, + } satisfies Json; + }); +} + +function summarizeTextGenerationOutput(output: Json | undefined): Json { + if (!isRecord(output)) { + return output ?? null; + } + + const generatedText = output.generated_text; + return { + ...output, + generated_text: + typeof generatedText === "string" ? "" : (generatedText ?? null), + } satisfies Json; +} + +function summarizeProviderSpan(event: CapturedLogEvent): Json { + const summary = summarizeWrapperContract(event, [ + "dimensions", + "endpointUrl", + "finish_reason", + "model", + "operation", + "provider", + "scenario", + ]) as Record; + + switch (event.span.name) { + case "huggingface.chat_completion": + case "huggingface.chat_completion_stream": + summary.output = summarizeChatOutput(event.output as Json); + break; + case "huggingface.text_generation": + case "huggingface.text_generation_stream": + summary.output = summarizeTextGenerationOutput(event.output as Json); + break; + case "huggingface.feature_extraction": + summary.output = (event.output as Json) ?? null; + break; + default: + break; + } + + return summary; +} + +function normalizeMetrics(value: Json): Json { + if (Array.isArray(value)) { + return value.map((entry) => normalizeMetrics(entry as Json)); + } + + if (!isRecord(value)) { + return value; + } + + const normalized: Record = {}; + for (const [key, entry] of Object.entries(value)) { + if ( + typeof entry === "number" && + ["end", "start", "time_to_first_token"].includes(key) + ) { + normalized[key] = ""; + continue; + } + + normalized[key] = normalizeMetrics(entry as Json); + } + return normalized; +} + +function normalizePayloadOutput(row: Json): Json { + if (!isRecord(row)) { + return row; + } + + return "output" in row + ? { + ...row, + output: normalizeLoggedOutput(row.output), + } + : row; +} + +function normalizeLoggedOutput(output: Json): Json { + if (Array.isArray(output)) { + return summarizeChatOutput(output); + } + + if (!isRecord(output)) { + return output; + } + + if ("generated_text" in output) { + return summarizeTextGenerationOutput(output); + } + + if (Array.isArray(output.choices)) { + return { + ...output, + choices: summarizeChatOutput(output.choices), + }; + } + + return output; +} + +function buildSpanSummary(events: CapturedLogEvent[]): Json { + const root = findLatestSpan(events, ROOT_NAME); + const chatOperation = findLatestSpan(events, "huggingface-chat-operation"); + const chatStreamOperation = findLatestSpan( + events, + "huggingface-chat-stream-operation", + ); + const textGenerationOperation = findLatestSpan( + events, + "huggingface-text-generation-operation", + ); + const textGenerationStreamOperation = findLatestSpan( + events, + "huggingface-text-generation-stream-operation", + ); + const featureExtractionOperation = findLatestSpan( + events, + "huggingface-feature-extraction-operation", + ); + + return [ + root ? summarizeWrapperContract(root, ["scenario"]) : null, + chatOperation + ? summarizeWrapperContract(chatOperation, ["operation"]) + : null, + chatOperation + ? summarizeProviderSpan( + findLatestChildSpan( + events, + "huggingface.chat_completion", + chatOperation.span.id, + )!, + ) + : null, + chatStreamOperation + ? summarizeWrapperContract(chatStreamOperation, ["operation"]) + : null, + chatStreamOperation + ? summarizeProviderSpan( + findLatestChildSpan( + events, + "huggingface.chat_completion_stream", + chatStreamOperation.span.id, + )!, + ) + : null, + textGenerationOperation + ? summarizeWrapperContract(textGenerationOperation, ["operation"]) + : null, + textGenerationOperation + ? summarizeProviderSpan( + findLatestChildSpan( + events, + "huggingface.text_generation", + textGenerationOperation.span.id, + )!, + ) + : null, + textGenerationStreamOperation + ? summarizeWrapperContract(textGenerationStreamOperation, ["operation"]) + : null, + textGenerationStreamOperation + ? summarizeProviderSpan( + findLatestChildSpan( + events, + "huggingface.text_generation_stream", + textGenerationStreamOperation.span.id, + )!, + ) + : null, + featureExtractionOperation + ? summarizeWrapperContract(featureExtractionOperation, ["operation"]) + : null, + featureExtractionOperation + ? summarizeProviderSpan( + findLatestChildSpan( + events, + "huggingface.feature_extraction", + featureExtractionOperation.span.id, + )!, + ) + : null, + ] satisfies Json; +} + +export function defineHuggingFaceInstrumentationAssertions(options: { + name: string; + runScenario: RunHuggingFaceScenario; + snapshotName: string; + testFileUrl: string; + timeoutMs: number; +}): void { + const spanSnapshotPath = resolveFileSnapshotPath( + options.testFileUrl, + `${options.snapshotName}.span-events.json`, + ); + const payloadSnapshotPath = resolveFileSnapshotPath( + options.testFileUrl, + `${options.snapshotName}.log-payloads.json`, + ); + + describe(options.name, () => { + let events: CapturedLogEvent[] = []; + let payloadRows: Json = []; + + beforeAll(async () => { + await withScenarioHarness(async (harness) => { + await options.runScenario(harness); + events = harness.events(); + + const root = findLatestSpan(events, ROOT_NAME); + payloadRows = payloadRowsForRootSpan(harness.payloads(), root?.span.id) + .map((row) => normalizePayloadOutput(normalizeMetrics(row as Json))) + .filter((row) => row !== null); + }); + }, options.timeoutMs); + + test( + "captures the root trace for the scenario", + { timeout: options.timeoutMs }, + () => { + const root = findLatestSpan(events, ROOT_NAME); + + expect(root).toBeDefined(); + expect(root?.row.metadata).toMatchObject({ + scenario: SCENARIO_NAME, + }); + }, + ); + + test( + "matches the span contract snapshot", + { timeout: options.timeoutMs }, + async ({ expect }) => { + await expect( + formatJsonFileSnapshot(buildSpanSummary(events)), + ).toMatchFileSnapshot(spanSnapshotPath); + }, + ); + + test( + "matches the log payload snapshot", + { timeout: options.timeoutMs }, + async ({ expect }) => { + await expect(formatJsonFileSnapshot(payloadRows)).toMatchFileSnapshot( + payloadSnapshotPath, + ); + }, + ); + }); +} diff --git a/e2e/scenarios/huggingface-instrumentation/package.json b/e2e/scenarios/huggingface-instrumentation/package.json new file mode 100644 index 000000000..528b06e0b --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/package.json @@ -0,0 +1,18 @@ +{ + "name": "@braintrust/e2e-huggingface-instrumentation", + "private": true, + "braintrustScenario": { + "canary": { + "dependencies": { + "@huggingface/inference": "@huggingface/inference@4", + "huggingface-inference-sdk-v2": "@huggingface/inference@2", + "huggingface-inference-sdk-v3": "@huggingface/inference@3" + } + } + }, + "dependencies": { + "@huggingface/inference": "4.13.15", + "huggingface-inference-sdk-v2": "npm:@huggingface/inference@2.8.1", + "huggingface-inference-sdk-v3": "npm:@huggingface/inference@3.15.0" + } +} diff --git a/e2e/scenarios/huggingface-instrumentation/pnpm-lock.yaml b/e2e/scenarios/huggingface-instrumentation/pnpm-lock.yaml new file mode 100644 index 000000000..39eaf4981 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/pnpm-lock.yaml @@ -0,0 +1,65 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + dependencies: + '@huggingface/inference': + specifier: 4.13.15 + version: 4.13.15 + huggingface-inference-sdk-v2: + specifier: npm:@huggingface/inference@2.8.1 + version: '@huggingface/inference@2.8.1' + huggingface-inference-sdk-v3: + specifier: npm:@huggingface/inference@3.15.0 + version: '@huggingface/inference@3.15.0' + +packages: + + '@huggingface/inference@2.8.1': + resolution: {integrity: sha512-EfsNtY9OR6JCNaUa5bZu2mrs48iqeTz0Gutwf+fU0Kypx33xFQB4DKMhp8u4Ee6qVbLbNWvTHuWwlppLQl4p4Q==} + engines: {node: '>=18'} + + '@huggingface/inference@3.15.0': + resolution: {integrity: sha512-C+Adt4fu4ztlq0Al9EOgEdK5Hl8ebV1eoDEWegJPdAJ97U8A1aqBbl1Sp4S4+wIy3nLApTrtcLuoizGZmLtDMA==} + engines: {node: '>=18'} + + '@huggingface/inference@4.13.15': + resolution: {integrity: sha512-V7B13KFDVhYkQqgx8vpMcmtEG+PoePlh65IUvpphTTItHAq6zRLcsS4torh1QavUaT+6nEwho4wFXzBQNGBwKQ==} + engines: {node: '>=18'} + + '@huggingface/jinja@0.5.6': + resolution: {integrity: sha512-MyMWyLnjqo+KRJYSH7oWNbsOn5onuIvfXYPcc0WOGxU0eHUV7oAYUoQTl2BMdu7ml+ea/bu11UM+EshbeHwtIA==} + engines: {node: '>=18'} + + '@huggingface/tasks@0.12.30': + resolution: {integrity: sha512-A1ITdxbEzx9L8wKR8pF7swyrTLxWNDFIGDLUWInxvks2ruQ8PLRBZe8r0EcjC3CDdtlj9jV1V4cgV35K/iy3GQ==} + + '@huggingface/tasks@0.19.90': + resolution: {integrity: sha512-nfV9luJbvwGQ/5oKXkKhCV9h4X7mwh1YaGG3ORd6UMLDSwr1OFSSatcBX0O9OtBtmNK19aGSjbLFqqgcIR6+IA==} + +snapshots: + + '@huggingface/inference@2.8.1': + dependencies: + '@huggingface/tasks': 0.12.30 + + '@huggingface/inference@3.15.0': + dependencies: + '@huggingface/jinja': 0.5.6 + '@huggingface/tasks': 0.19.90 + + '@huggingface/inference@4.13.15': + dependencies: + '@huggingface/jinja': 0.5.6 + '@huggingface/tasks': 0.19.90 + + '@huggingface/jinja@0.5.6': {} + + '@huggingface/tasks@0.12.30': {} + + '@huggingface/tasks@0.19.90': {} diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.mjs b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.mjs new file mode 100644 index 000000000..cb3f30be9 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.mjs @@ -0,0 +1,9 @@ +import * as huggingFace from "huggingface-inference-sdk-v2"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoHuggingFaceInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => + runAutoHuggingFaceInstrumentation(huggingFace, { + supportsLiveTextGeneration: false, + }), +); diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.ts b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.ts new file mode 100644 index 000000000..2cd4fa160 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v281.ts @@ -0,0 +1,9 @@ +import * as huggingFace from "huggingface-inference-sdk-v2"; +import { runMain } from "../../helpers/scenario-runtime"; +import { runWrappedHuggingFaceInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => + runWrappedHuggingFaceInstrumentation(huggingFace, { + supportsLiveTextGeneration: false, + }), +); diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.mjs b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.mjs new file mode 100644 index 000000000..1a7d6d191 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.mjs @@ -0,0 +1,5 @@ +import * as huggingFace from "huggingface-inference-sdk-v3"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoHuggingFaceInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => runAutoHuggingFaceInstrumentation(huggingFace)); diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.ts b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.ts new file mode 100644 index 000000000..8709de235 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.huggingface-v3150.ts @@ -0,0 +1,5 @@ +import * as huggingFace from "huggingface-inference-sdk-v3"; +import { runMain } from "../../helpers/scenario-runtime"; +import { runWrappedHuggingFaceInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => runWrappedHuggingFaceInstrumentation(huggingFace)); diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.impl.mjs b/e2e/scenarios/huggingface-instrumentation/scenario.impl.mjs new file mode 100644 index 000000000..e9aee92ac --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.impl.mjs @@ -0,0 +1,221 @@ +import { wrapHuggingFace } from "braintrust"; +import { + collectAsync, + runOperation, + runTracedScenario, +} from "../../helpers/provider-runtime.mjs"; + +const CHAT_MODEL = "meta-llama/Llama-3.1-8B-Instruct"; +const CHAT_PROVIDER = "featherless-ai"; +const EMBEDDING_MODEL = "thenlper/gte-large"; +const EMBEDDING_PROVIDER = "hf-inference"; +const ROOT_NAME = "huggingface-instrumentation-root"; +const SCENARIO_NAME = "huggingface-instrumentation"; +const TEXT_GENERATION_MODEL = "meta-llama/Llama-3.1-8B"; +const TEXT_GENERATION_PROVIDER = "featherless-ai"; +const HUGGINGFACE_SCENARIO_TIMEOUT_MS = 150_000; +const V2_CHAT_ENDPOINT_URL = "https://router.huggingface.co"; +const V2_FEATURE_EXTRACTION_ENDPOINT_URL = + "https://router.huggingface.co/hf-inference/models/thenlper/gte-large/pipeline/feature-extraction"; +const V2_TEXT_GENERATION_ENDPOINT_URL = + "https://router.huggingface.co/featherless-ai/v1/completions"; +const HUGGINGFACE_SCENARIO_SPECS = [ + { + autoEntry: "scenario.huggingface-v281.mjs", + dependencyName: "huggingface-inference-sdk-v2", + snapshotName: "huggingface-v281", + supportsLiveTextGeneration: false, + wrapperEntry: "scenario.huggingface-v281.ts", + }, + { + autoEntry: "scenario.huggingface-v3150.mjs", + dependencyName: "huggingface-inference-sdk-v3", + snapshotName: "huggingface-v3150", + wrapperEntry: "scenario.huggingface-v3150.ts", + }, + { + autoEntry: "scenario.mjs", + dependencyName: "@huggingface/inference", + snapshotName: "huggingface-v41315", + wrapperEntry: "scenario.ts", + }, +]; + +function getClientConstructor(sdk) { + return sdk.InferenceClient ?? sdk.HfInference; +} + +function getHuggingFaceApiKey() { + const apiKey = process.env.HUGGINGFACE_API_KEY; + if (!apiKey) { + throw new Error("HUGGINGFACE_API_KEY must be set for this scenario"); + } + return apiKey; +} + +function getScenarioCapabilities(options) { + return { + supportsLiveTextGeneration: options.supportsLiveTextGeneration !== false, + }; +} + +function createClient(Client, apiKey, options) { + const client = new Client(apiKey); + + if (options.supportsLiveTextGeneration === false) { + if (typeof client.endpoint !== "function") { + throw new Error("Expected HuggingFace v2 client to support endpoint()"); + } + + return client.endpoint(V2_CHAT_ENDPOINT_URL); + } + + return client; +} + +async function runHuggingFaceInstrumentationScenario(sdk, options = {}) { + const decoratedSDK = options.decorateSDK ? options.decorateSDK(sdk) : sdk; + const Client = getClientConstructor(decoratedSDK); + const apiKey = getHuggingFaceApiKey(); + const capabilities = getScenarioCapabilities(options); + + if (!Client) { + throw new Error("Could not resolve a HuggingFace client constructor"); + } + + const client = createClient(Client, apiKey, capabilities); + + await runTracedScenario({ + callback: async () => { + await runOperation("huggingface-chat-operation", "chat", async () => { + await client.chatCompletion({ + max_tokens: 16, + messages: [ + { + role: "user", + content: "Reply with exactly OK.", + }, + ], + model: CHAT_MODEL, + ...(capabilities.supportsLiveTextGeneration + ? { provider: CHAT_PROVIDER } + : {}), + temperature: 0, + }); + }); + + await runOperation( + "huggingface-chat-stream-operation", + "chat-stream", + async () => { + const stream = client.chatCompletionStream({ + max_tokens: 16, + messages: [ + { + role: "user", + content: "Reply with exactly OK.", + }, + ], + model: CHAT_MODEL, + ...(capabilities.supportsLiveTextGeneration + ? { provider: CHAT_PROVIDER } + : {}), + temperature: 0, + }); + await collectAsync(stream); + }, + ); + + if (capabilities.supportsLiveTextGeneration) { + await runOperation( + "huggingface-text-generation-operation", + "text-generation", + async () => { + await decoratedSDK.textGeneration({ + accessToken: apiKey, + inputs: "The capital of France is", + model: TEXT_GENERATION_MODEL, + parameters: { + do_sample: false, + max_new_tokens: 4, + return_full_text: false, + }, + provider: TEXT_GENERATION_PROVIDER, + }); + }, + ); + } + + await runOperation( + "huggingface-text-generation-stream-operation", + "text-generation-stream", + async () => { + const stream = decoratedSDK.textGenerationStream({ + ...(capabilities.supportsLiveTextGeneration + ? { + accessToken: apiKey, + inputs: "The capital of France is", + model: TEXT_GENERATION_MODEL, + parameters: { + do_sample: false, + max_new_tokens: 4, + return_full_text: false, + }, + provider: TEXT_GENERATION_PROVIDER, + } + : { + accessToken: apiKey, + endpointUrl: V2_TEXT_GENERATION_ENDPOINT_URL, + inputs: "The capital of France is", + max_tokens: 4, + model: TEXT_GENERATION_MODEL, + prompt: "The capital of France is", + }), + }); + await collectAsync(stream); + }, + ); + + await runOperation( + "huggingface-feature-extraction-operation", + "feature-extraction", + async () => { + await decoratedSDK.featureExtraction({ + accessToken: apiKey, + inputs: "Paris France", + model: EMBEDDING_MODEL, + ...(capabilities.supportsLiveTextGeneration + ? { provider: EMBEDDING_PROVIDER } + : { endpointUrl: V2_FEATURE_EXTRACTION_ENDPOINT_URL }), + }); + }, + ); + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-huggingface-instrumentation", + rootName: ROOT_NAME, + }); +} + +export async function runWrappedHuggingFaceInstrumentation(sdk, options = {}) { + await runHuggingFaceInstrumentationScenario(sdk, { + ...options, + decorateSDK: wrapHuggingFace, + }); +} + +export async function runAutoHuggingFaceInstrumentation(sdk, options = {}) { + await runHuggingFaceInstrumentationScenario(sdk, options); +} + +export { + CHAT_MODEL, + EMBEDDING_MODEL, + HUGGINGFACE_SCENARIO_SPECS, + HUGGINGFACE_SCENARIO_TIMEOUT_MS, + ROOT_NAME, + SCENARIO_NAME, + TEXT_GENERATION_MODEL, +}; diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.mjs b/e2e/scenarios/huggingface-instrumentation/scenario.mjs new file mode 100644 index 000000000..256e65ce9 --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.mjs @@ -0,0 +1,5 @@ +import * as huggingFace from "@huggingface/inference"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoHuggingFaceInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => runAutoHuggingFaceInstrumentation(huggingFace)); diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.test.ts b/e2e/scenarios/huggingface-instrumentation/scenario.test.ts new file mode 100644 index 000000000..7e9d8af1f --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.test.ts @@ -0,0 +1,60 @@ +import { describe } from "vitest"; +import { + prepareScenarioDir, + readInstalledPackageVersion, + resolveScenarioDir, +} from "../../helpers/scenario-harness"; +import { defineHuggingFaceInstrumentationAssertions } from "./assertions"; +import { + HUGGINGFACE_SCENARIO_SPECS, + HUGGINGFACE_SCENARIO_TIMEOUT_MS, +} from "./scenario.impl.mjs"; + +const scenarioDir = await prepareScenarioDir({ + scenarioDir: resolveScenarioDir(import.meta.url), +}); + +const huggingFaceScenarios = await Promise.all( + HUGGINGFACE_SCENARIO_SPECS.map(async (scenario) => ({ + ...scenario, + version: await readInstalledPackageVersion( + scenarioDir, + scenario.dependencyName, + ), + })), +); + +for (const scenario of huggingFaceScenarios) { + describe(`huggingface inference sdk ${scenario.version}`, () => { + defineHuggingFaceInstrumentationAssertions({ + name: "wrapped instrumentation", + runScenario: async ({ runScenarioDir }) => { + await runScenarioDir({ + entry: scenario.wrapperEntry, + runContext: { variantKey: scenario.snapshotName }, + scenarioDir, + timeoutMs: HUGGINGFACE_SCENARIO_TIMEOUT_MS, + }); + }, + snapshotName: scenario.snapshotName, + testFileUrl: import.meta.url, + timeoutMs: HUGGINGFACE_SCENARIO_TIMEOUT_MS, + }); + + defineHuggingFaceInstrumentationAssertions({ + name: "auto-hook instrumentation", + runScenario: async ({ runNodeScenarioDir }) => { + await runNodeScenarioDir({ + entry: scenario.autoEntry, + nodeArgs: ["--import", "braintrust/hook.mjs"], + runContext: { variantKey: scenario.snapshotName }, + scenarioDir, + timeoutMs: HUGGINGFACE_SCENARIO_TIMEOUT_MS, + }); + }, + snapshotName: scenario.snapshotName, + testFileUrl: import.meta.url, + timeoutMs: HUGGINGFACE_SCENARIO_TIMEOUT_MS, + }); + }); +} diff --git a/e2e/scenarios/huggingface-instrumentation/scenario.ts b/e2e/scenarios/huggingface-instrumentation/scenario.ts new file mode 100644 index 000000000..a4f217b2c --- /dev/null +++ b/e2e/scenarios/huggingface-instrumentation/scenario.ts @@ -0,0 +1,5 @@ +import * as huggingFace from "@huggingface/inference"; +import { runMain } from "../../helpers/scenario-runtime"; +import { runWrappedHuggingFaceInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => runWrappedHuggingFaceInstrumentation(huggingFace)); diff --git a/e2e/scripts/run-canary-tests-docker.mjs b/e2e/scripts/run-canary-tests-docker.mjs index 130980697..3035c524b 100644 --- a/e2e/scripts/run-canary-tests-docker.mjs +++ b/e2e/scripts/run-canary-tests-docker.mjs @@ -22,6 +22,7 @@ const ALLOWED_ENV_KEYS = [ "OPENAI_BASE_URL", "OPENROUTER_API_KEY", "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY", ]; function getAllowedEnv() { diff --git a/js/src/auto-instrumentations/bundler/plugin.ts b/js/src/auto-instrumentations/bundler/plugin.ts index c63ebd652..f899ee978 100644 --- a/js/src/auto-instrumentations/bundler/plugin.ts +++ b/js/src/auto-instrumentations/bundler/plugin.ts @@ -25,6 +25,7 @@ import { anthropicConfigs } from "../configs/anthropic"; import { aiSDKConfigs } from "../configs/ai-sdk"; import { claudeAgentSDKConfigs } from "../configs/claude-agent-sdk"; import { googleGenAIConfigs } from "../configs/google-genai"; +import { huggingFaceConfigs } from "../configs/huggingface"; import { openRouterAgentConfigs } from "../configs/openrouter-agent"; import { openRouterConfigs } from "../configs/openrouter"; import { mistralConfigs } from "../configs/mistral"; @@ -74,6 +75,7 @@ export const unplugin = createUnplugin((options = {}) => { ...aiSDKConfigs, ...claudeAgentSDKConfigs, ...googleGenAIConfigs, + ...huggingFaceConfigs, ...openRouterConfigs, ...openRouterAgentConfigs, ...mistralConfigs, diff --git a/js/src/auto-instrumentations/bundler/webpack-loader.ts b/js/src/auto-instrumentations/bundler/webpack-loader.ts index 1431a02cd..676e65175 100644 --- a/js/src/auto-instrumentations/bundler/webpack-loader.ts +++ b/js/src/auto-instrumentations/bundler/webpack-loader.ts @@ -34,6 +34,7 @@ import { anthropicConfigs } from "../configs/anthropic"; import { aiSDKConfigs } from "../configs/ai-sdk"; import { claudeAgentSDKConfigs } from "../configs/claude-agent-sdk"; import { googleGenAIConfigs } from "../configs/google-genai"; +import { huggingFaceConfigs } from "../configs/huggingface"; import { openRouterAgentConfigs } from "../configs/openrouter-agent"; import { openRouterConfigs } from "../configs/openrouter"; import { mistralConfigs } from "../configs/mistral"; @@ -69,6 +70,7 @@ function getMatcher(options: BundlerPluginOptions): InstrumentationMatcher { ...aiSDKConfigs, ...claudeAgentSDKConfigs, ...googleGenAIConfigs, + ...huggingFaceConfigs, ...openRouterConfigs, ...openRouterAgentConfigs, ...mistralConfigs, diff --git a/js/src/auto-instrumentations/configs/huggingface.ts b/js/src/auto-instrumentations/configs/huggingface.ts new file mode 100644 index 000000000..04ed77912 --- /dev/null +++ b/js/src/auto-instrumentations/configs/huggingface.ts @@ -0,0 +1,245 @@ +import type { InstrumentationConfig } from "@apm-js-collab/code-transformer"; +import { huggingFaceChannels } from "../../instrumentation/plugins/huggingface-channels"; + +export const huggingFaceConfigs: InstrumentationConfig[] = [ + { + channelName: huggingFaceChannels.chatCompletion.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.js", + }, + functionQuery: { + functionName: "chatCompletion", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.chatCompletion.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.cjs", + }, + functionQuery: { + functionName: "chatCompletion", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.chatCompletionStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.js", + }, + functionQuery: { + functionName: "chatCompletionStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.chatCompletionStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.cjs", + }, + functionQuery: { + functionName: "chatCompletionStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.textGeneration.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.js", + }, + functionQuery: { + functionName: "textGeneration", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.textGeneration.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.cjs", + }, + functionQuery: { + functionName: "textGeneration", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.textGenerationStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.js", + }, + functionQuery: { + functionName: "textGenerationStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.textGenerationStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.cjs", + }, + functionQuery: { + functionName: "textGenerationStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.featureExtraction.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.js", + }, + functionQuery: { + functionName: "featureExtraction", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.featureExtraction.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=2.0.0 <3.0.0", + filePath: "dist/index.cjs", + }, + functionQuery: { + functionName: "featureExtraction", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.chatCompletion.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/esm/tasks/nlp/chatCompletion.js", + }, + functionQuery: { + functionName: "chatCompletion", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.chatCompletion.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/commonjs/tasks/nlp/chatCompletion.js", + }, + functionQuery: { + functionName: "chatCompletion", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.chatCompletionStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/esm/tasks/nlp/chatCompletionStream.js", + }, + functionQuery: { + functionName: "chatCompletionStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.chatCompletionStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/commonjs/tasks/nlp/chatCompletionStream.js", + }, + functionQuery: { + functionName: "chatCompletionStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.textGeneration.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/esm/tasks/nlp/textGeneration.js", + }, + functionQuery: { + functionName: "textGeneration", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.textGeneration.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/commonjs/tasks/nlp/textGeneration.js", + }, + functionQuery: { + functionName: "textGeneration", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.textGenerationStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/esm/tasks/nlp/textGenerationStream.js", + }, + functionQuery: { + functionName: "textGenerationStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.textGenerationStream.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/commonjs/tasks/nlp/textGenerationStream.js", + }, + functionQuery: { + functionName: "textGenerationStream", + kind: "Sync", + }, + }, + { + channelName: huggingFaceChannels.featureExtraction.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/esm/tasks/nlp/featureExtraction.js", + }, + functionQuery: { + functionName: "featureExtraction", + kind: "Async", + }, + }, + { + channelName: huggingFaceChannels.featureExtraction.channelName, + module: { + name: "@huggingface/inference", + versionRange: ">=3.0.0 <5.0.0", + filePath: "dist/commonjs/tasks/nlp/featureExtraction.js", + }, + functionQuery: { + functionName: "featureExtraction", + kind: "Async", + }, + }, +]; diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index ec2083687..a80905352 100644 --- a/js/src/auto-instrumentations/hook.mts +++ b/js/src/auto-instrumentations/hook.mts @@ -19,6 +19,7 @@ import { anthropicConfigs } from "./configs/anthropic.js"; import { aiSDKConfigs } from "./configs/ai-sdk.js"; import { claudeAgentSDKConfigs } from "./configs/claude-agent-sdk.js"; import { googleGenAIConfigs } from "./configs/google-genai.js"; +import { huggingFaceConfigs } from "./configs/huggingface.js"; import { openRouterAgentConfigs } from "./configs/openrouter-agent.js"; import { openRouterConfigs } from "./configs/openrouter.js"; import { mistralConfigs } from "./configs/mistral.js"; @@ -39,6 +40,7 @@ const allConfigs = [ ...aiSDKConfigs, ...claudeAgentSDKConfigs, ...googleGenAIConfigs, + ...huggingFaceConfigs, ...openRouterConfigs, ...openRouterAgentConfigs, ...mistralConfigs, diff --git a/js/src/auto-instrumentations/index.ts b/js/src/auto-instrumentations/index.ts index 4177ae061..d81a79eec 100644 --- a/js/src/auto-instrumentations/index.ts +++ b/js/src/auto-instrumentations/index.ts @@ -33,6 +33,7 @@ export { anthropicConfigs } from "./configs/anthropic"; export { aiSDKConfigs } from "./configs/ai-sdk"; export { claudeAgentSDKConfigs } from "./configs/claude-agent-sdk"; export { googleGenAIConfigs } from "./configs/google-genai"; +export { huggingFaceConfigs } from "./configs/huggingface"; export { openRouterAgentConfigs } from "./configs/openrouter-agent"; export { openRouterConfigs } from "./configs/openrouter"; export { mistralConfigs } from "./configs/mistral"; diff --git a/js/src/exports.ts b/js/src/exports.ts index d4a1a61c8..9dbc592f3 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -175,6 +175,7 @@ export { wrapAnthropic } from "./wrappers/anthropic"; export { wrapMastraAgent } from "./wrappers/mastra"; export { wrapClaudeAgentSDK } from "./wrappers/claude-agent-sdk/claude-agent-sdk"; export { wrapGoogleGenAI } from "./wrappers/google-genai"; +export { wrapHuggingFace } from "./wrappers/huggingface"; export { wrapOpenRouterAgent } from "./wrappers/openrouter-agent"; export { wrapOpenRouter } from "./wrappers/openrouter"; export { wrapMistral } from "./wrappers/mistral"; diff --git a/js/src/instrumentation/braintrust-plugin.test.ts b/js/src/instrumentation/braintrust-plugin.test.ts index 07d91337d..1ec75ed9c 100644 --- a/js/src/instrumentation/braintrust-plugin.test.ts +++ b/js/src/instrumentation/braintrust-plugin.test.ts @@ -5,6 +5,7 @@ import { AnthropicPlugin } from "./plugins/anthropic-plugin"; import { AISDKPlugin } from "./plugins/ai-sdk-plugin"; import { ClaudeAgentSDKPlugin } from "./plugins/claude-agent-sdk-plugin"; import { GoogleGenAIPlugin } from "./plugins/google-genai-plugin"; +import { HuggingFacePlugin } from "./plugins/huggingface-plugin"; import { OpenRouterAgentPlugin } from "./plugins/openrouter-agent-plugin"; import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; import { MistralPlugin } from "./plugins/mistral-plugin"; @@ -46,6 +47,10 @@ vi.mock("./plugins/google-genai-plugin", () => ({ GoogleGenAIPlugin: createPluginClassMock(), })); +vi.mock("./plugins/huggingface-plugin", () => ({ + HuggingFacePlugin: createPluginClassMock(), +})); + vi.mock("./plugins/openrouter-plugin", () => ({ OpenRouterPlugin: createPluginClassMock(), })); @@ -110,6 +115,15 @@ describe("BraintrustPlugin", () => { expect(mockInstance.enable).toHaveBeenCalledTimes(1); }); + it("should create and enable HuggingFace plugin by default", () => { + const plugin = new BraintrustPlugin(); + plugin.enable(); + + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); + const mockInstance = vi.mocked(HuggingFacePlugin).mock.results[0].value; + expect(mockInstance.enable).toHaveBeenCalledTimes(1); + }); + it("should create and enable OpenRouter plugin by default", () => { const plugin = new BraintrustPlugin(); plugin.enable(); @@ -148,6 +162,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); @@ -162,6 +177,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); @@ -176,6 +192,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); @@ -195,6 +212,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); }); @@ -211,6 +229,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); }); @@ -227,6 +246,7 @@ describe("BraintrustPlugin", () => { expect(AnthropicPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); @@ -244,6 +264,24 @@ describe("BraintrustPlugin", () => { expect(AnthropicPlugin).toHaveBeenCalledTimes(1); expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); + expect(MistralPlugin).toHaveBeenCalledTimes(1); + }); + + it("should not create HuggingFace plugin when huggingface: false", () => { + const plugin = new BraintrustPlugin({ + integrations: { huggingface: false }, + }); + plugin.enable(); + + expect(HuggingFacePlugin).not.toHaveBeenCalled(); + expect(OpenAIPlugin).toHaveBeenCalledTimes(1); + expect(AnthropicPlugin).toHaveBeenCalledTimes(1); + expect(AISDKPlugin).toHaveBeenCalledTimes(1); + expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); + expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); @@ -279,6 +317,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -294,6 +333,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); }); @@ -315,6 +355,7 @@ describe("BraintrustPlugin", () => { aisdk: false, claudeAgentSDK: false, googleGenAI: false, + huggingface: false, openrouter: false, openrouterAgent: false, mistral: false, @@ -327,6 +368,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).not.toHaveBeenCalled(); expect(ClaudeAgentSDKPlugin).not.toHaveBeenCalled(); expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); + expect(HuggingFacePlugin).not.toHaveBeenCalled(); expect(OpenRouterPlugin).not.toHaveBeenCalled(); expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); expect(MistralPlugin).not.toHaveBeenCalled(); @@ -340,6 +382,7 @@ describe("BraintrustPlugin", () => { aisdk: false, claudeAgentSDK: true, googleGenAI: false, + huggingface: true, openrouter: true, mistral: false, }, @@ -348,6 +391,7 @@ describe("BraintrustPlugin", () => { expect(OpenAIPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(AnthropicPlugin).not.toHaveBeenCalled(); @@ -370,6 +414,7 @@ describe("BraintrustPlugin", () => { expect(AnthropicPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -386,6 +431,7 @@ describe("BraintrustPlugin", () => { expect(AnthropicPlugin).toHaveBeenCalledTimes(1); expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -457,6 +503,8 @@ describe("BraintrustPlugin", () => { vi.mocked(ClaudeAgentSDKPlugin).mock.results[0].value; const googleGenAIMock = vi.mocked(GoogleGenAIPlugin).mock.results[0].value; + const huggingFaceMock = + vi.mocked(HuggingFacePlugin).mock.results[0].value; const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock .results[0].value; @@ -467,6 +515,7 @@ describe("BraintrustPlugin", () => { expect(aiSDKMock.enable).toHaveBeenCalledTimes(1); expect(claudeAgentSDKMock.enable).toHaveBeenCalledTimes(1); expect(googleGenAIMock.enable).toHaveBeenCalledTimes(1); + expect(huggingFaceMock.enable).toHaveBeenCalledTimes(1); expect(openRouterMock.enable).toHaveBeenCalledTimes(1); expect(openRouterAgentMock.enable).toHaveBeenCalledTimes(1); expect(mistralMock.enable).toHaveBeenCalledTimes(1); @@ -483,6 +532,8 @@ describe("BraintrustPlugin", () => { vi.mocked(ClaudeAgentSDKPlugin).mock.results[0].value; const googleGenAIMock = vi.mocked(GoogleGenAIPlugin).mock.results[0].value; + const huggingFaceMock = + vi.mocked(HuggingFacePlugin).mock.results[0].value; const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock .results[0].value; @@ -495,6 +546,7 @@ describe("BraintrustPlugin", () => { expect(aiSDKMock.disable).toHaveBeenCalledTimes(1); expect(claudeAgentSDKMock.disable).toHaveBeenCalledTimes(1); expect(googleGenAIMock.disable).toHaveBeenCalledTimes(1); + expect(huggingFaceMock.disable).toHaveBeenCalledTimes(1); expect(openRouterMock.disable).toHaveBeenCalledTimes(1); expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); expect(mistralMock.disable).toHaveBeenCalledTimes(1); @@ -536,6 +588,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).not.toHaveBeenCalled(); expect(ClaudeAgentSDKPlugin).not.toHaveBeenCalled(); expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); + expect(HuggingFacePlugin).not.toHaveBeenCalled(); expect(OpenRouterPlugin).not.toHaveBeenCalled(); expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); expect(MistralPlugin).not.toHaveBeenCalled(); @@ -555,6 +608,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); + expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); @@ -568,6 +622,7 @@ describe("BraintrustPlugin", () => { aisdk: true, claudeAgentSDK: false, googleGenAI: true, + huggingface: true, openrouter: true, openrouterAgent: true, mistral: false, @@ -579,6 +634,8 @@ describe("BraintrustPlugin", () => { const aiSDKMock = vi.mocked(AISDKPlugin).mock.results[0].value; const googleGenAIMock = vi.mocked(GoogleGenAIPlugin).mock.results[0].value; + const huggingFaceMock = + vi.mocked(HuggingFacePlugin).mock.results[0].value; const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock .results[0].value; @@ -588,6 +645,7 @@ describe("BraintrustPlugin", () => { expect(openaiMock.disable).toHaveBeenCalledTimes(1); expect(aiSDKMock.disable).toHaveBeenCalledTimes(1); expect(googleGenAIMock.disable).toHaveBeenCalledTimes(1); + expect(huggingFaceMock.disable).toHaveBeenCalledTimes(1); expect(openRouterMock.disable).toHaveBeenCalledTimes(1); expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); expect(MistralPlugin).not.toHaveBeenCalled(); diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index cc6523de3..3a917c5ad 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -4,6 +4,7 @@ import { AnthropicPlugin } from "./plugins/anthropic-plugin"; import { AISDKPlugin } from "./plugins/ai-sdk-plugin"; import { ClaudeAgentSDKPlugin } from "./plugins/claude-agent-sdk-plugin"; import { GoogleGenAIPlugin } from "./plugins/google-genai-plugin"; +import { HuggingFacePlugin } from "./plugins/huggingface-plugin"; import { OpenRouterAgentPlugin } from "./plugins/openrouter-agent-plugin"; import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; import { MistralPlugin } from "./plugins/mistral-plugin"; @@ -16,6 +17,7 @@ export interface BraintrustPluginConfig { aisdk?: boolean; google?: boolean; googleGenAI?: boolean; + huggingface?: boolean; claudeAgentSDK?: boolean; openrouter?: boolean; openrouterAgent?: boolean; @@ -32,6 +34,7 @@ export interface BraintrustPluginConfig { * - Claude Agent SDK (agent interactions) * - Vercel AI SDK (generateText, streamText, etc.) * - Google GenAI SDK + * - HuggingFace Inference SDK * - Mistral SDK * * The plugin is automatically enabled when the Braintrust library is loaded. @@ -44,6 +47,7 @@ export class BraintrustPlugin extends BasePlugin { private aiSDKPlugin: AISDKPlugin | null = null; private claudeAgentSDKPlugin: ClaudeAgentSDKPlugin | null = null; private googleGenAIPlugin: GoogleGenAIPlugin | null = null; + private huggingFacePlugin: HuggingFacePlugin | null = null; private openRouterPlugin: OpenRouterPlugin | null = null; private openRouterAgentPlugin: OpenRouterAgentPlugin | null = null; private mistralPlugin: MistralPlugin | null = null; @@ -88,6 +92,11 @@ export class BraintrustPlugin extends BasePlugin { this.googleGenAIPlugin.enable(); } + if (integrations.huggingface !== false) { + this.huggingFacePlugin = new HuggingFacePlugin(); + this.huggingFacePlugin.enable(); + } + if (integrations.openrouter !== false) { this.openRouterPlugin = new OpenRouterPlugin(); this.openRouterPlugin.enable(); @@ -130,6 +139,11 @@ export class BraintrustPlugin extends BasePlugin { this.googleGenAIPlugin = null; } + if (this.huggingFacePlugin) { + this.huggingFacePlugin.disable(); + this.huggingFacePlugin = null; + } + if (this.openRouterPlugin) { this.openRouterPlugin.disable(); this.openRouterPlugin = null; diff --git a/js/src/instrumentation/plugins/huggingface-channels.ts b/js/src/instrumentation/plugins/huggingface-channels.ts new file mode 100644 index 000000000..9a7cc9c42 --- /dev/null +++ b/js/src/instrumentation/plugins/huggingface-channels.ts @@ -0,0 +1,57 @@ +import { channel, defineChannels } from "../core/channel-definitions"; +import type { + HuggingFaceChatCompletion, + HuggingFaceChatCompletionChunk, + HuggingFaceChatCompletionParams, + HuggingFaceFeatureExtractionOutput, + HuggingFaceFeatureExtractionParams, + HuggingFaceTextGenerationOutput, + HuggingFaceTextGenerationParams, + HuggingFaceTextGenerationStreamOutput, +} from "../../vendor-sdk-types/huggingface"; + +export const huggingFaceChannels = defineChannels("@huggingface/inference", { + chatCompletion: channel< + [HuggingFaceChatCompletionParams], + HuggingFaceChatCompletion + >({ + channelName: "chatCompletion", + kind: "async", + }), + + chatCompletionStream: channel< + [HuggingFaceChatCompletionParams], + AsyncIterable, + Record, + HuggingFaceChatCompletionChunk + >({ + channelName: "chatCompletionStream", + kind: "sync-stream", + }), + + textGeneration: channel< + [HuggingFaceTextGenerationParams], + HuggingFaceTextGenerationOutput + >({ + channelName: "textGeneration", + kind: "async", + }), + + textGenerationStream: channel< + [HuggingFaceTextGenerationParams], + AsyncIterable, + Record, + HuggingFaceTextGenerationStreamOutput + >({ + channelName: "textGenerationStream", + kind: "sync-stream", + }), + + featureExtraction: channel< + [HuggingFaceFeatureExtractionParams], + HuggingFaceFeatureExtractionOutput + >({ + channelName: "featureExtraction", + kind: "async", + }), +}); diff --git a/js/src/instrumentation/plugins/huggingface-plugin.ts b/js/src/instrumentation/plugins/huggingface-plugin.ts new file mode 100644 index 000000000..677c8ef03 --- /dev/null +++ b/js/src/instrumentation/plugins/huggingface-plugin.ts @@ -0,0 +1,513 @@ +import { + traceAsyncChannel, + traceSyncStreamChannel, + unsubscribeAll, +} from "../core/channel-tracing"; +import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; +import { BasePlugin } from "../core"; +import { SpanTypeAttribute, isObject } from "../../../util/index"; +import { getCurrentUnixTimestamp } from "../../util"; +import { parseMetricsFromUsage } from "../../openai-utils"; +import { huggingFaceChannels } from "./huggingface-channels"; +import type { + HuggingFaceChatCompletion, + HuggingFaceChatCompletionChunk, + HuggingFaceFeatureExtractionOutput, + HuggingFaceTextGenerationDetails, + HuggingFaceTextGenerationStreamOutput, +} from "../../vendor-sdk-types/huggingface"; +import type { Span } from "../../logger"; + +const REQUEST_METADATA_ALLOWLIST = new Set([ + "dimensions", + "encoding_format", + "endpointUrl", + "max_tokens", + "model", + "provider", + "seed", + "stop", + "stream", + "temperature", + "top_p", +]); + +const RESPONSE_METADATA_ALLOWLIST = new Set([ + "created", + "id", + "model", + "object", +]); + +export class HuggingFacePlugin extends BasePlugin { + protected onEnable(): void { + this.unsubscribers.push( + traceAsyncChannel(huggingFaceChannels.chatCompletion, { + name: "huggingface.chat_completion", + type: SpanTypeAttribute.LLM, + extractInput: extractChatInputWithMetadata, + extractOutput: (result) => result?.choices, + extractMetadata: (result) => extractResponseMetadata(result), + extractMetrics: (result) => parseMetricsFromUsage(result?.usage), + }), + traceSyncStreamChannel(huggingFaceChannels.chatCompletionStream, { + name: "huggingface.chat_completion_stream", + type: SpanTypeAttribute.LLM, + extractInput: extractChatInputWithMetadata, + patchResult: ({ result, span, startTime }) => + patchChatCompletionStream({ + result, + span, + startTime, + }), + }), + traceAsyncChannel(huggingFaceChannels.textGeneration, { + name: "huggingface.text_generation", + type: SpanTypeAttribute.LLM, + extractInput: extractTextGenerationInputWithMetadata, + extractOutput: (result) => + isObject(result) ? { generated_text: result.generated_text } : result, + extractMetadata: extractTextGenerationMetadata, + extractMetrics: (result) => + extractTextGenerationMetrics(result?.details ?? null), + }), + traceSyncStreamChannel(huggingFaceChannels.textGenerationStream, { + name: "huggingface.text_generation_stream", + type: SpanTypeAttribute.LLM, + extractInput: extractTextGenerationInputWithMetadata, + patchResult: ({ result, span, startTime }) => + patchTextGenerationStream({ + result, + span, + startTime, + }), + }), + traceAsyncChannel(huggingFaceChannels.featureExtraction, { + name: "huggingface.feature_extraction", + type: SpanTypeAttribute.LLM, + extractInput: extractFeatureExtractionInputWithMetadata, + extractOutput: summarizeFeatureExtractionOutput, + extractMetrics: () => ({}), + }), + ); + } + + protected onDisable(): void { + this.unsubscribers = unsubscribeAll(this.unsubscribers); + } +} + +function addProviderMetadata( + metadata: Record, +): Record { + return { + ...metadata, + provider: metadata.provider ?? "huggingface", + }; +} + +function normalizeArgs(args: unknown[] | unknown): unknown[] { + if (Array.isArray(args)) { + return args; + } + + if (isArrayLike(args)) { + return Array.from(args); + } + + return [args]; +} + +function isArrayLike(value: unknown): value is ArrayLike { + return ( + isObject(value) && + "length" in value && + typeof value.length === "number" && + Number.isInteger(value.length) && + value.length >= 0 + ); +} + +function getFirstObjectArg( + args: unknown[] | unknown, +): Record | undefined { + const firstObjectArg = normalizeArgs(args).find((arg) => isObject(arg)); + return isObject(firstObjectArg) ? firstObjectArg : undefined; +} + +function pickRequestMetadata( + params: Record | undefined, +): Record { + if (!params) { + return addProviderMetadata({}); + } + + const metadata: Record = {}; + for (const key of REQUEST_METADATA_ALLOWLIST) { + const value = params[key]; + if (value !== undefined) { + metadata[key] = value; + } + } + + if (isObject(params.parameters)) { + metadata.parameters = params.parameters; + } + + return addProviderMetadata(metadata); +} + +function extractChatInputWithMetadata(args: unknown[] | unknown): { + input: unknown; + metadata: Record; +} { + const params = getFirstObjectArg(args); + const { messages, ...rawMetadata } = params ?? {}; + return { + input: messages, + metadata: pickRequestMetadata(rawMetadata), + }; +} + +function extractTextGenerationInputWithMetadata(args: unknown[] | unknown): { + input: unknown; + metadata: Record; +} { + const params = getFirstObjectArg(args); + const { inputs, ...rawMetadata } = params ?? {}; + return { + input: inputs, + metadata: pickRequestMetadata(rawMetadata), + }; +} + +function extractFeatureExtractionInputWithMetadata(args: unknown[] | unknown): { + input: unknown; + metadata: Record; +} { + const params = getFirstObjectArg(args); + const { inputs, ...rawMetadata } = params ?? {}; + return { + input: inputs, + metadata: pickRequestMetadata(rawMetadata), + }; +} + +function extractResponseMetadata( + result: HuggingFaceChatCompletion | undefined, +): Record | undefined { + if (!isObject(result)) { + return undefined; + } + + const metadata: Record = {}; + for (const key of RESPONSE_METADATA_ALLOWLIST) { + const value = result[key]; + if (value !== undefined) { + metadata[key] = value; + } + } + + return Object.keys(metadata).length > 0 ? metadata : undefined; +} + +function extractTextGenerationMetrics( + details: HuggingFaceTextGenerationDetails | null | undefined, +): Record { + if (!isObject(details)) { + return {}; + } + + const promptTokens = Array.isArray(details.prefill) + ? details.prefill.length + : undefined; + const completionTokens = + typeof details.generated_tokens === "number" + ? details.generated_tokens + : Array.isArray(details.tokens) + ? details.tokens.length + : undefined; + + const metrics: Record = {}; + if (promptTokens !== undefined) { + metrics.prompt_tokens = promptTokens; + } + if (completionTokens !== undefined) { + metrics.completion_tokens = completionTokens; + } + if (promptTokens !== undefined || completionTokens !== undefined) { + metrics.tokens = (promptTokens ?? 0) + (completionTokens ?? 0); + } + return metrics; +} + +function extractTextGenerationMetadata(result: { + details?: HuggingFaceTextGenerationDetails | null; +}): Record | undefined { + if (!isObject(result?.details)) { + return undefined; + } + + return typeof result.details.finish_reason === "string" + ? { + finish_reason: result.details.finish_reason, + } + : undefined; +} + +function summarizeFeatureExtractionOutput( + result: HuggingFaceFeatureExtractionOutput, +): Record | undefined { + if (!Array.isArray(result)) { + return undefined; + } + + const first = result[0]; + if (typeof first === "number") { + return { embedding_length: result.length }; + } + + if ( + Array.isArray(first) && + first.every((value) => typeof value === "number") + ) { + return { + embedding_count: result.length, + embedding_length: first.length, + }; + } + + if ( + Array.isArray(first) && + first.length > 0 && + Array.isArray(first[0]) && + first[0].every((value) => typeof value === "number") + ) { + return { + embedding_batch_count: result.length, + embedding_count: first.length, + embedding_length: first[0].length, + }; + } + + return undefined; +} + +function patchChatCompletionStream(args: { + result: AsyncIterable; + span: Span; + startTime: number; +}): boolean { + const { result, span, startTime } = args; + if (!result || !isAsyncIterable(result)) { + return false; + } + + let firstChunkTime: number | undefined; + patchStreamIfNeeded(result, { + onChunk: () => { + if (firstChunkTime === undefined) { + firstChunkTime = getCurrentUnixTimestamp(); + } + }, + onComplete: (chunks) => { + const lastChunk = chunks.at(-1); + const metrics = { + ...parseMetricsFromUsage(lastChunk?.usage), + ...(firstChunkTime !== undefined + ? { time_to_first_token: firstChunkTime - startTime } + : {}), + }; + + span.log({ + output: aggregateChatCompletionChunks(chunks), + ...(extractResponseMetadata(lastChunk) + ? { metadata: extractResponseMetadata(lastChunk) } + : {}), + metrics, + }); + span.end(); + }, + onError: (error) => { + span.log({ + error: error.message, + }); + span.end(); + }, + }); + + return true; +} + +function patchTextGenerationStream(args: { + result: AsyncIterable; + span: Span; + startTime: number; +}): boolean { + const { result, span, startTime } = args; + if (!result || !isAsyncIterable(result)) { + return false; + } + + let firstChunkTime: number | undefined; + patchStreamIfNeeded(result, { + onChunk: () => { + if (firstChunkTime === undefined) { + firstChunkTime = getCurrentUnixTimestamp(); + } + }, + onComplete: (chunks) => { + const lastChunk = chunks.at(-1); + span.log({ + output: aggregateTextGenerationStreamChunks(chunks), + ...(extractTextGenerationStreamMetadata(chunks) + ? { + metadata: extractTextGenerationStreamMetadata(chunks), + } + : {}), + metrics: { + ...extractTextGenerationMetrics(lastChunk?.details ?? null), + ...parseMetricsFromUsage(lastChunk?.usage), + ...(firstChunkTime !== undefined + ? { time_to_first_token: firstChunkTime - startTime } + : {}), + }, + }); + span.end(); + }, + onError: (error) => { + span.log({ + error: error.message, + }); + span.end(); + }, + }); + + return true; +} + +function aggregateChatCompletionChunks( + chunks: HuggingFaceChatCompletionChunk[], +): { choices: Array> } | undefined { + if (chunks.length === 0) { + return undefined; + } + + const aggregatedChoices = new Map< + number, + { + content: string; + finish_reason?: string | null; + role?: string; + } + >(); + + for (const chunk of chunks) { + for (const choice of chunk.choices ?? []) { + const index = typeof choice.index === "number" ? choice.index : 0; + const existing = aggregatedChoices.get(index) ?? { content: "" }; + const delta = isObject(choice.delta) ? choice.delta : undefined; + const message = isObject(choice.message) ? choice.message : undefined; + + if (typeof delta?.content === "string") { + existing.content += delta.content; + } else if (typeof message?.content === "string") { + existing.content = message.content; + } + + if (typeof delta?.role === "string") { + existing.role = delta.role; + } else if (typeof message?.role === "string") { + existing.role = message.role; + } + + if (choice.finish_reason !== undefined) { + existing.finish_reason = choice.finish_reason; + } + + aggregatedChoices.set(index, existing); + } + } + + return { + choices: [...aggregatedChoices.entries()].map(([index, choice]) => ({ + index, + message: { + content: choice.content, + role: choice.role ?? "assistant", + }, + ...(choice.finish_reason !== undefined + ? { finish_reason: choice.finish_reason } + : {}), + })), + }; +} + +function aggregateTextGenerationStreamChunks( + chunks: HuggingFaceTextGenerationStreamOutput[], +): { generated_text: string; finish_reason?: string | null } | undefined { + if (chunks.length === 0) { + return undefined; + } + + let generatedText = ""; + let finishReason: string | null | undefined; + for (const chunk of chunks) { + if (typeof chunk.generated_text === "string") { + generatedText = chunk.generated_text; + } else if (typeof chunk.token?.text === "string" && !chunk.token.special) { + generatedText += chunk.token.text; + } else if (Array.isArray(chunk.choices)) { + for (const choice of chunk.choices) { + if (typeof choice.text === "string") { + generatedText += choice.text; + } + + if (choice.finish_reason !== undefined) { + finishReason = choice.finish_reason; + } + } + } + + if ( + isObject(chunk.details) && + typeof chunk.details.finish_reason === "string" + ) { + finishReason = chunk.details.finish_reason; + } + } + + return { + generated_text: generatedText, + ...(finishReason !== undefined ? { finish_reason: finishReason } : {}), + }; +} + +function extractTextGenerationStreamMetadata( + chunks: HuggingFaceTextGenerationStreamOutput[], +): Record | undefined { + for (let index = chunks.length - 1; index >= 0; index--) { + const chunk = chunks[index]; + if ( + isObject(chunk?.details) && + typeof chunk.details.finish_reason === "string" + ) { + return { + finish_reason: chunk.details.finish_reason, + }; + } + + if (!Array.isArray(chunk?.choices)) { + continue; + } + + const finishReason = chunk.choices.findLast( + (choice) => choice.finish_reason !== undefined, + )?.finish_reason; + if (finishReason !== undefined) { + return { finish_reason: finishReason }; + } + } + + return undefined; +} diff --git a/js/src/instrumentation/registry.test.ts b/js/src/instrumentation/registry.test.ts index fc2c36bd3..f5dd6a2ac 100644 --- a/js/src/instrumentation/registry.test.ts +++ b/js/src/instrumentation/registry.test.ts @@ -118,6 +118,7 @@ describe("configureInstrumentation API", () => { integrations: { openai: false, anthropic: true, + huggingface: true, openrouter: false, mistral: false, }, diff --git a/js/src/instrumentation/registry.ts b/js/src/instrumentation/registry.ts index d506dcc0b..0aeb7a82c 100644 --- a/js/src/instrumentation/registry.ts +++ b/js/src/instrumentation/registry.ts @@ -19,6 +19,7 @@ export interface InstrumentationConfig { vercel?: boolean; aisdk?: boolean; google?: boolean; + huggingface?: boolean; claudeAgentSDK?: boolean; openrouter?: boolean; openrouterAgent?: boolean; @@ -107,6 +108,7 @@ class PluginRegistry { vercel: true, aisdk: true, google: true, + huggingface: true, claudeAgentSDK: true, openrouter: true, openrouterAgent: true, diff --git a/js/src/vendor-sdk-types/huggingface.ts b/js/src/vendor-sdk-types/huggingface.ts new file mode 100644 index 000000000..1b4c46c07 --- /dev/null +++ b/js/src/vendor-sdk-types/huggingface.ts @@ -0,0 +1,175 @@ +export interface HuggingFaceRequestOptions { + retry_on_error?: boolean; + fetch?: typeof fetch; + signal?: AbortSignal; + includeCredentials?: string | boolean; + billTo?: string; + [key: string]: unknown; +} + +export interface HuggingFaceClientConstructor { + new (...args: unknown[]): HuggingFaceClient; +} + +export interface HuggingFaceModule { + InferenceClient?: HuggingFaceClientConstructor; + InferenceClientEndpoint?: HuggingFaceClientConstructor; + HfInference?: HuggingFaceClientConstructor; + HfInferenceEndpoint?: HuggingFaceClientConstructor; + chatCompletion?: ( + params: HuggingFaceChatCompletionParams, + options?: HuggingFaceRequestOptions, + ) => Promise; + chatCompletionStream?: ( + params: HuggingFaceChatCompletionParams, + options?: HuggingFaceRequestOptions, + ) => AsyncIterable; + textGeneration?: ( + params: HuggingFaceTextGenerationParams, + options?: HuggingFaceRequestOptions, + ) => Promise; + textGenerationStream?: ( + params: HuggingFaceTextGenerationParams, + options?: HuggingFaceRequestOptions, + ) => AsyncIterable; + featureExtraction?: ( + params: HuggingFaceFeatureExtractionParams, + options?: HuggingFaceRequestOptions, + ) => Promise; + [key: string]: unknown; +} + +export interface HuggingFaceClient { + chatCompletion: ( + params: HuggingFaceChatCompletionParams, + options?: HuggingFaceRequestOptions, + ) => Promise; + chatCompletionStream: ( + params: HuggingFaceChatCompletionParams, + options?: HuggingFaceRequestOptions, + ) => AsyncIterable; + textGeneration: ( + params: HuggingFaceTextGenerationParams, + options?: HuggingFaceRequestOptions, + ) => Promise; + textGenerationStream: ( + params: HuggingFaceTextGenerationParams, + options?: HuggingFaceRequestOptions, + ) => AsyncIterable; + featureExtraction: ( + params: HuggingFaceFeatureExtractionParams, + options?: HuggingFaceRequestOptions, + ) => Promise; + endpoint?: (endpointUrl: string) => HuggingFaceClient; + [key: string]: unknown; +} + +export interface HuggingFaceChatCompletionParams { + messages?: unknown; + model?: string; + provider?: string; + endpointUrl?: string; + stream?: boolean; + [key: string]: unknown; +} + +export interface HuggingFaceTextGenerationParams { + inputs?: unknown; + model?: string; + provider?: string; + endpointUrl?: string; + stream?: boolean; + parameters?: Record; + [key: string]: unknown; +} + +export interface HuggingFaceFeatureExtractionParams { + inputs?: unknown; + model?: string; + provider?: string; + endpointUrl?: string; + dimensions?: number | null; + encoding_format?: "float" | "base64"; + [key: string]: unknown; +} + +export interface HuggingFaceUsage { + prompt_tokens?: number; + completion_tokens?: number; + total_tokens?: number; + [key: string]: unknown; +} + +export interface HuggingFaceChatMessage { + role?: string; + content?: string | null | unknown[]; + tool_calls?: unknown; + [key: string]: unknown; +} + +export interface HuggingFaceChatCompletionChoice { + index?: number; + message?: HuggingFaceChatMessage; + delta?: HuggingFaceChatMessage; + finish_reason?: string | null; + [key: string]: unknown; +} + +export interface HuggingFaceChatCompletion { + id?: string; + object?: string; + model?: string; + created?: number; + usage?: HuggingFaceUsage; + choices?: HuggingFaceChatCompletionChoice[]; + [key: string]: unknown; +} + +export type HuggingFaceChatCompletionChunk = HuggingFaceChatCompletion; + +export interface HuggingFaceTextGenerationToken { + id?: number; + text?: string; + logprob?: number; + special?: boolean; + [key: string]: unknown; +} + +export interface HuggingFaceTextGenerationDetails { + finish_reason?: string | null; + generated_tokens?: number; + prefill?: HuggingFaceTextGenerationToken[]; + tokens?: HuggingFaceTextGenerationToken[]; + [key: string]: unknown; +} + +export interface HuggingFaceTextGenerationOutput { + generated_text?: string | null; + details?: HuggingFaceTextGenerationDetails | null; + [key: string]: unknown; +} + +export interface HuggingFaceTextGenerationChoice { + index?: number; + text?: string; + finish_reason?: string | null; + [key: string]: unknown; +} + +export interface HuggingFaceTextGenerationStreamOutput { + index?: number; + token?: HuggingFaceTextGenerationToken; + choices?: HuggingFaceTextGenerationChoice[]; + generated_text?: string | null; + details?: HuggingFaceTextGenerationDetails | null; + usage?: HuggingFaceUsage; + model?: string; + object?: string; + created?: number; + [key: string]: unknown; +} + +export type HuggingFaceFeatureExtractionOutput = + | number[] + | number[][] + | number[][][]; diff --git a/js/src/wrappers/huggingface.test.ts b/js/src/wrappers/huggingface.test.ts new file mode 100644 index 000000000..dc0805cb9 --- /dev/null +++ b/js/src/wrappers/huggingface.test.ts @@ -0,0 +1,398 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, +} from "vitest"; +import { configureNode } from "../node/config"; +import { + _exportsForTestingOnly, + initLogger, + Logger, + TestBackgroundLogger, +} from "../logger"; +import { wrapHuggingFace } from "./huggingface"; + +try { + configureNode(); +} catch { + // Shared process setup in tests can run this more than once. +} + +function buildFakeClientClass() { + return class FakeInferenceClient { + public constructor( + _accessToken = "", + _defaultOptions: Record = {}, + ) {} + + public chatCompletion = async (params: { + model?: string; + messages?: unknown; + }) => ({ + id: "chatcmpl_test", + object: "chat.completion", + model: params.model, + usage: { + prompt_tokens: 6, + completion_tokens: 1, + total_tokens: 7, + }, + choices: [ + { + index: 0, + message: { + role: "assistant", + content: "Paris", + }, + finish_reason: "stop", + }, + ], + }); + + public chatCompletionStream = (params: { model?: string }) => + (async function* () { + yield { + id: "chatcmpl_stream", + model: params.model, + choices: [ + { + index: 0, + delta: { + role: "assistant", + content: "Par", + }, + }, + ], + }; + yield { + id: "chatcmpl_stream", + model: params.model, + usage: { + prompt_tokens: 6, + completion_tokens: 1, + total_tokens: 7, + }, + choices: [ + { + index: 0, + delta: { + content: "is", + }, + finish_reason: "stop", + }, + ], + }; + })(); + + public textGeneration = async (params: { model?: string }) => ({ + generated_text: "Paris", + details: { + finish_reason: "eos_token", + generated_tokens: 1, + prefill: [{ id: 1, text: "Prompt" }], + }, + model: params.model, + }); + + public textGenerationStream = () => + (async function* () { + yield { + token: { + id: 1, + text: "Par", + logprob: 0, + special: false, + }, + generated_text: null, + details: null, + }; + yield { + token: { + id: 2, + text: "is", + logprob: 0, + special: false, + }, + generated_text: "Paris", + details: { + finish_reason: "stop_sequence", + generated_tokens: 2, + prefill: [{ id: 0, text: "Prompt" }], + tokens: [ + { id: 1, text: "Par", logprob: 0, special: false }, + { id: 2, text: "is", logprob: 0, special: false }, + ], + }, + }; + })(); + + public featureExtraction = async () => [[0.1, 0.2, 0.3]]; + + public endpoint(_endpointUrl: string) { + return new FakeInferenceClient(); + } + }; +} + +function buildModernModule() { + const FakeInferenceClient = buildFakeClientClass(); + const fakeClient = new FakeInferenceClient(); + + return { + InferenceClient: FakeInferenceClient, + chatCompletion: fakeClient.chatCompletion, + chatCompletionStream: fakeClient.chatCompletionStream, + textGeneration: fakeClient.textGeneration, + textGenerationStream: fakeClient.textGenerationStream, + featureExtraction: fakeClient.featureExtraction, + }; +} + +function buildLegacyModule() { + const FakeInferenceClient = buildFakeClientClass(); + const fakeClient = new FakeInferenceClient(); + + return { + HfInference: FakeInferenceClient, + featureExtraction: fakeClient.featureExtraction, + }; +} + +function buildCompletionStyleTextGenerationModule() { + const FakeInferenceClient = buildFakeClientClass(); + const fakeClient = new FakeInferenceClient(); + + return { + InferenceClient: FakeInferenceClient, + textGenerationStream: () => + (async function* () { + yield { + id: "textgen_stream", + object: "text_completion", + model: "meta-llama/Meta-Llama-3.1-8B", + choices: [ + { + index: 0, + text: "Par", + finish_reason: null, + }, + ], + }; + yield { + id: "textgen_stream", + object: "text_completion", + model: "meta-llama/Meta-Llama-3.1-8B", + choices: [ + { + index: 0, + text: "is", + finish_reason: "stop", + }, + ], + }; + yield { + id: "textgen_stream", + object: "text_completion", + model: "meta-llama/Meta-Llama-3.1-8B", + choices: [], + usage: { + prompt_tokens: 5, + completion_tokens: 2, + total_tokens: 7, + }, + }; + })(), + chatCompletion: fakeClient.chatCompletion, + chatCompletionStream: fakeClient.chatCompletionStream, + textGeneration: fakeClient.textGeneration, + featureExtraction: fakeClient.featureExtraction, + }; +} + +async function collectAsync(records: AsyncIterable) { + const items: T[] = []; + for await (const record of records) { + items.push(record); + } + return items; +} + +describe("wrapHuggingFace", () => { + let backgroundLogger: TestBackgroundLogger; + let _logger: Logger; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + _logger = initLogger({ + projectName: "huggingface.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + test("wraps InferenceClient chatCompletion calls", async () => { + const { InferenceClient } = wrapHuggingFace(buildModernModule()); + const client = new InferenceClient("hf_test"); + + await client.chatCompletion({ + model: "Qwen/Qwen3-32B", + messages: [{ role: "user", content: "Reply with exactly PARIS." }], + temperature: 0, + }); + + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + span_attributes: { + name: "huggingface.chat_completion", + type: "llm", + }, + metadata: expect.objectContaining({ + model: "Qwen/Qwen3-32B", + provider: "huggingface", + }), + metrics: expect.objectContaining({ + prompt_tokens: 6, + completion_tokens: 1, + tokens: 7, + }), + output: [ + expect.objectContaining({ + message: expect.objectContaining({ + content: "Paris", + }), + }), + ], + }); + }); + + test("wraps direct textGenerationStream exports", async () => { + const huggingFace = wrapHuggingFace(buildModernModule()); + + const stream = huggingFace.textGenerationStream!( + { + model: "mistralai/Mixtral-8x7B-v0.1", + inputs: "Reply with exactly PARIS.", + }, + {}, + ); + const chunks = await collectAsync(stream); + + expect(chunks).toHaveLength(2); + + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + span_attributes: { + name: "huggingface.text_generation_stream", + type: "llm", + }, + metadata: expect.objectContaining({ + model: "mistralai/Mixtral-8x7B-v0.1", + provider: "huggingface", + finish_reason: "stop_sequence", + }), + metrics: expect.objectContaining({ + prompt_tokens: 1, + completion_tokens: 2, + tokens: 3, + time_to_first_token: expect.any(Number), + }), + output: expect.objectContaining({ + generated_text: "Paris", + }), + }); + }); + + test("supports legacy HfInference exports and feature extraction", async () => { + const huggingFace = wrapHuggingFace(buildLegacyModule()); + const client = new huggingFace.HfInference!("hf_test"); + + await client.endpoint("https://example.invalid").chatCompletion({ + model: "Qwen/Qwen3-32B", + messages: [{ role: "user", content: "Reply with exactly PARIS." }], + }); + await huggingFace.featureExtraction!( + { + model: "sentence-transformers/distilbert-base-nli-mean-tokens", + inputs: "Paris France", + dimensions: 3, + }, + {}, + ); + + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(2); + expect(spans[0]).toMatchObject({ + span_attributes: { + name: "huggingface.chat_completion", + }, + }); + expect(spans[1]).toMatchObject({ + span_attributes: { + name: "huggingface.feature_extraction", + type: "llm", + }, + metadata: expect.objectContaining({ + dimensions: 3, + model: "sentence-transformers/distilbert-base-nli-mean-tokens", + provider: "huggingface", + }), + output: expect.objectContaining({ + embedding_count: 1, + embedding_length: 3, + }), + }); + }); + + test("wraps completion-style textGenerationStream chunks", async () => { + const huggingFace = wrapHuggingFace( + buildCompletionStyleTextGenerationModule(), + ); + + const stream = huggingFace.textGenerationStream!( + { + model: "meta-llama/Llama-3.1-8B", + inputs: "The capital of France is", + }, + {}, + ); + const chunks = await collectAsync(stream); + + expect(chunks).toHaveLength(3); + + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + span_attributes: { + name: "huggingface.text_generation_stream", + type: "llm", + }, + metadata: expect.objectContaining({ + model: "meta-llama/Llama-3.1-8B", + provider: "huggingface", + finish_reason: "stop", + }), + metrics: expect.objectContaining({ + prompt_tokens: 5, + completion_tokens: 2, + tokens: 7, + time_to_first_token: expect.any(Number), + }), + output: expect.objectContaining({ + generated_text: "Paris", + finish_reason: "stop", + }), + }); + }); +}); diff --git a/js/src/wrappers/huggingface.ts b/js/src/wrappers/huggingface.ts new file mode 100644 index 000000000..b357099b8 --- /dev/null +++ b/js/src/wrappers/huggingface.ts @@ -0,0 +1,315 @@ +import { huggingFaceChannels } from "../instrumentation/plugins/huggingface-channels"; +import type { + HuggingFaceChatCompletion, + HuggingFaceChatCompletionChunk, + HuggingFaceChatCompletionParams, + HuggingFaceClient, + HuggingFaceClientConstructor, + HuggingFaceFeatureExtractionOutput, + HuggingFaceFeatureExtractionParams, + HuggingFaceModule, + HuggingFaceRequestOptions, + HuggingFaceTextGenerationOutput, + HuggingFaceTextGenerationParams, + HuggingFaceTextGenerationStreamOutput, +} from "../vendor-sdk-types/huggingface"; + +const HUGGINGFACE_CONSTRUCTOR_KEYS = [ + "InferenceClient", + "InferenceClientEndpoint", + "HfInference", + "HfInferenceEndpoint", +] as const; +const HUGGINGFACE_CONSTRUCTOR_KEY_SET: ReadonlySet = new Set( + HUGGINGFACE_CONSTRUCTOR_KEYS, +); + +/** + * Wrap a HuggingFace Inference SDK module or client with Braintrust tracing. + * + * Supports the LLM and embeddings APIs we intentionally instrument: + * - chatCompletion + * - chatCompletionStream + * - textGeneration + * - textGenerationStream + * - featureExtraction + */ +export function wrapHuggingFace( + huggingFace: HuggingFaceModule, +): HuggingFaceModule; +export function wrapHuggingFace( + huggingFace: HuggingFaceClient, +): HuggingFaceClient; +export function wrapHuggingFace(huggingFace: T): T; +export function wrapHuggingFace(huggingFace: unknown): unknown { + if (isSupportedHuggingFaceModule(huggingFace)) { + return moduleProxy(huggingFace); + } + + if (isSupportedHuggingFaceClient(huggingFace)) { + return clientProxy(huggingFace); + } + + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.warn("Unsupported HuggingFace Inference SDK. Not wrapping."); + return huggingFace; +} + +function isHuggingFaceConstructorKey( + value: string, +): value is (typeof HUGGINGFACE_CONSTRUCTOR_KEYS)[number] { + return HUGGINGFACE_CONSTRUCTOR_KEY_SET.has(value); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function hasFunction(value: unknown, methodName: string): boolean { + return ( + isRecord(value) && + methodName in value && + typeof value[methodName] === "function" + ); +} + +function isSupportedHuggingFaceModule( + value: unknown, +): value is HuggingFaceModule { + if (!isRecord(value)) { + return false; + } + + return ( + HUGGINGFACE_CONSTRUCTOR_KEYS.some( + (key) => key in value && typeof value[key] === "function", + ) || isSupportedHuggingFaceClient(value) + ); +} + +function isSupportedHuggingFaceClient( + value: unknown, +): value is HuggingFaceClient { + return ( + hasFunction(value, "chatCompletion") && + hasFunction(value, "chatCompletionStream") && + hasFunction(value, "textGeneration") && + hasFunction(value, "textGenerationStream") && + hasFunction(value, "featureExtraction") + ); +} + +function moduleProxy(module: HuggingFaceModule): HuggingFaceModule { + const shadowTarget = Object.create(module); + return new Proxy(shadowTarget, { + get(target, prop, receiver) { + if (typeof prop === "string" && isHuggingFaceConstructorKey(prop)) { + const value = Reflect.get(module, prop, receiver); + return typeof value === "function" + ? wrapClientConstructor(value) + : value; + } + + switch (prop) { + case "chatCompletion": + return target.chatCompletion + ? wrapChatCompletion(target.chatCompletion.bind(target)) + : target.chatCompletion; + case "chatCompletionStream": + return target.chatCompletionStream + ? wrapChatCompletionStream(target.chatCompletionStream.bind(target)) + : target.chatCompletionStream; + case "textGeneration": + return target.textGeneration + ? wrapTextGeneration(target.textGeneration.bind(target)) + : target.textGeneration; + case "textGenerationStream": + return target.textGenerationStream + ? wrapTextGenerationStream(target.textGenerationStream.bind(target)) + : target.textGenerationStream; + case "featureExtraction": + return target.featureExtraction + ? wrapFeatureExtraction(target.featureExtraction.bind(target)) + : target.featureExtraction; + default: + return Reflect.get(module, prop, receiver); + } + }, + }); +} + +function wrapClientConstructor( + constructor: HuggingFaceClientConstructor, +): HuggingFaceClientConstructor { + return new Proxy(constructor, { + construct(target, args) { + const instance: HuggingFaceClient = Reflect.construct(target, args); + return clientProxy(instance); + }, + }); +} + +function clientProxy(client: HuggingFaceClient): HuggingFaceClient { + return clientProxyWithContext(client); +} + +function clientProxyWithContext( + client: HuggingFaceClient, + endpointUrl?: string, +): HuggingFaceClient { + const shadowTarget = Object.create(client); + return new Proxy(shadowTarget, { + get(_target, prop, receiver) { + switch (prop) { + case "chatCompletion": + return wrapChatCompletion( + client.chatCompletion.bind(client), + endpointUrl, + ); + case "chatCompletionStream": + return wrapChatCompletionStream( + client.chatCompletionStream.bind(client), + endpointUrl, + ); + case "textGeneration": + return wrapTextGeneration( + client.textGeneration.bind(client), + endpointUrl, + ); + case "textGenerationStream": + return wrapTextGenerationStream( + client.textGenerationStream.bind(client), + endpointUrl, + ); + case "featureExtraction": + return wrapFeatureExtraction( + client.featureExtraction.bind(client), + endpointUrl, + ); + case "endpoint": + if (!client.endpoint) { + return client.endpoint; + } + { + const endpoint = client.endpoint.bind(client); + return (nextEndpointUrl: string) => + clientProxyWithContext( + endpoint(nextEndpointUrl), + nextEndpointUrl, + ); + } + default: + return Reflect.get(client, prop, receiver); + } + }, + }); +} + +function withEndpointUrl>( + params: T, + endpointUrl?: string, +): T { + if (!endpointUrl || params.endpointUrl !== undefined) { + return params; + } + + return { + ...params, + endpointUrl, + }; +} + +function wrapChatCompletion( + original: ( + params: HuggingFaceChatCompletionParams, + options?: HuggingFaceRequestOptions, + ) => Promise, + endpointUrl?: string, +): HuggingFaceClient["chatCompletion"] { + return (params, options) => { + const traceParams = withEndpointUrl(params, endpointUrl); + const context: Parameters< + typeof huggingFaceChannels.chatCompletion.tracePromise + >[1] = { + arguments: [traceParams], + }; + return huggingFaceChannels.chatCompletion.tracePromise( + () => original(params, options), + context, + ); + }; +} + +function wrapChatCompletionStream( + original: ( + params: HuggingFaceChatCompletionParams, + options?: HuggingFaceRequestOptions, + ) => AsyncIterable, + endpointUrl?: string, +): HuggingFaceClient["chatCompletionStream"] { + return (params, options) => + huggingFaceChannels.chatCompletionStream.traceSync( + () => original(params, options), + { + arguments: [withEndpointUrl(params, endpointUrl)], + }, + ); +} + +function wrapTextGeneration( + original: ( + params: HuggingFaceTextGenerationParams, + options?: HuggingFaceRequestOptions, + ) => Promise, + endpointUrl?: string, +): HuggingFaceClient["textGeneration"] { + return (params, options) => { + const traceParams = withEndpointUrl(params, endpointUrl); + const context: Parameters< + typeof huggingFaceChannels.textGeneration.tracePromise + >[1] = { + arguments: [traceParams], + }; + return huggingFaceChannels.textGeneration.tracePromise( + () => original(params, options), + context, + ); + }; +} + +function wrapTextGenerationStream( + original: ( + params: HuggingFaceTextGenerationParams, + options?: HuggingFaceRequestOptions, + ) => AsyncIterable, + endpointUrl?: string, +): HuggingFaceClient["textGenerationStream"] { + return (params, options) => + huggingFaceChannels.textGenerationStream.traceSync( + () => original(params, options), + { + arguments: [withEndpointUrl(params, endpointUrl)], + }, + ); +} + +function wrapFeatureExtraction( + original: ( + params: HuggingFaceFeatureExtractionParams, + options?: HuggingFaceRequestOptions, + ) => Promise, + endpointUrl?: string, +): HuggingFaceClient["featureExtraction"] { + return (params, options) => { + const traceParams = withEndpointUrl(params, endpointUrl); + const context: Parameters< + typeof huggingFaceChannels.featureExtraction.tracePromise + >[1] = { + arguments: [traceParams], + }; + return huggingFaceChannels.featureExtraction.tracePromise( + () => original(params, options), + context, + ); + }; +} diff --git a/turbo.json b/turbo.json index 42046bc5d..446b3e171 100644 --- a/turbo.json +++ b/turbo.json @@ -6,7 +6,8 @@ "ANTHROPIC_API_KEY", "GEMINI_API_KEY", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "tasks": { "build": { @@ -26,7 +27,8 @@ "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "dependsOn": ["^build"], "outputs": [] @@ -41,7 +43,8 @@ "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "dependsOn": ["^build"], "outputs": [] @@ -61,7 +64,8 @@ "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "dependsOn": [], "outputs": [] @@ -76,7 +80,8 @@ "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "dependsOn": ["^build"], "outputs": [] @@ -105,7 +110,8 @@ "GEMINI_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "dependsOn": ["^build", "build"] }, @@ -117,7 +123,8 @@ "GEMINI_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "dependsOn": ["^build", "build"] }, @@ -129,7 +136,8 @@ "GEMINI_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY", - "MISTRAL_API_KEY" + "MISTRAL_API_KEY", + "HUGGINGFACE_API_KEY" ], "dependsOn": ["^build", "build"] } From 3c7f7b901a437bcb98dd8653a6e740b766c98b6c Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 13 Apr 2026 16:10:37 +0200 Subject: [PATCH 2/3] ts --- .../instrumentation/plugins/huggingface-plugin.ts | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/js/src/instrumentation/plugins/huggingface-plugin.ts b/js/src/instrumentation/plugins/huggingface-plugin.ts index 677c8ef03..b47992e8c 100644 --- a/js/src/instrumentation/plugins/huggingface-plugin.ts +++ b/js/src/instrumentation/plugins/huggingface-plugin.ts @@ -501,11 +501,15 @@ function extractTextGenerationStreamMetadata( continue; } - const finishReason = chunk.choices.findLast( - (choice) => choice.finish_reason !== undefined, - )?.finish_reason; - if (finishReason !== undefined) { - return { finish_reason: finishReason }; + for ( + let choiceIndex = chunk.choices.length - 1; + choiceIndex >= 0; + choiceIndex-- + ) { + const choice = chunk.choices[choiceIndex]; + if (choice?.finish_reason !== undefined) { + return { finish_reason: choice.finish_reason }; + } } } From 9dbe505d10cf097294a9d26565bb0690b90a1c0c Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 13 Apr 2026 17:21:23 +0200 Subject: [PATCH 3/3] Review comments --- .agents/skills/instrumentation/SKILL.md | 1 + js/src/instrumentation/plugins/huggingface-plugin.ts | 12 ++++-------- js/src/wrappers/huggingface.ts | 9 +++------ 3 files changed, 8 insertions(+), 14 deletions(-) diff --git a/.agents/skills/instrumentation/SKILL.md b/.agents/skills/instrumentation/SKILL.md index c59a7859b..4d094d761 100644 --- a/.agents/skills/instrumentation/SKILL.md +++ b/.agents/skills/instrumentation/SKILL.md @@ -28,6 +28,7 @@ Map the change before editing: - Support both auto-instrumentation and manual instrumentation. Auto-instrumentation does not cover every environment, loader, or framework. - For orchestrion auto-instrumentation, prefer targeting public API functions. Instrumenting internal helpers is more likely to break across library versions. - Auto and manual paths should share logic. Prefer both paths emitting the same tracing-channel events, with provider plugins converting those events into spans/logs/errors. Manual wrappers should not directly emit observability data. +- Reuse shared repo utilities before introducing local helpers. Check `js/util/index.ts`, neighboring instrumentation files, and existing plugins/wrappers for utilities like `isObject`, merge helpers, and sanitizers before adding ad hoc replacements. - If a public instrumentation surface changes, check whether the export surface also needs updates in `js/src/instrumentation/index.ts` or `js/src/exports.ts`. - Preserve async context propagation. Changes around tracing channels, stream patching, or loader hooks must keep the current span context across awaits and stream consumption. - Maintain isomorphic behavior. Node and browser/bundled paths must use compatible channel implementations and avoid channel-registry mismatches. diff --git a/js/src/instrumentation/plugins/huggingface-plugin.ts b/js/src/instrumentation/plugins/huggingface-plugin.ts index b47992e8c..20e06823e 100644 --- a/js/src/instrumentation/plugins/huggingface-plugin.ts +++ b/js/src/instrumentation/plugins/huggingface-plugin.ts @@ -312,6 +312,7 @@ function patchChatCompletionStream(args: { }, onComplete: (chunks) => { const lastChunk = chunks.at(-1); + const responseMetadata = extractResponseMetadata(lastChunk); const metrics = { ...parseMetricsFromUsage(lastChunk?.usage), ...(firstChunkTime !== undefined @@ -321,9 +322,7 @@ function patchChatCompletionStream(args: { span.log({ output: aggregateChatCompletionChunks(chunks), - ...(extractResponseMetadata(lastChunk) - ? { metadata: extractResponseMetadata(lastChunk) } - : {}), + ...(responseMetadata ? { metadata: responseMetadata } : {}), metrics, }); span.end(); @@ -358,13 +357,10 @@ function patchTextGenerationStream(args: { }, onComplete: (chunks) => { const lastChunk = chunks.at(-1); + const streamMetadata = extractTextGenerationStreamMetadata(chunks); span.log({ output: aggregateTextGenerationStreamChunks(chunks), - ...(extractTextGenerationStreamMetadata(chunks) - ? { - metadata: extractTextGenerationStreamMetadata(chunks), - } - : {}), + ...(streamMetadata ? { metadata: streamMetadata } : {}), metrics: { ...extractTextGenerationMetrics(lastChunk?.details ?? null), ...parseMetricsFromUsage(lastChunk?.usage), diff --git a/js/src/wrappers/huggingface.ts b/js/src/wrappers/huggingface.ts index b357099b8..6ff98a259 100644 --- a/js/src/wrappers/huggingface.ts +++ b/js/src/wrappers/huggingface.ts @@ -1,4 +1,5 @@ import { huggingFaceChannels } from "../instrumentation/plugins/huggingface-channels"; +import { isObject } from "../../util"; import type { HuggingFaceChatCompletion, HuggingFaceChatCompletionChunk, @@ -61,13 +62,9 @@ function isHuggingFaceConstructorKey( return HUGGINGFACE_CONSTRUCTOR_KEY_SET.has(value); } -function isRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null; -} - function hasFunction(value: unknown, methodName: string): boolean { return ( - isRecord(value) && + isObject(value) && methodName in value && typeof value[methodName] === "function" ); @@ -76,7 +73,7 @@ function hasFunction(value: unknown, methodName: string): boolean { function isSupportedHuggingFaceModule( value: unknown, ): value is HuggingFaceModule { - if (!isRecord(value)) { + if (!isObject(value)) { return false; }