Merge branch 'tensorzero:main' into main

Author: pycoder49

.github/workflows/merge-queue.yml

@@ -275,6 +275,7 @@ jobs:
       OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
       FIREWORKS_ACCOUNT_ID: ${{ secrets.FIREWORKS_ACCOUNT_ID }}
       FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
+      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
 
   # See 'ci/README.md' at the repository root for more details.
   check-all-live-tests-passed:

.github/workflows/slash-command-regen-fixtures.yml

@@ -9,6 +9,7 @@ env:
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
   AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
   AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+  ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
 
 permissions:
   pull-requests: write # For doing the emoji reaction on a PR comment
@@ -42,6 +43,7 @@ jobs:
           echo "FIREWORKS_API_KEY=not_used" >> ui/fixtures/.env
           echo "FIREWORKS_BASE_URL=http://mock-inference-provider:3030/fireworks/" >> ui/fixtures/.env
           echo "OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }}" >> ui/fixtures/.env
+          echo "ANTHROPIC_API_KEY=${{ secrets.ANTHROPIC_API_KEY }}" >> ui/fixtures/.env
           echo "OPENAI_BASE_URL=http://mock-inference-provider:3030/openai/" >> ui/fixtures/.env
           echo "FIREWORKS_BASE_URL=http://mock-inference-provider:3030/fireworks/" >> ui/fixtures/.env
           echo "S3_ACCESS_KEY_ID=${{ secrets.AWS_ACCESS_KEY_ID }}" >> ui/fixtures/.env

.github/workflows/ui-tests-e2e-model-inference-cache.yml

@@ -13,6 +13,8 @@ on:
         required: false
       FIREWORKS_API_KEY:
         required: false
+      ANTHROPIC_API_KEY:
+        required: false
     inputs:
       regen_cache:
         required: true
@@ -95,6 +97,7 @@ jobs:
           echo "FIREWORKS_ACCOUNT_ID=${{ secrets.FIREWORKS_ACCOUNT_ID }}" >> fixtures/.env-gateway
           echo "FIREWORKS_API_KEY=${{ secrets.FIREWORKS_API_KEY }}" >> fixtures/.env-gateway
           echo "OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }}" >> fixtures/.env-gateway
+          echo "ANTHROPIC_API_KEY=${{ secrets.ANTHROPIC_API_KEY }}" >> fixtures/.env-gateway
           echo "S3_ACCESS_KEY_ID=${{ secrets.S3_ACCESS_KEY_ID }}" >> fixtures/.env-gateway
           echo "S3_SECRET_ACCESS_KEY=${{ secrets.S3_SECRET_ACCESS_KEY }}" >> fixtures/.env-gateway
           ./fixtures/regenerate-model-inference-cache.sh
@@ -124,6 +127,7 @@ jobs:
           echo "FIREWORKS_ACCOUNT_ID=${{ secrets.FIREWORKS_ACCOUNT_ID || 'not_used' }}" >> fixtures/.env-gateway
           echo "FIREWORKS_API_KEY=${{ secrets.FIREWORKS_API_KEY || 'not_used' }}" >> fixtures/.env-gateway
           echo "OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY || 'not_used' }}" >> fixtures/.env-gateway
+          echo "ANTHROPIC_API_KEY=${{ secrets.ANTHROPIC_API_KEY || 'not_used' }}" >> fixtures/.env-gateway
           echo "S3_ACCESS_KEY_ID=${{ secrets.S3_ACCESS_KEY_ID }}" >> fixtures/.env-gateway
           echo "S3_SECRET_ACCESS_KEY=${{ secrets.S3_SECRET_ACCESS_KEY }}" >> fixtures/.env-gateway
           docker compose -f fixtures/docker-compose.e2e.yml up --no-build -d

.github/workflows/ui-tests-e2e.yml

