langwatch 0.1.0 → 0.1.3

package/src/langchain.ts

@@ -0,0 +1,557 @@
+import type { AgentAction, AgentFinish } from "@langchain/core/agents";
+import { BaseCallbackHandler } from "@langchain/core/callbacks/base";
+import { type DocumentInterface } from "@langchain/core/documents";
+import type { Serialized } from "@langchain/core/load/serializable";
+import {
+  AIMessage,
+  AIMessageChunk,
+  FunctionMessage,
+  FunctionMessageChunk,
+  HumanMessage,
+  HumanMessageChunk,
+  SystemMessage,
+  SystemMessageChunk,
+  ToolMessage,
+  ToolMessageChunk,
+  mapChatMessagesToStoredMessages,
+  type BaseMessage,
+  type StoredMessage,
+} from "@langchain/core/messages";
+import type { ChatGeneration, LLMResult } from "@langchain/core/outputs";
+import type { ChainValues } from "@langchain/core/utils/types";
+import { stringify } from "javascript-stringify";
+import {
+  type LangWatchRAGSpan,
+  type LangWatchSpan,
+  type LangWatchTrace,
+} from ".";
+import {
+  type RAGSpan,
+  type BaseSpan,
+  type ChatMessage,
+  type ChatRichContent,
+  type SpanInputOutput,
+} from "./types";
+
+export class LangWatchCallbackHandler extends BaseCallbackHandler {
+  name = "LangWatchCallbackHandler";
+  trace: LangWatchTrace;
+  spans: Record<string, LangWatchSpan> = {};
+
+  constructor({ trace }: { trace: LangWatchTrace }) {
+    super();
+    this.trace = trace;
+  }
+
+  async handleLLMStart(
+    llm: Serialized,
+    prompts: string[],
+    runId: string,
+    parentRunId?: string | undefined,
+    extraParams?: Record<string, unknown> | undefined,
+    _tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    this.spans[runId] = this.buildLLMSpan({
+      llm,
+      runId,
+      parentRunId,
+      input: {
+        type: "json",
+        value: prompts,
+      },
+      extraParams,
+      metadata,
+      name,
+    });
+  }
+
+  private buildLLMSpan({
+    llm,
+    runId,
+    parentRunId,
+    input,
+    extraParams,
+    metadata,
+    name,
+  }: {
+    llm: Serialized;
+    runId: string;
+    parentRunId?: string | undefined;
+    input: SpanInputOutput;
+    extraParams?: Record<string, unknown> | undefined;
+    metadata?: Record<string, unknown> | undefined;
+    name?: string | undefined;
+  }) {
+    try {
+      const parent = this.getParent(parentRunId);
+
+      const vendor = metadata?.ls_provider ?? llm.id.at(-2)?.toString();
+      const model =
+        metadata?.ls_model_name ?? (llm as any).kwargs?.model ?? "unknown";
+
+      const span = parent.startLLMSpan({
+        spanId: runId,
+        name: name ?? llm.id.at(-1)?.toString(),
+        input,
+        model: [vendor, model].filter((x) => x).join("/"),
+        params: {
+          temperature: (extraParams?.invocation_params as any)?.temperature,
+          ...((extraParams?.invocation_params as any)?.functions
+            ? { functions: (extraParams?.invocation_params as any)?.functions }
+            : {}),
+        },
+      });
+
+      return span;
+    } catch (e) {
+      this.trace.client.emit("error", e);
+      throw e;
+    }
+  }
+
+  async handleChatModelStart(
+    llm: Serialized,
+    messages: BaseMessage[][],
+    runId: string,
+    parentRunId?: string | undefined,
+    extraParams?: Record<string, unknown> | undefined,
+    tags?: string[] | undefined,
+    metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    this.spans[runId] = this.buildLLMSpan({
+      name,
+      llm,
+      runId,
+      parentRunId,
+      input: {
+        type: "chat_messages",
+        value: messages.flatMap(convertFromLangChainMessages),
+      },
+      extraParams,
+      metadata,
+    });
+  }
+
+  async handleNewToken(_token: string, runId: string): Promise<void> {
+    const span = this.spans[runId];
+    if (runId && span && !span.timestamps.firstTokenAt) {
+      span.update({
+        timestamps: { ...span.timestamps, firstTokenAt: Date.now() },
+      });
+    }
+  }
+
+  async handleLLMEnd(
+    response: LLMResult,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    try {
+      const span = this.spans[runId];
+      if (!span) {
+        return;
+      }
+
+      const outputs: SpanInputOutput[] = [];
+      for (const generation of response.generations) {
+        // TODO: again, why the twice loop? Can OpenAI generate multiple chat outputs?
+        for (const generation_ of generation) {
+          if ("message" in generation_) {
+            outputs.push({
+              type: "chat_messages",
+              value: convertFromLangChainMessages([
+                (generation_ as ChatGeneration).message,
+              ]),
+            });
+          } else if ("text" in generation_) {
+            outputs.push({
+              type: "text",
+              value: generation_.text,
+            });
+          } else {
+            outputs.push({
+              type: "text",
+              value: JSON.stringify(generation_),
+            });
+          }
+        }
+      }
+
+      const output: SpanInputOutput | undefined =
+        outputs.length === 1
+          ? outputs[0]
+          : { type: "list", value: outputs as any };
+
+      // Commenting it out because LangChain.js prompt and completion tokens is broken, this one doesn't work as it should with python,
+      // and response_metadata.prompt and response_metadata.completion is there but it's always 0. Better let our server count.
+      // const metrics = response.llmOutput?.token_usage
+      //   ? {
+      //       promptTokens: response.llmOutput.token_usage.prompt_tokens,
+      //       completionTokens: response.llmOutput.token_usage.completion_tokens,
+      //     }
+      //   : undefined;
+
+      span.end({
+        output,
+        // ...(metrics ? { metrics } : {}),
+      });
+    } catch (e) {
+      this.trace.client.emit("error", e);
+      throw e;
+    }
+  }
+
+  async handleLLMError(
+    err: Error,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    this.errorSpan({ runId, error: err });
+  }
+
+  async handleChainStart(
+    chain: Serialized,
+    inputs: ChainValues,
+    runId: string,
+    parentRunId?: string | undefined,
+    _tags?: string[] | undefined,
+    _metadata?: Record<string, unknown> | undefined,
+    _runType?: string,
+    name?: string
+  ): Promise<void> {
+    this.spans[runId] = this.buildSpan({
+      type: "chain",
+      serialized: chain,
+      runId,
+      parentRunId,
+      input: inputs,
+      name,
+    });
+  }
+
+  async handleChainEnd(
+    output: ChainValues,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    this.endSpan({
+      runId,
+      output,
+    });
+  }
+
+  async handleChainError(
+    err: Error,
+    runId: string,
+    _parentRunId?: string | undefined,
+    _tags?: string[] | undefined,
+    _kwargs?: { inputs?: Record<string, unknown> | undefined } | undefined
+  ): Promise<void> {
+    this.errorSpan({ runId, error: err });
+  }
+
+  async handleToolStart(
+    tool: Serialized,
+    input: string,
+    runId: string,
+    parentRunId?: string | undefined,
+    _tags?: string[] | undefined,
+    _metadata?: Record<string, unknown> | undefined,
+    name?: string
+  ): Promise<void> {
+    this.spans[runId] = this.buildSpan({
+      type: "tool",
+      serialized: tool,
+      runId,
+      parentRunId,
+      input,
+      name,
+    });
+  }
+
+  async handleToolEnd(
+    output: string,
+    runId: string,
+    _parentRunId?: string | undefined
+  ): Promise<void> {
+    this.endSpan({ runId, output });
+  }
+
+  async handleToolError(
+    err: Error,
+    runId: string,
+    _parentRunId?: string | undefined,
+    _tags?: string[] | undefined
+  ): Promise<void> {
+    this.errorSpan({ runId, error: err });
+  }
+
+  async handleRetrieverStart(
+    retriever: Serialized,
+    query: string,
+    runId: string,
+    parentRunId?: string | undefined,
+    _tags?: string[] | undefined,
+    _metadata?: Record<string, unknown> | undefined,
+    name?: string | undefined
+  ) {
+    try {
+      const parent = this.getParent(parentRunId);
+
+      this.spans[runId] = parent.startRAGSpan({
+        spanId: runId,
+        name: name ?? retriever.name ?? retriever.id.at(-1)?.toString(),
+        input: this.autoconvertTypedValues(query),
+      });
+    } catch (e) {
+      this.trace.client.emit("error", e);
+      throw e;
+    }
+  }
+
+  async handleRetrieverEnd(
+    documents: DocumentInterface<Record<string, any>>[],
+    runId: string,
+    _parentRunId?: string | undefined,
+    _tags?: string[] | undefined
+  ) {
+    try {
+      const contexts: RAGSpan["contexts"] = documents.map((doc) => ({
+        content: doc.pageContent,
+        ...(doc.metadata.source ? { documentId: doc.metadata.source } : {}),
+      }));
+
+      const span = this.spans[runId] as LangWatchRAGSpan;
+      if (!span) {
+        return;
+      }
+
+      span.end({
+        contexts,
+        output: this.autoconvertTypedValues(documents),
+      });
+    } catch (e) {
+      this.trace.client.emit("error", e);
+      throw e;
+    }
+  }
+
+  async handleRetrieverError(
+    err: Error,
+    runId: string,
+    _parentRunId?: string | undefined,
+    _tags?: string[] | undefined
+  ) {
+    this.errorSpan({ runId, error: err });
+  }
+
+  async handleAgentAction(
+    _action: AgentAction,
+    runId: string,
+    _parentRunId?: string | undefined,
+    _tags?: string[] | undefined
+  ): Promise<void> {
+    const span = this.spans[runId];
+    if (!span) {
+      return;
+    }
+
+    span.update({
+      type: "agent",
+    });
+  }
+
+  async handleAgentEnd(
+    action: AgentFinish,
+    runId: string,
+    _parentRunId?: string | undefined,
+    _tags?: string[] | undefined
+  ): Promise<void> {
+    this.endSpan({
+      runId,
+      output: action.returnValues,
+    });
+  }
+
+  private buildSpan({
+    type,
+    serialized,
+    runId,
+    parentRunId,
+    input,
+    name,
+  }: {
+    type: BaseSpan["type"];
+    serialized: Serialized;
+    runId: string;
+    parentRunId?: string | undefined;
+    input: unknown;
+    name?: string | undefined;
+  }) {
+    try {
+      const parent = this.getParent(parentRunId);
+
+      const span = parent.startSpan({
+        spanId: runId,
+        type,
+        name: name ?? serialized.name ?? serialized.id.at(-1)?.toString(),
+        input: this.autoconvertTypedValues(input),
+      });
+
+      return span;
+    } catch (e) {
+      this.trace.client.emit("error", e);
+      throw e;
+    }
+  }
+
+  private endSpan({ runId, output }: { runId: string; output: unknown }): void {
+    try {
+      const span = this.spans[runId];
+      if (!span) {
+        return;
+      }
+
+      span.end({
+        output: this.autoconvertTypedValues(output),
+      });
+    } catch (e) {
+      this.trace.client.emit("error", e);
+      throw e;
+    }
+  }
+
+  private errorSpan({ runId, error }: { runId: string; error: Error }): void {
+    const span = this.spans[runId];
+    if (!span) {
+      return;
+    }
+
+    span.end({
+      error,
+    });
+  }
+
+  private autoconvertTypedValues(value: any): SpanInputOutput | undefined {
+    if (
+      !value ||
+      (typeof value === "object" && Object.keys(value).length === 0)
+    ) {
+      return undefined;
+    }
+    if (typeof value === "string") {
+      return { type: "text", value };
+    }
+    try {
+      JSON.stringify(value);
+      return { type: "json", value };
+    } catch (e) {
+      return { type: "text", value: stringify(value) ?? value.toString() };
+    }
+  }
+
+  private getParent(
+    parentRunId?: string | undefined
+  ): LangWatchSpan | LangWatchTrace {
+    return (
+      (parentRunId
+        ? this.spans[parentRunId]
+        : this.spans[Object.keys(this.spans).at(-1) ?? ""]) ?? this.trace
+    );
+  }
+}
+
+export const convertFromLangChainMessages = (
+  messages: BaseMessage[]
+): ChatMessage[] => {
+  const chatMessages: ChatMessage[] = [];
+  for (const message of messages) {
+    chatMessages.push(convertFromLangChainMessage(message));
+  }
+  return chatMessages;
+};
+
+const convertFromLangChainMessage = (
+  message: BaseMessage & { id?: string[] }
+): ChatMessage => {
+  let role: ChatMessage["role"] = "user";
+
+  const message_: (BaseMessage | StoredMessage) & {
+    id?: string[];
+    type?: string;
+  } = message.lc_serializable
+    ? mapChatMessagesToStoredMessages([message])[0]!
+    : message;
+
+  // Dang this is so hard, langchain.js has 3 ways of representing the same thing...
+  if (
+    message_ instanceof HumanMessage ||
+    message_ instanceof HumanMessageChunk ||
+    message_.id?.at(-1) === "HumanMessage" ||
+    message_.id?.at(-1) === "HumanMessageChunk" ||
+    message_.type === "human"
+  ) {
+    role = "user";
+  } else if (
+    message instanceof AIMessage ||
+    message instanceof AIMessageChunk ||
+    message.id?.at(-1) === "AIMessage" ||
+    message.id?.at(-1) === "AIMessageChunk" ||
+    message_.type === "ai"
+  ) {
+    role = "assistant";
+  } else if (
+    message instanceof SystemMessage ||
+    message instanceof SystemMessageChunk ||
+    message.id?.at(-1) === "SystemMessage" ||
+    message.id?.at(-1) === "SystemMessageChunk" ||
+    message_.type === "system"
+  ) {
+    role = "system";
+  } else if (
+    message instanceof FunctionMessage ||
+    message instanceof FunctionMessageChunk ||
+    message.id?.at(-1) === "FunctionMessage" ||
+    message.id?.at(-1) === "FunctionMessageChunk" ||
+    message_.type === "function"
+  ) {
+    role = "function";
+  } else if (
+    message instanceof ToolMessage ||
+    message instanceof ToolMessageChunk ||
+    message.id?.at(-1) === "ToolMessage" ||
+    message.id?.at(-1) === "ToolMessageChunk" ||
+    message_.type === "tool"
+  ) {
+    role = "tool";
+  }
+
+  const content =
+    typeof message.content === "string"
+      ? message.content
+      : message.content.map(
+          (content): ChatRichContent =>
+            content.type === "text"
+              ? { type: "text", text: content.text }
+              : content.type == "image_url"
+              ? { type: "image_url", image_url: content.image_url }
+              : { type: "text", text: JSON.stringify(content) }
+        );
+
+  const functionCall = message.additional_kwargs as
+    | ChatMessage["function_call"]
+    | undefined;
+
+  return {
+    role,
+    content,
+    ...(functionCall &&
+    typeof functionCall === "object" &&
+    Object.keys(functionCall).length > 0
+      ? { function_call: functionCall }
+      : {}),
+  };
+};

