fix: prompt improvements

Author: dosco

src/ax/dsp/prompt.ts

@@ -35,7 +35,7 @@ export class AxPromptTemplate {
     const inArgs = this.renderDescFields(this.sig.getInputFields())
     const outArgs = this.renderDescFields(this.sig.getOutputFields())
     const task = [
-      `#Task\nGiven the fields ${inArgs}, produce the fields ${outArgs}.`,
+      `You will be provided with the following fields: ${inArgs}. Your task is to generate two new fields: ${outArgs}.`,
     ]
 
     const fnNames = functions?.map((f) => {
@@ -45,40 +45,36 @@ export class AxPromptTemplate {
       return f.name
     })
 
-    const funcList = fnNames?.map((fname) => `'${fname}'`).join(', ')
+    const funcList = fnNames?.map((fname) => `\`${fname}\``).join(', ')
 
     if (funcList && funcList.length > 0) {
       task.push(
-        `Use the following functions ${funcList} to complete the task. The functions must be used to resolve the output field values.`
+        `Complete the task, using the following functions as needed: ${funcList}. Refer to the function descriptions for proper usage. The output field values may be generated by applying these functions if appropriate for the task.`
       )
     }
 
     const desc = this.sig.getDescription()
     if (desc) {
-      task.push(desc.endsWith('.') ? desc : desc + '.')
+      task.push(capitalizeFirstLetter(desc.endsWith('.') ? desc : desc + '.'))
     }
 
-    task.push(
-      'Ensure the output strictly follows a plain text format, `key: value` separated by a new line.'
-    )
-
     const inputFields = this.renderFields(this.sig.getInputFields())
+    task.push(`## Input Fields\n${inputFields}`)
+
     const outputFields = this.renderFields(this.sig.getOutputFields())
+    task.push(`## Output Fields\n${outputFields}`)
 
     task.push(
-      [
-        '\n',
-        '## Input Fields',
-        inputFields,
-        '\n',
-        '## Output Fields',
-        outputFields,
-      ].join('\n')
+      'Output must be in plain text, with each `key: value` pair on a new line. The format of each `value` should strictly adhere to the formatting instructions for its corresponding `key`, as defined earlier in this prompt.'
+    )
+
+    task.push(
+      'No additional text or formatting is permitted. The specific formatting rules for each key are provided above; ensure the corresponding values follow those rules precisely.'
     )
 
     this.task = {
       type: 'text' as const,
-      text: task.join(' '),
+      text: task.join('\n\n'),
     }
   }
 
@@ -95,7 +91,7 @@ export class AxPromptTemplate {
   ): AxChatRequest['chatPrompt'] => {
     const renderedExamples = examples
       ? [
-          { type: 'text' as const, text: 'Examples:\n' },
+          { type: 'text' as const, text: '\n## Examples:\n' },
           ...this.renderExamples(examples),
         ]
       : []
@@ -400,22 +396,20 @@ export class AxPromptTemplate {
     list.map((v) => `\`${v.title}\``).join(', ')
 
   private renderFields = (fields: readonly AxField[]) => {
-    // Header
-    const header = 'Field Name | Field Type | Required/Optional | Description'
-    const separator = '|'
-
     // Transform each field into table row
     const rows = fields.map((field) => {
       const name = field.title
       const type = field.type?.name ? toFieldType(field.type) : 'string'
       const required = field.isOptional ? 'optional' : 'required'
-      const description = field.description ?? ''
+      const description = field.description
+        ? `: ${capitalizeFirstLetter(field.description)}`
+        : ''
 
-      return [name, type, required, description].join(` ${separator} `).trim()
+      // Eg. - `Conversation` (string, optional): The conversation context.
+      return `- \`${name}\` (${type}, ${required})${description}`.trim()
     })
 
-    // Combine header and rows
-    return [header, ...rows].join('\n')
+    return rows.join('\n')
   }
 }
 
@@ -519,3 +513,10 @@ const isEmptyValue = (
   }
   return false
 }
+
+function capitalizeFirstLetter(str: string) {
+  if (str.length === 0) {
+    return ''
+  }
+  return `${str.charAt(0).toUpperCase()}${str.slice(1)}`
+}

src/ax/package.json

@@ -25,7 +25,7 @@
     "fix": "run-s fix:*",
     "fix:lint": "eslint --fix",
     "fix:format": "prettier --write \"**/*.{ts,json,md}\"",
-    "doc:build:markdown": "typedoc",
+    "doc:build:markdown": "typedoc --readme none",
     "coverage": "c8 ava",
     "prepare": "husky install",
     "tsx": "node --env-file=.env --import=tsx",

src/docs/README.md

@@ -7,7 +7,7 @@ Building intelligent agents is a breeze with the Ax framework, inspired by the p
 
 Large language models (LLMs) are becoming really powerful and have reached a point where they can work as the backend for your entire product. However, there's still a lot of complexity to manage from using the correct prompts, models, streaming, function calls, error correction, and much more. We aim help manage this complexity via this easy-to-use library that can work with all state-of-the-art LLMs. Additionally, we are using the latest research to add new capabilities like DSPy to the library.
 
-## Install
+### Install
 
 With NPM
 
@@ -37,7 +37,7 @@ yarn add @ax-llm/ax
 - Lite weight, zero-dependencies
 
 
-## Quick Start
+### Quick Start
 
 1. Pick an AI to work with
 
@@ -69,6 +69,6 @@ const res = await ai.chat([
 ]);
 ```
 
-## Reach out
+### Reach out
 
 https://twitter.com/dosco

src/docs/src/content/docs/01-start/01-about.md

@@ -9,7 +9,7 @@ Efficient type-safe prompts are auto-generated from a simple signature. A prompt
 
 You can have multiple input and output fields, and each field can be of the types `string`, `number`, `boolean`, `date`, `datetime`, `class "class1, class2"`, `JSON`, or an array of any of these, e.g., `string[]`. When a type is not defined, it defaults to `string`. The underlying AI is encouraged to generate the correct JSON when the `JSON` type is used.
 
-## Output Field Types
+### Output Field Types
 
 | Type                      | Description                       | Usage                      | Example Output                                     |
 |---------------------------|-----------------------------------|----------------------------|----------------------------------------------------|

src/docs/src/content/docs/01-start/02-signatures.md

@@ -4,7 +4,7 @@ description: A more detailed guide to using Ax
 ---
 
 
-## Pick an LLM
+### Pick an LLM
 
 Ax is a zero-dependency framework. Every LLM API integration we build is solid, works well with Ax, and supports all required features, such as function calling, multi-modal, JSON, streaming, etc.
 
@@ -19,7 +19,7 @@ const ai = new AxAI({
 
 The LLMs are pre-configured with sensible defaults such as models and other conifgurations such as topK, temperature, etc
 
-## Prompting
+### Prompting
 
 Prompts are usually stressful and complex. You never know what the right prompt is, and blobs of text in your code are hard to deal with. We fix this by adopting the prompt signatures from the popular Stanford DSPy paper.
 
@@ -46,7 +46,7 @@ question:string, animalImage:image -> answer:string
 question:string, answers:string[] -> rating:number[]
 ```
 
-## Putting it all together
+### Putting it all together
 
 Use the above AI and a prompt to build an LLM-powered program to summarize the text.
 
@@ -80,7 +80,7 @@ tsx example.ts
 }
 ```
 
-## Build your first agent
+### Build your first agent
 
 Ax makes it really simple to build agents. An agent requires a `name`, `description` and `signature`. it can optionally use `functions` and other `agents`.
 

src/docs/src/content/docs/01-start/04-quick.md

@@ -0,0 +1,70 @@
+---
+title: RAG & Vector DBs
+description: A guide on working with vector databases and Retrieval Augmented Generation (RAG) in ax.
+---
+
+Vector databases are critical to building LLM workflows. We have clean abstractions over popular vector databases and our own quick in-memory vector database.
+
+| Provider   | Tested  |
+| ---------- | ------- |
+| In Memory  | 🟢 100% |
+| Weaviate   | 🟢 100% |
+| Cloudflare | 🟡 50%  |
+| Pinecone   | 🟡 50%  |
+
+```typescript
+// Create embeddings from text using an LLM
+const ret = await this.ai.embed({ texts: 'hello world' });
+
+// Create an in memory vector db
+const db = new axDB('memory');
+
+// Insert into vector db
+await this.db.upsert({
+  id: 'abc',
+  table: 'products',
+  values: ret.embeddings[0]
+});
+
+// Query for similar entries using embeddings
+const matches = await this.db.query({
+  table: 'products',
+  values: embeddings[0]
+});
+```
+
+Alternatively you can use the `AxDBManager` which handles smart chunking, embedding and querying everything
+for you, it makes things almost too easy.
+
+```typescript
+const manager = new AxDBManager({ ai, db });
+await manager.insert(text);
+
+const matches = await manager.query(
+  'John von Neumann on human intelligence and singularity.'
+);
+console.log(matches);
+```
+
+### RAG Documents
+
+Using documents like PDF, DOCX, PPT, XLS, etc., with LLMs is a huge pain. We make it easy with Apache Tika, an open-source document processing engine.
+
+Launch Apache Tika
+
+```shell
+docker run -p 9998:9998 apache/tika
+```
+
+Convert documents to text and embed them for retrieval using the `AxDBManager`, which also supports a reranker and query rewriter. Two default implementations, `AxDefaultResultReranker` and `AxDefaultQueryRewriter`, are available.
+
+```typescript
+const tika = new AxApacheTika();
+const text = await tika.convert('/path/to/document.pdf');
+
+const manager = new AxDBManager({ ai, db });
+await manager.insert(text);
+
+const matches = await manager.query('Find some text');
+console.log(matches);
+```

src/docs/src/content/docs/01-start/05-rag.md

@@ -0,0 +1,34 @@
+---
+title: Multi-modal DSPy
+description: Using multi-modal inputs like images and audio with DSPy pipelines and LLMs
+---
+
+When using models like `GPT-4o` and `Gemini` that support multi-modal prompts, we support using image fields, and this works with the whole DSP pipeline.
+
+```typescript
+const image = fs
+  .readFileSync('./src/examples/assets/kitten.jpeg')
+  .toString('base64');
+
+const gen = new AxChainOfThought(`question, animalImage:image -> answer`);
+
+const res = await gen.forward(ai, {
+  question: 'What family does this animal belong to?',
+  animalImage: { mimeType: 'image/jpeg', data: image }
+});
+```
+
+When using models like `gpt-4o-audio-preview` that support multi-modal prompts with audio support, we support using audio fields, and this works with the whole DSP pipeline.
+
+```typescript
+const audio = fs
+  .readFileSync('./src/examples/assets/comment.wav')
+  .toString('base64');
+
+const gen = new AxGen(`question, commentAudio:audio -> answer`);
+
+const res = await gen.forward(ai, {
+  question: 'What family does this animal belong to?',
+  commentAudio: { format: 'wav', data: audio }
+});
+```