@gammatech/aijsx 0.10.1-dev.2024-06-07 → 0.11.0-dev.2024-06-17

package/dist/index.d.mts

@@ -6,8 +6,8 @@ export { OpenAI as OpenAIClient } from 'openai';
 import { ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam, ChatCompletionAssistantMessageParam, ChatCompletionCreateParams } from 'openai/resources';
 import AnthropicClient from '@anthropic-ai/sdk';
 export { default as AnthropicClient } from '@anthropic-ai/sdk';
-import { PromptFeedback, GenerateContentRequest, VertexAI } from '@google-cloud/vertexai';
-export { VertexAI } from '@google-cloud/vertexai';
+import { PromptFeedback, GenerateContentRequest, VertexAI, HarmCategory, HarmBlockThreshold } from '@google-cloud/vertexai';
+export { HarmBlockThreshold as GoogleHarmBlockThreshold, HarmCategory as GoogleHarmCategory, VertexAI } from '@google-cloud/vertexai';
 
 declare class ChatCompletionError extends Error {
     readonly chatCompletionRequest: LogChatCompletionRequest;
@@ -43,6 +43,10 @@ type ChatCompletionClientAndProvider<K> = {
         completion: number;
     }) => number;
 };
+type GetChatCompletionClientAndProvider<Model, Client> = (model: Model, args: {
+    retryCount: number;
+    lastError?: Error | null;
+}) => Promise<ChatCompletionClientAndProvider<Client>>;
 
 type CreateRenderContextOptions = {
     logger?: LogImplementation;
@@ -52,15 +56,25 @@ type CreateRenderContextOptions = {
 };
 declare function createRenderContext({ logger, traceId, processor, contextValues, }?: CreateRenderContextOptions): RenderContext;
 
-declare const RetryCountContext: Context<number>;
+type ExcludeNumber<E extends number, T = number> = T extends E ? never : T;
+type RetryCountContextValue = {
+    retryCount: 0;
+    lastError: null;
+} | {
+    retryCount: ExcludeNumber<0>;
+    lastError: Error;
+};
+declare const RetryCountContext: Context<RetryCountContextValue>;
+declare const RetryLastErrorContext: Context<Error | null>;
 declare const DefaultMaxRetriesContext: Context<number>;
 type RetryProps = {
     shouldRetry: (error: Error) => boolean;
     retries?: number;
+    lastError?: Error;
     maxRetries?: number;
     children: AINode;
 };
-declare function Retry({ shouldRetry, retries, maxRetries, children }: RetryProps, ctx: RenderContext): AsyncGenerator<string, void, unknown>;
+declare function Retry({ shouldRetry, retries, maxRetries, lastError, children }: RetryProps, ctx: RenderContext): AsyncGenerator<string, void, unknown>;
 
 type FallbackProps = {
     fallback: AINode;
@@ -241,7 +255,7 @@ declare module '@gammatech/aijsx' {
 }
 type ValidOpenAIVisionModel = 'gpt-4o' | 'gpt-4o-2024-05-13' | 'gpt-4-turbo-2024-04-09' | 'gpt-4-turbo' | 'gpt-4-vision-preview';
 type ValidOpenAIChatModel = ValidOpenAIVisionModel | 'gpt-4' | 'gpt-4-0314' | 'gpt-4-0613' | 'gpt-4-32k' | 'gpt-4-32k-0314' | 'gpt-4-32k-0613' | 'gpt-4-1106-preview' | 'gpt-4-0125-preview' | 'gpt-3.5-turbo' | 'gpt-3.5-turbo-0301' | 'gpt-3.5-turbo-0613' | 'gpt-3.5-turbo-16k' | 'gpt-3.5-turbo-16k-0613' | 'gpt-3.5-turbo-1106' | 'gpt-3.5-turbo-0125';
-declare const OpenAIClientContext: Context<() => ChatCompletionClientAndProvider<OpenAI>>;
+declare const OpenAIClientContext: Context<GetChatCompletionClientAndProvider<ValidOpenAIChatModel, OpenAI>>;
 type OpenAIChatCompletionProps = {
     model: ValidOpenAIChatModel;
     maxTokens?: number;
@@ -266,7 +280,7 @@ declare module '@gammatech/aijsx' {
  * @see https://docs.anthropic.com/claude/reference/selecting-a-model
  */
 type ValidAnthropicChatModel = 'claude-instant-1.2' | 'claude-2.1' | 'claude-3-opus-20240229' | 'claude-3-sonnet-20240229' | 'claude-3-haiku-20240307';
-declare const AnthropicClientContext: Context<() => ChatCompletionClientAndProvider<AnthropicClient>>;
+declare const AnthropicClientContext: Context<GetChatCompletionClientAndProvider<ValidAnthropicChatModel, AnthropicClient>>;
 type AnthropicChatCompletionProps = {
     model: ValidAnthropicChatModel;
     maxTokens?: number;
@@ -291,7 +305,7 @@ declare module '@gammatech/aijsx' {
     }
 }
 type ValidGoogleChatModel = 'gemini-1.5-pro' | 'gemini-1.5-flash';
-declare const GoogleClientContext: Context<() => ChatCompletionClientAndProvider<VertexAI>>;
+declare const GoogleClientContext: Context<GetChatCompletionClientAndProvider<ValidGoogleChatModel, VertexAI>>;
 type GoogleChatCompletionProps = {
     model: ValidGoogleChatModel;
     maxTokens?: number;
@@ -299,7 +313,11 @@ type GoogleChatCompletionProps = {
     stop?: string | string[];
     maxRetries?: number;
     children: AINode;
+    safetySettings?: {
+        category: HarmCategory;
+        threshold: HarmBlockThreshold;
+    }[];
 };
 declare function GoogleChatCompletion(props: GoogleChatCompletionProps, ctx: RenderContext): JSX.Element;
 
-export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EnrichingSpanProcessor, EvaluatorFn, EvaluatorResult, Fallback, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };
+export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EnrichingSpanProcessor, EvaluatorFn, EvaluatorResult, Fallback, type GetChatCompletionClientAndProvider, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, RetryLastErrorContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };

package/dist/index.d.ts

@@ -6,8 +6,8 @@ export { OpenAI as OpenAIClient } from 'openai';
 import { ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam, ChatCompletionAssistantMessageParam, ChatCompletionCreateParams } from 'openai/resources';
 import AnthropicClient from '@anthropic-ai/sdk';
 export { default as AnthropicClient } from '@anthropic-ai/sdk';
-import { PromptFeedback, GenerateContentRequest, VertexAI } from '@google-cloud/vertexai';
-export { VertexAI } from '@google-cloud/vertexai';
+import { PromptFeedback, GenerateContentRequest, VertexAI, HarmCategory, HarmBlockThreshold } from '@google-cloud/vertexai';
+export { HarmBlockThreshold as GoogleHarmBlockThreshold, HarmCategory as GoogleHarmCategory, VertexAI } from '@google-cloud/vertexai';
 
 declare class ChatCompletionError extends Error {
     readonly chatCompletionRequest: LogChatCompletionRequest;
@@ -43,6 +43,10 @@ type ChatCompletionClientAndProvider<K> = {
         completion: number;
     }) => number;
 };
+type GetChatCompletionClientAndProvider<Model, Client> = (model: Model, args: {
+    retryCount: number;
+    lastError?: Error | null;
+}) => Promise<ChatCompletionClientAndProvider<Client>>;
 
 type CreateRenderContextOptions = {
     logger?: LogImplementation;
@@ -52,15 +56,25 @@ type CreateRenderContextOptions = {
 };
 declare function createRenderContext({ logger, traceId, processor, contextValues, }?: CreateRenderContextOptions): RenderContext;
 
-declare const RetryCountContext: Context<number>;
+type ExcludeNumber<E extends number, T = number> = T extends E ? never : T;
+type RetryCountContextValue = {
+    retryCount: 0;
+    lastError: null;
+} | {
+    retryCount: ExcludeNumber<0>;
+    lastError: Error;
+};
+declare const RetryCountContext: Context<RetryCountContextValue>;
+declare const RetryLastErrorContext: Context<Error | null>;
 declare const DefaultMaxRetriesContext: Context<number>;
 type RetryProps = {
     shouldRetry: (error: Error) => boolean;
     retries?: number;
+    lastError?: Error;
     maxRetries?: number;
     children: AINode;
 };
-declare function Retry({ shouldRetry, retries, maxRetries, children }: RetryProps, ctx: RenderContext): AsyncGenerator<string, void, unknown>;
+declare function Retry({ shouldRetry, retries, maxRetries, lastError, children }: RetryProps, ctx: RenderContext): AsyncGenerator<string, void, unknown>;
 
 type FallbackProps = {
     fallback: AINode;
@@ -241,7 +255,7 @@ declare module '@gammatech/aijsx' {
 }
 type ValidOpenAIVisionModel = 'gpt-4o' | 'gpt-4o-2024-05-13' | 'gpt-4-turbo-2024-04-09' | 'gpt-4-turbo' | 'gpt-4-vision-preview';
 type ValidOpenAIChatModel = ValidOpenAIVisionModel | 'gpt-4' | 'gpt-4-0314' | 'gpt-4-0613' | 'gpt-4-32k' | 'gpt-4-32k-0314' | 'gpt-4-32k-0613' | 'gpt-4-1106-preview' | 'gpt-4-0125-preview' | 'gpt-3.5-turbo' | 'gpt-3.5-turbo-0301' | 'gpt-3.5-turbo-0613' | 'gpt-3.5-turbo-16k' | 'gpt-3.5-turbo-16k-0613' | 'gpt-3.5-turbo-1106' | 'gpt-3.5-turbo-0125';
-declare const OpenAIClientContext: Context<() => ChatCompletionClientAndProvider<OpenAI>>;
+declare const OpenAIClientContext: Context<GetChatCompletionClientAndProvider<ValidOpenAIChatModel, OpenAI>>;
 type OpenAIChatCompletionProps = {
     model: ValidOpenAIChatModel;
     maxTokens?: number;
@@ -266,7 +280,7 @@ declare module '@gammatech/aijsx' {
  * @see https://docs.anthropic.com/claude/reference/selecting-a-model
  */
 type ValidAnthropicChatModel = 'claude-instant-1.2' | 'claude-2.1' | 'claude-3-opus-20240229' | 'claude-3-sonnet-20240229' | 'claude-3-haiku-20240307';
-declare const AnthropicClientContext: Context<() => ChatCompletionClientAndProvider<AnthropicClient>>;
+declare const AnthropicClientContext: Context<GetChatCompletionClientAndProvider<ValidAnthropicChatModel, AnthropicClient>>;
 type AnthropicChatCompletionProps = {
     model: ValidAnthropicChatModel;
     maxTokens?: number;
@@ -291,7 +305,7 @@ declare module '@gammatech/aijsx' {
     }
 }
 type ValidGoogleChatModel = 'gemini-1.5-pro' | 'gemini-1.5-flash';
-declare const GoogleClientContext: Context<() => ChatCompletionClientAndProvider<VertexAI>>;
+declare const GoogleClientContext: Context<GetChatCompletionClientAndProvider<ValidGoogleChatModel, VertexAI>>;
 type GoogleChatCompletionProps = {
     model: ValidGoogleChatModel;
     maxTokens?: number;
@@ -299,7 +313,11 @@ type GoogleChatCompletionProps = {
     stop?: string | string[];
     maxRetries?: number;
     children: AINode;
+    safetySettings?: {
+        category: HarmCategory;
+        threshold: HarmBlockThreshold;
+    }[];
 };
 declare function GoogleChatCompletion(props: GoogleChatCompletionProps, ctx: RenderContext): JSX.Element;
 
-export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EnrichingSpanProcessor, EvaluatorFn, EvaluatorResult, Fallback, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };
+export { AIComponent, AINode, AISpanAttributes, AISpanProcessor, AnthropicChatCompletion, type AnthropicChatCompletionRequest, AnthropicClientContext, AssistantMessage, type ChatCompletionClientAndProvider, ChatCompletionError, type ChatCompletionRequestPayloads, ChatMessage, Context, type CostFn, DebugMessage, DefaultMaxRetriesContext, EnrichingSpanProcessor, EvaluatorFn, EvaluatorResult, Fallback, type GetChatCompletionClientAndProvider, GoogleChatCompletion, type GoogleChatCompletionRequest, GoogleClientContext, LogChatCompletionRequest, LogImplementation, OpenAIChatCompletion, type OpenAIChatCompletionRequest, type OpenAIChatMessage, OpenAIClientContext, ParseVariablesError, ProcessedAISpanAttributes, type Prompt, PromptInvalidOutputError, PromptParsed, ReadableSpan, RenderContext, Retry, RetryCountContext, RetryLastErrorContext, SpanAttributes, SpanExporter, SpanProcessor, SystemMessage, type TokenizerFn, Trace, Tracer, UserMessage, type ValidAnthropicChatModel, type ValidGoogleChatModel, type ValidOpenAIChatModel, type ValidOpenAIVisionModel, anthropicTokenizer, computeUsage, createPrompt, createRenderContext, evaluatePrompt, isPromptParsed, openaiTokenizer, tracing };

package/dist/index.js

@@ -44,6 +44,8 @@ __export(src_exports, {
   Fallback: () => Fallback,
   GoogleChatCompletion: () => GoogleChatCompletion,
   GoogleClientContext: () => GoogleClientContext,
+  GoogleHarmBlockThreshold: () => import_vertexai2.HarmBlockThreshold,
+  GoogleHarmCategory: () => import_vertexai2.HarmCategory,
   ImagePart: () => ImagePart,
   LogImplementation: () => LogImplementation,
   NoopLogImplementation: () => NoopLogImplementation,
@@ -54,6 +56,7 @@ __export(src_exports, {
   PromptInvalidOutputError: () => PromptInvalidOutputError,
   Retry: () => Retry,
   RetryCountContext: () => RetryCountContext,
+  RetryLastErrorContext: () => RetryLastErrorContext,
   SystemMessage: () => SystemMessage,
   Trace: () => Trace,
   UserMessage: () => UserMessage,
@@ -1536,26 +1539,41 @@ function renderCloseTag(element) {
 }
 
 // src/retry.tsx
-var RetryCountContext = createContext(0);
+var RetryCountContext = createContext({
+  retryCount: 0,
+  lastError: null
+});
+var RetryLastErrorContext = createContext(null);
 var DefaultMaxRetriesContext = createContext(0);
-async function* Retry({ shouldRetry: shouldRetry4, retries = 0, maxRetries = 3, children }, ctx) {
+async function* Retry({ shouldRetry: shouldRetry4, retries = 0, maxRetries = 3, lastError, children }, ctx) {
   const { render } = ctx;
   let hasYieldedData = false;
   try {
+    const ctxValue = retries === 0 ? { retryCount: 0, lastError: null } : {
+      retryCount: retries,
+      lastError
+    };
     const result = render(
-      /* @__PURE__ */ jsx(RetryCountContext.Provider, { value: retries, children })
+      /* @__PURE__ */ jsx(RetryCountContext.Provider, { value: ctxValue, children })
     );
     for await (const value of result) {
       hasYieldedData = true;
       yield value;
     }
   } catch (e) {
-    if (hasYieldedData || retries >= maxRetries || !shouldRetry4(e)) {
+    const err = castToError(e);
+    if (hasYieldedData || retries >= maxRetries || !shouldRetry4(err)) {
       throw e;
     }
     await backoff(retries);
     yield* Retry(
-      { shouldRetry: shouldRetry4, retries: retries + 1, maxRetries, children },
+      {
+        shouldRetry: shouldRetry4,
+        retries: retries + 1,
+        maxRetries,
+        lastError: err,
+        children
+      },
       ctx
     );
   }
@@ -1565,6 +1583,15 @@ var backoff = (retries) => {
   const waitTime = BASE_BACKOFF * Math.pow(4, retries);
   return new Promise((resolve) => setTimeout(resolve, waitTime));
 };
+var castToError = (e) => {
+  if (e instanceof Error) {
+    return e;
+  }
+  if (typeof e === "string") {
+    return new Error(e);
+  }
+  return new Error("Unknown error");
+};
 
 // src/fallback.tsx
 async function* Fallback({ shouldFallback = () => true, fallback, children }, ctx) {
@@ -2049,7 +2076,7 @@ function getEnvVar(name, shouldThrow = true) {
 
 // src/lib/openai/OpenAI.tsx
 var defaultClient = null;
-var OpenAIClientContext = createContext(() => {
+var OpenAIClientContext = createContext(async () => {
   if (defaultClient) {
     return defaultClient;
   }
@@ -2121,9 +2148,16 @@ function OpenAIChatCompletion(props, ctx) {
 async function* OpenAIChatCompletionInner(props, ctx) {
   const startTime = performance.now();
   const { logger, tracer, getContext } = ctx;
-  const retryCount = getContext(RetryCountContext);
+  const { retryCount, lastError } = getContext(RetryCountContext);
   const span = tracer.getActiveSpan();
-  const { client, provider, providerRegion, costFn } = getContext(OpenAIClientContext)();
+  const getClientFn = getContext(OpenAIClientContext);
+  const { client, provider, providerRegion, costFn } = await getClientFn(
+    props.model,
+    {
+      retryCount,
+      lastError
+    }
+  );
   if (!client) {
     throw new Error("[OpenAI] must supply OpenAI model via context");
   }
@@ -2293,7 +2327,7 @@ var anthropicTokenizer = (message) => {
 
 // src/lib/anthropic/Anthropic.tsx
 var defaultClient2 = null;
-var AnthropicClientContext = createContext(() => {
+var AnthropicClientContext = createContext(async () => {
   if (defaultClient2) {
     return defaultClient2;
   }
@@ -2411,11 +2445,16 @@ function AnthropicChatCompletion(props, ctx) {
 async function* AnthropicChatCompletionInner(props, ctx) {
   const startTime = performance.now();
   const { logger, tracer, getContext } = ctx;
-  const retryCount = getContext(RetryCountContext);
+  const { retryCount, lastError } = getContext(RetryCountContext);
   const span = tracer.getActiveSpan();
-  const { client, provider, providerRegion, costFn } = getContext(
-    AnthropicClientContext
-  )();
+  const getClientFn = getContext(AnthropicClientContext);
+  const { client, provider, providerRegion, costFn } = await getClientFn(
+    props.model,
+    {
+      retryCount,
+      lastError
+    }
+  );
   if (!client) {
     throw new Error(
       "[AnthropicChatCompletion] must supply AnthropicClient via context"
@@ -2600,7 +2639,7 @@ var errorToChatCompletionError = (error, requestData) => {
 };
 
 // src/lib/google/Google.tsx
-var SAFETY_SETTINGS = [
+var DEFAULT_SAFETY_SETTINGS = [
   {
     category: import_vertexai.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
     threshold: import_vertexai.HarmBlockThreshold.BLOCK_ONLY_HIGH
@@ -2622,7 +2661,7 @@ var SAFETY_SETTINGS = [
     threshold: import_vertexai.HarmBlockThreshold.BLOCK_ONLY_HIGH
   }
 ];
-var GoogleClientContext = createContext(() => {
+var GoogleClientContext = createContext(async () => {
   if (defaultClient3) {
     return defaultClient3;
   }
@@ -2703,9 +2742,16 @@ function GoogleChatCompletion(props, ctx) {
 async function* GoogleChatCompletionInner(props, ctx) {
   const startTime = performance.now();
   const { logger, tracer, getContext } = ctx;
-  const retryCount = getContext(RetryCountContext);
+  const { retryCount, lastError } = getContext(RetryCountContext);
   const span = tracer.getActiveSpan();
-  const { client, provider, providerRegion, costFn } = getContext(GoogleClientContext)();
+  const getClientFn = getContext(GoogleClientContext);
+  const { client, provider, providerRegion, costFn } = await getClientFn(
+    props.model,
+    {
+      retryCount,
+      lastError
+    }
+  );
   if (!client) {
     throw new Error(
       "[GoogleChatCompletion] must supply GoogleClient via context"
@@ -2754,7 +2800,7 @@ async function* GoogleChatCompletionInner(props, ctx) {
   });
   const model = client.getGenerativeModel({
     model: props.model,
-    safetySettings: SAFETY_SETTINGS
+    safetySettings: props.safetySettings || DEFAULT_SAFETY_SETTINGS
   });
   let response;
   try {
@@ -2882,6 +2928,8 @@ var import_vertexai2 = require("@google-cloud/vertexai");
   Fallback,
   GoogleChatCompletion,
   GoogleClientContext,
+  GoogleHarmBlockThreshold,
+  GoogleHarmCategory,
   ImagePart,
   LogImplementation,
   NoopLogImplementation,
@@ -2892,6 +2940,7 @@ var import_vertexai2 = require("@google-cloud/vertexai");
   PromptInvalidOutputError,
   Retry,
   RetryCountContext,
+  RetryLastErrorContext,
   SystemMessage,
   Trace,
   UserMessage,

package/dist/index.mjs

@@ -1436,26 +1436,41 @@ function renderCloseTag(element) {
 }
 
 // src/retry.tsx
-var RetryCountContext = createContext(0);
+var RetryCountContext = createContext({
+  retryCount: 0,
+  lastError: null
+});
+var RetryLastErrorContext = createContext(null);
 var DefaultMaxRetriesContext = createContext(0);
-async function* Retry({ shouldRetry: shouldRetry4, retries = 0, maxRetries = 3, children }, ctx) {
+async function* Retry({ shouldRetry: shouldRetry4, retries = 0, maxRetries = 3, lastError, children }, ctx) {
   const { render } = ctx;
   let hasYieldedData = false;
   try {
+    const ctxValue = retries === 0 ? { retryCount: 0, lastError: null } : {
+      retryCount: retries,
+      lastError
+    };
     const result = render(
-      /* @__PURE__ */ jsx(RetryCountContext.Provider, { value: retries, children })
+      /* @__PURE__ */ jsx(RetryCountContext.Provider, { value: ctxValue, children })
     );
     for await (const value of result) {
       hasYieldedData = true;
       yield value;
     }
   } catch (e) {
-    if (hasYieldedData || retries >= maxRetries || !shouldRetry4(e)) {
+    const err = castToError(e);
+    if (hasYieldedData || retries >= maxRetries || !shouldRetry4(err)) {
       throw e;
     }
     await backoff(retries);
     yield* Retry(
-      { shouldRetry: shouldRetry4, retries: retries + 1, maxRetries, children },
+      {
+        shouldRetry: shouldRetry4,
+        retries: retries + 1,
+        maxRetries,
+        lastError: err,
+        children
+      },
       ctx
     );
   }
@@ -1465,6 +1480,15 @@ var backoff = (retries) => {
   const waitTime = BASE_BACKOFF * Math.pow(4, retries);
   return new Promise((resolve) => setTimeout(resolve, waitTime));
 };
+var castToError = (e) => {
+  if (e instanceof Error) {
+    return e;
+  }
+  if (typeof e === "string") {
+    return new Error(e);
+  }
+  return new Error("Unknown error");
+};
 
 // src/fallback.tsx
 async function* Fallback({ shouldFallback = () => true, fallback, children }, ctx) {
@@ -1949,7 +1973,7 @@ function getEnvVar(name, shouldThrow = true) {
 
 // src/lib/openai/OpenAI.tsx
 var defaultClient = null;
-var OpenAIClientContext = createContext(() => {
+var OpenAIClientContext = createContext(async () => {
   if (defaultClient) {
     return defaultClient;
   }
@@ -2021,9 +2045,16 @@ function OpenAIChatCompletion(props, ctx) {
 async function* OpenAIChatCompletionInner(props, ctx) {
   const startTime = performance.now();
   const { logger, tracer, getContext } = ctx;
-  const retryCount = getContext(RetryCountContext);
+  const { retryCount, lastError } = getContext(RetryCountContext);
   const span = tracer.getActiveSpan();
-  const { client, provider, providerRegion, costFn } = getContext(OpenAIClientContext)();
+  const getClientFn = getContext(OpenAIClientContext);
+  const { client, provider, providerRegion, costFn } = await getClientFn(
+    props.model,
+    {
+      retryCount,
+      lastError
+    }
+  );
   if (!client) {
     throw new Error("[OpenAI] must supply OpenAI model via context");
   }
@@ -2193,7 +2224,7 @@ var anthropicTokenizer = (message) => {
 
 // src/lib/anthropic/Anthropic.tsx
 var defaultClient2 = null;
-var AnthropicClientContext = createContext(() => {
+var AnthropicClientContext = createContext(async () => {
   if (defaultClient2) {
     return defaultClient2;
   }
@@ -2311,11 +2342,16 @@ function AnthropicChatCompletion(props, ctx) {
 async function* AnthropicChatCompletionInner(props, ctx) {
   const startTime = performance.now();
   const { logger, tracer, getContext } = ctx;
-  const retryCount = getContext(RetryCountContext);
+  const { retryCount, lastError } = getContext(RetryCountContext);
   const span = tracer.getActiveSpan();
-  const { client, provider, providerRegion, costFn } = getContext(
-    AnthropicClientContext
-  )();
+  const getClientFn = getContext(AnthropicClientContext);
+  const { client, provider, providerRegion, costFn } = await getClientFn(
+    props.model,
+    {
+      retryCount,
+      lastError
+    }
+  );
   if (!client) {
     throw new Error(
       "[AnthropicChatCompletion] must supply AnthropicClient via context"
@@ -2473,9 +2509,9 @@ import AnthropicClient2 from "@anthropic-ai/sdk";
 
 // src/lib/google/Google.tsx
 import {
-  VertexAI,
+  HarmBlockThreshold,
   HarmCategory,
-  HarmBlockThreshold
+  VertexAI
 } from "@google-cloud/vertexai";
 
 // src/lib/google/errors.ts
@@ -2504,7 +2540,7 @@ var errorToChatCompletionError = (error, requestData) => {
 };
 
 // src/lib/google/Google.tsx
-var SAFETY_SETTINGS = [
+var DEFAULT_SAFETY_SETTINGS = [
   {
     category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
     threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH
@@ -2526,7 +2562,7 @@ var SAFETY_SETTINGS = [
     threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH
   }
 ];
-var GoogleClientContext = createContext(() => {
+var GoogleClientContext = createContext(async () => {
   if (defaultClient3) {
     return defaultClient3;
   }
@@ -2607,9 +2643,16 @@ function GoogleChatCompletion(props, ctx) {
 async function* GoogleChatCompletionInner(props, ctx) {
   const startTime = performance.now();
   const { logger, tracer, getContext } = ctx;
-  const retryCount = getContext(RetryCountContext);
+  const { retryCount, lastError } = getContext(RetryCountContext);
   const span = tracer.getActiveSpan();
-  const { client, provider, providerRegion, costFn } = getContext(GoogleClientContext)();
+  const getClientFn = getContext(GoogleClientContext);
+  const { client, provider, providerRegion, costFn } = await getClientFn(
+    props.model,
+    {
+      retryCount,
+      lastError
+    }
+  );
   if (!client) {
     throw new Error(
       "[GoogleChatCompletion] must supply GoogleClient via context"
@@ -2658,7 +2701,7 @@ async function* GoogleChatCompletionInner(props, ctx) {
   });
   const model = client.getGenerativeModel({
     model: props.model,
-    safetySettings: SAFETY_SETTINGS
+    safetySettings: props.safetySettings || DEFAULT_SAFETY_SETTINGS
   });
   let response;
   try {
@@ -2768,7 +2811,11 @@ function cleanChatCompletionRequest3(chatCompletionRequest) {
 }
 
 // src/lib/google/index.ts
-import { VertexAI as VertexAI2 } from "@google-cloud/vertexai";
+import {
+  VertexAI as VertexAI2,
+  HarmBlockThreshold as HarmBlockThreshold2,
+  HarmCategory as HarmCategory2
+} from "@google-cloud/vertexai";
 export {
   AIFragment,
   AISpanProcessor,
@@ -2785,6 +2832,8 @@ export {
   Fallback,
   GoogleChatCompletion,
   GoogleClientContext,
+  HarmBlockThreshold2 as GoogleHarmBlockThreshold,
+  HarmCategory2 as GoogleHarmCategory,
   ImagePart,
   LogImplementation,
   NoopLogImplementation,
@@ -2795,6 +2844,7 @@ export {
   PromptInvalidOutputError,
   Retry,
   RetryCountContext,
+  RetryLastErrorContext,
   SystemMessage,
   Trace,
   UserMessage,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@gammatech/aijsx",
-  "version": "0.10.1-dev.2024-06-07",
+  "version": "0.11.0-dev.2024-06-17",
   "description": "Rewrite of aijsx",
   "author": "Jordan Garcia",
   "license": "MIT",