@botpress/zai 1.0.1 → 1.1.0

package/src/adapters/botpress-table.ts

@@ -0,0 +1,210 @@
+import { type Client } from '@botpress/client'
+import { z } from '@bpinternal/zui'
+
+import { BotpressClient, GenerationMetadata } from '../utils'
+import { Adapter, GetExamplesProps, SaveExampleProps } from './adapter'
+
+const CRITICAL_TAGS = {
+  system: 'true',
+  'schema-purpose': 'active-learning',
+  'schema-version': 'Oct-2024'
+} as const
+
+const OPTIONAL_TAGS = {
+  'x-studio-title': 'Active Learning',
+  'x-studio-description': 'Table for storing active learning tasks and examples',
+  'x-studio-readonly': 'true',
+  'x-studio-icon': 'lucide://atom',
+  'x-studio-color': 'green'
+} as const
+
+const FACTOR = 30
+
+const Props = z.object({
+  client: BotpressClient,
+  tableName: z
+    .string()
+    .regex(
+      /^[a-zA-Z0-9_]{1,45}Table$/,
+      'Table name must be lowercase and contain only letters, numbers and underscores'
+    )
+})
+
+export type TableSchema = z.input<typeof TableSchema>
+const TableSchema = z.object({
+  taskType: z.string().describe('The type of the task (filter, extract, etc.)'),
+  taskId: z.string(),
+  key: z.string().describe('A unique key for the task (e.g. a hash of the input, taskId, taskType and instructions)'),
+  instructions: z.string(),
+  input: z.object({}).passthrough().describe('The input to the task'),
+  output: z.object({}).passthrough().describe('The expected output'),
+  explanation: z.string().nullable(),
+  metadata: GenerationMetadata,
+  status: z.enum(['pending', 'rejected', 'approved']),
+  feedback: z
+    .object({
+      rating: z.enum(['very-bad', 'bad', 'good', 'very-good']),
+      comment: z.string().nullable()
+    })
+    .nullable()
+    .default(null)
+})
+
+const searchableColumns = ['input'] as const satisfies Array<keyof typeof TableSchema.shape> as string[]
+
+const TableJsonSchema = Object.entries(TableSchema.shape).reduce((acc, [key, value]) => {
+  acc[key] = value.toJsonSchema()
+  acc[key]['x-zui'] ??= {}
+  acc[key]['x-zui'].searchable = searchableColumns.includes(key)
+  return acc
+}, {})
+
+export class TableAdapter extends Adapter {
+  private client: Client
+  private tableName: string
+
+  private status: 'initialized' | 'ready' | 'error'
+
+  constructor(props: z.input<typeof Props>) {
+    super()
+    props = Props.parse(props)
+    this.client = props.client
+    this.tableName = props.tableName
+    this.status = 'ready'
+  }
+
+  public async getExamples<TInput, TOutput>({ taskType, taskId, input }: GetExamplesProps<TInput>) {
+    await this.assertTableExists()
+
+    const { rows } = await this.client
+      .findTableRows({
+        table: this.tableName,
+        search: JSON.stringify({ value: input }).substring(0, 1023), // Search is limited to 1024 characters
+        limit: 10, // TODO
+        filter: {
+          // Proximity match of approved examples
+          taskType,
+          taskId,
+          status: 'approved'
+        } satisfies Partial<TableSchema>
+      })
+      .catch((err) => {
+        // TODO: handle error
+        console.error(`Error fetching examples: ${err.message}`)
+        return { rows: [] }
+      })
+
+    return rows.map((row) => ({
+      key: row.key,
+      input: row.input.value as TInput,
+      output: row.output.value as TOutput,
+      explanation: row.explanation,
+      similarity: row.similarity ?? 0
+    }))
+  }
+
+  public async saveExample<TInput, TOutput>({
+    key,
+    taskType,
+    taskId,
+    instructions,
+    input,
+    output,
+    explanation,
+    metadata,
+    status = 'pending'
+  }: SaveExampleProps<TInput, TOutput>) {
+    await this.assertTableExists()
+
+    await this.client
+      .upsertTableRows({
+        table: this.tableName,
+        keyColumn: 'key',
+        rows: [
+          {
+            key,
+            taskType,
+            taskId,
+            instructions,
+            input: { value: input },
+            output: { value: output },
+            explanation: explanation ?? null,
+            status,
+            metadata
+          } satisfies TableSchema
+        ]
+      })
+      .catch(() => {
+        // TODO: handle error
+      })
+  }
+
+  private async assertTableExists() {
+    if (this.status !== 'ready') {
+      return
+    }
+
+    const { table, created } = await this.client
+      .getOrCreateTable({
+        table: this.tableName,
+        factor: FACTOR,
+        frozen: true,
+        isComputeEnabled: false,
+        tags: {
+          ...CRITICAL_TAGS,
+          ...OPTIONAL_TAGS
+        },
+        schema: TableJsonSchema
+      })
+      .catch(() => {
+        this.status = 'error'
+        return { table: null, created: false }
+      })
+
+    if (!table) {
+      return
+    }
+
+    if (!created) {
+      const issues: string[] = []
+
+      if (table.factor !== FACTOR) {
+        issues.push(`Factor is ${table.factor} instead of ${FACTOR}`)
+      }
+
+      if (table.frozen !== true) {
+        issues.push('Table is not frozen')
+      }
+
+      for (const [key, value] of Object.entries(CRITICAL_TAGS)) {
+        if (table.tags?.[key] !== value) {
+          issues.push(`Tag ${key} is ${table.tags?.[key]} instead of ${value}`)
+        }
+      }
+
+      for (const key of Object.keys(TableJsonSchema)) {
+        const column = table.schema?.properties[key]
+        const expected = TableJsonSchema[key] as { type: string }
+
+        if (!column) {
+          issues.push(`Column ${key} is missing`)
+          continue
+        }
+
+        if (column.type !== expected.type) {
+          issues.push(`Column ${key} has type ${column.type} instead of ${expected.type}`)
+        }
+
+        if (expected['x-zui'].searchable && !column['x-zui'].searchable) {
+          issues.push(`Column ${key} is not searchable but should be`)
+        }
+      }
+
+      if (issues.length) {
+        this.status = 'error'
+      }
+    }
+
+    this.status = 'initialized'
+  }
+}