@@ -113,6 +113,7 @@ jobs:
           # The 'ui-tests-e2e' job tests that the UI container starts without some of these variables set,
           echo "FIREWORKS_ACCOUNT_ID=fake_fireworks_account" >> fixtures/.env
           echo "FIREWORKS_API_KEY=not_used" >> fixtures/.env
+          echo "ANTHROPIC_API_KEY=not_used" >> fixtures/.env
           echo "FIREWORKS_BASE_URL=http://mock-inference-provider:3030/fireworks/" >> fixtures/.env
           echo "OPENAI_API_KEY=not_used" >> fixtures/.env
           echo "OPENAI_BASE_URL=http://mock-inference-provider:3030/openai/" >> fixtures/.env

tensorzero-core/src/endpoints/inference.rs

@@ -239,7 +239,10 @@ pub async fn inference(
         &clickhouse_connection_info,
     )
     .await?;
-    tracing::Span::current().record("episode_id", episode_id.to_string());
+    // Record the episode id if we didn't already have one
+    if params.episode_id.is_none() {
+        tracing::Span::current().record("episode_id", episode_id.to_string());
+    }
 
     let (function, function_name) = find_function(&params, &config)?;
     let mut candidate_variants: BTreeMap<String, Arc<VariantInfo>> =

tensorzero-core/src/observability.rs

@@ -93,9 +93,13 @@ pub fn build_opentelemetry_layer<T: SpanExporter + 'static>(
             // We only expose spans that explicitly contain field prefixed with "http." or "otel."
             // For example, `#[instrument(fields(otel.name = "my_otel_name"))]` will be exported
             let filter = filter::filter_fn(|metadata| {
-                metadata.fields().iter().any(|field| {
-                    field.name().starts_with("http.") || field.name().starts_with("otel.")
-                })
+                if metadata.is_event() {
+                    matches!(metadata.level(), &tracing::Level::ERROR)
+                } else {
+                    metadata.fields().iter().any(|field| {
+                        field.name().starts_with("http.") || field.name().starts_with("otel.")
+                    })
+                }
             });
 
             reload_handle

tensorzero-core/tests/e2e/otel.rs

@@ -3,7 +3,7 @@ use std::{
     sync::{Arc, Mutex},
 };
 
