diff --git a/README.md b/README.md
index 7bf7590..25df22d 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,7 @@
 # Ax - Build LLMs Powered Agents (Typescript)
 
-Build intelligent agents quickly — inspired by the power of "Agentic workflows" and the Stanford DSPy paper. Seamlessly integrates with multiple LLMs and VectorDBs to build RAG pipelines or collaborative agents that can solve complex problems. Advanced features streaming validation, multi-modal DSPy, etc.
+Use Ax and get a streaming, multi-modal DSPy framework with agents and typed signatures. Works with all LLMs. Ax is always streaming and handles output type validation while streaming for faster responses and lower token usage.
+
 
 [![NPM Package](https://img.shields.io/npm/v/@ax-llm/ax?style=for-the-badge&color=green)](https://www.npmjs.com/package/@ax-llm/ax)
 [![Discord Chat](https://dcbadge.vercel.app/api/server/DSHg3dU7dW?style=for-the-badge)](https://discord.gg/DSHg3dU7dW)
@@ -8,9 +9,9 @@ Build intelligent agents quickly — inspired by the power of "Agentic workflows
 
 ![image](https://github.com/ax-llm/ax/assets/832235/3a250031-591c-42e0-b4fc-06afb8c351c4)
 
-## Our focus on agents
+## Focus on agents
 
-We've renamed from "llmclient" to "ax" to highlight our focus on powering agentic workflows. We agree with many experts like "Andrew Ng" that agentic workflows are the key to unlocking the true power of large language models and what can be achieved with in-context learning. Also, we are big fans of the Stanford DSPy paper, and this library is the result of all of this coming together to build a powerful framework for you to build with.
+While we're focused on building agents, Ax has all the tools needed to quickly build powerful and production ready workflows with LLMs
 
 ![image](https://github.com/ax-llm/ax/assets/832235/801b8110-4cba-4c50-8ec7-4d5859121fe5)
 
@@ -95,11 +96,10 @@ const ai = new AxAI({
 });
 
 const gen = new AxChainOfThought(
-  ai,
   `textToSummarize -> textType:class "note, email, reminder", shortSummary "summarize in 5 to 10 words"`
 );
 
-const res = await gen.forward({ textToSummarize });
+const res = await gen.forward(ai, { textToSummarize });
 
 console.log('>', res);
 ```
@@ -111,26 +111,26 @@ Use the agent prompt (framework) to build agents that work with other agents to
 ```typescript
 # npm run tsx ./src/examples/agent.ts
 
-const researcher = new AxAgent(ai, {
+const researcher = new AxAgent({
   name: 'researcher',
   description: 'Researcher agent',
   signature: `physicsQuestion "physics questions" -> answer "reply in bullet points"`
 });
 
-const summarizer = new AxAgent(ai, {
+const summarizer = new AxAgent({
   name: 'summarizer',
   description: 'Summarizer agent',
   signature: `text "text so summarize" -> shortSummary "summarize in 5 to 10 words"`
 });
 
-const agent = new AxAgent(ai, {
+const agent = new AxAgent({
   name: 'agent',
   description: 'A an agent to research complex topics',
   signature: `question -> answer`,
   agents: [researcher, summarizer]
 });
 
-agent.forward({ questions: "How many atoms are there in the universe" })
+agent.forward(ai, { questions: "How many atoms are there in the universe" })
 ```
 
 ## Vector DBs Supported
@@ -210,9 +210,9 @@ const image = fs
   .readFileSync('./src/examples/assets/kitten.jpeg')
   .toString('base64');
 
-const gen = new AxChainOfThought(ai, `question, animalImage:image -> answer`);
+const gen = new AxChainOfThought(`question, animalImage:image -> answer`);
 
-const res = await gen.forward({
+const res = await gen.forward(ai, {
   question: 'What family does this animal belong to?',
   animalImage: { mimeType: 'image/jpeg', data: image }
 });
@@ -329,7 +329,7 @@ const ai = new AxAI({
 // Create a model using the provider
 const model = new AxAIProvider(ai);
 
-export const foodAgent = new AxAgent(ai, {
+export const foodAgent = new AxAgent({
   name: 'food-search',
   description:
     'Use this agent to find restaurants based on what the customer wants',
@@ -341,7 +341,7 @@ export const foodAgent = new AxAgent(ai, {
 const aiState = getMutableAIState()
 
 // Create an agent for a specific task
-const foodAgent = new AxAgentProvider({
+const foodAgent = new AxAgentProvider(ai, {
     agent: foodAgent,
     updateState: (state) => {
          aiState.done({ ...aiState.get(), state })
diff --git a/src/ax/ai/types.ts b/src/ax/ai/types.ts
index 488599f..454702d 100644
--- a/src/ax/ai/types.ts
+++ b/src/ax/ai/types.ts
@@ -174,6 +174,7 @@ export type AxAIServiceOptions = {
 };
 
 export type AxAIServiceActionOptions = {
+  ai?: Readonly<AxAIService>;
   sessionId?: string;
   traceId?: string;
 };
diff --git a/src/ax/docs/manager.ts b/src/ax/docs/manager.ts
index f196625..e9de3df 100644
--- a/src/ax/docs/manager.ts
+++ b/src/ax/docs/manager.ts
@@ -114,7 +114,7 @@ export class AxDBManager {
 
     if (typeof texts[0] === 'string' && this.rewriter) {
       for (const [i, text] of texts.entries()) {
-        const { rewrittenQuery } = await this.rewriter.forward({
+        const { rewrittenQuery } = await this.rewriter.forward(this.ai, {
           query: text
         });
         texts[i] = rewrittenQuery;
@@ -147,7 +147,7 @@ export class AxDBManager {
       const resultItems = tp ? getTopInPercent(m, tp) : m;
 
       if (this.reranker) {
-        const { rankedItems } = await this.reranker.forward({
+        const { rankedItems } = await this.reranker.forward(this.ai, {
           query: texts[0] as string,
           items: resultItems.map((item) => item.text)
         });
diff --git a/src/ax/docs/reranker.ts b/src/ax/docs/reranker.ts
index 22fc683..69b0396 100644
--- a/src/ax/docs/reranker.ts
+++ b/src/ax/docs/reranker.ts
@@ -1,29 +1,30 @@
 import type { AxAIService } from '../ai/types.js';
 import {
-  AxGenerate,
-  type AxGenerateOptions,
+  AxGen,
+  type AxGenOptions,
   type AxProgramForwardOptions,
   axStringUtil
 } from '../dsp/index.js';
 
 import type { AxRerankerIn, AxRerankerOut } from './manager.js';
 
-export class AxDefaultResultReranker extends AxGenerate<
+export class AxDefaultResultReranker extends AxGen<
   AxRerankerIn,
   AxRerankerOut
 > {
-  constructor(ai: AxAIService, options?: Readonly<AxGenerateOptions>) {
+  constructor(options?: Readonly<AxGenOptions>) {
     const signature = `"You are a re-ranker assistant tasked with evaluating a set of content items in relation to a specific question. Your role involves critically analyzing each content item to determine its relevance to the question and re-ranking them accordingly. This process includes assigning a relevance score from 0 to 10 to each content item based on how well it answers the question, its coverage of the topic, and the reliability of its information. This re-ranked list should start with the content item that is most relevant to the question and end with the least relevant. Output only the list."
     query: string, items: string[] -> rankedItems: string[] "list of id, 5-words Rationale, relevance score"`;
 
-    super(ai, signature, options);
+    super(signature, options);
   }
 
   public override forward = async (
+    ai: Readonly<AxAIService>,
     input: Readonly<AxRerankerIn>,
     options?: Readonly<AxProgramForwardOptions>
   ): Promise<AxRerankerOut> => {
-    const { rankedItems } = await super.forward(input, options);
+    const { rankedItems } = await super.forward(ai, input, options);
 
     const sortedIndexes: number[] = rankedItems.map((item) => {
       const { id: index } = axStringUtil.extractIdAndText(item);
diff --git a/src/ax/docs/rewriter.ts b/src/ax/docs/rewriter.ts
index 4d44787..cc5181b 100644
--- a/src/ax/docs/rewriter.ts
+++ b/src/ax/docs/rewriter.ts
@@ -1,19 +1,15 @@
 import {
-  type AxAIService,
-  AxGenerate,
-  type AxGenerateOptions,
+  AxGen,
+  type AxGenOptions,
   type AxRewriteIn,
   type AxRewriteOut
 } from '../index.js';
 
-export class AxDefaultQueryRewriter extends AxGenerate<
-  AxRewriteIn,
-  AxRewriteOut
-> {
-  constructor(ai: AxAIService, options?: Readonly<AxGenerateOptions>) {
+export class AxDefaultQueryRewriter extends AxGen<AxRewriteIn, AxRewriteOut> {
+  constructor(options?: Readonly<AxGenOptions>) {
     const signature = `"You are a query rewriter assistant tasked with rewriting a given query to improve its clarity, specificity, and relevance. Your role involves analyzing the query to identify any ambiguities, generalizations, or irrelevant information and then rephrasing it to make it more focused and precise. The rewritten query should be concise, easy to understand, and directly related to the original query. Output only the rewritten query."
     query: string -> rewrittenQuery: string`;
 
-    super(ai, signature, options);
+    super(signature, options);
   }
 }
diff --git a/src/ax/dsp/evaluate.ts b/src/ax/dsp/evaluate.ts
index ae845f2..d41dc40 100644
--- a/src/ax/dsp/evaluate.ts
+++ b/src/ax/dsp/evaluate.ts
@@ -1,8 +1,11 @@
+import type { AxAIService } from '../ai/types.js';
+
 import type { AxExample, AxMetricFn } from './optimize.js';
 import type { AxGenIn, AxGenOut, AxProgram } from './program.js';
 import { updateProgressBar } from './util.js';
 
 export type AxEvaluateArgs<IN extends AxGenIn, OUT extends AxGenOut> = {
+  ai: AxAIService;
   program: Readonly<AxProgram<IN, OUT>>;
   examples: Readonly<AxExample[]>;
 };
@@ -11,13 +14,19 @@ export class AxTestPrompt<
   IN extends AxGenIn = AxGenIn,
   OUT extends AxGenOut = AxGenOut
 > {
+  private ai: AxAIService;
   private program: Readonly<AxProgram<IN, OUT>>;
   private examples: Readonly<AxExample[]>;
 
-  constructor({ program, examples = [] }: Readonly<AxEvaluateArgs<IN, OUT>>) {
+  constructor({
+    ai,
+    program,
+    examples = []
+  }: Readonly<AxEvaluateArgs<IN, OUT>>) {
     if (examples.length == 0) {
       throw new Error('No examples found');
     }
+    this.ai = ai;
     this.program = program;
     this.examples = examples;
   }
@@ -33,7 +42,7 @@ export class AxTestPrompt<
         throw new Error('Invalid example');
       }
 
-      const res = await this.program.forward(ex as IN);
+      const res = await this.program.forward(this.ai, ex as IN);
       const success = metricFn({ prediction: res, example: ex });
       if (success) {
         successCount++;
diff --git a/src/ax/dsp/functions.ts b/src/ax/dsp/functions.ts
index 01f2ced..8e50115 100644
--- a/src/ax/dsp/functions.ts
+++ b/src/ax/dsp/functions.ts
@@ -41,7 +41,8 @@ export class AxFunctionProcessor {
     const opt = options
       ? {
           sessionId: options.sessionId,
-          traceId: options.traceId
+          traceId: options.traceId,
+          ai: options.ai
         }
       : undefined;
 
diff --git a/src/ax/dsp/gen2.ts b/src/ax/dsp/gen2.ts
deleted file mode 100644
index c269988..0000000
--- a/src/ax/dsp/gen2.ts
+++ /dev/null
@@ -1,30 +0,0 @@
-import type { AxAIService } from '../ai/types.js';
-
-import {
-  AxGenerate,
-  type AxGenerateOptions,
-  type AxGenerateResult
-} from './generate.js';
-import type { AxGenIn, AxGenOut, AxProgramForwardOptions } from './program.js';
-import type { AxSignature } from './sig.js';
-
-export class AxGen<
-  IN extends AxGenIn = AxGenIn,
-  OUT extends AxGenerateResult<AxGenOut> = AxGenerateResult<AxGenOut>
-> extends AxGenerate<IN, OUT> {
-  constructor(
-    signature: Readonly<AxSignature | string>,
-    options?: Readonly<AxGenerateOptions>
-  ) {
-    super(null as unknown as AxAIService, signature, options);
-  }
-
-  // @ts-expect-error changing the over ridden function type
-  async forward(
-    ai: Readonly<AxAIService>,
-    values: IN,
-    options?: Readonly<AxProgramForwardOptions>
-  ): Promise<OUT> {
-    return await super.forward(values, { ...options, ai });
-  }
-}
diff --git a/src/ax/dsp/generate.ts b/src/ax/dsp/generate.ts
index b02a00b..eb885fb 100644
--- a/src/ax/dsp/generate.ts
+++ b/src/ax/dsp/generate.ts
@@ -40,7 +40,7 @@ import {
 import { AxPromptTemplate } from './prompt.js';
 import { AxSignature } from './sig.js';
 
-export interface AxGenerateOptions {
+export interface AxGenOptions {
   maxCompletions?: number;
   maxRetries?: number;
   maxSteps?: number;
@@ -49,6 +49,7 @@ export interface AxGenerateOptions {
   rateLimiter?: AxRateLimiterFunction;
   stream?: boolean;
   debug?: boolean;
+  instructions?: string;
 
   functions?: AxFunction[] | { toFunction: () => AxFunction }[];
   functionCall?: AxChatRequest['functionCall'];
@@ -62,6 +63,8 @@ export type AxGenerateResult<OUT extends AxGenOut> = OUT & {
 };
 
 export interface AxResponseHandlerArgs<T> {
+  ai: Readonly<AxAIService>;
+  sig: Readonly<AxSignature>;
   res: T;
   usageInfo: { ai: string; model: string };
   mem: AxAIMemory;
@@ -69,26 +72,24 @@ export interface AxResponseHandlerArgs<T> {
   traceId?: string;
 }
 
-export class AxGenerate<
+export class AxGen<
   IN extends AxGenIn = AxGenIn,
   OUT extends AxGenerateResult<AxGenOut> = AxGenerateResult<AxGenOut>
 > extends AxProgramWithSignature<IN, OUT> {
-  private ai: AxAIService;
   private pt: AxPromptTemplate;
   private asserts: AxAssertion[];
   private streamingAsserts: AxStreamingAssertion[];
-  private options?: Omit<AxGenerateOptions, 'functions'>;
+  private options?: Omit<AxGenOptions, 'functions'>;
 
   private functions?: AxFunction[];
   private funcProc?: AxFunctionProcessor;
   private functionList?: string;
 
   constructor(
-    ai: AxAIService,
     signature: Readonly<AxSignature | string>,
-    options?: Readonly<AxGenerateOptions>
+    options?: Readonly<AxGenOptions>
   ) {
-    super(signature);
+    super(signature, { instructions: options?.instructions });
 
     this.functions = options?.functions?.map((f) => {
       if ('toFunction' in f) {
@@ -97,7 +98,6 @@ export class AxGenerate<
       return f;
     });
 
-    this.ai = ai;
     this.options = options;
     this.pt = new (options?.promptTemplate ?? AxPromptTemplate)(this.signature);
     this.asserts = this.options?.asserts ?? [];
@@ -107,29 +107,32 @@ export class AxGenerate<
 
     if (this.functions) {
       this.funcProc = new AxFunctionProcessor(this.functions);
-      this.updateSigForFunctions();
     }
   }
 
-  private updateSigForFunctions = () => {
+  private updateSigForFunctions = (ai: AxAIService) => {
     // AI supports function calling natively so
     // no need to add fields for function call
-    if (this.ai.getFeatures().functions) {
+    if (ai.getFeatures().functions) {
       return;
     }
 
+    const sig = new AxSignature(this.signature);
+
     // These are the fields for the function call only needed when the underlying LLM API does not support function calling natively in the API.
-    this.signature.addOutputField({
+    sig.addOutputField({
       name: 'functionName',
       description: 'Name of function to call',
       isOptional: true
     });
 
-    this.signature.addOutputField({
+    sig.addOutputField({
       name: 'functionArguments',
       description: 'Arguments of function to call',
       isOptional: true
     });
+
+    return sig;
   };
 
   public addAssert = (
@@ -201,6 +204,7 @@ export class AxGenerate<
   }
 
   private async forwardCore({
+    sig,
     mem,
     sessionId,
     traceId,
@@ -210,7 +214,8 @@ export class AxGenerate<
     stream = false
   }: Readonly<
     Omit<AxProgramForwardOptions, 'ai' | 'mem'> & {
-      ai: AxAIService;
+      sig: Readonly<AxSignature>;
+      ai: Readonly<AxAIService>;
       mem: AxAIMemory;
     }
   >): Promise<OUT> {
@@ -231,6 +236,8 @@ export class AxGenerate<
 
     if (res instanceof ReadableStream) {
       return (await this.processSteamingResponse({
+        ai,
+        sig,
         res,
         usageInfo,
         mem,
@@ -240,6 +247,8 @@ export class AxGenerate<
     }
 
     return (await this.processResponse({
+      ai,
+      sig,
       res,
       usageInfo,
       mem,
@@ -249,6 +258,8 @@ export class AxGenerate<
   }
 
   private async processSteamingResponse({
+    ai,
+    sig,
     res,
     usageInfo,
     mem,
@@ -281,7 +292,7 @@ export class AxGenerate<
             xstate,
             content
           );
-          streamingExtractValues(this.signature, values, xstate, content);
+          streamingExtractValues(sig, values, xstate, content);
           assertAssertions(this.asserts, values);
         }
 
@@ -300,9 +311,9 @@ export class AxGenerate<
       }
     }
 
-    const funcs = parseFunctions(this.ai, functionCalls, values);
+    const funcs = parseFunctions(ai, functionCalls, values);
     if (funcs) {
-      await this.processFunctions(funcs, mem, sessionId, traceId);
+      await this.processFunctions(ai, funcs, mem, sessionId, traceId);
     }
 
     streamingExtractFinalValue(values, xstate, content);
@@ -312,6 +323,7 @@ export class AxGenerate<
   }
 
   private async processResponse({
+    ai,
     res,
     usageInfo,
     mem,
@@ -333,7 +345,7 @@ export class AxGenerate<
       }
 
       if (result.functionCalls) {
-        const funcs = parseFunctions(this.ai, result.functionCalls, values);
+        const funcs = parseFunctions(ai, result.functionCalls, values);
 
         if (funcs) {
           await this.processFunctions(funcs, mem, sessionId, traceId);
@@ -349,11 +361,12 @@ export class AxGenerate<
   }
 
   private async _forward(
+    ai: Readonly<AxAIService>,
+    sig: Readonly<AxSignature>,
     values: IN,
     options?: Readonly<AxProgramForwardOptions>,
     span?: AxSpan
   ): Promise<OUT> {
-    const ai = options?.ai ?? this.ai;
     const maxRetries = options?.maxRetries ?? this.options?.maxRetries ?? 5;
     const maxSteps = options?.maxSteps ?? this.options?.maxSteps ?? 10;
     const mem = options?.mem ?? this.options?.mem ?? new AxMemory();
@@ -361,9 +374,9 @@ export class AxGenerate<
 
     let err: ValidationError | AxAssertionError | undefined;
 
-    if (this.sigHash !== this.signature.hash()) {
+    if (this.sigHash !== sig.hash()) {
       const promptTemplate = this.options?.promptTemplate ?? AxPromptTemplate;
-      this.pt = new promptTemplate(this.signature);
+      this.pt = new promptTemplate(sig);
     }
 
     const prompt = this.pt.render<IN>(values, {
@@ -386,6 +399,7 @@ export class AxGenerate<
 
           const output = await this.forwardCore({
             ai,
+            sig,
             mem,
             sessionId,
             traceId,
@@ -401,7 +415,7 @@ export class AxGenerate<
             continue multiStepLoop;
           }
 
-          assertRequiredFields(this.signature, output);
+          assertRequiredFields(sig, output);
           this.trace = { ...output };
           return output;
         } catch (e) {
@@ -440,17 +454,20 @@ export class AxGenerate<
   }
 
   public override async forward(
+    ai: Readonly<AxAIService>,
     values: IN,
     options?: Readonly<AxProgramForwardOptions>
   ): Promise<OUT> {
+    const sig = this.updateSigForFunctions(ai) ?? this.signature;
+
     const tracer = this.options?.tracer ?? options?.tracer;
 
     if (!tracer) {
-      return await this._forward(values, options);
+      return await this._forward(ai, sig, values, options);
     }
 
     const attributes = {
-      ['generate.signature']: this.signature.toString(),
+      ['generate.signature']: sig.toString(),
       ['generate.functions']: this.functionList ?? 'none'
     };
 
@@ -461,7 +478,7 @@ export class AxGenerate<
         attributes
       },
       async (span) => {
-        const res = this._forward(values, options, span);
+        const res = this._forward(ai, sig, values, options, span);
         span.end();
         return res;
       }
@@ -469,6 +486,7 @@ export class AxGenerate<
   }
 
   public processFunctions = async (
+    ai: Readonly<AxAIService>,
     functionCalls: readonly AxChatResponseFunctionCall[],
     mem: Readonly<AxMemory>,
     sessionId?: string,
@@ -476,7 +494,7 @@ export class AxGenerate<
   ) => {
     // Map each function call to a promise that resolves to the function result or null
     const promises = functionCalls.map((func) =>
-      this.funcProc?.execute(func, { sessionId, traceId }).then((fres) => {
+      this.funcProc?.execute(func, { sessionId, traceId, ai }).then((fres) => {
         if (fres?.id) {
           return {
             role: 'function' as const,
diff --git a/src/ax/dsp/index.ts b/src/ax/dsp/index.ts
index 4a0bf2c..830b750 100644
--- a/src/ax/dsp/index.ts
+++ b/src/ax/dsp/index.ts
@@ -8,7 +8,6 @@ export * from './loader.js';
 export * from './strutil.js';
 export * from './router.js';
 export * from './loader.js';
-export * from './gen2.js';
 
 export type { AxAssertion, AxStreamingAssertion } from './asserts.js';
 export type { AxPromptTemplate, AxFieldTemplateFn } from './prompt.js';
diff --git a/src/ax/dsp/optimize.ts b/src/ax/dsp/optimize.ts
index 8aaf30f..e9cfc28 100644
--- a/src/ax/dsp/optimize.ts
+++ b/src/ax/dsp/optimize.ts
@@ -1,3 +1,5 @@
+import type { AxAIService } from '../ai/types.js';
+
 import type {
   AxFieldValue,
   AxGenIn,
@@ -17,6 +19,7 @@ export type AxMetricFn = <T extends AxGenOut = AxGenOut>(
 export type AxMetricFnArgs = Parameters<AxMetricFn>[0];
 
 export type AxOptimizerArgs<IN extends AxGenIn, OUT extends AxGenOut> = {
+  ai: AxAIService;
   program: Readonly<AxProgram<IN, OUT>>;
   examples: Readonly<AxExample[]>;
   options?: { maxRounds?: number; maxExamples?: number; maxDemos?: number };
@@ -26,6 +29,7 @@ export class AxBootstrapFewShot<
   IN extends AxGenIn = AxGenIn,
   OUT extends AxGenOut = AxGenOut
 > {
+  private ai: AxAIService;
   private program: Readonly<AxProgram<IN, OUT>>;
   private examples: Readonly<AxExample[]>;
   private maxRounds: number;
@@ -34,6 +38,7 @@ export class AxBootstrapFewShot<
   private traces: AxProgramTrace[] = [];
 
   constructor({
+    ai,
     program,
     examples = [],
     options
@@ -45,6 +50,7 @@ export class AxBootstrapFewShot<
     this.maxDemos = options?.maxDemos ?? 4;
     this.maxExamples = options?.maxExamples ?? 16;
 
+    this.ai = ai;
     this.program = program;
     this.examples = examples;
   }
@@ -71,7 +77,7 @@ export class AxBootstrapFewShot<
       const exList = [...examples.slice(0, i), ...examples.slice(i + 1)];
       this.program.setExamples(exList);
 
-      const res = await this.program.forward(ex as IN, aiOpt);
+      const res = await this.program.forward(this.ai, ex as IN, aiOpt);
       const success = metricFn({ prediction: res, example: ex });
       if (success) {
         this.traces = [...this.traces, ...this.program.getTraces()];
diff --git a/src/ax/dsp/program.ts b/src/ax/dsp/program.ts
index 39a32b8..895c00b 100644
--- a/src/ax/dsp/program.ts
+++ b/src/ax/dsp/program.ts
@@ -17,6 +17,7 @@ export type AxFieldValue =
   | number
   | boolean
   | object
+  | null
   | { mimeType: string; data: string }
   | { mimeType: string; data: string }[];
 
@@ -71,6 +72,10 @@ export type AxProgramUsage = AxChatResponse['modelUsage'] & {
   model: string;
 };
 
+export interface AxProgramWithSignatureOptions {
+  instructions?: string;
+}
+
 export class AxProgramWithSignature<IN extends AxGenIn, OUT extends AxGenOut>
   implements AxTunable, AxUsable
 {
@@ -85,11 +90,18 @@ export class AxProgramWithSignature<IN extends AxGenIn, OUT extends AxGenOut>
   private key: { id: string; custom?: boolean };
   private children: AxInstanceRegistry<Readonly<AxTunable & AxUsable>>;
 
-  constructor(signature: Readonly<AxSignature | string>) {
+  constructor(
+    signature: Readonly<AxSignature | string>,
+    options?: Readonly<AxProgramWithSignatureOptions>
+  ) {
     this.signature = new AxSignature(signature);
     this.sigHash = this.signature?.hash();
     this.children = new AxInstanceRegistry();
     this.key = { id: this.constructor.name };
+
+    if (options?.instructions) {
+      this.signature.setDescription(options.instructions);
+    }
   }
 
   public getSignature() {
@@ -105,7 +117,9 @@ export class AxProgramWithSignature<IN extends AxGenIn, OUT extends AxGenOut>
 
   public async forward(
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    _arg0: IN,
+    _ai: Readonly<AxAIService>,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _values: IN,
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     _options?: Readonly<AxProgramForwardOptions>
   ): Promise<OUT> {
@@ -232,7 +246,9 @@ export class AxProgram<IN extends AxGenIn, OUT extends AxGenOut>
 
   public async forward(
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    _arg0: IN,
+    _ai: Readonly<AxAIService>,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _values: IN,
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     _options?: Readonly<AxProgramForwardOptions>
   ): Promise<OUT> {
diff --git a/src/ax/dsp/prompt.ts b/src/ax/dsp/prompt.ts
index 976abf5..0be01a2 100644
--- a/src/ax/dsp/prompt.ts
+++ b/src/ax/dsp/prompt.ts
@@ -108,6 +108,9 @@ export class AxPromptTemplate {
 
     if (extraFields && extraFields.length > 0) {
       extraFields.forEach((field) => {
+        // if (!field.isOptional && !field.value) {
+        //   throw new Error(`Value for field '${field.name}' is required.`);
+        // }
         const fn =
           this.fieldTemplates?.[field.name] ?? this.defaultRenderInField;
         if (!field.description || field.description.length === 0) {
@@ -218,28 +221,25 @@ export class AxPromptTemplate {
     values: Readonly<Record<string, AxFieldValue>>,
     skipMissing?: boolean
   ) => {
-    const textFieldFn: AxFieldTemplateFn =
-      this.fieldTemplates?.[field.name] ?? this.defaultRenderInField;
     const value = values[field.name];
 
     if (skipMissing && !value) {
       return;
     }
 
-    if (
-      !value ||
-      ((Array.isArray(value) || typeof value === 'string') &&
-        value.length === 0)
-    ) {
-      if (field.isOptional) {
-        return;
-      }
-      throw new Error(`Value for input field '${field.name}' is required.`);
+    if (isEmptyValue(field, value)) {
+      return;
     }
+
     if (field.type) {
-      validateValue(field, value);
+      validateValue(field, value!);
     }
-    const processedValue = processValue(field, value);
+
+    const processedValue = processValue(field, value!);
+
+    const textFieldFn: AxFieldTemplateFn =
+      this.fieldTemplates?.[field.name] ?? this.defaultRenderInField;
+
     return textFieldFn(field, processedValue);
   };
 
@@ -251,6 +251,10 @@ export class AxPromptTemplate {
       const validateImage = (
         value: Readonly<AxFieldValue>
       ): { mimeType: string; data: string } => {
+        if (!value) {
+          throw new Error('Image field value is required.');
+        }
+
         if (typeof value !== 'object') {
           throw new Error('Image field value must be an object.');
         }
@@ -407,3 +411,19 @@ function combineConsecutiveStrings(separator: string) {
     return acc;
   };
 }
+
+const isEmptyValue = (
+  field: Readonly<AxField>,
+  value?: Readonly<AxFieldValue>
+) => {
+  if (
+    !value ||
+    ((Array.isArray(value) || typeof value === 'string') && value.length === 0)
+  ) {
+    if (field.isOptional) {
+      return true;
+    }
+    throw new Error(`Value for input field '${field.name}' is required.`);
+  }
+  return false;
+};
diff --git a/src/ax/dsp/util.ts b/src/ax/dsp/util.ts
index b67efec..aee78ab 100644
--- a/src/ax/dsp/util.ts
+++ b/src/ax/dsp/util.ts
@@ -23,7 +23,7 @@ export const updateProgressBar = (
     elapsedTime > 0 ? (current / elapsedTime).toFixed(2) : '0.00';
 
   process.stdout.write(
-    `\r${msg}: ${current} / ${total}  (${colorLog.yellow(percentage)}%): 100%|${filledBar}${emptyBar}| Success: ${success}/${total} [${colorLog.red(elapsedTime.toFixed(2))}, ${itemsPerSecond}it/s]`
+    `\r${msg}: ${current} / ${total} (${colorLog.yellow(percentage)}%): 100%|${filledBar}${emptyBar}| Success: ${success}/${total} [${colorLog.red(elapsedTime.toFixed(2))}, ${itemsPerSecond}it/s]`
   );
 };
 
@@ -54,7 +54,12 @@ export const validateValue = (
   };
 
   const validImage = (val: Readonly<AxFieldValue>): boolean => {
-    if (typeof val !== 'object' || !('mimeType' in val) || !('data' in val)) {
+    if (
+      !val ||
+      typeof val !== 'object' ||
+      !('mimeType' in val) ||
+      !('data' in val)
+    ) {
       return false;
     }
     return true;
@@ -100,9 +105,7 @@ export const validateValue = (
 
   if (!isValid) {
     throw new Error(
-      `Validation failed: Expected '${field.name}' to be a ${ft.isArray ? 'an array of ' : ''}${
-        ft.name
-      } instead got '${value}'`
+      `Validation failed: Expected '${field.name}' to be a ${field.type?.isArray ? 'an array of ' : ''}${ft.name} instead got '${typeof value}' (${value})`
     );
   }
 };
diff --git a/src/ax/funcs/code.ts b/src/ax/funcs/code.ts
index 5891399..5b79ca9 100644
--- a/src/ax/funcs/code.ts
+++ b/src/ax/funcs/code.ts
@@ -68,7 +68,7 @@ export class AxJSInterpreter {
     return {
       name: 'javascriptInterpreter',
       description:
-        'Use this function to run Javascript code and get any expected return value.',
+        'Use this function to run Javascript code and get any expected return value',
       parameters: {
         type: 'object',
         properties: {
diff --git a/src/ax/prompts/agent.ts b/src/ax/prompts/agent.ts
index 826bafd..324624c 100644
--- a/src/ax/prompts/agent.ts
+++ b/src/ax/prompts/agent.ts
@@ -1,5 +1,5 @@
 import type { AxAIService, AxFunction } from '../ai/types.js';
-import { type AxGenerateOptions, AxSignature } from '../dsp/index.js';
+import { type AxGenOptions, AxSignature } from '../dsp/index.js';
 import {
   type AxGenIn,
   type AxGenOut,
@@ -19,16 +19,16 @@ export interface AxAgentic extends AxTunable, AxUsable {
   getFunction(): AxFunction;
 }
 
-export type AxAgentOptions = Omit<
-  AxGenerateOptions,
-  'functions' | 'functionCall'
->;
+export type AxAgentOptions = Omit<AxGenOptions, 'functions' | 'functionCall'>;
 
 export class AxAgent<IN extends AxGenIn, OUT extends AxGenOut>
   implements AxAgentic
 {
+  private ai?: AxAIService;
   private signature: AxSignature;
   private program: AxProgramWithSignature<IN, OUT>;
+  private functions?: AxFunction[];
+  private agents?: AxAgentic[];
 
   private name: string;
   private description: string;
@@ -36,14 +36,15 @@ export class AxAgent<IN extends AxGenIn, OUT extends AxGenOut>
   private func: AxFunction;
 
   constructor(
-    ai: AxAIService,
     {
+      ai,
       name,
       description,
       signature,
       agents,
       functions
     }: Readonly<{
+      ai?: Readonly<AxAIService>;
       name: string;
       description: string;
       signature: AxSignature | string;
@@ -52,22 +53,22 @@ export class AxAgent<IN extends AxGenIn, OUT extends AxGenOut>
     }>,
     options?: Readonly<AxAgentOptions>
   ) {
+    this.ai = ai;
     this.signature = new AxSignature(signature);
+    this.functions = functions;
+    this.agents = agents;
 
     const funcs: AxFunction[] = [
       ...(functions ?? []),
       ...(agents?.map((a) => a.getFunction()) ?? [])
     ];
 
-    const opt = {
-      ...options,
-      functions: funcs
-    };
+    const opt = { ...options, functions: funcs };
 
     this.program =
       funcs.length > 0
-        ? new AxReAct<IN, OUT>(ai, this.signature, opt)
-        : new AxChainOfThought<IN, OUT>(ai, this.signature, opt);
+        ? new AxReAct<IN, OUT>(this.signature, opt)
+        : new AxChainOfThought<IN, OUT>(this.signature, opt);
 
     if (!name || name.length < 5) {
       throw new Error(
@@ -128,18 +129,34 @@ export class AxAgent<IN extends AxGenIn, OUT extends AxGenOut>
 
   public getFunction(): AxFunction {
     const boundFunc = this.forward.bind(this);
+
+    // Create a wrapper function that excludes the 'ai' parameter
+    const wrappedFunc = (
+      values: IN,
+      options?: Readonly<AxProgramForwardOptions>
+    ) => {
+      const ai = this.ai ?? options?.ai;
+      if (!ai) {
+        throw new Error('AI service is required to run the agent');
+      }
+      return boundFunc(ai, values, options);
+    };
+
     return {
       ...this.func,
-      func: boundFunc
+      func: wrappedFunc
     };
   }
 
   public async forward(
+    ai: Readonly<AxAIService>,
     values: IN,
     options?: Readonly<AxProgramForwardOptions>
   ): Promise<OUT> {
+    const _ai = this.ai ?? ai;
+
     if (!options?.tracer) {
-      return await this.program.forward(values, options);
+      return await this.program.forward(_ai, values, options);
     }
 
     const attributes = {
@@ -155,7 +172,7 @@ export class AxAgent<IN extends AxGenIn, OUT extends AxGenOut>
         attributes
       },
       async (span) => {
-        const res = await this.program.forward(values, options);
+        const res = await this.program.forward(_ai, values, options);
         span.end();
         return res;
       }
diff --git a/src/ax/prompts/cot.ts b/src/ax/prompts/cot.ts
index 6c61ef5..2fb7141 100644
--- a/src/ax/prompts/cot.ts
+++ b/src/ax/prompts/cot.ts
@@ -1,16 +1,14 @@
-import type { AxAIService } from '../ai/types.js';
-import { AxGenerate, type AxGenerateOptions } from '../dsp/generate.js';
+import { AxGen, type AxGenOptions } from '../dsp/generate.js';
 import type { AxGenIn, AxGenOut } from '../dsp/program.js';
 import { AxSignature } from '../dsp/sig.js';
 
 export class AxChainOfThought<
   IN extends AxGenIn = AxGenIn,
   OUT extends AxGenOut = AxGenOut
-> extends AxGenerate<IN, OUT & { reason: string }> {
+> extends AxGen<IN, OUT & { reason: string }> {
   constructor(
-    ai: AxAIService,
     signature: Readonly<AxSignature | string>,
-    options?: Readonly<AxGenerateOptions>
+    options?: Readonly<AxGenOptions>
   ) {
     const sig = new AxSignature(signature);
     const description = `Let's work this out in a step by step way in order to ensure we have the right answer.`;
@@ -23,6 +21,6 @@ export class AxChainOfThought<
       ...sig.getOutputFields()
     ]);
 
-    super(ai, sig, options);
+    super(sig, options);
   }
 }
diff --git a/src/ax/prompts/prompts.test.ts b/src/ax/prompts/prompts.test.ts
index ed6e2fd..135f425 100644
--- a/src/ax/prompts/prompts.test.ts
+++ b/src/ax/prompts/prompts.test.ts
@@ -54,12 +54,11 @@ test('generate prompt', async (t) => {
   // const ai = new AxAI({ name: 'ollama', config: { model: 'nous-hermes2' } });
 
   const gen = new AxChainOfThought(
-    ai,
     `someText -> shortSummary "summarize in 5 to 10 words"`
   );
   gen.setExamples(examples);
 
-  const res = await gen.forward({ someText }, { stream: false });
+  const res = await gen.forward(ai, { someText });
 
   t.deepEqual(res, {
     reason: 'Blah blah blah',
diff --git a/src/ax/prompts/rag.ts b/src/ax/prompts/rag.ts
index 59cd4a9..cce0ce3 100644
--- a/src/ax/prompts/rag.ts
+++ b/src/ax/prompts/rag.ts
@@ -1,11 +1,11 @@
-import type { AxAIService } from '../ai/types.js';
 import {
-  AxGenerate,
-  type AxGenerateOptions,
+  AxGen,
+  type AxGenOptions,
   AxSignature,
   axStringUtil
 } from '../dsp/index.js';
 import { type AxProgramForwardOptions } from '../dsp/program.js';
+import type { AxAIService } from '../index.js';
 
 import { AxChainOfThought } from './cot.js';
 
@@ -13,7 +13,7 @@ export class AxRAG extends AxChainOfThought<
   { context: string[]; question: string },
   { answer: string }
 > {
-  private genQuery: AxGenerate<
+  private genQuery: AxGen<
     { context: string[]; question: string },
     { query: string }
   >;
@@ -21,28 +21,28 @@ export class AxRAG extends AxChainOfThought<
   private maxHops: number;
 
   constructor(
-    ai: AxAIService,
     queryFn: (query: string) => Promise<string>,
-    options: Readonly<AxGenerateOptions & { maxHops?: number }>
+    options: Readonly<AxGenOptions & { maxHops?: number }>
   ) {
     const sig =
       '"Answer questions with short factoid answers." context:string[] "may contain relevant facts", question -> answer';
-    super(ai, sig, options);
+    super(sig, options);
 
     this.maxHops = options?.maxHops ?? 3;
 
     const qsig = new AxSignature(
       '"Write a simple search query that will help answer a complex question." context?:string[] "may contain relevant facts", question -> query "question to further our understanding"'
     );
-    this.genQuery = new AxGenerate<
+    this.genQuery = new AxGen<
       { context: string[]; question: string },
       { query: string }
-    >(ai, qsig);
+    >(qsig);
     this.queryFn = queryFn;
     this.register(this.genQuery);
   }
 
   public override async forward(
+    ai: Readonly<AxAIService>,
     { question }: Readonly<{ question: string }>,
     options?: Readonly<AxProgramForwardOptions>
   ): Promise<{ answer: string; reason: string }> {
@@ -50,6 +50,7 @@ export class AxRAG extends AxChainOfThought<
 
     for (let i = 0; i < this.maxHops; i++) {
       const { query } = await this.genQuery.forward(
+        ai,
         {
           context,
           question
@@ -60,6 +61,6 @@ export class AxRAG extends AxChainOfThought<
       context = axStringUtil.dedup([...context, val]);
     }
 
-    return super.forward({ context, question }, options);
+    return super.forward(ai, { context, question }, options);
   }
 }
diff --git a/src/ax/prompts/react.ts b/src/ax/prompts/react.ts
index 01276fb..57db04e 100644
--- a/src/ax/prompts/react.ts
+++ b/src/ax/prompts/react.ts
@@ -1,5 +1,4 @@
-import type { AxAIService } from '../ai/index.js';
-import type { AxGenerateOptions } from '../dsp/generate.js';
+import type { AxGenOptions } from '../dsp/generate.js';
 import type { AxGenIn, AxGenOut } from '../dsp/program.js';
 import { AxSignature } from '../dsp/sig.js';
 
@@ -10,9 +9,8 @@ export class AxReAct<
   OUT extends AxGenOut = AxGenOut
 > extends AxChainOfThought<IN, OUT> {
   constructor(
-    ai: AxAIService,
     signature: Readonly<AxSignature | string>,
-    options: Readonly<AxGenerateOptions>
+    options: Readonly<AxGenOptions>
   ) {
     if (!options?.functions || options.functions.length === 0) {
       throw new Error('No functions provided');
@@ -43,6 +41,6 @@ export class AxReAct<
       sig.getOutputFields().map((v) => ({ ...v, isOptional: true }))
     );
 
-    super(ai, sig, options);
+    super(sig, options);
   }
 }
diff --git a/src/ax/util/apicall.ts b/src/ax/util/apicall.ts
index e493152..b67e1d3 100644
--- a/src/ax/util/apicall.ts
+++ b/src/ax/util/apicall.ts
@@ -6,8 +6,8 @@ import {
 
 import type { AxSpan } from '../trace/index.js';
 
+import { SSEParser } from './sse.js';
 import { TextDecoderStreamPolyfill } from './stream.js';
-import { JSONStringifyStream } from './transform.js';
 
 /**
  * Util: API details
@@ -75,7 +75,7 @@ export const apiCall = async <TRequest = unknown, TResponse = unknown>(
 
     const st = res.body
       .pipeThrough(new textDecoderStream())
-      .pipeThrough(new JSONStringifyStream<TResponse>());
+      .pipeThrough(new SSEParser<TResponse>());
 
     return st;
   } catch (e) {
diff --git a/src/ax/util/sse.ts b/src/ax/util/sse.ts
new file mode 100644
index 0000000..bf8d343
--- /dev/null
+++ b/src/ax/util/sse.ts
@@ -0,0 +1,113 @@
+import { TransformStream, TransformStreamDefaultController } from 'stream/web';
+
+interface CurrentEventState {
+  event?: string;
+  rawData: string;
+  id?: string;
+  retry?: number;
+}
+
+// eslint-disable-next-line @typescript-eslint/naming-convention
+export class SSEParser<T = unknown> extends TransformStream<string, T> {
+  private buffer: string = '';
+  private currentEvent: CurrentEventState = { rawData: '' };
+
+  constructor(private readonly dataParser: (data: string) => T = JSON.parse) {
+    super({
+      transform: (chunk, controller) => this.handleChunk(chunk, controller),
+      flush: (controller) => this.handleFlush(controller)
+    });
+  }
+
+  private handleChunk(
+    chunk: string,
+    controller: TransformStreamDefaultController<T>
+  ): void {
+    this.buffer += chunk;
+    this.processBuffer(controller);
+  }
+
+  private handleFlush(controller: TransformStreamDefaultController<T>): void {
+    this.processBuffer(controller);
+    if (this.currentEvent.rawData) {
+      this.emitEvent(controller);
+    }
+  }
+
+  private processBuffer(controller: TransformStreamDefaultController<T>): void {
+    const lines = this.buffer.split(/\r\n|\r|\n/);
+    this.buffer = lines.pop() || '';
+
+    for (const line of lines) {
+      if (line.trim() === '') {
+        this.emitEvent(controller);
+      } else {
+        this.parseLine(line);
+      }
+    }
+  }
+
+  private parseLine(line: string): void {
+    const colonIndex = line.indexOf(':');
+    if (colonIndex === -1) {
+      // If there's no colon, treat the whole line as data
+      this.currentEvent.rawData += this.currentEvent.rawData
+        ? '\n' + line.trim()
+        : line.trim();
+      return;
+    }
+
+    const field = line.slice(0, colonIndex).trim();
+    const value = line.slice(colonIndex + 1).trim();
+
+    switch (field) {
+      case 'event':
+        this.currentEvent.event = value;
+        break;
+      case 'data':
+        this.currentEvent.rawData += this.currentEvent.rawData
+          ? '\n' + value
+          : value;
+        break;
+      case 'id':
+        this.currentEvent.id = value;
+        break;
+      case 'retry': {
+        const retryValue = parseInt(value, 10);
+        if (!isNaN(retryValue)) {
+          this.currentEvent.retry = retryValue;
+        }
+        break;
+      }
+    }
+  }
+
+  private emitEvent(controller: TransformStreamDefaultController<T>): void {
+    if (this.currentEvent.rawData) {
+      // Check for special "[DONE]" message or other non-JSON data
+      if (
+        this.currentEvent.rawData.trim() === '[DONE]' ||
+        this.currentEvent.rawData.trim().startsWith('[')
+      ) {
+        return;
+      } else {
+        try {
+          // Attempt to parse the data using the provided dataParser
+          const parsedData: T = this.dataParser(this.currentEvent.rawData);
+          // Emit only if we successfully parsed the data
+          controller.enqueue(parsedData);
+        } catch (e) {
+          // If parsing fails, log the error without emitting
+          console.warn('Failed to parse event data:', e);
+          console.log(
+            'Raw data that failed to parse:',
+            this.currentEvent.rawData
+          );
+        }
+      }
+
+      // Reset the current event
+      this.currentEvent = { rawData: '' };
+    }
+  }
+}
diff --git a/src/ax/util/transform.ts b/src/ax/util/transform.ts
index 3126172..d02882f 100644
--- a/src/ax/util/transform.ts
+++ b/src/ax/util/transform.ts
@@ -4,32 +4,6 @@ import {
   type TransformStreamDefaultController
 } from 'stream/web';
 
-class JSONTransformer<O> implements Transformer<string, O> {
-  async transform(
-    obj: string,
-    controller: TransformStreamDefaultController<O>
-  ) {
-    (obj.split('\n') ?? [])
-      .map(extractJson)
-      .map<O>((v) => {
-        try {
-          return v && v.length > 0 ? JSON.parse(v as string) : null;
-        } catch (e: unknown) {
-          return null;
-        }
-      })
-      .filter((v) => v)
-      .forEach((v) => v && controller.enqueue(v));
-  }
-}
-
-// eslint-disable-next-line @typescript-eslint/naming-convention
-export class JSONStringifyStream<O> extends TransformStream<string, O> {
-  constructor() {
-    super(new JSONTransformer<O>());
-  }
-}
-
 class TypeTransformer<I, O> implements Transformer<I, O> {
   private buffer?: O[];
   private doneCallback?: (args0: readonly O[]) => Promise<void>;
@@ -67,14 +41,3 @@ export class RespTransformStream<I, O> extends TransformStream<I, O> {
     super(new TypeTransformer<I, O>(transformFn, doneCallback));
   }
 }
-
-const extractJson = (str: string): string | null => {
-  const startIndex = str.indexOf('{');
-  const endIndex = str.lastIndexOf('}');
-
-  if (startIndex !== -1 && endIndex !== -1) {
-    return str.substring(startIndex, endIndex + 1);
-  }
-
-  return null;
-};
diff --git a/src/docs/src/content/docs/guides/dsp.md b/src/docs/src/content/docs/guides/dsp.md
index 696df23..53d4521 100644
--- a/src/docs/src/content/docs/guides/dsp.md
+++ b/src/docs/src/content/docs/guides/dsp.md
@@ -22,7 +22,7 @@ There are various prompts available in Ax, pick one based on your needs.
 A signature defines the task you want to do, the inputs you’ll provide, and the outputs you expect the LLM to generate.
 
 ```typescript
-const prompt = new Generate(ai, 
+const prompt = new AxGen(
 `"Extract customer query details" customerMessage:string -> customerName, customerIssue, ,productName:string, troubleshootingAttempted?:string`)
 ```
 
@@ -62,7 +62,7 @@ You are now ready to use this prompt in your workflows.
 const ai = new AxAI("openai", { apiKey: process.env.OPENAI_APIKEY })
 
 # Execute the prompt
-const { customerName, productName, troubleshootingAttempted } = prompt.forward({ customerMessage })
+const { customerName, productName, troubleshootingAttempted } = prompt.forward(ai, { customerMessage })
 ```
 
 Easy enough! this is all you need
@@ -95,7 +95,7 @@ const examples = await hf.getRows<{ question: string; answer: string }>({
 
 ```typescript
 // Create your prompt
-const prompt = new Generate(ai, `question -> answer`)
+const prompt = new AxGen(`question -> answer`)
 ```
 
 ```typescript
diff --git a/src/docs/src/content/docs/guides/functions-1.md b/src/docs/src/content/docs/guides/functions-1.md
index 93e1103..e635f1b 100644
--- a/src/docs/src/content/docs/guides/functions-1.md
+++ b/src/docs/src/content/docs/guides/functions-1.md
@@ -78,11 +78,11 @@ class GoogleSearch {
 Just set the function on the prompt
 
 ```typescript
-const prompt = new Generate(ai, 'inputs -> output', { functions: [ googleSearch ] })
+const prompt = new AxGen('inputs -> output', { functions: [ googleSearch ] })
 ```
 
 Or in the case of function classes
 
 ```typescript
-const prompt = new Generate(ai, 'inputs -> output', { functions: [ new GoogleSearch(apiKey) ] })
+const prompt = new AxGen('inputs -> output', { functions: [ new GoogleSearch(apiKey) ] })
 ```
\ No newline at end of file
diff --git a/src/docs/src/content/docs/guides/functions-2.md b/src/docs/src/content/docs/guides/functions-2.md
index 6246bac..e5b4a8c 100644
--- a/src/docs/src/content/docs/guides/functions-2.md
+++ b/src/docs/src/content/docs/guides/functions-2.md
@@ -170,17 +170,16 @@ const ai = new Ax({
   apiKey: process.env.OPENAI_APIKEY as string
 });
 
-const agent = new AxAgent(ai, {
+const agent = new AxAgent({
   name: 'Restaurant search agent'
   description:
     'Search for restaurants to dine at based on the weather and food preferences',
   signature:
     `customerQuery:string  -> restaurant:string, priceRange:string "use $ signs to indicate price range"`
     functions,
-
 });
 
-const res = await agent.forward({ customerQuery });
+const res = await agent.forward(ai, { customerQuery });
 console.log(res);
 ```
 
diff --git a/src/docs/src/content/docs/guides/production-deploy.md b/src/docs/src/content/docs/guides/production-deploy.md
index 003de58..d850599 100644
--- a/src/docs/src/content/docs/guides/production-deploy.md
+++ b/src/docs/src/content/docs/guides/production-deploy.md
@@ -94,7 +94,7 @@ const app = new Hono();
 // Endpoint to get questions and return answers
 app.post('/get-answer', async (c) => {
   // RAG instance for processing queries using recursive answers generation
-  const rag = new AxRAG(ai, fetchFromVectorDB, { maxHops: 1 });
+  const rag = new AxRAG(fetchFromVectorDB, { maxHops: 1 });
 
   // Extract question from request body
   const question = c.req.body.question;
@@ -106,7 +106,7 @@ app.post('/get-answer', async (c) => {
 
   try {
     // Process the question using RAG
-    const answer = await rag.forward({ question });
+    const answer = await rag.forward(ai, { question });
 
     // Return the answer in JSON format
     return c.json({ answer });
diff --git a/src/docs/src/content/docs/start/quick.md b/src/docs/src/content/docs/start/quick.md
index 2b316ab..5b92fd5 100644
--- a/src/docs/src/content/docs/start/quick.md
+++ b/src/docs/src/content/docs/start/quick.md
@@ -80,10 +80,8 @@ const ai = new AxAI({
   apiKey: process.env.OPENAI_APIKEY as string
 });
 
-const prompt = `textToSummarize -> shortSummary "summarize in 5 to 10 words"`;
-
-const gen = new AxChainOfThought(ai, prompt);
-const res = await gen.forward({ textToSummarize });
+const gen = new AxChainOfThought(`textToSummarize -> shortSummary "summarize in 5 to 10 words"`);
+const res = await gen.forward(ai, { textToSummarize });
 
 console.log(res);
 ```
@@ -107,7 +105,7 @@ The Stock Analyst Agent is an advanced AI-powered tool that provides comprehensi
 This is only an example, but it highlights the power of agentic workflows, where you can build agents who work with agents to handle complex tasks.
 
 ```typescript title="Stock Analyst Agent"
-const agent = new AxAgent(ai, {
+const agent = new AxAgent({
   name: 'Stock Analyst',
   description:
     'An AI agent specialized in analyzing stocks, market trends, and providing financial insights.',
@@ -140,19 +138,19 @@ const agent = new AxAgent(ai, {
 ```typescript
 // ./src/examples/agent.ts
 
-const researcher = new AxAgent(ai, {
+const researcher = new AxAgent({
   name: 'researcher',
   description: 'Researcher agent',
   signature: `physicsQuestion "physics questions" -> answer "reply in bullet points"`
 });
 
-const summarizer = new AxAgent(ai, {
+const summarizer = new AxAgent({
   name: 'summarizer',
   description: 'Summarizer agent',
   signature: `text "text so summarize" -> shortSummary "summarize in 5 to 10 words"`
 });
 
-const agent = new AxAgent(ai, {
+const agent = new AxAgent({
   name: 'agent',
   description: 'A an agent to research complex topics',
   signature: `question -> answer`,
diff --git a/src/examples/agent.ts b/src/examples/agent.ts
index f1cedd8..d69c800 100644
--- a/src/examples/agent.ts
+++ b/src/examples/agent.ts
@@ -1,37 +1,20 @@
 import { AxAgent, AxAI } from '@ax-llm/ax';
 
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-});
-
-// const ai = new AxAI({
-//   name: 'google-gemini',
-//   apiKey: process.env.GOOGLE_APIKEY as string
-// });
-
-// const ai = new AxAI({
-//   name: 'groq',
-//   apiKey: process.env.GROQ_APIKEY as string
-// });
-
-// ai.setOptions({ debug: true });
-
-const researcher = new AxAgent(ai, {
+const researcher = new AxAgent({
   name: 'Physics Researcher',
   description:
     'Researcher for physics questions can answer questions about advanced physics',
   signature: `physicsQuestion "physics questions" -> answer "reply in bullet points"`
 });
 
-const summarizer = new AxAgent(ai, {
+const summarizer = new AxAgent({
   name: 'Science Summarizer',
   description:
     'Summarizer can write short summaries of advanced science topics',
   signature: `answer "bullet points to summarize" -> shortSummary "summarize in 10 to 20 words"`
 });
 
-const agent = new AxAgent(ai, {
+const agent = new AxAgent({
   name: 'Scientist',
   description: 'An agent that can answer advanced science questions',
   signature: `question -> answer`,
@@ -46,5 +29,22 @@ const question = `
   Include the summary and bullet points in the answer.
   Return this only if you can answer all the questions.`;
 
-const res = await agent.forward({ question });
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string
+});
+
+// const ai = new AxAI({
+//   name: 'google-gemini',
+//   apiKey: process.env.GOOGLE_APIKEY as string
+// });
+
+// const ai = new AxAI({
+//   name: 'groq',
+//   apiKey: process.env.GROQ_APIKEY as string
+// });
+
+// ai.setOptions({ debug: true });
+
+const res = await agent.forward(ai, { question });
 console.log('>', res);
diff --git a/src/examples/balancer.ts b/src/examples/balancer.ts
index 1ed0729..a63c0c3 100644
--- a/src/examples/balancer.ts
+++ b/src/examples/balancer.ts
@@ -29,12 +29,10 @@ const ai2 = new AxAI({
   }
 });
 
-const ai = new AxBalancer([ai1, ai2]);
-
 const gen = new AxChainOfThought(
-  ai,
   `textToSummarize -> shortSummary "summarize in 5 to 10 words"`
 );
 
-const res = await gen.forward({ textToSummarize }, { model: 'chill' });
+const ai = new AxBalancer([ai1, ai2]);
+const res = await gen.forward(ai, { textToSummarize }, { model: 'chill' });
 console.log('>', res);
diff --git a/src/examples/chain-of-thought.ts b/src/examples/chain-of-thought.ts
index 2f5b0cd..3a41ce3 100644
--- a/src/examples/chain-of-thought.ts
+++ b/src/examples/chain-of-thought.ts
@@ -1,14 +1,6 @@
 import { AxAI, AxChainOfThought } from '@ax-llm/ax';
 
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-});
-
-ai.setOptions({ debug: true });
-
 const cot = new AxChainOfThought(
-  ai,
   `
   context:string[] "Information to answer the question",
   question:string
@@ -24,5 +16,11 @@ const values = {
   ]
 };
 
-const res = await cot.forward(values);
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string
+});
+ai.setOptions({ debug: true });
+
+const res = await cot.forward(ai, values);
 console.log(res);
diff --git a/src/examples/customer-support.ts b/src/examples/customer-support.ts
index 658459c..645337f 100644
--- a/src/examples/customer-support.ts
+++ b/src/examples/customer-support.ts
@@ -1,12 +1,6 @@
-import { AxAI, AxGenerate } from '@ax-llm/ax';
+import { AxAI, AxGen } from '@ax-llm/ax';
 
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-});
-
-const gen = new AxGenerate(
-  ai,
+const gen = new AxGen(
   `customerEmail:string  -> productName:string "The name of the product",
 issueDescription:string "A description of the issue",
 issueSummary:string "A summary of the issue",
@@ -27,4 +21,9 @@ Best regards,
 John Doe.
   `;
 
-console.log(await gen.forward({ customerEmail: customerMessage }));
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string
+});
+
+console.log(await gen.forward(ai, { customerEmail: customerMessage }));
diff --git a/src/examples/docker.ts b/src/examples/docker.ts
index a625397..eab5f87 100644
--- a/src/examples/docker.ts
+++ b/src/examples/docker.ts
@@ -1,4 +1,4 @@
-import { AxAI, AxDockerSession, AxGenerate } from '@ax-llm/ax';
+import { AxAI, AxDockerSession, AxGen } from '@ax-llm/ax';
 
 // Initialize Docker session
 const dockerSession = new AxDockerSession();
@@ -9,22 +9,21 @@ await dockerSession.findOrCreateContainer({
   tag: 'ax:example'
 });
 
-// Initialize the AI instance with your API key
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-});
-
 // Define the task for generating a command sequence
-const prompt = new AxGenerate(
-  ai,
+const prompt = new AxGen(
   `"Find requested file and display top 3 lines of its content and a hash of the file."
   fileQuery:string -> content:string, hash:string`,
   { functions: [dockerSession] }
 );
 
+// Initialize the AI instance with your API key
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string
+});
+
 // Execute the task
-const res = await prompt.forward({
+const res = await prompt.forward(ai, {
   fileQuery: 'config file for current shell'
 });
 
diff --git a/src/examples/dope-or-nope.ts b/src/examples/dope-or-nope.ts
index 6bf740c..0f538df 100644
--- a/src/examples/dope-or-nope.ts
+++ b/src/examples/dope-or-nope.ts
@@ -29,7 +29,6 @@ const ai = new AxAI({
 
 // Setup the program to tune
 const program = new AxChainOfThought<{ question: string }, { answer: string }>(
-  ai,
   `question -> answer:number "numerical rating from 1 to 4"`
 );
 
@@ -40,6 +39,7 @@ const optimize = new AxBootstrapFewShot<
   { question: string },
   { answer: string }
 >({
+  ai,
   program,
   examples
 });
diff --git a/src/examples/extract.ts b/src/examples/extract.ts
index 2aaeb20..99038f7 100644
--- a/src/examples/extract.ts
+++ b/src/examples/extract.ts
@@ -1,4 +1,4 @@
-import { AxAI, AxAIOpenAIModel, AxGenerate } from '@ax-llm/ax';
+import { AxAI, AxAIOpenAIModel, AxGen } from '@ax-llm/ax';
 
 const chatMessage = `Hello Mike, How are you set for a call tomorrow or Friday? I have a few things to discuss with you. Also the ticket number is 300. Let me know what time works best for you. Thanks!`;
 
@@ -15,14 +15,10 @@ const ai = new AxAI({
 });
 ai.setOptions({ debug: true });
 
-const gen = new AxGenerate(
-  ai,
+const gen = new AxGen(
   `chatMessage, currentDate:datetime -> subject, foundMeeting:boolean, ticketNumber?:number, customerNumber?:number, datesMentioned:datetime[], messageType:class "reminder, follow-up, meeting, other"`
 );
 
-const res = await gen.forward(
-  { chatMessage, currentDate },
-  { modelConfig: { stream: true } }
-);
+const res = await gen.forward(ai, { chatMessage, currentDate });
 
 console.log('>', res);
diff --git a/src/examples/fibonacci.ts b/src/examples/fibonacci.ts
index 32ba4bc..68a472e 100644
--- a/src/examples/fibonacci.ts
+++ b/src/examples/fibonacci.ts
@@ -4,16 +4,16 @@ const sig = new AxSignature(
   `numberSeriesTask:string  -> fibonacciSeries:number[]`
 );
 
+const gen = new AxReAct(sig, {
+  functions: [new AxJSInterpreter()]
+});
+
 const ai = new AxAI({
   name: 'openai',
   apiKey: process.env.OPENAI_APIKEY as string
 });
 
-const gen = new AxReAct(ai, sig, {
-  functions: [new AxJSInterpreter()]
-});
-
-const res = await gen.forward({
+const res = await gen.forward(ai, {
   numberSeriesTask: 'Use code to calculate the fibonacci series of 10'
 });
 
diff --git a/src/examples/food-search.ts b/src/examples/food-search.ts
index 7dd0ebe..007713c 100644
--- a/src/examples/food-search.ts
+++ b/src/examples/food-search.ts
@@ -150,8 +150,11 @@ const signature = new AxSignature(
 const ai = new AxAI({
   name: 'openai',
   apiKey: process.env.OPENAI_APIKEY as string
+  //   config: { model: AxAIOpenAIModel.GPT4OMini }
 });
 
+// ai.setOptions({ debug: true });
+
 // const ai = new AxAI({
 //   name: 'groq',
 //   apiKey: process.env.GROQ_APIKEY as string
@@ -174,7 +177,7 @@ const ai = new AxAI({
 
 // ai.setOptions({ debug: true });
 
-const gen = new AxAgent(ai, {
+const gen = new AxAgent({
   name: 'food-search',
   description:
     'Use this agent to find restaurants based on what the customer wants',
@@ -182,6 +185,6 @@ const gen = new AxAgent(ai, {
   functions
 });
 
-const res = await gen.forward({ customerQuery }, { stream: false });
+const res = await gen.forward(ai, { customerQuery }, { stream: true });
 
 console.log('>', res);
diff --git a/src/examples/marketing.ts b/src/examples/marketing.ts
index 18d8519..5c509d1 100644
--- a/src/examples/marketing.ts
+++ b/src/examples/marketing.ts
@@ -1,9 +1,4 @@
-import { AxAI, AxGenerate } from '@ax-llm/ax';
-
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-});
+import { AxAI, AxGen } from '@ax-llm/ax';
 
 const product = {
   name: 'Acme Toilet Cleaning',
@@ -22,12 +17,16 @@ const messageGuidelines = [
   'Employs emojis and friendly language'
 ];
 
-const gen = new AxGenerate(
-  ai,
+const gen = new AxGen(
   `productName, productDescription, toName, toDescription, messageGuidelines -> message`
 );
 
-const res = await gen.forward({
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string
+});
+
+const res = await gen.forward(ai, {
   productName: product.name,
   productDescription: product.description,
   toName: to.name,
diff --git a/src/examples/multi-modal.ts b/src/examples/multi-modal.ts
index c9bac0c..1939a11 100644
--- a/src/examples/multi-modal.ts
+++ b/src/examples/multi-modal.ts
@@ -2,19 +2,19 @@ import fs from 'node:fs';
 
 import { AxAI, AxAIOpenAIModel, AxChainOfThought } from '@ax-llm/ax';
 
+const gen = new AxChainOfThought(`question, animalImage:image -> answer`);
+
+const image = fs
+  .readFileSync('./src/examples/assets/kitten.jpeg')
+  .toString('base64');
+
 const ai = new AxAI({
   name: 'openai',
   apiKey: process.env.OPENAI_APIKEY as string,
   config: { model: AxAIOpenAIModel.GPT4O }
 });
 
-const gen = new AxChainOfThought(ai, `question, animalImage:image -> answer`);
-
-const image = fs
-  .readFileSync('./src/examples/assets/kitten.jpeg')
-  .toString('base64');
-
-const res = await gen.forward({
+const res = await gen.forward(ai, {
   question: 'What family does this animal belong to?',
   animalImage: { mimeType: 'image/jpeg', data: image }
 });
diff --git a/src/examples/qna-tune.ts b/src/examples/qna-tune.ts
index 32f0fd5..ba1bc27 100644
--- a/src/examples/qna-tune.ts
+++ b/src/examples/qna-tune.ts
@@ -26,31 +26,30 @@ const examples = await hf.getRows<{ question: string; answer: string }>({
   renameMap: { query: 'question', answer: 'answer' }
 });
 
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string,
-  config: { model: AxAIOpenAIModel.GPT4O, maxTokens: 3000 }
-});
-
-ai.setOptions({ debug: true });
-
 const fetchFromVectorDB = async (query: string) => {
   const cot = new AxChainOfThought<{ query: string }, { answer: string }>(
-    ai,
     'query -> answer:string "answer to the query"'
   );
-  const { answer } = await cot.forward({ query });
+  const { answer } = await cot.forward(ai, { query });
   return answer;
 };
 
 // Setup the program to tune
-const program = new AxRAG(ai, fetchFromVectorDB, { maxHops: 1 });
+const program = new AxRAG(fetchFromVectorDB, { maxHops: 1 });
+
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string,
+  config: { model: AxAIOpenAIModel.GPT4O, maxTokens: 3000 }
+});
+ai.setOptions({ debug: true });
 
 // Setup a Bootstrap Few Shot optimizer to tune the above program
 const optimize = new AxBootstrapFewShot<
   { question: string },
   { answer: string }
 >({
+  ai,
   program,
   examples
 });
diff --git a/src/examples/qna-use-tuned.ts b/src/examples/qna-use-tuned.ts
index 29e3599..da6fc6d 100644
--- a/src/examples/qna-use-tuned.ts
+++ b/src/examples/qna-use-tuned.ts
@@ -9,13 +9,7 @@ import {
   AxTestPrompt
 } from '@ax-llm/ax';
 
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-});
-
 const program = new AxChainOfThought<{ question: string }, { answer: string }>(
-  ai,
   `question -> answer "in short 2 or 3 words"`
 );
 
@@ -55,5 +49,10 @@ const metricFn: AxMetricFn = ({ prediction, example }) => {
   );
 };
 
-const ev = new AxTestPrompt({ program, examples });
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string
+});
+
+const ev = new AxTestPrompt({ ai, program, examples });
 await ev.run(metricFn);
diff --git a/src/examples/rag.ts b/src/examples/rag.ts
index 66a81f9..1f2c435 100644
--- a/src/examples/rag.ts
+++ b/src/examples/rag.ts
@@ -3,21 +3,20 @@ import { AxAI, AxChainOfThought, AxRAG } from '@ax-llm/ax';
 // simulated vector db call using an llm
 const fetchFromVectorDB = async (query: string) => {
   const cot = new AxChainOfThought<{ query: string }, { answer: string }>(
-    ai,
     'query -> answer'
   );
-  const { answer } = await cot.forward({ query });
+  const { answer } = await cot.forward(ai, { query });
   return answer;
 };
 
+const rag = new AxRAG(fetchFromVectorDB, { maxHops: 3 });
+
 const ai = new AxAI({
   name: 'openai',
   apiKey: process.env.OPENAI_APIKEY as string
 });
 
-const rag = new AxRAG(ai, fetchFromVectorDB, { maxHops: 3 });
-
-const res = await rag.forward({
+const res = await rag.forward(ai, {
   question: 'List 3 of the top most important work done by Michael Stonebraker?'
 });
 
diff --git a/src/examples/react.ts b/src/examples/react.ts
index 646c49a..bb4137f 100644
--- a/src/examples/react.ts
+++ b/src/examples/react.ts
@@ -31,14 +31,14 @@ const functions = [
   }
 ];
 
+const cot = new AxReAct(`question:string -> answer:string`, {
+  functions
+});
+
 const ai = new AxAI({
   name: 'openai',
   apiKey: process.env.OPENAI_APIKEY as string
 });
 
-const cot = new AxReAct(ai, `question:string -> answer:string`, {
-  functions
-});
-
-const res = await cot.forward(values);
+const res = await cot.forward(ai, values);
 console.log(res);
diff --git a/src/examples/smart-home.ts b/src/examples/smart-home.ts
index 7ec6ac3..8605b54 100644
--- a/src/examples/smart-home.ts
+++ b/src/examples/smart-home.ts
@@ -35,7 +35,7 @@ const ai = new AxAI({
   apiKey: process.env.OPENAI_APIKEY as string
 });
 
-const agent = new AxAgent(ai, {
+const agent = new AxAgent({
   name: 'lares',
   description: 'Lares smart home assistant',
   signature: `instruction -> room:string "the room where the dog is found"`,
@@ -50,7 +50,10 @@ const agent = new AxAgent(ai, {
         },
         required: ['room']
       } as AxFunctionJSONSchema,
-      func: async (args: Readonly<{ room: string }>) => {
+      func: async (args) => {
+        if (!args?.room) {
+          throw new Error('Missing required parameter: room');
+        }
         const roomState = state.rooms[args.room];
         if (roomState) {
           roomState.light = !roomState.light;
@@ -120,5 +123,5 @@ const instruction = `
     The initial state is: ${JSON.stringify({ ...state, dogLocation: 'unknown' })}.
   `;
 
-const res = await agent.forward({ instruction });
+const res = await agent.forward(ai, { instruction });
 console.log('Response:', res);
diff --git a/src/examples/streaming1.ts b/src/examples/streaming1.ts
index 459aae1..499eb97 100644
--- a/src/examples/streaming1.ts
+++ b/src/examples/streaming1.ts
@@ -5,14 +5,8 @@ import { AxAI, AxChainOfThought } from '@ax-llm/ax';
 //   apiKey: process.env.OPENAI_APIKEY as string
 // });
 
-const ai = new AxAI({
-  name: 'google-gemini',
-  apiKey: process.env.GOOGLE_APIKEY as string
-});
-
 // setup the prompt program
 const gen = new AxChainOfThought(
-  ai,
   `startNumber:number -> next10Numbers:number[]`
 );
 
@@ -21,10 +15,12 @@ gen.addAssert(({ next10Numbers }: Readonly<{ next10Numbers: number[] }>) => {
   return next10Numbers ? !next10Numbers.includes(5) : undefined;
 }, 'Numbers 5 is not allowed');
 
+const ai = new AxAI({
+  name: 'google-gemini',
+  apiKey: process.env.GOOGLE_APIKEY as string
+});
+
 // run the program with streaming enabled
-const res = await gen.forward(
-  { startNumber: 1 },
-  { stream: true, debug: true }
-);
+const res = await gen.forward(ai, { startNumber: 1 }, { debug: true });
 
 console.log('>', res);
diff --git a/src/examples/streaming2.ts b/src/examples/streaming2.ts
index 4218969..1ff165d 100644
--- a/src/examples/streaming2.ts
+++ b/src/examples/streaming2.ts
@@ -1,20 +1,12 @@
 import { AxAI, AxChainOfThought } from '@ax-llm/ax';
 
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-});
-
 // const ai = new AxAI({
 //     name: 'anthropic'
 //     apiKey: process.env.ANTHROPIC_APIKEY as string
 // });
 
 // setup the prompt program
-const gen = new AxChainOfThought(
-  ai,
-  `question:string -> answerInPoints:string`
-);
+const gen = new AxChainOfThought(`question:string -> answerInPoints:string`);
 
 // add a assertion to ensure all lines start with a number and a dot.
 gen.addStreamingAssert(
@@ -33,12 +25,14 @@ gen.addStreamingAssert(
   'Lines must start with a number and a dot. Eg: 1. This is a line.'
 );
 
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string
+});
+
 // run the program with streaming enabled
-const res = await gen.forward(
-  {
-    question: 'Provide a list of 3 optimizations to speedup LLM inference.'
-  },
-  { stream: true }
-);
+const res = await gen.forward(ai, {
+  question: 'Provide a list of 3 optimizations to speedup LLM inference.'
+});
 
 console.log('>', res);
diff --git a/src/examples/summarize.ts b/src/examples/summarize.ts
index 0956f7e..43d8712 100644
--- a/src/examples/summarize.ts
+++ b/src/examples/summarize.ts
@@ -2,16 +2,6 @@ import { AxAI, AxAIOpenAIModel, AxChainOfThought } from '@ax-llm/ax';
 
 const noteText = `The technological singularity—or simply the singularity[1]—is a hypothetical future point in time at which technological growth becomes uncontrollable and irreversible, resulting in unforeseeable changes to human civilization.[2][3] According to the most popular version of the singularity hypothesis, I.J. Good's intelligence explosion model, an upgradable intelligent agent will eventually enter a "runaway reaction" of self-improvement cycles, each new and more intelligent generation appearing more and more rapidly, causing an "explosion" in intelligence and resulting in a powerful superintelligence that qualitatively far surpasses all human intelligence.[4]`;
 
-// Example with OpenAI using custom labels in place of model names
-const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string,
-  config: { model: 'model-a' },
-  modelMap: {
-    'model-a': AxAIOpenAIModel.GPT4OMini
-  }
-});
-
 // const ai = new AxAI({ name: 'ollama', model: 'nous-hermes2' });
 
 const gen = new AxChainOfThought(
@@ -38,6 +28,20 @@ gen.addAssert(({ reason }: Readonly<{ reason: string }>) => {
   return !reason.includes('goat');
 }, 'Reason should not contain "the"');
 
-const res = await gen.forward({ noteText }, { modelConfig: { stream: true } });
+// Example with OpenAI using custom labels in place of model names
+const ai = new AxAI({
+  name: 'openai',
+  apiKey: process.env.OPENAI_APIKEY as string,
+  config: { model: 'model-a' },
+  modelMap: {
+    'model-a': AxAIOpenAIModel.GPT4OMini
+  }
+});
+
+const res = await gen.forward(
+  ai,
+  { noteText },
+  { modelConfig: { stream: true } }
+);
 
 console.log('>', res);
diff --git a/src/web-api/package.json b/src/web-api/package.json
index 7864302..ae54b35 100644
--- a/src/web-api/package.json
+++ b/src/web-api/package.json
@@ -6,7 +6,7 @@
     "lint": "npx eslint src --ext .ts,.tsx",
     "lint:fix": "npx eslint src --ext .ts,.tsx --fix",
     "build": "rm -rf dist/* && npx tsc && npx tsc-alias && npm run build:web",
-    "build:web": "(cd ../web-ui && npm run build && cp -r dist/ ../api/dist/public)",
+    "build:web": "(cd ../web-ui && npm run build && cp -r dist/ ../web-api/dist/public)",
     "start": "npm run tsx ./dist/index.js",
     "dev": "npm run tsx -- --watch ./src/index.ts",
     "tsx": "node --env-file=.env --import=tsx --no-warnings --trace-warnings  --trace-deprecation"
diff --git a/src/web-api/src/api/util.ts b/src/web-api/src/api/util.ts
index 2e1f300..a10f0a8 100644
--- a/src/web-api/src/api/util.ts
+++ b/src/web-api/src/api/util.ts
@@ -1,5 +1,3 @@
-import type { HandlerContext } from '@/util';
-
 import { ObjectId } from 'mongodb';
 import sharp from 'sharp';
 
@@ -79,30 +77,3 @@ export const getFiles = async (form: FormData): Promise<GetFilesResult> => {
 
   return { docs, images };
 };
-
-export const createAI = async (hc: Readonly<HandlerContext>) => {
-  let args: AxAIArgs | undefined;
-
-  if (aiType === 'big') {
-    const apiKey = (await decryptKey(hc, agent.aiBigModel.apiKey)) ?? '';
-    args = {
-      apiKey,
-      config: { model: agent.aiBigModel.model },
-      name: agent.aiBigModel.id
-    } as AxAIArgs;
-  }
-
-  if (aiType === 'small') {
-    const apiKey = (await decryptKey(hc, agent.aiSmallModel.apiKey)) ?? '';
-    args = {
-      apiKey,
-      config: { model: agent.aiSmallModel.model },
-      name: agent.aiSmallModel.id
-    } as AxAIArgs;
-  }
-  if (!args) {
-    throw new Error('Invalid AI type: ' + aiType);
-  }
-
-  return new AxAI(args);
-};