package/src/adapters/memory.ts

@@ -0,0 +1,13 @@
+import { Adapter } from './adapter'
+
+export class MemoryAdapter extends Adapter {
+  constructor(public examples: any[]) {
+    super()
+  }
+
+  async getExamples() {
+    return this.examples
+  }
+
+  async saveExample() {}
+}

package/src/index.ts

@@ -0,0 +1,11 @@
+import { Zai } from './zai'
+
+import './operations/text'
+import './operations/rewrite'
+import './operations/summarize'
+import './operations/check'
+import './operations/filter'
+import './operations/extract'
+import './operations/label'
+
+export { Zai }

package/src/models.ts

@@ -0,0 +1,394 @@
+
+// This file is generated. Do not edit it manually.
+// See 'scripts/update-models.ts'
+
+/* eslint-disable */
+/* tslint:disable */
+
+export const Models = [
+  {
+    "id": "anthropic__claude-3-haiku-20240307",
+    "name": "Claude 3 Haiku",
+    "integration": "anthropic",
+    "input": {
+      "maxTokens": 200000
+    },
+    "output": {
+      "maxTokens": 4096
+    }
+  },
+  {
+    "id": "anthropic__claude-3-5-sonnet-20240620",
+    "name": "Claude 3.5 Sonnet",
+    "integration": "anthropic",
+    "input": {
+      "maxTokens": 200000
+    },
+    "output": {
+      "maxTokens": 4096
+    }
+  },
+  {
+    "id": "cerebras__llama3.1-70b",
+    "name": "Llama 3.1 70B",
+    "integration": "cerebras",
+    "input": {
+      "maxTokens": 8192
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "cerebras__llama3.1-8b",
+    "name": "Llama 3.1 8B",
+    "integration": "cerebras",
+    "input": {
+      "maxTokens": 8192
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/deepseek-coder-v2-instruct",
+    "name": "DeepSeek Coder V2 Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 131072
+    },
+    "output": {
+      "maxTokens": 131072
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/deepseek-coder-v2-lite-instruct",
+    "name": "DeepSeek Coder V2 Lite",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 163840
+    },
+    "output": {
+      "maxTokens": 163840
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/firellava-13b",
+    "name": "FireLLaVA-13B",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 4096
+    },
+    "output": {
+      "maxTokens": 4096
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/firefunction-v2",
+    "name": "Firefunction V2",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 8192
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/gemma2-9b-it",
+    "name": "Gemma 2 9B Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 8192
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/llama-v3p1-405b-instruct",
+    "name": "Llama 3.1 405B Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 131072
+    },
+    "output": {
+      "maxTokens": 131072
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/llama-v3p1-70b-instruct",
+    "name": "Llama 3.1 70B Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 131072
+    },
+    "output": {
+      "maxTokens": 131072
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/llama-v3p1-8b-instruct",
+    "name": "Llama 3.1 8B Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 131072
+    },
+    "output": {
+      "maxTokens": 131072
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/mixtral-8x22b-instruct",
+    "name": "Mixtral MoE 8x22B Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 65536
+    },
+    "output": {
+      "maxTokens": 65536
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/mixtral-8x7b-instruct",
+    "name": "Mixtral MoE 8x7B Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 32768
+    },
+    "output": {
+      "maxTokens": 32768
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/mythomax-l2-13b",
+    "name": "MythoMax L2 13b",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 4096
+    },
+    "output": {
+      "maxTokens": 4096
+    }
+  },
+  {
+    "id": "fireworks-ai__accounts/fireworks/models/qwen2-72b-instruct",
+    "name": "Qwen2 72b Instruct",
+    "integration": "fireworks-ai",
+    "input": {
+      "maxTokens": 32768
+    },
+    "output": {
+      "maxTokens": 32768
+    }
+  },
+  {
+    "id": "groq__gemma2-9b-it",
+    "name": "Gemma2 9B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 8192
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama3-70b-8192",
+    "name": "LLaMA 3 70B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 8192
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama3-8b-8192",
+    "name": "LLaMA 3 8B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 8192
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama-3.1-70b-versatile",
+    "name": "LLaMA 3.1 70B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama-3.1-8b-instant",
+    "name": "LLaMA 3.1 8B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama-3.2-11b-vision-preview",
+    "name": "LLaMA 3.2 11B Vision",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama-3.2-1b-preview",
+    "name": "LLaMA 3.2 1B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama-3.2-3b-preview",
+    "name": "LLaMA 3.2 3B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama-3.2-90b-vision-preview",
+    "name": "LLaMA 3.2 90B Vision",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 8192
+    }
+  },
+  {
+    "id": "groq__llama-3.3-70b-versatile",
+    "name": "LLaMA 3.3 70B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 32768
+    }
+  },
+  {
+    "id": "groq__mixtral-8x7b-32768",
+    "name": "Mixtral 8x7B",
+    "integration": "groq",
+    "input": {
+      "maxTokens": 32768
+    },
+    "output": {
+      "maxTokens": 32768
+    }
+  },
+  {
+    "id": "openai__o1-2024-12-17",
+    "name": "GPT o1",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 200000
+    },
+    "output": {
+      "maxTokens": 100000
+    }
+  },
+  {
+    "id": "openai__o1-mini-2024-09-12",
+    "name": "GPT o1-mini",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 65536
+    }
+  },
+  {
+    "id": "openai__gpt-3.5-turbo-0125",
+    "name": "GPT-3.5 Turbo",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 4096
+    }
+  },
+  {
+    "id": "openai__gpt-4-turbo-2024-04-09",
+    "name": "GPT-4 Turbo",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 4096
+    }
+  },
+  {
+    "id": "openai__gpt-4o-2024-08-06",
+    "name": "GPT-4o (August 2024)",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 16384
+    }
+  },
+  {
+    "id": "openai__gpt-4o-2024-05-13",
+    "name": "GPT-4o (May 2024)",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 4096
+    }
+  },
+  {
+    "id": "openai__gpt-4o-2024-11-20",
+    "name": "GPT-4o (November 2024)",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 16384
+    }
+  },
+  {
+    "id": "openai__gpt-4o-mini-2024-07-18",
+    "name": "GPT-4o Mini",
+    "integration": "openai",
+    "input": {
+      "maxTokens": 128000
+    },
+    "output": {
+      "maxTokens": 16384
+    }
+  }
+] as const