-use opentelemetry::{KeyValue, SpanId, Value};
+use opentelemetry::{trace::Status, KeyValue, SpanId, Value};
 use opentelemetry_sdk::{
     error::OTelSdkResult,
     trace::{SpanData, SpanExporter},
@@ -15,6 +15,7 @@ use tensorzero::{
 use tensorzero_core::inference::types::TextKind;
 use tensorzero_core::observability::build_opentelemetry_layer;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+use uuid::Uuid;
 
 use crate::providers::common::make_embedded_gateway_no_config;
 
@@ -50,6 +51,7 @@ impl CapturingOtelExporter {
     }
 }
 
+#[derive(Debug)]
 pub struct SpanMap {
     pub root_spans: Vec<SpanData>,
     pub span_children: HashMap<SpanId, Vec<SpanData>>,
@@ -135,6 +137,7 @@ pub async fn test_capture_simple_inference_spans() {
     // Since we're using the embedded gateway, the root span will be `function_inference`
     // (we won't have a top-level HTTP span)
     assert_eq!(root_span.name, "function_inference");
+    assert_eq!(root_span.status, Status::Unset);
     let root_attr_map = attrs_to_map(&root_span.attributes);
     assert_eq!(root_attr_map["model_name"], "dummy::good".into());
     assert_eq!(
@@ -154,6 +157,7 @@ pub async fn test_capture_simple_inference_spans() {
     };
 
     assert_eq!(variant_span.name, "variant_inference");
+    assert_eq!(variant_span.status, Status::Unset);
     let variant_attr_map = attrs_to_map(&variant_span.attributes);
     assert_eq!(
         variant_attr_map["function_name"],
@@ -168,6 +172,7 @@ pub async fn test_capture_simple_inference_spans() {
     };
 
     assert_eq!(model_span.name, "model_inference");
+    assert_eq!(model_span.status, Status::Unset);
     let model_attr_map = attrs_to_map(&model_span.attributes);
     assert_eq!(model_attr_map["model_name"], "dummy::good".into());
     assert_eq!(model_attr_map["stream"], false.into());
@@ -177,6 +182,7 @@ pub async fn test_capture_simple_inference_spans() {
         panic!("Expected one child span: {model_children:#?}");
     };
     assert_eq!(model_provider_span.name, "model_provider_inference");
+    assert_eq!(model_provider_span.status, Status::Unset);
     let model_provider_attr_map = attrs_to_map(&model_provider_span.attributes);
     assert_eq!(model_provider_attr_map["provider_name"], "dummy".into());
     assert_eq!(
@@ -199,3 +205,135 @@ pub async fn test_capture_simple_inference_spans() {
 
     assert_eq!(num_spans, 4);
 }
+
+#[tokio::test]
+pub async fn test_capture_model_error() {
+    let episode_uuid = Uuid::now_v7();
+    let exporter = install_capturing_otel_exporter();
+
+    let client = make_embedded_gateway_no_config().await;
+    let _err = client
+        .inference(ClientInferenceParams {
+            episode_id: Some(episode_uuid),
+            model_name: Some("openai::missing-model-name".to_string()),
+            input: ClientInput {
+                system: None,
+                messages: vec![ClientInputMessage {
+                    role: Role::User,
+                    content: vec![ClientInputMessageContent::Text(TextKind::Text {
+                        text: "What is your name?".to_string(),
+                    })],
+                }],
+            },
+            ..Default::default()
+        })
+        .await
+        .unwrap_err();
+
+    let all_spans = exporter.take_spans();
+    let num_spans = all_spans.len();
+    let spans = build_span_map(all_spans);
+
+    let [root_span] = spans.root_spans.as_slice() else {
+        panic!("Expected one root span: {:#?}", spans.root_spans);
+    };
+    // Since we're using the embedded gateway, the root span will be `function_inference`
+    // (we won't have a top-level HTTP span)
+    assert_eq!(root_span.name, "function_inference");
+    assert_eq!(
+        root_span.status,
+        Status::Error {
+            description: "".into()
+        }
+    );
+    let root_attr_map = attrs_to_map(&root_span.attributes);
+    assert_eq!(
+        root_attr_map["model_name"],
+        "openai::missing-model-name".into()
+    );
+    assert_eq!(root_attr_map["episode_id"], episode_uuid.to_string().into());
+    assert_eq!(root_attr_map.get("function_name"), None);
+    assert_eq!(root_attr_map.get("variant_name"), None);
+
+    let root_children = &spans.span_children[&root_span.span_context.span_id()];
+    let [variant_span] = root_children.as_slice() else {
+        panic!("Expected one child span: {root_children:#?}");
+    };
+
+    assert_eq!(variant_span.name, "variant_inference");
+    assert_eq!(variant_span.status, Status::Unset);
+    let variant_attr_map = attrs_to_map(&variant_span.attributes);
+    assert_eq!(
+        variant_attr_map["function_name"],
+        "tensorzero::default".into()
+    );
+    assert_eq!(
+        variant_attr_map["variant_name"],
+        "openai::missing-model-name".into()
+    );
+    assert_eq!(variant_attr_map["stream"], false.into());
+
+    let variant_children = &spans.span_children[&variant_span.span_context.span_id()];
+    let [model_span] = variant_children.as_slice() else {
+        panic!("Expected one child span: {variant_children:#?}");
+    };
+
+    assert_eq!(model_span.name, "model_inference");
+    assert_eq!(
+        model_span.status,
+        Status::Error {
+            description: "".into()
+        }
+    );
+    let model_attr_map = attrs_to_map(&model_span.attributes);
+    assert_eq!(
+        model_attr_map["model_name"],
+        "openai::missing-model-name".into()
+    );
+    assert_eq!(model_attr_map["stream"], false.into());
+
+    let model_children = &spans.span_children[&model_span.span_context.span_id()];
+    let [model_provider_span] = model_children.as_slice() else {
+        panic!("Expected one child span: {model_children:#?}");
+    };
+    assert_eq!(model_provider_span.name, "model_provider_inference");
+    assert_eq!(
+        model_provider_span.status,
+        Status::Error {
+            description: "".into()
+        }
+    );
+    assert_eq!(
+        model_provider_span.events.len(),
+        1,
+        "Unexpected number of events: {model_provider_span:#?}",
+    );
+    assert!(
+        model_provider_span.events[0]
+            .name
+            .starts_with("Error from openai server:"),
+        "Unexpected span event: {:?}",
+        model_provider_span.events[0]
+    );
+    let model_provider_attr_map = attrs_to_map(&model_provider_span.attributes);
+    assert_eq!(model_provider_attr_map["provider_name"], "openai".into());
+    assert_eq!(
+        model_provider_attr_map["gen_ai.operation.name"],
+        "chat".into()
+    );
+    assert_eq!(model_provider_attr_map["gen_ai.system"], "openai".into());
+    assert_eq!(
+        model_provider_attr_map["gen_ai.request.model"],
+        "missing-model-name".into()
+    );
+    assert_eq!(model_attr_map["stream"], false.into());
+
+    assert_eq!(
+        spans
+            .span_children
+            .get(&model_provider_span.span_context.span_id()),
+        None
+    );
+
+    assert_eq!(num_spans, 4);
+}

ui/app/entry.server.tsx

@@ -0,0 +1,71 @@
+import { PassThrough } from "node:stream";
+
+import type { AppLoadContext, EntryContext } from "react-router";
+import { createReadableStreamFromReadable } from "@react-router/node";
+import { ServerRouter } from "react-router";
+import { isbot } from "isbot";
+import type { RenderToPipeableStreamOptions } from "react-dom/server";
+import { renderToPipeableStream } from "react-dom/server";
+
+export const streamTimeout = 30_000;
+
+export default function handleRequest(
+  request: Request,
+  responseStatusCode: number,
+  responseHeaders: Headers,
+  routerContext: EntryContext,
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  loadContext: AppLoadContext,
+  // If you have middleware enabled:
+  // loadContext: unstable_RouterContextProvider
+) {
+  return new Promise((resolve, reject) => {
+    let shellRendered = false;
+    const userAgent = request.headers.get("user-agent");
+
+    // Ensure requests from bots and SPA Mode renders wait for all content to load before responding
+    // https://react.dev/reference/react-dom/server/renderToPipeableStream#waiting-for-all-content-to-load-for-crawlers-and-static-generation
+    const readyOption: keyof RenderToPipeableStreamOptions =
+      (userAgent && isbot(userAgent)) || routerContext.isSpaMode
+        ? "onAllReady"
+        : "onShellReady";
+
+    const { pipe, abort } = renderToPipeableStream(
+      <ServerRouter context={routerContext} url={request.url} />,
+      {
+        [readyOption]() {
+          shellRendered = true;
+          const body = new PassThrough();
+          const stream = createReadableStreamFromReadable(body);
+
+          responseHeaders.set("Content-Type", "text/html");
+
+          resolve(
+            new Response(stream, {
+              headers: responseHeaders,
+              status: responseStatusCode,
+            }),
+          );
+
+          pipe(body);
+        },
+        onShellError(error: unknown) {
+          reject(error);
+        },
+        onError(error: unknown) {
+          responseStatusCode = 500;
+          // Log streaming rendering errors from inside the shell.  Don't log
+          // errors encountered during initial shell rendering since they'll
+          // reject and get logged in handleDocumentRequest.
+          if (shellRendered) {
+            console.error(error);
+          }
+        },
+      },
+    );
+
+    // Abort the rendering stream after the `streamTimeout` so it has time to
+    // flush down the rejected boundaries
+    setTimeout(abort, streamTimeout + 1000);
+  });
+}