feat: new api allows for more flexibility

README.md

@@ -1,16 +1,17 @@
 # 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)
 [![Twitter](https://img.shields.io/twitter/follow/dosco?style=for-the-badge&color=red)](https://twitter.com/dosco)
 
 ![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 })

src/ax/ai/types.ts

@@ -174,6 +174,7 @@ export type AxAIServiceOptions = {
 };
 
 export type AxAIServiceActionOptions = {
+  ai?: Readonly<AxAIService>;
   sessionId?: string;
   traceId?: string;
 };

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)
         });

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);

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);
   }
 }

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++;

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;