package/src/{helpers.ts → typeUtils.ts}

@@ -1,3 +1,8 @@
+import {
+  reservedSpanParamsSchema,
+  reservedTraceMetadataSchema
+} from "./server/types/tracer.generated";
+
 export type Strict<T> = T & { [K in Exclude<keyof any, keyof T>]: never };
 
 type SnakeToCamelCase<S extends string> = S extends `${infer T}_${infer U}`
@@ -44,7 +49,10 @@ function camelToSnakeCase(str: string): string {
   return str.replace(/[A-Z]/g, (letter) => `_${letter.toLowerCase()}`);
 }
 
-export function camelToSnakeCaseNested<T>(obj: T): CamelToSnakeCaseNested<T> {
+export function camelToSnakeCaseNested<T>(
+  obj: T,
+  parentKey?: string
+): CamelToSnakeCaseNested<T> {
   if (Array.isArray(obj)) {
     return obj.map((item) =>
       camelToSnakeCaseNested(item)
@@ -54,7 +62,17 @@ export function camelToSnakeCaseNested<T>(obj: T): CamelToSnakeCaseNested<T> {
     for (const key in obj) {
       if (obj.hasOwnProperty(key)) {
         const newKey = camelToSnakeCase(key);
-        newObj[newKey] = camelToSnakeCaseNested((obj as any)[key]);
+        // Keep arbitrary keys the same
+        if (
+          (parentKey === "metadata" &&
+            !Object.keys(reservedTraceMetadataSchema.shape).includes(newKey)) ||
+          (parentKey === "params" &&
+            !Object.keys(reservedSpanParamsSchema.shape).includes(newKey))
+        ) {
+          newObj[key] = (obj as any)[key];
+        } else {
+          newObj[newKey] = camelToSnakeCaseNested((obj as any)[key], newKey);
+        }
       }
     }
     return newObj as CamelToSnakeCaseNested<T>;

package/src/types.ts

@@ -1,6 +1,5 @@
-import type modelPrices from "llm-cost/model_prices_and_context_window.json";
 import type { OpenAI } from "openai";
-import { type SnakeToCamelCaseNested } from "./helpers";
+import { type SnakeToCamelCaseNested } from "./typeUtils";
 import {
   type BaseSpan as ServerBaseSpan,
   type ChatMessage as ServerChatMessage,
@@ -10,6 +9,7 @@ import {
   type SpanInputOutput as ServerSpanInputOutput,
   type TypedValueChatMessages,
   type Trace,
+  type RESTEvaluation as ServerRESTEvaluation,
 } from "./server/types/tracer";
 
 export type Metadata = SnakeToCamelCaseNested<Trace["metadata"]>;
@@ -63,8 +63,12 @@ export type PendingBaseSpan = PendingSpan<BaseSpan>;
 // vendor is deprecated, and we try to force the available models here
 export type LLMSpan = ConvertServerSpan<
   Omit<ServerLLMSpan, "vendor" | "model">
-> & { model: keyof typeof modelPrices | (string & NonNullable<unknown>) };
+> & { model: string };
 export type PendingLLMSpan = PendingSpan<LLMSpan>;
 
 export type RAGSpan = ConvertServerSpan<ServerRAGSpan>;
 export type PendingRAGSpan = PendingSpan<RAGSpan>;
+
+export type RESTEvaluation = SnakeToCamelCaseNested<
+  Omit<ServerRESTEvaluation, "error">
+> & { error?: ServerRESTEvaluation["error"] };

package/src/utils.ts

@@ -1,7 +1,9 @@
 import { convertUint8ArrayToBase64 } from "@ai-sdk/provider-utils";
-import { type ImagePart, type CoreMessage } from "ai";
-import { type ChatMessage } from "./types";
+import { type CoreMessage, type ImagePart } from "ai";
+import { z } from "zod";
 import { type ErrorCapture } from "./server/types/tracer";
+import { chatMessageSchema } from "./server/types/tracer.generated";
+import { type ChatMessage, type SpanInputOutput } from "./types";
 
 const convertImageToUrl = (
   image: ImagePart["image"],
@@ -177,3 +179,24 @@ export const captureError = (error: unknown): ErrorCapture => {
     };
   }
 };
+
+export const autoconvertTypedValues = (value: unknown): SpanInputOutput => {
+  if (typeof value === "string") {
+    return { type: "text", value };
+  }
+
+  const chatMessages = z.array(chatMessageSchema).safeParse(value);
+  if (Array.isArray(value) && chatMessages.success) {
+    return {
+      type: "chat_messages",
+      value: chatMessages.data,
+    };
+  }
+
+  try {
+    JSON.stringify(value);
+    return { type: "json", value: value as object };
+  } catch (e) {
+    return { type: "raw", value: value as any };
+  }
+};

package/ts-to-zod.config.js

@@ -14,5 +14,7 @@ module.exports = {
       "TraceCheckJob",
       "AnalyticsMetric",
       "NewDatasetEntries",
+      "EvaluationRESTParams",
+      "EvaluationRESTResult",
     ].includes(name),
 };