Initial commit
This commit is contained in:
Riley Tomasek
@@ -0,0 +1,8 @@
|
||||
# Changesets
|
||||
|
||||
Hello and welcome! This folder has been automatically generated by `@changesets/cli`, a build tool that works
|
||||
with multi-package repos, or single-package repos to help you version and publish your code. You can
|
||||
find the full documentation for it [in our repository](https://github.com/changesets/changesets)
|
||||
|
||||
We have a quick list of common questions to get you started engaging with this project in
|
||||
[our documentation](https://github.com/changesets/changesets/blob/main/docs/common-questions.md)
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"$schema": "https://unpkg.com/@changesets/config@2.3.0/schema.json",
|
||||
"changelog": "@changesets/cli/changelog",
|
||||
"commit": false,
|
||||
"fixed": [],
|
||||
"linked": [],
|
||||
"access": "restricted",
|
||||
"baseBranch": "master",
|
||||
"updateInternalDependencies": "patch",
|
||||
"ignore": ["examples-basic"]
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
root = true
|
||||
|
||||
[*]
|
||||
indent_style = space
|
||||
indent_size = 2
|
||||
end_of_line = lf
|
||||
charset = utf-8
|
||||
trim_trailing_whitespace = true
|
||||
insert_final_newline = true
|
||||
@@ -0,0 +1 @@
|
||||
* text=auto eol=lf
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
.cache
|
||||
.DS_Store
|
||||
.idea
|
||||
*.log
|
||||
*.tgz
|
||||
coverage
|
||||
dist
|
||||
lib-cov
|
||||
logs
|
||||
node_modules
|
||||
temp
|
||||
todo.md
|
||||
.env*
|
||||
.dev*
|
||||
.env.local
|
||||
.local-files
|
||||
@@ -0,0 +1,8 @@
|
||||
*.env
|
||||
*.log
|
||||
*.jsonl
|
||||
**/node_modules
|
||||
**/dist
|
||||
**/build
|
||||
**/dbschema
|
||||
**/public
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"singleQuote": true
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
{
|
||||
"$schema": "https://json.schemastore.org/tsconfig",
|
||||
"display": "Node ESM",
|
||||
"compilerOptions": {
|
||||
"moduleResolution": "node16",
|
||||
"module": "node16",
|
||||
"lib": ["ESNext", "DOM"],
|
||||
"jsx": "preserve",
|
||||
"outDir": "dist",
|
||||
"allowJs": true,
|
||||
"importHelpers": true,
|
||||
"esModuleInterop": true,
|
||||
"allowSyntheticDefaultImports": true,
|
||||
"strict": true,
|
||||
"noUnusedLocals": true,
|
||||
"noUnusedParameters": true,
|
||||
"noImplicitReturns": true,
|
||||
"noFallthroughCasesInSwitch": true,
|
||||
"skipLibCheck": true,
|
||||
"resolveJsonModule": true,
|
||||
"declaration": true,
|
||||
"sourceMap": true,
|
||||
"checkJs": true,
|
||||
"noEmit": false,
|
||||
"forceConsistentCasingInFileNames": true,
|
||||
"isolatedModules": true,
|
||||
"incremental": false,
|
||||
"skipDefaultLibCheck": true
|
||||
},
|
||||
"exclude": ["**/node_modules", "**/.*/"]
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"name": "tsconfig",
|
||||
"version": "0.0.0",
|
||||
"private": true,
|
||||
"files": [
|
||||
"node-esm.json"
|
||||
],
|
||||
"publishConfig": {
|
||||
"access": "public"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"extends": ["@dexaai/eslint-config", "@dexaai/eslint-config/node"]
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"name": "examples-basic",
|
||||
"type": "module",
|
||||
"version": "0.0.0",
|
||||
"private": true,
|
||||
"description": "",
|
||||
"scripts": {
|
||||
"clean": "rm -rf .turbo node_modules dist",
|
||||
"typecheck": "tsc --noEmit"
|
||||
},
|
||||
"dependencies": {
|
||||
"@dexaai/agent": "workspace:*",
|
||||
"@dexaai/datastore": "workspace:*",
|
||||
"@dexaai/model": "workspace:*",
|
||||
"tslib": "^2.5.3",
|
||||
"workgpt": "^0.1.1",
|
||||
"zod": "^3.21.4"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@dexaai/eslint-config": "^0.2.0",
|
||||
"@types/node": "^20.3.1",
|
||||
"dotenv": "^16.3.1",
|
||||
"eslint": "^8.43.0",
|
||||
"tsconfig": "workspace:*",
|
||||
"type-fest": "3.12.0",
|
||||
"typescript": "^5.1.3"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
# Example Usage of LLM Tools
|
||||
@@ -0,0 +1,30 @@
|
||||
import 'dotenv/config';
|
||||
import { Datastore } from '@dexaai/datastore/pinecone';
|
||||
import { EmbeddingModel } from '@dexaai/model/openai';
|
||||
|
||||
// Base Datastore
|
||||
(async () => {
|
||||
// Embedding model with cache
|
||||
const embeddingModel = new EmbeddingModel({
|
||||
params: { model: 'text-embedding-ada-002' },
|
||||
context: { test: 'test' },
|
||||
debug: true,
|
||||
}).addCache();
|
||||
await embeddingModel.run({ input: ['cat'] });
|
||||
// This should be cached
|
||||
await embeddingModel.run({ input: ['cat'] });
|
||||
|
||||
// Datastore with cache
|
||||
const store = new Datastore({
|
||||
namespace: 'test',
|
||||
embeddingModel,
|
||||
debug: true,
|
||||
}).addCache();
|
||||
await store.upsert([
|
||||
{ id: '1', metadata: { content: 'cat' } },
|
||||
{ id: '2', metadata: { content: 'dog' } },
|
||||
]);
|
||||
await store.query({ query: 'cat' });
|
||||
// This should be cached
|
||||
await store.query({ query: 'cat' });
|
||||
})();
|
||||
@@ -0,0 +1,85 @@
|
||||
import 'dotenv/config';
|
||||
import { Datastore, HybridDatastore } from '@dexaai/datastore/pinecone';
|
||||
import { EmbeddingModel } from '@dexaai/model/openai';
|
||||
import { SpladeModel } from '@dexaai/model/custom';
|
||||
|
||||
// Base Datastore
|
||||
(async () => {
|
||||
const embeddingModel = new EmbeddingModel({
|
||||
params: { model: 'text-embedding-ada-002' },
|
||||
context: { test: 'test' },
|
||||
debug: true,
|
||||
});
|
||||
|
||||
const store = new Datastore({
|
||||
namespace: 'test',
|
||||
embeddingModel,
|
||||
debug: true,
|
||||
});
|
||||
|
||||
await store.upsert([
|
||||
{ id: '1', metadata: { content: 'cat' } },
|
||||
{ id: '2', metadata: { content: 'dog' } },
|
||||
{ id: '3', metadata: { content: 'whale' } },
|
||||
{ id: '4', metadata: { content: 'shark' } },
|
||||
{ id: '5', metadata: { content: 'computer' } },
|
||||
{ id: '6', metadata: { content: 'laptop' } },
|
||||
{ id: '7', metadata: { content: 'phone' } },
|
||||
{ id: '8', metadata: { content: 'tablet' } },
|
||||
]);
|
||||
|
||||
const result = await store.query({
|
||||
query: 'dolphin',
|
||||
});
|
||||
console.log(JSON.stringify(result, null, 2));
|
||||
|
||||
const result2 = await store.query({
|
||||
query: 'fox',
|
||||
});
|
||||
console.log(JSON.stringify(result2, null, 2));
|
||||
|
||||
const result3 = await store.query({
|
||||
query: 'keyboard',
|
||||
});
|
||||
console.log(JSON.stringify(result3, null, 2));
|
||||
|
||||
// Hybrid Datastore
|
||||
const spladeModel = new SpladeModel({
|
||||
params: { model: 'naver/splade-cocondenser-ensembledistil' },
|
||||
debug: true,
|
||||
});
|
||||
|
||||
const hStore = new HybridDatastore({
|
||||
namespace: 'test-hybrid',
|
||||
embeddingModel,
|
||||
spladeModel,
|
||||
context: { baz: 'qux' },
|
||||
debug: true,
|
||||
});
|
||||
|
||||
await hStore.upsert([
|
||||
{ id: '1', metadata: { content: 'cat' } },
|
||||
{ id: '2', metadata: { content: 'dog' } },
|
||||
{ id: '3', metadata: { content: 'whale' } },
|
||||
{ id: '4', metadata: { content: 'shark' } },
|
||||
{ id: '5', metadata: { content: 'computer' } },
|
||||
{ id: '6', metadata: { content: 'laptop' } },
|
||||
{ id: '7', metadata: { content: 'phone' } },
|
||||
{ id: '8', metadata: { content: 'tablet' } },
|
||||
]);
|
||||
|
||||
await hStore.query({
|
||||
query: 'dolphin',
|
||||
});
|
||||
// console.log(JSON.stringify(r, null, 2));
|
||||
|
||||
await hStore.query({
|
||||
query: 'fox',
|
||||
});
|
||||
// console.log(JSON.stringify(r2, null, 2));
|
||||
|
||||
await hStore.query({
|
||||
query: 'keyboard',
|
||||
});
|
||||
// console.log(JSON.stringify(r3, null, 2));
|
||||
})();
|
||||
@@ -0,0 +1,13 @@
|
||||
import 'dotenv/config';
|
||||
import { Calculator } from 'workgpt/apis/calculator';
|
||||
import { FactApi } from 'workgpt/apis/fact';
|
||||
import { createRunner } from '@dexaai/agent';
|
||||
|
||||
const apis = await Promise.all([new Calculator(), new FactApi()]);
|
||||
const runner = createRunner(apis, undefined, true);
|
||||
|
||||
await runner.runWithDirective(
|
||||
'What is the age of obama multiplied by the age of trump?'
|
||||
);
|
||||
|
||||
console.log('--- Done ---');
|
||||
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"extends": "tsconfig/node-esm.json",
|
||||
"include": ["./src/**/*"],
|
||||
"exclude": ["**/node_modules", "**/.*/"],
|
||||
"compilerOptions": {
|
||||
"experimentalDecorators": true,
|
||||
"emitDecoratorMetadata": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"name": "llm-tools",
|
||||
"type": "module",
|
||||
"private": true,
|
||||
"version": "0.0.0",
|
||||
"description": "",
|
||||
"repository": "dexaai/llm-tools",
|
||||
"scripts": {
|
||||
"build": "turbo run build",
|
||||
"clean": "turbo run clean && rm -rf node_modules",
|
||||
"dev": "turbo run dev --parallel",
|
||||
"format": "prettier --write \"**/*.{ts,tsx}\"",
|
||||
"lint": "turbo run lint",
|
||||
"test": "turbo run test",
|
||||
"typecheck": "turbo run typecheck"
|
||||
},
|
||||
"dependencies": {},
|
||||
"devDependencies": {
|
||||
"@changesets/cli": "^2.26.1",
|
||||
"prettier": "^2.8.8",
|
||||
"tsx": "^3.12.7",
|
||||
"turbo": "^1.10.5"
|
||||
},
|
||||
"packageManager": "pnpm@7.25.0",
|
||||
"engines": {
|
||||
"node": ">= 18"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"extends": ["@dexaai/eslint-config", "@dexaai/eslint-config/node"]
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
{
|
||||
"name": "@dexaai/agent",
|
||||
"type": "module",
|
||||
"version": "0.0.0",
|
||||
"description": "",
|
||||
"repository": "dexaai/llm-tools",
|
||||
"scripts": {
|
||||
"build": "tsup",
|
||||
"clean": "rm -rf .turbo node_modules dist",
|
||||
"dev": "tsup --watch",
|
||||
"lint": "eslint --cache --cache-location ./node_modules/.cache/eslint .",
|
||||
"test": "vitest",
|
||||
"typecheck": "tsc --noEmit"
|
||||
},
|
||||
"sideEffects": false,
|
||||
"files": [
|
||||
"dist"
|
||||
],
|
||||
"types": "./dist/index.d.ts",
|
||||
"main": "./dist/index.cjs",
|
||||
"module": "./dist/index.js",
|
||||
"exports": {
|
||||
"./package.json": "./package.json",
|
||||
".": {
|
||||
"types": "./dist/index.d.ts",
|
||||
"require": "./dist/index.cjs",
|
||||
"import": "./dist/index.js",
|
||||
"default": "./dist/index.js"
|
||||
}
|
||||
},
|
||||
"dependencies": {
|
||||
"@dexaai/model": "workspace:*",
|
||||
"workgpt": "^0.1.1"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@dexaai/eslint-config": "^0.2.0",
|
||||
"@types/node": "^20.3.1",
|
||||
"eslint": "^8.43.0",
|
||||
"tsconfig": "workspace:*",
|
||||
"tsup": "^7.0.0",
|
||||
"typescript": "^5.1.3",
|
||||
"vite": "^4.3.9",
|
||||
"vitest": "^0.32.2"
|
||||
},
|
||||
"publishConfig": {
|
||||
"access": "restricted"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
# Agent
|
||||
|
||||
Standardized agents for interacting with OpenAI functions.
|
||||
@@ -0,0 +1 @@
|
||||
export { createRunner } from './runner.js';
|
||||
@@ -0,0 +1,58 @@
|
||||
import { WorkGptRunner } from 'workgpt/runners/workgpt';
|
||||
import type { Api } from 'workgpt/apis/base';
|
||||
import { ChatModel } from '@dexaai/model/openai';
|
||||
import { ChatAgent } from 'workgpt/chat-agents/base';
|
||||
import type { ChatRequest, ChatResponse } from 'workgpt/chat-agents/types';
|
||||
import type { ChatMessage } from '@dexaai/model';
|
||||
|
||||
/** Create a WorkGptRunner using an LLM ChatModel. */
|
||||
export function createRunner(
|
||||
apis: Api[],
|
||||
chatModel?: ChatModel,
|
||||
verbose: boolean = false
|
||||
): WorkGptRunner {
|
||||
const model =
|
||||
chatModel ??
|
||||
new ChatModel({
|
||||
params: { model: 'gpt-3.5-turbo-0613' },
|
||||
});
|
||||
const agent = new LlmChatAgent(model, verbose);
|
||||
return new WorkGptRunner({ apis, agent });
|
||||
}
|
||||
|
||||
/** Agent that uses the LLM ChatModel to make requests */
|
||||
class LlmChatAgent extends ChatAgent {
|
||||
private readonly chatModel: ChatModel;
|
||||
|
||||
constructor(chatModel: ChatModel, verbose: boolean = false) {
|
||||
super({ verbose });
|
||||
if (verbose) {
|
||||
chatModel.addHooks({ afterApiResponse: agentLogger });
|
||||
}
|
||||
this.chatModel = chatModel;
|
||||
}
|
||||
|
||||
async call(args: ChatRequest): Promise<ChatResponse> {
|
||||
const { message } = await this.chatModel.run({
|
||||
messages: args.messages as ChatMessage[],
|
||||
functions: args.functions,
|
||||
});
|
||||
return message as ChatResponse;
|
||||
}
|
||||
}
|
||||
|
||||
/** Logging function for hooks to debug agents */
|
||||
function agentLogger(evt: any) {
|
||||
const { choices } = evt.response as {
|
||||
choices: { message: ChatMessage }[];
|
||||
};
|
||||
const req = evt.params.messages[evt.params.messages.length - 1];
|
||||
const name = req.name ? ` (${req.name})` : '';
|
||||
const message = choices[0].message;
|
||||
console.log(
|
||||
`>> ${req.role}: ${name} ${
|
||||
req.content || JSON.stringify(req.function_call || '')
|
||||
}`
|
||||
);
|
||||
console.log(`<< ${message.role}: ${message.content}`, message.function_call);
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
describe('should', () => {
|
||||
it('exported', () => {
|
||||
expect(1).toEqual(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"extends": "tsconfig/node-esm.json",
|
||||
"include": ["./src/**/*"],
|
||||
"exclude": ["**/node_modules", "**/.*/"],
|
||||
"compilerOptions": {
|
||||
"types": ["vitest/globals", "node"]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
import { defineConfig } from 'tsup';
|
||||
|
||||
export default defineConfig({
|
||||
entry: ['src/index.ts'],
|
||||
format: ['esm', 'cjs'],
|
||||
clean: true,
|
||||
dts: true,
|
||||
target: 'es2022',
|
||||
});
|
||||
@@ -0,0 +1,12 @@
|
||||
import { defineConfig } from 'vite';
|
||||
|
||||
export default defineConfig(() => {
|
||||
return {
|
||||
mode: 'test',
|
||||
|
||||
test: {
|
||||
environment: 'node',
|
||||
globals: true,
|
||||
},
|
||||
};
|
||||
});
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"extends": ["@dexaai/eslint-config", "@dexaai/eslint-config/node"]
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
{
|
||||
"name": "@dexaai/datastore",
|
||||
"type": "module",
|
||||
"version": "0.0.0",
|
||||
"description": "",
|
||||
"repository": "dexaai/llm-tools",
|
||||
"scripts": {
|
||||
"build": "tsup",
|
||||
"clean": "rm -rf .turbo node_modules dist",
|
||||
"dev": "tsup --watch",
|
||||
"lint": "eslint --cache --cache-location ./node_modules/.cache/eslint .",
|
||||
"test": "vitest",
|
||||
"typecheck": "tsc --noEmit"
|
||||
},
|
||||
"sideEffects": false,
|
||||
"files": [
|
||||
"dist"
|
||||
],
|
||||
"types": "./dist/index.d.ts",
|
||||
"main": "./dist/index.cjs",
|
||||
"module": "./dist/index.js",
|
||||
"exports": {
|
||||
"./package.json": "./package.json",
|
||||
".": {
|
||||
"types": "./dist/index.d.ts",
|
||||
"require": "./dist/index.cjs",
|
||||
"import": "./dist/index.js",
|
||||
"default": "./dist/index.js"
|
||||
},
|
||||
"./pinecone": {
|
||||
"types": "./dist/pinecone/index.d.ts",
|
||||
"require": "./dist/pinecone/index.cjs",
|
||||
"import": "./dist/pinecone/index.js",
|
||||
"default": "./dist/pinecone/index.js"
|
||||
}
|
||||
},
|
||||
"typesVersions": {
|
||||
"*": {
|
||||
"pinecone": [
|
||||
"./dist/pinecone/index.d.ts"
|
||||
]
|
||||
}
|
||||
},
|
||||
"dependencies": {
|
||||
"@dexaai/model": "workspace:*",
|
||||
"keyv": "^4.5.2",
|
||||
"ky": "^0.33.3",
|
||||
"p-map": "^6.0.0",
|
||||
"p-retry": "^5.1.2",
|
||||
"p-throttle": "^5.1.0",
|
||||
"pinecone-client": "^1.1.0",
|
||||
"tkv-cache": "^0.1.2",
|
||||
"ts-extras": "^0.11.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@dexaai/eslint-config": "^0.2.0",
|
||||
"@types/node": "^20.3.1",
|
||||
"eslint": "^8.43.0",
|
||||
"tsconfig": "workspace:*",
|
||||
"tsup": "^7.0.0",
|
||||
"type-fest": "3.12.0",
|
||||
"typescript": "^5.1.3",
|
||||
"vite": "^4.3.9",
|
||||
"vitest": "^0.32.2"
|
||||
},
|
||||
"publishConfig": {
|
||||
"access": "restricted"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
# Datastore
|
||||
|
||||
Standardized classes for interacting with vector databases.
|
||||
@@ -0,0 +1,126 @@
|
||||
import Keyv from 'keyv';
|
||||
import { KvCache } from 'tkv-cache';
|
||||
import type { IEmbeddingModel, ISparseModel } from '@dexaai/model';
|
||||
import type {
|
||||
BaseMeta,
|
||||
Ctx,
|
||||
DatastoreOpts,
|
||||
DatastoreProvider,
|
||||
DatastoreType,
|
||||
Doc,
|
||||
Hooks,
|
||||
HDatastoreOpts,
|
||||
IDatastore,
|
||||
Query,
|
||||
QueryCache,
|
||||
QueryResult,
|
||||
} from './types.js';
|
||||
|
||||
export abstract class BaseDatastore<DocMeta extends BaseMeta>
|
||||
implements IDatastore<DocMeta>
|
||||
{
|
||||
protected abstract runQuery(
|
||||
query: Query<DocMeta>,
|
||||
context?: Ctx
|
||||
): Promise<QueryResult<DocMeta>>;
|
||||
abstract upsert(docs: Doc<DocMeta>[], context?: Ctx): Promise<void>;
|
||||
abstract delete(docIds: string[]): Promise<void>;
|
||||
abstract deleteAll(): Promise<void>;
|
||||
|
||||
abstract datastoreType: DatastoreType;
|
||||
abstract datastoreProvider: DatastoreProvider;
|
||||
|
||||
protected namespace: string;
|
||||
protected embeddingModel: IEmbeddingModel;
|
||||
protected cache?: QueryCache<DocMeta>;
|
||||
protected hooks: Hooks<DocMeta>;
|
||||
protected context: Ctx;
|
||||
|
||||
constructor(args: DatastoreOpts<DocMeta>) {
|
||||
this.namespace = args.namespace;
|
||||
this.embeddingModel = args.embeddingModel;
|
||||
this.cache = args.cache;
|
||||
this.context = args.context ?? {};
|
||||
this.hooks = args.debug
|
||||
? {
|
||||
afterApiResponse: console.log,
|
||||
afterCacheHit: console.log,
|
||||
beforeError: console.error,
|
||||
...args.hooks,
|
||||
}
|
||||
: args.hooks ?? {};
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a cache to the datastore for storing query responses.
|
||||
* Defaults to an in-memory Keyv cache.
|
||||
*/
|
||||
addCache(args?: {
|
||||
/** A Keyv instance to use for caching */
|
||||
keyv?: Keyv<QueryResult<DocMeta>>;
|
||||
/** Keyv options to create a new Keyv instance */
|
||||
keyvOpts?: Keyv.Options<QueryResult<DocMeta>>;
|
||||
/** A function to normalize the cache key */
|
||||
normalizeKey?: (params: Query<DocMeta>) => Partial<Query<DocMeta>>;
|
||||
errorHandler?: (error: unknown) => void;
|
||||
}): this {
|
||||
const { keyvOpts, normalizeKey, errorHandler } = args ?? {};
|
||||
const keyv = args?.keyv ?? new Keyv<QueryResult<DocMeta>>(keyvOpts);
|
||||
this.cache = new KvCache(keyv, normalizeKey, errorHandler);
|
||||
return this;
|
||||
}
|
||||
|
||||
async query(
|
||||
query: Query<DocMeta>,
|
||||
context?: Ctx
|
||||
): Promise<QueryResult<DocMeta>> {
|
||||
// Return cached response if available
|
||||
const cached = await this?.cache?.get(query);
|
||||
if (cached) {
|
||||
const mergedContext = { ...this.context, ...context };
|
||||
await this.hooks.afterCacheHit?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
datastoreType: this.datastoreType,
|
||||
datastoreProvider: this.datastoreProvider,
|
||||
query,
|
||||
response: cached,
|
||||
context: mergedContext,
|
||||
});
|
||||
return cached;
|
||||
}
|
||||
|
||||
try {
|
||||
// Run the query
|
||||
const response = await this.runQuery(query, context);
|
||||
|
||||
// Update the cache
|
||||
await this?.cache?.set(query, response);
|
||||
|
||||
return response;
|
||||
} catch (error) {
|
||||
const mergedContext = { ...this.context, ...context };
|
||||
await this.hooks.beforeError?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
datastoreType: this.datastoreType,
|
||||
datastoreProvider: this.datastoreProvider,
|
||||
query,
|
||||
error,
|
||||
context: mergedContext,
|
||||
});
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export abstract class BaseHybridDatastore<DocMeta extends BaseMeta>
|
||||
extends BaseDatastore<DocMeta>
|
||||
implements IDatastore<DocMeta>
|
||||
{
|
||||
protected spladeModel: ISparseModel;
|
||||
|
||||
constructor(args: HDatastoreOpts<DocMeta>) {
|
||||
const { spladeModel, ...rest } = args;
|
||||
super(rest);
|
||||
this.spladeModel = args.spladeModel;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
export { BaseDatastore, BaseHybridDatastore } from './datastore.js';
|
||||
export { VectorUtils } from './utils/vectors.js';
|
||||
export type {
|
||||
BaseMeta,
|
||||
Ctx,
|
||||
DatastoreOpts,
|
||||
Doc,
|
||||
HDatastoreOpts,
|
||||
IDatastore,
|
||||
Prettify,
|
||||
Query,
|
||||
QueryResult,
|
||||
} from './types.js';
|
||||
@@ -0,0 +1,13 @@
|
||||
import { PineconeClient } from 'pinecone-client';
|
||||
import type { JsonObject } from 'type-fest';
|
||||
|
||||
type PineconeConfig = ConstructorParameters<typeof PineconeClient>[0];
|
||||
|
||||
/** Create a new Pinecone client instance. */
|
||||
export function createPineconeClient<Meta extends JsonObject = {}>(
|
||||
config: PineconeConfig
|
||||
): PineconeClient<Meta> {
|
||||
return new PineconeClient<Meta>(config);
|
||||
}
|
||||
|
||||
export type { PineconeClient };
|
||||
@@ -0,0 +1,271 @@
|
||||
import { BaseDatastore, BaseHybridDatastore } from '../datastore.js';
|
||||
import type {
|
||||
BaseMeta,
|
||||
Ctx,
|
||||
DatastoreOpts,
|
||||
Doc,
|
||||
HDatastoreOpts,
|
||||
IDatastore,
|
||||
Prettify,
|
||||
Query,
|
||||
QueryResult,
|
||||
} from '../types.js';
|
||||
import type { PineconeClient } from './client.js';
|
||||
import { createPineconeClient } from './client.js';
|
||||
|
||||
export class Datastore<DocMeta extends BaseMeta>
|
||||
extends BaseDatastore<DocMeta>
|
||||
implements IDatastore<DocMeta>
|
||||
{
|
||||
datastoreType = 'embedding' as const;
|
||||
datastoreProvider = 'pinecone' as const;
|
||||
private readonly pinecone: PineconeClient<DocMeta>;
|
||||
|
||||
constructor(
|
||||
args: Prettify<
|
||||
DatastoreOpts<DocMeta> & {
|
||||
pinecone?: PineconeClient<DocMeta>;
|
||||
}
|
||||
>
|
||||
) {
|
||||
const { pinecone, ...rest } = args;
|
||||
super(rest);
|
||||
this.pinecone =
|
||||
pinecone ||
|
||||
createPineconeClient<DocMeta>({
|
||||
namespace: this.namespace,
|
||||
});
|
||||
}
|
||||
|
||||
async runQuery(
|
||||
query: Query<DocMeta>,
|
||||
context?: Ctx
|
||||
): Promise<QueryResult<DocMeta>> {
|
||||
const mergedContext = { ...this.context, ...context };
|
||||
|
||||
// Get the query embedding
|
||||
let queryEmbedding = query.embedding;
|
||||
|
||||
// If no query embedding is provided, run the query embedder
|
||||
if (!queryEmbedding) {
|
||||
const {
|
||||
embeddings: [embedding],
|
||||
} = await this.embeddingModel.run(
|
||||
{
|
||||
input: [query.query],
|
||||
},
|
||||
mergedContext
|
||||
);
|
||||
queryEmbedding = embedding;
|
||||
}
|
||||
|
||||
// Query Pinecone
|
||||
const start = Date.now();
|
||||
const response = await this.pinecone.query({
|
||||
topK: query.topK ?? 10,
|
||||
...(typeof query.minScore === 'number'
|
||||
? { minScore: query.minScore }
|
||||
: {}),
|
||||
...(query.filter && { filter: query.filter }),
|
||||
includeValues: query.includeValues ?? false,
|
||||
includeMetadata: true,
|
||||
vector: queryEmbedding,
|
||||
});
|
||||
const latency = Date.now() - start;
|
||||
await this.hooks.afterApiResponse?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
datastoreType: this.datastoreType,
|
||||
datastoreProvider: this.datastoreProvider,
|
||||
query,
|
||||
response,
|
||||
context: mergedContext,
|
||||
latency,
|
||||
});
|
||||
|
||||
const queryResult: QueryResult<DocMeta> = {
|
||||
query: query.query,
|
||||
docs: response.matches,
|
||||
};
|
||||
|
||||
return queryResult;
|
||||
}
|
||||
|
||||
async upsert(docs: Doc<DocMeta>[], context?: Ctx): Promise<void> {
|
||||
const mergedContext = { ...this.context, ...context };
|
||||
try {
|
||||
// Get the text content from the docs
|
||||
const texts = docs.map((doc) => doc.metadata.content);
|
||||
|
||||
// Create the embeddings for the docs
|
||||
// This relies on the classes to handle batching and throttling
|
||||
const embeddingRes = await this.embeddingModel.run(
|
||||
{ input: texts },
|
||||
mergedContext
|
||||
);
|
||||
|
||||
// Combine the results into Pinecones vector format and upsert
|
||||
// The Pinecone client will handle batching and throttling
|
||||
return this.pinecone.upsert({
|
||||
vectors: docs.map((doc, i) => ({
|
||||
id: doc.id,
|
||||
values: embeddingRes.embeddings[i],
|
||||
metadata: doc.metadata,
|
||||
})),
|
||||
});
|
||||
} catch (error) {
|
||||
await this.hooks.beforeError?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
datastoreType: this.datastoreType,
|
||||
datastoreProvider: this.datastoreProvider,
|
||||
upsert: docs,
|
||||
error,
|
||||
context: mergedContext,
|
||||
});
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async delete(docIds: string[]): Promise<void> {
|
||||
return this.pinecone.delete({ ids: docIds });
|
||||
}
|
||||
|
||||
async deleteAll(): Promise<void> {
|
||||
return this.pinecone.delete({ deleteAll: true });
|
||||
}
|
||||
}
|
||||
|
||||
export class HybridDatastore<DocMeta extends BaseMeta>
|
||||
extends BaseHybridDatastore<DocMeta>
|
||||
implements IDatastore<DocMeta>
|
||||
{
|
||||
datastoreType = 'hybrid' as const;
|
||||
datastoreProvider = 'pinecone' as const;
|
||||
private readonly pinecone: PineconeClient<DocMeta>;
|
||||
|
||||
constructor(
|
||||
args: Prettify<
|
||||
HDatastoreOpts<DocMeta> & {
|
||||
pinecone?: PineconeClient<DocMeta>;
|
||||
}
|
||||
>
|
||||
) {
|
||||
const { pinecone, ...rest } = args;
|
||||
super(rest);
|
||||
this.pinecone =
|
||||
pinecone ||
|
||||
createPineconeClient<DocMeta>({
|
||||
namespace: this.namespace,
|
||||
});
|
||||
}
|
||||
|
||||
async runQuery(
|
||||
query: Query<DocMeta>,
|
||||
context?: Ctx
|
||||
): Promise<QueryResult<DocMeta>> {
|
||||
const mergedContext = { ...this.context, ...context };
|
||||
|
||||
// Get the query embedding and sparse vector
|
||||
const queryEmbedding = query.embedding;
|
||||
const querySparseVector = query.sparseVector;
|
||||
const [
|
||||
{
|
||||
embeddings: [embedding],
|
||||
},
|
||||
{
|
||||
vectors: [sparseVector],
|
||||
},
|
||||
] = await Promise.all([
|
||||
queryEmbedding
|
||||
? { embeddings: [queryEmbedding] }
|
||||
: this.embeddingModel.run(
|
||||
{
|
||||
input: [query.query],
|
||||
},
|
||||
mergedContext
|
||||
),
|
||||
querySparseVector
|
||||
? { vectors: [querySparseVector] }
|
||||
: this.spladeModel.run(
|
||||
{
|
||||
input: [query.query],
|
||||
},
|
||||
mergedContext
|
||||
),
|
||||
]);
|
||||
|
||||
// Query Pinecone
|
||||
const start = Date.now();
|
||||
const response = await this.pinecone.query({
|
||||
topK: query.topK ?? 10,
|
||||
...(typeof query.minScore === 'number'
|
||||
? { minScore: query.minScore }
|
||||
: {}),
|
||||
...(query.filter && { filter: query.filter }),
|
||||
includeValues: query.includeValues ?? false,
|
||||
includeMetadata: true,
|
||||
vector: embedding,
|
||||
sparseVector: sparseVector,
|
||||
});
|
||||
const latency = Date.now() - start;
|
||||
await this.hooks.afterApiResponse?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
|
||||
datastoreType: this.datastoreType,
|
||||
datastoreProvider: this.datastoreProvider,
|
||||
query,
|
||||
response,
|
||||
context: mergedContext,
|
||||
latency,
|
||||
});
|
||||
|
||||
const queryResult: QueryResult<DocMeta> = {
|
||||
query: query.query,
|
||||
docs: response.matches,
|
||||
};
|
||||
|
||||
return queryResult;
|
||||
}
|
||||
|
||||
async upsert(docs: Doc<DocMeta>[], context?: Ctx): Promise<void> {
|
||||
const mergedContext = { ...this.context, ...context };
|
||||
try {
|
||||
// Get the text content from the docs
|
||||
const texts = docs.map((doc) => doc.metadata.content);
|
||||
|
||||
// Create the embeddings and sparse vectors
|
||||
// This relies on the classes to handle batching and throttling
|
||||
const [embeddingRes, spladeRes] = await Promise.all([
|
||||
this.embeddingModel.run({ input: texts }),
|
||||
this.spladeModel.run({ input: texts }),
|
||||
]);
|
||||
|
||||
// Combine the results into Pinecones vector format and upsert
|
||||
return this.pinecone.upsert({
|
||||
vectors: docs.map((doc, i) => ({
|
||||
id: doc.id,
|
||||
values: embeddingRes.embeddings[i],
|
||||
sparseValues: spladeRes.vectors[i],
|
||||
metadata: doc.metadata,
|
||||
})),
|
||||
});
|
||||
} catch (error) {
|
||||
await this.hooks.beforeError?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
datastoreType: this.datastoreType,
|
||||
datastoreProvider: this.datastoreProvider,
|
||||
upsert: docs,
|
||||
error,
|
||||
context: mergedContext,
|
||||
});
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async delete(docIds: string[]): Promise<void> {
|
||||
return this.pinecone.delete({ ids: docIds });
|
||||
}
|
||||
|
||||
async deleteAll(): Promise<void> {
|
||||
return this.pinecone.delete({ deleteAll: true });
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
export { Datastore, HybridDatastore } from './datastore.js';
|
||||
|
||||
export type { Doc, IDatastore } from '../types.js';
|
||||
@@ -0,0 +1,138 @@
|
||||
import type {
|
||||
IEmbeddingModel,
|
||||
ISparseModel,
|
||||
SparseValues,
|
||||
} from '@dexaai/model';
|
||||
|
||||
export type DatastoreType = 'embedding' | 'hybrid';
|
||||
export type DatastoreProvider = 'pinecone' | 'custom';
|
||||
|
||||
export type BaseMeta = { content: string };
|
||||
|
||||
/** Generic metadata object. */
|
||||
export type Ctx = { [key: string]: any };
|
||||
|
||||
export interface QueryCache<DocMeta extends BaseMeta> {
|
||||
get(key: Query<DocMeta>): Promise<QueryResult<DocMeta> | null>;
|
||||
set(key: Query<DocMeta>, value: QueryResult<DocMeta>): Promise<boolean>;
|
||||
}
|
||||
|
||||
export interface Hooks<DocMeta extends BaseMeta> {
|
||||
afterApiResponse?: (event: {
|
||||
timestamp: string;
|
||||
datastoreType: DatastoreType;
|
||||
datastoreProvider: DatastoreProvider;
|
||||
query: Query<DocMeta>;
|
||||
response: unknown;
|
||||
latency: number;
|
||||
context: Ctx;
|
||||
}) => void | Promise<void>;
|
||||
afterCacheHit?: (event: {
|
||||
timestamp: string;
|
||||
datastoreType: DatastoreType;
|
||||
datastoreProvider: DatastoreProvider;
|
||||
query: Query<DocMeta>;
|
||||
response: unknown;
|
||||
context: Ctx;
|
||||
}) => void | Promise<void>;
|
||||
beforeError?: (event: {
|
||||
timestamp: string;
|
||||
datastoreType: DatastoreType;
|
||||
datastoreProvider: DatastoreProvider;
|
||||
query?: Query<DocMeta>;
|
||||
upsert?: Doc<DocMeta>[];
|
||||
error: unknown;
|
||||
context: Ctx;
|
||||
}) => void | Promise<void>;
|
||||
}
|
||||
|
||||
export interface IDatastore<DocMeta extends BaseMeta> {
|
||||
/** Query the DataStore for documents. */
|
||||
query(query: Query<DocMeta>, context?: Ctx): Promise<QueryResult<DocMeta>>;
|
||||
|
||||
/** Insert or update documents in the DataStore. */
|
||||
upsert(docs: Doc<DocMeta>[], context?: Ctx): Promise<void>;
|
||||
|
||||
/** Delete documents by ID from the DataStore. */
|
||||
delete(docIds: string[]): Promise<void>;
|
||||
|
||||
/** Delete all documents from the DataStore. */
|
||||
deleteAll(): Promise<void>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Options for creating a Datastore instance.
|
||||
*/
|
||||
export interface DatastoreOpts<DocMeta extends BaseMeta> {
|
||||
namespace: string;
|
||||
embeddingModel: IEmbeddingModel;
|
||||
cache?: QueryCache<DocMeta>;
|
||||
hooks?: Hooks<DocMeta>;
|
||||
context?: Ctx;
|
||||
debug?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Options for creating a hybrid Datastore instance (with Splade).
|
||||
*/
|
||||
export interface HDatastoreOpts<DocMeta extends BaseMeta>
|
||||
extends DatastoreOpts<DocMeta> {
|
||||
/**
|
||||
* Splade instance for creating sparse vectors
|
||||
*/
|
||||
spladeModel: ISparseModel;
|
||||
}
|
||||
|
||||
export interface Doc<Meta extends BaseMeta = BaseMeta> {
|
||||
id: string;
|
||||
metadata: Meta;
|
||||
embedding?: number[];
|
||||
sparseVector?: SparseValues;
|
||||
score?: number;
|
||||
}
|
||||
export interface ScoredDoc<Meta extends BaseMeta = BaseMeta> extends Doc<Meta> {
|
||||
score: number;
|
||||
}
|
||||
|
||||
export interface Query<Meta extends BaseMeta> {
|
||||
query: string;
|
||||
embedding?: number[];
|
||||
sparseVector?: SparseValues;
|
||||
topK?: number;
|
||||
minScore?: number;
|
||||
filter?: QueryFilter<Meta>;
|
||||
includeValues?: boolean;
|
||||
}
|
||||
export interface QueryResult<Meta extends BaseMeta> {
|
||||
query: string;
|
||||
docs: ScoredDoc<Meta>[];
|
||||
}
|
||||
|
||||
/**
|
||||
* The possible leaf values for filter objects.
|
||||
* @note Null values aren't supported in metadata for filters, but are allowed here and automatically removed for convenience.
|
||||
*/
|
||||
type FilterValue = string | number | boolean | null | string[] | number[];
|
||||
type FilterOperator =
|
||||
| '$eq'
|
||||
| '$ne'
|
||||
| '$gt'
|
||||
| '$gte'
|
||||
| '$lt'
|
||||
| '$lte'
|
||||
| '$in'
|
||||
| '$nin';
|
||||
/**
|
||||
* An object of metadata filters.
|
||||
* @see https://www.pinecone.io/docs/metadata-filtering/
|
||||
*/
|
||||
export type QueryFilter<Metadata extends BaseMeta> = {
|
||||
[key in keyof Metadata | FilterOperator]?:
|
||||
| FilterValue
|
||||
| {
|
||||
[key in keyof Metadata | FilterOperator]?: FilterValue;
|
||||
};
|
||||
};
|
||||
|
||||
/** Improve preview of union types in autocomplete. */
|
||||
export type Prettify<T> = { [K in keyof T]: T[K] } & {};
|
||||
@@ -0,0 +1,80 @@
|
||||
/** Utilities for working with vectors/embeddings. */
|
||||
export class VectorUtils {
|
||||
/** Calculate the cosine similarity between two vectors. */
|
||||
static cosineSimilarity(a: number[], b: number[]): number {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error('Vectors must have the same length.');
|
||||
}
|
||||
|
||||
let dotProduct = 0;
|
||||
let aMagnitudeSquared = 0;
|
||||
let bMagnitudeSquared = 0;
|
||||
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
const aValue = a[i];
|
||||
const bValue = b[i];
|
||||
|
||||
dotProduct += aValue * bValue;
|
||||
aMagnitudeSquared += aValue * aValue;
|
||||
bMagnitudeSquared += bValue * bValue;
|
||||
}
|
||||
|
||||
// Check for non-zero magnitude after the loop to avoid multiple divisions by zero.
|
||||
if (aMagnitudeSquared === 0 || bMagnitudeSquared === 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const magnitudeProduct = aMagnitudeSquared * bMagnitudeSquared;
|
||||
const magnitudeProductRoot = Math.sqrt(magnitudeProduct);
|
||||
|
||||
return dotProduct / magnitudeProductRoot;
|
||||
}
|
||||
|
||||
/** Calculate the dot product of two vectors */
|
||||
static dotProduct(a: number[], b: number[]): number {
|
||||
if (a.length !== b.length) {
|
||||
throw new Error('Vectors must have the same length.');
|
||||
}
|
||||
|
||||
let dotProduct = 0;
|
||||
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
const aValue = a[i];
|
||||
const bValue = b[i];
|
||||
|
||||
dotProduct += aValue * bValue;
|
||||
}
|
||||
|
||||
return dotProduct;
|
||||
}
|
||||
|
||||
/**
|
||||
* Find the nearest neighbors of a vector in a set of documents with embeddings.
|
||||
* @param vector The vector to find neighbors for.
|
||||
* @param docs The set of documents with a vector/embedding to search.
|
||||
* @param topK The number of neighbors to find.
|
||||
* @param distanceFunction The distance function to use.
|
||||
* @returns The k nearest neighbors of the vector with the similarity score added (sorted by similarity).
|
||||
*/
|
||||
static nearestNeighbors<D extends { embedding: number[] }>(args: {
|
||||
vector: number[];
|
||||
docs: D[];
|
||||
topK: number;
|
||||
distanceFunction?: 'cosineSimilarity' | 'dotProduct';
|
||||
}): (D & { score: number })[] {
|
||||
const { vector, docs, topK, distanceFunction } = args;
|
||||
const distance =
|
||||
distanceFunction === 'cosineSimilarity'
|
||||
? VectorUtils.cosineSimilarity
|
||||
: VectorUtils.dotProduct;
|
||||
const distances: (D & { score: number })[] = docs.map((doc) => ({
|
||||
...doc,
|
||||
score: distance(vector, doc.embedding),
|
||||
}));
|
||||
|
||||
// Sort the distances in descending order of similarity (higher similarity means closer)
|
||||
distances.sort((a, b) => b.score - a.score);
|
||||
|
||||
return distances.slice(0, topK);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
describe('should', () => {
|
||||
it('exported', () => {
|
||||
expect(1).toEqual(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"extends": "tsconfig/node-esm.json",
|
||||
"include": ["./src/**/*"],
|
||||
"exclude": ["**/node_modules", "**/.*/"],
|
||||
"compilerOptions": {
|
||||
"types": ["vitest/globals", "node"]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
import { defineConfig } from 'tsup';
|
||||
|
||||
export default defineConfig({
|
||||
entry: ['src/index.ts', 'src/pinecone/index.ts'],
|
||||
format: ['esm', 'cjs'],
|
||||
clean: true,
|
||||
dts: true,
|
||||
target: 'es2022',
|
||||
});
|
||||
@@ -0,0 +1,12 @@
|
||||
import { defineConfig } from 'vite';
|
||||
|
||||
export default defineConfig(() => {
|
||||
return {
|
||||
mode: 'test',
|
||||
|
||||
test: {
|
||||
environment: 'node',
|
||||
globals: true,
|
||||
},
|
||||
};
|
||||
});
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"extends": ["@dexaai/eslint-config", "@dexaai/eslint-config/node"]
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
{
|
||||
"name": "@dexaai/model",
|
||||
"type": "module",
|
||||
"version": "0.0.0",
|
||||
"description": "",
|
||||
"repository": "dexaai/llm-tools",
|
||||
"scripts": {
|
||||
"build": "tsup",
|
||||
"clean": "rm -rf .turbo node_modules dist",
|
||||
"dev": "tsup --watch",
|
||||
"lint": "eslint --cache --cache-location ./node_modules/.cache/eslint .",
|
||||
"test": "vitest",
|
||||
"typecheck": "tsc --noEmit"
|
||||
},
|
||||
"sideEffects": false,
|
||||
"files": [
|
||||
"dist"
|
||||
],
|
||||
"types": "./dist/index.d.ts",
|
||||
"main": "./dist/index.cjs",
|
||||
"module": "./dist/index.js",
|
||||
"exports": {
|
||||
"./package.json": "./package.json",
|
||||
".": {
|
||||
"types": "./dist/index.d.ts",
|
||||
"require": "./dist/index.cjs",
|
||||
"import": "./dist/index.js",
|
||||
"default": "./dist/index.js"
|
||||
},
|
||||
"./openai": {
|
||||
"types": "./dist/openai/index.d.ts",
|
||||
"require": "./dist/openai/index.cjs",
|
||||
"import": "./dist/openai/index.js",
|
||||
"default": "./dist/openai/index.js"
|
||||
},
|
||||
"./custom": {
|
||||
"types": "./dist/custom/index.d.ts",
|
||||
"require": "./dist/custom/index.cjs",
|
||||
"import": "./dist/custom/index.js",
|
||||
"default": "./dist/custom/index.js"
|
||||
}
|
||||
},
|
||||
"typesVersions": {
|
||||
"*": {
|
||||
"openai": [
|
||||
"./dist/openai/index.d.ts"
|
||||
],
|
||||
"custom": [
|
||||
"./dist/custom/index.d.ts"
|
||||
],
|
||||
"tokenizer": [
|
||||
"./dist/tokenizer/index.d.ts"
|
||||
]
|
||||
}
|
||||
},
|
||||
"dependencies": {
|
||||
"@dqbd/tiktoken": "^1.0.7",
|
||||
"froebel": "^0.23.2",
|
||||
"keyv": "^4.5.2",
|
||||
"ky": "^0.33.3",
|
||||
"openai-fetch": "^1.6.3",
|
||||
"p-map": "^6.0.0",
|
||||
"p-retry": "^5.1.2",
|
||||
"p-throttle": "^5.1.0",
|
||||
"parse-json": "^7.0.0",
|
||||
"tkv-cache": "^0.1.2",
|
||||
"ts-extras": "^0.11.0",
|
||||
"uuid": "^9.0.0",
|
||||
"zod": "^3.21.4"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@dexaai/eslint-config": "^0.2.0",
|
||||
"@types/common-tags": "^1.8.1",
|
||||
"@types/node": "^20.3.1",
|
||||
"@types/uuid": "^9.0.2",
|
||||
"eslint": "^8.43.0",
|
||||
"tsconfig": "workspace:*",
|
||||
"tsup": "^7.0.0",
|
||||
"type-fest": "3.12.0",
|
||||
"typescript": "^5.1.3",
|
||||
"vite": "^4.3.9",
|
||||
"vitest": "^0.32.2"
|
||||
},
|
||||
"publishConfig": {
|
||||
"access": "restricted"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">= 18"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
# LLM Models
|
||||
|
||||
Standardized classes for interacting with large language models.
|
||||
@@ -0,0 +1,3 @@
|
||||
export type { ISpladeModel } from './splade.js';
|
||||
|
||||
export { SpladeModel } from './splade.js';
|
||||
@@ -0,0 +1,112 @@
|
||||
import ky from 'ky';
|
||||
import pThrottle from 'p-throttle';
|
||||
import pMap from 'p-map';
|
||||
import type {
|
||||
Ctx,
|
||||
ISparseModel,
|
||||
SparseConfig,
|
||||
SparseResponse,
|
||||
SparseRun,
|
||||
SparseValues,
|
||||
} from '../types.js';
|
||||
import { AbstractModel } from '../model.js';
|
||||
|
||||
const SPLADE_MODELS = ['naver/splade-cocondenser-ensembledistil'] as const;
|
||||
type SpladeModelName = (typeof SPLADE_MODELS)[number];
|
||||
|
||||
interface CSparseConfig extends SparseConfig {
|
||||
model: SpladeModelName;
|
||||
}
|
||||
|
||||
export type ISpladeModel = ISparseModel<
|
||||
CSparseConfig,
|
||||
SparseRun,
|
||||
SparseResponse
|
||||
>;
|
||||
|
||||
interface CMultiSparseConfig extends CSparseConfig {
|
||||
concurrency?: number;
|
||||
throttleLimit?: number;
|
||||
throttleInterval?: number;
|
||||
}
|
||||
|
||||
export class SpladeModel
|
||||
extends AbstractModel<CMultiSparseConfig, SparseRun, SparseResponse>
|
||||
implements ISpladeModel
|
||||
{
|
||||
modelType = 'sparse-vector' as const;
|
||||
modelProvider = 'custom' as const;
|
||||
|
||||
protected async runModel(
|
||||
params: SparseRun & CMultiSparseConfig,
|
||||
context: Ctx
|
||||
): Promise<SparseResponse> {
|
||||
const interval = params.throttleInterval ?? 1000 * 60; // 1 minute
|
||||
const limit = params.throttleLimit ?? 600;
|
||||
const concurrency = params.concurrency ?? 10;
|
||||
|
||||
// Create a throttled version of the function for a single request
|
||||
const throttled = pThrottle({ limit, interval })(
|
||||
async (params: { input: string; model: SpladeModelName }) =>
|
||||
this.runSingle(params, context)
|
||||
);
|
||||
|
||||
// Run the requests in parallel, respecting the maxConcurrentRequests value
|
||||
const inputs = params.input.map((input) => ({
|
||||
input,
|
||||
model: params.model,
|
||||
}));
|
||||
const responses = await pMap(inputs, throttled, { concurrency });
|
||||
|
||||
return {
|
||||
vectors: responses.map((r) => r.vector),
|
||||
tokens: { prompt: 0, completion: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
protected async runSingle(
|
||||
params: { input: string } & CSparseConfig,
|
||||
context: Ctx
|
||||
) {
|
||||
const start = Date.now();
|
||||
const vector = await createSpladeVector(params);
|
||||
const latency = Date.now() - start;
|
||||
|
||||
// Don't need tokens for this model
|
||||
const tokens = { prompt: 0, completion: 0, total: 0 } as const;
|
||||
await this.hooks.afterApiResponse?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
modelType: this.modelType,
|
||||
modelProvider: this.modelProvider,
|
||||
// @ts-ignore
|
||||
params,
|
||||
response: vector,
|
||||
tokens,
|
||||
context,
|
||||
latency,
|
||||
});
|
||||
return { vector, tokens };
|
||||
}
|
||||
}
|
||||
|
||||
async function createSpladeVector(params: {
|
||||
input: string;
|
||||
model: SpladeModelName;
|
||||
}): Promise<SparseValues> {
|
||||
try {
|
||||
const SERVICE_URL = process.env['SPLADE_SERVICE_URL'];
|
||||
if (!SERVICE_URL) {
|
||||
throw new Error('Missing process.env.SPLADE_SERVICE_URL');
|
||||
}
|
||||
const sparseValues = await ky
|
||||
.post(SERVICE_URL, {
|
||||
timeout: 1000 * 60,
|
||||
json: { text: params.input },
|
||||
})
|
||||
.json<SparseValues>();
|
||||
return sparseValues;
|
||||
} catch (e) {
|
||||
// @ts-ignore: TODO: add custom Error class that handles this
|
||||
throw new Error('Failed to create splade vector', { cause: e });
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
export type {
|
||||
ChatMessage,
|
||||
IChatModel,
|
||||
ICompletionModel,
|
||||
IEmbeddingModel,
|
||||
ISparseModel,
|
||||
ITokenizer,
|
||||
SparseValues,
|
||||
} from './types.js';
|
||||
|
||||
export { AbstractModel } from './model.js';
|
||||
@@ -0,0 +1,131 @@
|
||||
import Keyv from 'keyv';
|
||||
import { KvCache } from 'tkv-cache';
|
||||
import type {
|
||||
Ctx,
|
||||
Hooks,
|
||||
IModel,
|
||||
ModelCache,
|
||||
ModelConfig,
|
||||
ModelProvider,
|
||||
ModelResponse,
|
||||
ModelRun,
|
||||
ModelType,
|
||||
Prettify,
|
||||
} from './types.js';
|
||||
|
||||
export interface ModelArgs<
|
||||
MConfig extends ModelConfig,
|
||||
MRun extends ModelRun,
|
||||
MResponse extends ModelResponse
|
||||
> {
|
||||
cache?: ModelCache<MRun & MConfig, MResponse>;
|
||||
context?: Ctx;
|
||||
params: MConfig & Partial<MRun>;
|
||||
hooks?: Hooks<MRun & MConfig, MResponse>;
|
||||
debug?: boolean;
|
||||
}
|
||||
|
||||
export abstract class AbstractModel<
|
||||
MConfig extends ModelConfig,
|
||||
MRun extends ModelRun,
|
||||
MResponse extends ModelResponse,
|
||||
AResponse extends any = any
|
||||
> implements IModel<MConfig, MRun, MResponse>
|
||||
{
|
||||
protected abstract runModel(
|
||||
params: Prettify<MRun & MConfig>,
|
||||
context: Ctx
|
||||
): Promise<MResponse>;
|
||||
|
||||
abstract modelType: ModelType;
|
||||
abstract modelProvider: ModelProvider;
|
||||
protected cache?: ModelCache<MRun & MConfig, MResponse>;
|
||||
protected context: Ctx;
|
||||
protected params: MConfig & Partial<MRun>;
|
||||
protected hooks: Hooks<MRun & MConfig, MResponse, AResponse>;
|
||||
|
||||
constructor(args: ModelArgs<MConfig, MRun, MResponse>) {
|
||||
this.cache = args.cache;
|
||||
this.context = args.context ?? {};
|
||||
this.params = args.params;
|
||||
this.hooks = args.debug
|
||||
? {
|
||||
afterApiResponse: console.log,
|
||||
afterCacheHit: console.log,
|
||||
beforeError: console.error,
|
||||
...args.hooks,
|
||||
}
|
||||
: args.hooks || {};
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a cache to the model for storing responses
|
||||
* Defaults to an in-memory Keyv cache.
|
||||
*/
|
||||
addCache(args?: {
|
||||
/** A Keyv instance to use for caching */
|
||||
keyv?: Keyv<MResponse>;
|
||||
/** Keyv options to create a new Keyv instance */
|
||||
keyvOpts?: Keyv.Options<MResponse>;
|
||||
/** A function to normalize the cache key */
|
||||
normalizeKey?: (params: MRun & MConfig) => Partial<MRun & MConfig>;
|
||||
errorHandler?: (error: unknown) => void;
|
||||
}): this {
|
||||
const { keyvOpts, normalizeKey, errorHandler } = args ?? {};
|
||||
const keyv = args?.keyv ?? new Keyv<MResponse>(keyvOpts);
|
||||
this.cache = new KvCache(keyv, normalizeKey, errorHandler);
|
||||
return this;
|
||||
}
|
||||
|
||||
/** Add hooks to the model for logging and debugging */
|
||||
addHooks(hooks: Hooks<MRun & MConfig, MResponse, AResponse>): this {
|
||||
this.hooks = { ...this.hooks, ...hooks };
|
||||
return this;
|
||||
}
|
||||
|
||||
async run(
|
||||
params: Prettify<MRun & Partial<MConfig>>,
|
||||
context?: Ctx
|
||||
): Promise<MResponse> {
|
||||
// Merge the default params with the provided params
|
||||
const mergedParams: MRun & MConfig = {
|
||||
...this.params,
|
||||
...params,
|
||||
};
|
||||
const mergedContext = { ...this.context, ...context };
|
||||
|
||||
// Return cached response if available
|
||||
const cached = await this?.cache?.get(mergedParams);
|
||||
if (cached) {
|
||||
await this.hooks.afterCacheHit?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
modelType: this.modelType,
|
||||
modelProvider: this.modelProvider,
|
||||
params: mergedParams,
|
||||
response: cached,
|
||||
context: mergedContext,
|
||||
});
|
||||
return cached;
|
||||
}
|
||||
|
||||
try {
|
||||
// Run the model (e.g. make the API request)
|
||||
const response = await this.runModel(mergedParams, mergedContext);
|
||||
|
||||
// Update the cache
|
||||
await this?.cache?.set(mergedParams, response);
|
||||
|
||||
return response;
|
||||
} catch (error) {
|
||||
await this.hooks.beforeError?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
modelType: this.modelType,
|
||||
modelProvider: this.modelProvider,
|
||||
params: mergedParams,
|
||||
error,
|
||||
context: mergedContext,
|
||||
});
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
import type {
|
||||
ChatConfig,
|
||||
ChatResponse,
|
||||
ChatRun,
|
||||
Ctx,
|
||||
IChatModel,
|
||||
Prettify,
|
||||
TokenCounts,
|
||||
} from '../types.js';
|
||||
import type { ModelArgs } from '../model.js';
|
||||
import { AbstractModel } from '../model.js';
|
||||
import { createOpenAIClient, extractTokens } from './client.js';
|
||||
import { createTokenizer } from './tokenizer.js';
|
||||
import type {
|
||||
ChatMessage,
|
||||
OpenAIChatModel,
|
||||
OpenAIChatParams,
|
||||
OpenAIChatResponse,
|
||||
OpenAIClient,
|
||||
} from './client.js';
|
||||
|
||||
interface OChatConfig
|
||||
extends ChatConfig,
|
||||
Omit<OpenAIChatParams, 'messages' | 'user'> {
|
||||
/** Handle new chunks from streaming requests. */
|
||||
handleUpdate?: (chunk: string) => void;
|
||||
model: OpenAIChatModel;
|
||||
}
|
||||
|
||||
export type IOChatModel = IChatModel<OChatConfig, ChatRun, ChatResponse>;
|
||||
|
||||
export class OChatModel
|
||||
extends AbstractModel<OChatConfig, ChatRun, ChatResponse, OpenAIChatResponse>
|
||||
implements IOChatModel
|
||||
{
|
||||
modelType = 'chat' as const;
|
||||
modelProvider = 'openai' as const;
|
||||
openaiClient: OpenAIClient;
|
||||
|
||||
constructor(
|
||||
args: Prettify<
|
||||
ModelArgs<OChatConfig, ChatRun, ChatResponse> & {
|
||||
openaiClient?: OpenAIClient;
|
||||
}
|
||||
>
|
||||
) {
|
||||
const { openaiClient, ...rest } = args;
|
||||
super(rest);
|
||||
this.openaiClient = openaiClient || createOpenAIClient();
|
||||
}
|
||||
|
||||
async runModel(
|
||||
params: ChatRun & OChatConfig,
|
||||
context: Ctx
|
||||
): Promise<ChatResponse> {
|
||||
const { handleUpdate, ...restParams } = params;
|
||||
// Use non-streaming API if no handler is provided
|
||||
if (!handleUpdate) {
|
||||
const start = Date.now();
|
||||
const { message, response } =
|
||||
await this.openaiClient.createChatCompletion({
|
||||
user: typeof context.user === 'string' ? context.user : '',
|
||||
...restParams,
|
||||
});
|
||||
const latency = Date.now() - start;
|
||||
const tokens = extractTokens(response.usage);
|
||||
await this.hooks.afterApiResponse?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
modelType: this.modelType,
|
||||
modelProvider: this.modelProvider,
|
||||
params: restParams,
|
||||
response,
|
||||
tokens,
|
||||
context,
|
||||
latency,
|
||||
});
|
||||
return { message, tokens };
|
||||
} else {
|
||||
const start = Date.now();
|
||||
const stream = await this.openaiClient.streamChatCompletion({
|
||||
user: typeof context.user === 'string' ? context.user : '',
|
||||
...restParams,
|
||||
});
|
||||
|
||||
// Keep track of the stream's output
|
||||
let text = '';
|
||||
let role = '';
|
||||
let response = {};
|
||||
|
||||
// Get a reader from the stream
|
||||
const reader = stream.getReader();
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
|
||||
if (done) {
|
||||
// If the stream is done, break out of the loop and save the conversation
|
||||
// to the cache before returning.
|
||||
break;
|
||||
}
|
||||
|
||||
// Only set role once
|
||||
if (!role && value?.message?.role) {
|
||||
role = value.message.role;
|
||||
}
|
||||
|
||||
response = { ...response, ...value?.response };
|
||||
|
||||
// Send an update to the caller
|
||||
const messageContent = value?.message?.content;
|
||||
if (typeof messageContent === 'string') {
|
||||
text = `${text}${messageContent}`;
|
||||
try {
|
||||
// @TODO: Should this be throttled?
|
||||
handleUpdate(messageContent);
|
||||
} catch (err) {
|
||||
console.error('Error handling update', err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Once the stream is done, release the reader
|
||||
reader.releaseLock();
|
||||
const latency = Date.now() - start;
|
||||
|
||||
// Construct the complete message
|
||||
const message: ChatMessage = {
|
||||
role: role as ChatMessage['role'],
|
||||
content: text,
|
||||
};
|
||||
|
||||
// Streamed responses don't include token usage so we have to add it
|
||||
const tokenizer = createTokenizer(params.model);
|
||||
const promptTokens = tokenizer.countTokens(params.messages);
|
||||
const completionTokens = tokenizer.countTokens(message);
|
||||
const tokens: TokenCounts = {
|
||||
prompt: promptTokens,
|
||||
completion: completionTokens,
|
||||
total: promptTokens + completionTokens,
|
||||
};
|
||||
const responseUsage = {
|
||||
prompt_tokens: tokens.prompt,
|
||||
completion_tokens: tokens.completion,
|
||||
total_tokens: tokens.total,
|
||||
};
|
||||
response = { ...response, usage: responseUsage };
|
||||
|
||||
await this.hooks.afterApiResponse?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
modelType: this.modelType,
|
||||
modelProvider: this.modelProvider,
|
||||
params: restParams,
|
||||
response: response as OpenAIChatResponse,
|
||||
tokens,
|
||||
context,
|
||||
latency,
|
||||
});
|
||||
return { message, tokens };
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
import type { ChatMessage } from 'openai-fetch';
|
||||
import { OpenAIClient } from 'openai-fetch';
|
||||
import type { TokenCounts } from '../types.js';
|
||||
|
||||
let cachedClient: OpenAIClient | undefined;
|
||||
|
||||
/** Create a new openai-fetch OpenAIClient. */
|
||||
export function createOpenAIClient(
|
||||
/** Options to pass to the OpenAI client. */
|
||||
opts?: ConstructorParameters<typeof OpenAIClient>[0],
|
||||
/** Force a new client to be created. */
|
||||
forceNew = false
|
||||
) {
|
||||
// Only use a cached client if no options are passed.
|
||||
if (opts === undefined && cachedClient && !forceNew) {
|
||||
return cachedClient;
|
||||
}
|
||||
return new OpenAIClient(opts);
|
||||
}
|
||||
|
||||
export { OpenAIApiError, OpenAIClient } from 'openai-fetch';
|
||||
|
||||
export type {
|
||||
ChatCompletionParams as OpenAIChatParams,
|
||||
ChatCompletionResponse as OpenAIChatResponse,
|
||||
ChatMessageFunction as OpenAIChatFunction,
|
||||
ChatMessage,
|
||||
CompletionParams as OpenAICompletionParams,
|
||||
CompletionResponse as OpenAICompletionResponse,
|
||||
EmbeddingParams as OpenAIEmbeddingParams,
|
||||
EmbeddingResponse as OpenAIEmbeddingResponse,
|
||||
} from 'openai-fetch';
|
||||
|
||||
/** Extract tokens from an OpenAI API response */
|
||||
export function extractTokens(
|
||||
usage: Record<string, number | undefined> | undefined
|
||||
) {
|
||||
const tokens: TokenCounts = {
|
||||
prompt: usage?.['prompt_tokens'] ?? 0,
|
||||
completion: usage?.['completion_tokens'] ?? 0,
|
||||
total: usage?.['total_tokens'] ?? 0,
|
||||
};
|
||||
return tokens;
|
||||
}
|
||||
|
||||
/** Reformat name to adhere to OpenAI's naming restrictions: /^[a-zA-Z0-9_-]{1,64}$/ */
|
||||
export function formatName<Msg extends ChatMessage>(message: Msg): Msg {
|
||||
const { name, role, content } = message;
|
||||
|
||||
// Remove the name key if it's empty.
|
||||
if (!name) return { role, content } as Msg;
|
||||
|
||||
// Reformat to meet OpenAI's naming restrictions.
|
||||
const formattedName = name
|
||||
.replace(/\s/g, '-')
|
||||
.replace(/[^a-zA-Z0-9_]/g, '')
|
||||
.toLowerCase();
|
||||
|
||||
return { role, name: formattedName, content } as Msg;
|
||||
}
|
||||
|
||||
const CHAT_MODELS = [
|
||||
'gpt-3.5-turbo',
|
||||
'gpt-3.5-turbo-0301',
|
||||
'gpt-4',
|
||||
'gpt-4-0314',
|
||||
'gpt-4-32k',
|
||||
'gpt-4-32k-0314',
|
||||
'gpt-4-0613',
|
||||
'gpt-4-32k-0613',
|
||||
'gpt-3.5-turbo-0613',
|
||||
'gpt-3.5-turbo-16k',
|
||||
] as const;
|
||||
export type OpenAIChatModel = (typeof CHAT_MODELS)[number];
|
||||
|
||||
const EMBEDDING_MODELS = [
|
||||
'text-embedding-ada-002',
|
||||
'text-similarity-davinci-001',
|
||||
'text-similarity-curie-001',
|
||||
'text-similarity-babbage-001',
|
||||
'text-similarity-ada-001',
|
||||
'text-search-davinci-doc-001',
|
||||
'text-search-curie-doc-001',
|
||||
'text-search-babbage-doc-001',
|
||||
'text-search-ada-doc-001',
|
||||
'code-search-babbage-code-001',
|
||||
'code-search-ada-code-001',
|
||||
] as const;
|
||||
export type OpenAIEmbeddingModel = (typeof EMBEDDING_MODELS)[number];
|
||||
@@ -0,0 +1,93 @@
|
||||
import type {
|
||||
CompletionConfig,
|
||||
CompletionResponse,
|
||||
CompletionRun,
|
||||
Ctx,
|
||||
ICompletionModel,
|
||||
Prettify,
|
||||
} from '../types.js';
|
||||
import type { ModelArgs } from '../model.js';
|
||||
import { AbstractModel } from '../model.js';
|
||||
import type {
|
||||
OpenAIClient,
|
||||
OpenAICompletionParams,
|
||||
OpenAICompletionResponse,
|
||||
} from './client.js';
|
||||
import { createOpenAIClient, extractTokens } from './client.js';
|
||||
|
||||
interface OCompletionConfig
|
||||
extends CompletionConfig,
|
||||
Omit<OpenAICompletionParams, 'prompt' | 'user'> {}
|
||||
|
||||
interface OCompletionsResponse extends CompletionResponse {
|
||||
completions: {
|
||||
completion: string;
|
||||
logprobs?: OpenAICompletionResponse['choices'][0]['logprobs'];
|
||||
}[];
|
||||
}
|
||||
|
||||
export type IOCompletionModel = ICompletionModel<
|
||||
OCompletionConfig,
|
||||
CompletionRun,
|
||||
OCompletionsResponse
|
||||
>;
|
||||
|
||||
export class OCompletionModel
|
||||
extends AbstractModel<
|
||||
OCompletionConfig,
|
||||
CompletionRun,
|
||||
OCompletionsResponse,
|
||||
OpenAICompletionResponse
|
||||
>
|
||||
implements IOCompletionModel
|
||||
{
|
||||
modelType = 'completion' as const;
|
||||
modelProvider = 'openai' as const;
|
||||
openaiClient: OpenAIClient;
|
||||
|
||||
constructor(
|
||||
args: Prettify<
|
||||
ModelArgs<OCompletionConfig, CompletionRun, OCompletionsResponse> & {
|
||||
openaiClient?: OpenAIClient;
|
||||
}
|
||||
>
|
||||
) {
|
||||
const { openaiClient, ...rest } = args;
|
||||
super(rest);
|
||||
this.openaiClient = openaiClient || createOpenAIClient();
|
||||
}
|
||||
|
||||
async runModel(
|
||||
params: CompletionRun & OCompletionConfig,
|
||||
context: Ctx
|
||||
): Promise<OCompletionsResponse> {
|
||||
const start = Date.now();
|
||||
const { response: apiResponse } = await this.openaiClient.createCompletions(
|
||||
{
|
||||
user: typeof context.user === 'string' ? context.user : '',
|
||||
...params,
|
||||
}
|
||||
);
|
||||
const latency = Date.now() - start;
|
||||
const tokens = extractTokens(apiResponse.usage);
|
||||
await this.hooks.afterApiResponse?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
modelType: this.modelType,
|
||||
modelProvider: this.modelProvider,
|
||||
params,
|
||||
response: apiResponse,
|
||||
tokens,
|
||||
context,
|
||||
latency,
|
||||
});
|
||||
|
||||
const response: OCompletionsResponse = {
|
||||
tokens,
|
||||
completions: apiResponse.choices.map((choice) => ({
|
||||
completion: choice.text || '',
|
||||
logprobs: choice.logprobs,
|
||||
})),
|
||||
};
|
||||
return response;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
import pThrottle from 'p-throttle';
|
||||
import pRetry, { AbortError } from 'p-retry';
|
||||
import pMap from 'p-map';
|
||||
import type {
|
||||
EmbeddingConfig,
|
||||
Ctx,
|
||||
EmbeddingRun,
|
||||
EmbeddingResponse,
|
||||
IEmbeddingModel,
|
||||
ITokenizer,
|
||||
TokenCounts,
|
||||
} from '../types.js';
|
||||
import type { ModelArgs } from '../model.js';
|
||||
import { AbstractModel } from '../model.js';
|
||||
import { OpenAIApiError, createOpenAIClient, extractTokens } from './client.js';
|
||||
import type {
|
||||
OpenAIClient,
|
||||
OpenAIEmbeddingModel,
|
||||
OpenAIEmbeddingParams,
|
||||
OpenAIEmbeddingResponse,
|
||||
} from './client.js';
|
||||
import { createTokenizer } from './tokenizer.js';
|
||||
|
||||
interface BatchOptions {
|
||||
maxTokensPerBatch: number;
|
||||
maxBatchSize: number;
|
||||
}
|
||||
interface ThrottleOptions {
|
||||
maxRequestsPerMin: number;
|
||||
maxConcurrentRequests: number;
|
||||
}
|
||||
interface RetryOptions {
|
||||
maxRetries: number;
|
||||
jitter: number;
|
||||
initialDelay: number;
|
||||
exponentialBase: number;
|
||||
}
|
||||
interface OEmbeddingConfig
|
||||
extends EmbeddingConfig,
|
||||
Omit<OpenAIEmbeddingParams, 'input' | 'user'> {
|
||||
model: OpenAIEmbeddingModel;
|
||||
batch?: Partial<BatchOptions>;
|
||||
throttle?: Partial<ThrottleOptions>;
|
||||
retry?: Partial<RetryOptions>;
|
||||
}
|
||||
|
||||
export type IOEmbeddingModel = IEmbeddingModel<
|
||||
OEmbeddingConfig,
|
||||
EmbeddingRun,
|
||||
EmbeddingResponse
|
||||
>;
|
||||
|
||||
type BulkEmbedder = (
|
||||
params: EmbeddingRun & OEmbeddingConfig,
|
||||
context: Ctx
|
||||
) => Promise<EmbeddingResponse>;
|
||||
|
||||
const DEFAULTS = {
|
||||
/** The interval that the OpenAI API rate limit resets */
|
||||
throttleInterval: 1000 * 60, // 1 minute
|
||||
maxBatchSize: 10,
|
||||
maxTokensPerBatch: 20000,
|
||||
maxConcurrentRequests: 1,
|
||||
maxRequestsPerMin: 3500,
|
||||
maxRetries: 3,
|
||||
jitter: 0.3,
|
||||
initialDelay: 1000,
|
||||
exponentialBase: 3,
|
||||
} as const;
|
||||
|
||||
export class OEmbeddingModel
|
||||
extends AbstractModel<
|
||||
OEmbeddingConfig,
|
||||
EmbeddingRun,
|
||||
EmbeddingResponse,
|
||||
OpenAIEmbeddingResponse
|
||||
>
|
||||
implements IOEmbeddingModel
|
||||
{
|
||||
modelType = 'embedding' as const;
|
||||
modelProvider = 'openai' as const;
|
||||
openaiClient: OpenAIClient;
|
||||
throttledModel: BulkEmbedder;
|
||||
tokenizer: ITokenizer;
|
||||
|
||||
/** Doesn't accept OpenAIClient because retry needs to be handled at the model level. */
|
||||
constructor(
|
||||
args: ModelArgs<OEmbeddingConfig, EmbeddingRun, EmbeddingResponse>
|
||||
) {
|
||||
const { params, ...rest } = args;
|
||||
super({ params, ...rest });
|
||||
this.openaiClient = createOpenAIClient();
|
||||
this.tokenizer = createTokenizer(params.model);
|
||||
const interval = DEFAULTS.throttleInterval;
|
||||
const limit =
|
||||
this.params.throttle?.maxRequestsPerMin || DEFAULTS.maxRequestsPerMin;
|
||||
// Create the throttled function
|
||||
this.throttledModel = pThrottle({ limit, interval })(
|
||||
async (params: EmbeddingRun & OEmbeddingConfig, context: Ctx) => {
|
||||
try {
|
||||
const start = Date.now();
|
||||
// Make the request to OpenAI API to generate embeddings
|
||||
const { embeddings, response } =
|
||||
await this.openaiClient.createEmbeddings({
|
||||
model: params.model,
|
||||
input: params.input,
|
||||
});
|
||||
const latency = Date.now() - start;
|
||||
const tokens = extractTokens(response.usage);
|
||||
await this.hooks.afterApiResponse?.({
|
||||
timestamp: new Date().toISOString(),
|
||||
modelType: this.modelType,
|
||||
modelProvider: this.modelProvider,
|
||||
params,
|
||||
response,
|
||||
tokens,
|
||||
context,
|
||||
latency,
|
||||
});
|
||||
return { embeddings, tokens };
|
||||
} catch (error) {
|
||||
const codesToRetry = [408, 413, 429, 500, 502, 503, 504];
|
||||
if (
|
||||
error instanceof OpenAIApiError &&
|
||||
error.status &&
|
||||
codesToRetry.includes(error.status)
|
||||
) {
|
||||
// If rate limit error, throw an error that can be retried
|
||||
throw error;
|
||||
}
|
||||
// Otherwise, throw an error that will not be retried
|
||||
throw new AbortError(error as Error);
|
||||
}
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
async runModel(
|
||||
params: EmbeddingRun,
|
||||
context: Ctx
|
||||
): Promise<EmbeddingResponse> {
|
||||
// Batch the inputs for the requests
|
||||
const batches = batchInputs({
|
||||
input: params.input,
|
||||
tokenizer: this.tokenizer,
|
||||
options: this.params.batch,
|
||||
}).map((batch) => batch.map((input) => input.text));
|
||||
|
||||
const mergedContext = { ...this.context, ...context } as Ctx;
|
||||
|
||||
// Make the requests in parallel, respecting concurrency setting
|
||||
const embeddingBatches = await pMap(
|
||||
batches,
|
||||
async (batch: string[]) => {
|
||||
return pRetry<EmbeddingResponse>(
|
||||
async () => {
|
||||
const response = await this.throttledModel(
|
||||
{
|
||||
input: batch,
|
||||
model: this.params.model,
|
||||
},
|
||||
mergedContext
|
||||
);
|
||||
return response;
|
||||
},
|
||||
{
|
||||
retries: this.params.retry?.maxRetries || DEFAULTS.maxRetries,
|
||||
onFailedAttempt: async (error) => {
|
||||
const delay: number = getBackoffTime({
|
||||
attemptNumber: error.attemptNumber,
|
||||
...this.params.retry,
|
||||
});
|
||||
await new Promise((resolve) => setTimeout(resolve, delay));
|
||||
},
|
||||
}
|
||||
);
|
||||
},
|
||||
{
|
||||
concurrency:
|
||||
this.params.throttle?.maxConcurrentRequests ||
|
||||
DEFAULTS.maxConcurrentRequests,
|
||||
}
|
||||
);
|
||||
|
||||
// Flatten the batches of embeddings into a single array
|
||||
const embeddings = embeddingBatches.map((batch) => batch.embeddings).flat();
|
||||
const tokens = embeddingBatches
|
||||
.map((batch) => batch.tokens)
|
||||
.reduce((acc, curr) => {
|
||||
return {
|
||||
total: acc.total + curr.total,
|
||||
prompt: acc.prompt + curr.prompt,
|
||||
completion: 0,
|
||||
};
|
||||
}, {} as TokenCounts);
|
||||
|
||||
return { embeddings, tokens };
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the time in milliseconds to wait before making another request after an error.
|
||||
* Use an exponential backoff strategy with a random jitter.
|
||||
* Adapted from: https://platform.openai.com/docs/guides/rate-limits/retrying-with-exponential-backoff
|
||||
*/
|
||||
function getBackoffTime(
|
||||
args: { attemptNumber: number } & Partial<RetryOptions>
|
||||
): number {
|
||||
const {
|
||||
exponentialBase = DEFAULTS.exponentialBase,
|
||||
initialDelay = DEFAULTS.initialDelay,
|
||||
jitter = DEFAULTS.jitter,
|
||||
} = args;
|
||||
|
||||
// Calculate the delay based on exponential backoff strategy
|
||||
const delay = initialDelay * exponentialBase ** (args.attemptNumber - 1);
|
||||
|
||||
// Calculate the random jitter (randomness within a certain range to avoid simultaneous retries)
|
||||
const randomJitter = delay * jitter * numberBetween(-1, 1);
|
||||
|
||||
// The final backoff time includes the jitter
|
||||
const backoffTime = delay + randomJitter;
|
||||
|
||||
return backoffTime;
|
||||
}
|
||||
|
||||
/** Get a random number between the specified range [min, max]. */
|
||||
function numberBetween(min: number, max: number): number {
|
||||
return Math.random() * (max - min) + min;
|
||||
}
|
||||
|
||||
type InputBatch = { text: string; tokenCount: number }[];
|
||||
|
||||
/**
|
||||
* Split text inputs into batches, respecting token limits and batch size.
|
||||
* @throws {Error} If an input exceeds the max tokens per batch.
|
||||
*/
|
||||
function batchInputs(args: {
|
||||
input: string[];
|
||||
tokenizer: ITokenizer;
|
||||
options?: Partial<BatchOptions>;
|
||||
}): InputBatch[] {
|
||||
const { input: inputs, tokenizer, options } = args;
|
||||
const {
|
||||
maxTokensPerBatch = DEFAULTS.maxTokensPerBatch,
|
||||
maxBatchSize = DEFAULTS.maxBatchSize,
|
||||
} = options || {};
|
||||
|
||||
// State for constructing batches
|
||||
const batches: InputBatch[] = [];
|
||||
let currentBatch: InputBatch = [];
|
||||
let currentBatchTokens = 0;
|
||||
|
||||
for (let input of inputs) {
|
||||
const tokenCount = tokenizer.countTokens(input);
|
||||
|
||||
// Ensure that the input does not exceed the max tokens per batch
|
||||
if (tokenCount > maxTokensPerBatch) {
|
||||
throw new Error(
|
||||
`Input exceeds max tokens per batch: ${tokenCount} > ${maxTokensPerBatch}`
|
||||
);
|
||||
}
|
||||
|
||||
// Add the first input to the batch
|
||||
if (currentBatch.length === 0) {
|
||||
currentBatch.push({ text: input, tokenCount });
|
||||
currentBatchTokens += tokenCount;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check if the input would exceed the max tokens per batch or max batch size
|
||||
const overTokenLimit = currentBatchTokens + tokenCount > maxTokensPerBatch;
|
||||
const overBatchSize = currentBatch.length + 1 > maxBatchSize;
|
||||
|
||||
// Start a new batch if either limit would be exceeded
|
||||
if (overTokenLimit || overBatchSize) {
|
||||
batches.push(currentBatch);
|
||||
currentBatch = [{ text: input, tokenCount }];
|
||||
currentBatchTokens = tokenCount;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Add the input to the current batch
|
||||
currentBatch.push({ text: input, tokenCount });
|
||||
currentBatchTokens += tokenCount;
|
||||
}
|
||||
|
||||
// Add the last batch
|
||||
if (currentBatch.length > 0) {
|
||||
batches.push(currentBatch);
|
||||
}
|
||||
|
||||
return batches;
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
export type { IOChatModel as IChatModel } from './chat.js';
|
||||
export type { IOCompletionModel as ICompletionModel } from './completion.js';
|
||||
export type { IOEmbeddingModel as IEmbeddingModel } from './embedding.js';
|
||||
export type { ITokenizer } from '../types.js';
|
||||
|
||||
export { OChatModel as ChatModel } from './chat.js';
|
||||
export { OCompletionModel as CompletionModel } from './completion.js';
|
||||
export { OEmbeddingModel as EmbeddingModel } from './embedding.js';
|
||||
export { createTokenizer } from './tokenizer.js';
|
||||
@@ -0,0 +1,128 @@
|
||||
import { encoding_for_model } from '@dqbd/tiktoken';
|
||||
import type { Tiktoken, TiktokenModel } from '@dqbd/tiktoken';
|
||||
import type { ChatMessage } from 'openai-fetch';
|
||||
import type { ITokenizer, CreateTokenizer } from '../types.js';
|
||||
|
||||
// Store instances of the tokenizer to avoid re-creating them for the same model
|
||||
const tokenizerCache = new Map<string, Tokenizer>();
|
||||
|
||||
/** Create a tokenizer for a specific model */
|
||||
export const createTokenizer: CreateTokenizer = (model: string) => {
|
||||
// export function createTokenizer(model: ModelName): ITokenizer {
|
||||
if (tokenizerCache.has(model)) {
|
||||
return tokenizerCache.get(model)!;
|
||||
} else {
|
||||
const tokenizer = new Tokenizer(model);
|
||||
tokenizerCache.set(model, tokenizer);
|
||||
return tokenizer;
|
||||
}
|
||||
};
|
||||
|
||||
const GPT_4_MODELS = [
|
||||
'gpt-4',
|
||||
'gpt-4-0314',
|
||||
'gpt-4-32k',
|
||||
'gpt-4-32k-0314',
|
||||
] as const;
|
||||
type Gpt4ModelName = (typeof GPT_4_MODELS)[number];
|
||||
|
||||
class Tokenizer implements ITokenizer {
|
||||
model: string;
|
||||
tiktoken: Tiktoken;
|
||||
|
||||
constructor(model: string) {
|
||||
this.model = model;
|
||||
this.tiktoken = encoding_for_model(model as TiktokenModel);
|
||||
}
|
||||
|
||||
/** Encode text to tokens */
|
||||
encode(text: string): Uint32Array {
|
||||
return this.tiktoken.encode(text);
|
||||
}
|
||||
|
||||
/** Decode tokens to text */
|
||||
decode(tokens: number[] | Uint32Array): string {
|
||||
const toDecode = Array.isArray(tokens) ? new Uint32Array(tokens) : tokens;
|
||||
return new TextDecoder().decode(this.tiktoken.decode(toDecode));
|
||||
}
|
||||
|
||||
/**
|
||||
* Count the number of tokens in a string or ChatMessage(s)
|
||||
* A single message is counted as a completion and an array as a prompt
|
||||
**/
|
||||
countTokens(input?: string | ChatMessage | ChatMessage[]): number {
|
||||
if (!input) return 0;
|
||||
if (typeof input === 'string') {
|
||||
return this.tiktoken.encode(input).length;
|
||||
} else if (Array.isArray(input)) {
|
||||
// This is copied from OpenAI's Python implementation: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_format_inputs_to_ChatGPT_models.ipynb
|
||||
// NOTE: The array version assumes this is for the prompt and adds more tokens.
|
||||
// Use the single ChatMessage for completion token count.
|
||||
// The tokensPerName and tokensPerMessage will need to be updated over time
|
||||
const isGpt4 = this.isGpt4Model(this.model);
|
||||
const tokensPerMessage = isGpt4 ? 3 : 4;
|
||||
let numTokens = 0;
|
||||
|
||||
for (const message of input) {
|
||||
numTokens += tokensPerMessage;
|
||||
if (message.content) {
|
||||
numTokens += this.countTokens(message.content);
|
||||
}
|
||||
// The name+role are handled differently for GPT-3.5 vs GPT-4
|
||||
if (isGpt4) {
|
||||
// For 4, the name and role are included
|
||||
// Details here: https://github.com/openai/openai-python/blob/main/chatml.md
|
||||
numTokens += 1; // role
|
||||
if (message.name) {
|
||||
// No idea why this, but tested with many examples and it works...
|
||||
numTokens += this.countTokens(`${message.name}`) + 1;
|
||||
}
|
||||
} else {
|
||||
// For 3.5, the name replaces the role if it's present
|
||||
numTokens += this.countTokens(message.name || message.role);
|
||||
}
|
||||
}
|
||||
|
||||
// Every reply is primed with assistant
|
||||
numTokens += 3;
|
||||
|
||||
return numTokens;
|
||||
} else {
|
||||
return this.countTokens(input.content || '');
|
||||
}
|
||||
}
|
||||
|
||||
/** Truncate text to a maximum number of tokens */
|
||||
truncate(args: {
|
||||
/** Text to truncate */
|
||||
text: string;
|
||||
/** Maximum number of tokens to keep (inclusive) */
|
||||
max: number;
|
||||
/** Truncate from the start or end of the text */
|
||||
from?: 'start' | 'end';
|
||||
}): string {
|
||||
const { text, max, from = 'start' } = args;
|
||||
const tokens = this.encode(text);
|
||||
|
||||
if (tokens.length <= max) {
|
||||
return text;
|
||||
}
|
||||
|
||||
const truncatedTokens =
|
||||
from === 'start'
|
||||
? tokens.slice(0, max)
|
||||
: tokens.slice(tokens.length - max);
|
||||
|
||||
let truncatedText = this.decode(truncatedTokens);
|
||||
|
||||
// Handle edge case where the last token is part of a multi-byte character sequence
|
||||
truncatedText = truncatedText.replace(/[\uDC00-\uDFFF]$/, '');
|
||||
|
||||
return truncatedText;
|
||||
}
|
||||
|
||||
/** Check if the given model is a GPT-4 variant. */
|
||||
private isGpt4Model(model: string): model is Gpt4ModelName {
|
||||
return GPT_4_MODELS.includes(model as any);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
import type { ChatMessage } from 'openai-fetch';
|
||||
|
||||
export type { ChatMessage } from 'openai-fetch';
|
||||
|
||||
/** Generic metadata object. */
|
||||
export type Ctx = { [key: string]: any };
|
||||
|
||||
/** Token counts for model request. */
|
||||
export type TokenCounts = { prompt: number; completion: number; total: number };
|
||||
|
||||
/** Sparse vectors from SPLADE models. */
|
||||
export type SparseValues = {
|
||||
indices: number[];
|
||||
values: number[];
|
||||
};
|
||||
|
||||
/** Improve preview of union types in autocomplete. */
|
||||
export type Prettify<T> = { [K in keyof T]: T[K] } & {};
|
||||
|
||||
export type ModelType =
|
||||
| 'base'
|
||||
| 'completion'
|
||||
| 'chat'
|
||||
| 'embedding'
|
||||
| 'sparse-vector';
|
||||
export type ModelProvider = 'openai' | 'custom';
|
||||
|
||||
/**
|
||||
* Base model
|
||||
*/
|
||||
export interface ModelConfig {
|
||||
model: string;
|
||||
}
|
||||
export interface ModelRun {}
|
||||
interface ModelParams extends ModelConfig, ModelRun {}
|
||||
export interface ModelResponse {
|
||||
tokens: TokenCounts;
|
||||
}
|
||||
export interface IModel<
|
||||
MConfig extends ModelConfig,
|
||||
MRun extends ModelRun,
|
||||
ModelResp extends ModelResponse
|
||||
> {
|
||||
modelType: ModelType;
|
||||
modelProvider: ModelProvider;
|
||||
run(params: MRun & Partial<MConfig>, context?: Ctx): Promise<ModelResp>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Chat model
|
||||
*/
|
||||
export interface ChatRun extends ModelRun {
|
||||
messages: ChatMessage[];
|
||||
}
|
||||
export interface ChatConfig extends ModelConfig {}
|
||||
export interface ChatResponse extends ModelResponse {
|
||||
message: ChatMessage;
|
||||
}
|
||||
export interface IChatModel<
|
||||
CConfig extends ChatConfig,
|
||||
CRun extends ChatRun,
|
||||
CResponse extends ChatResponse
|
||||
> extends IModel<CConfig, CRun, CResponse> {
|
||||
modelType: 'chat';
|
||||
}
|
||||
|
||||
/**
|
||||
* Completion model
|
||||
*/
|
||||
export interface CompletionRun extends ModelRun {
|
||||
prompt: string[];
|
||||
}
|
||||
export interface CompletionConfig extends ModelConfig {}
|
||||
export interface CompletionResponse extends ModelResponse {
|
||||
completions: { completion: string }[];
|
||||
}
|
||||
export interface ICompletionModel<
|
||||
CConfig extends CompletionConfig,
|
||||
CRun extends CompletionRun,
|
||||
CResponse extends CompletionResponse
|
||||
> extends IModel<CConfig, CRun, CResponse> {
|
||||
modelType: 'completion';
|
||||
}
|
||||
|
||||
/**
|
||||
* Embedding model
|
||||
*/
|
||||
export interface EmbeddingRun extends ModelRun {
|
||||
input: string[];
|
||||
}
|
||||
export interface EmbeddingConfig extends ModelConfig {}
|
||||
export interface EmbeddingResponse extends ModelResponse {
|
||||
embeddings: number[][];
|
||||
}
|
||||
export interface IEmbeddingModel<
|
||||
EConfig extends EmbeddingConfig = EmbeddingConfig,
|
||||
ERun extends EmbeddingRun = EmbeddingRun,
|
||||
EResponse extends EmbeddingResponse = EmbeddingResponse
|
||||
> extends IModel<EConfig, ERun, EResponse> {
|
||||
modelType: 'embedding';
|
||||
}
|
||||
|
||||
/**
|
||||
* Sparse vector model (SPLADE)
|
||||
*/
|
||||
export interface SparseRun extends ModelRun {
|
||||
input: string[];
|
||||
}
|
||||
export interface SparseConfig extends ModelConfig {}
|
||||
export interface SparseResponse extends ModelResponse {
|
||||
vectors: SparseValues[];
|
||||
}
|
||||
export interface ISparseModel<
|
||||
SConfig extends SparseConfig = SparseConfig,
|
||||
SRun extends SparseRun = SparseRun,
|
||||
SResponse extends SparseResponse = SparseResponse
|
||||
> extends IModel<SConfig, SRun, SResponse> {
|
||||
modelType: 'sparse-vector';
|
||||
}
|
||||
|
||||
/**
|
||||
* Hooks for logging and debugging
|
||||
*/
|
||||
export interface Hooks<
|
||||
MParams extends ModelParams,
|
||||
MResponse extends ModelResponse,
|
||||
AResponse extends any = any
|
||||
> {
|
||||
afterApiResponse?: (event: {
|
||||
timestamp: string;
|
||||
modelType: ModelType;
|
||||
modelProvider: ModelProvider;
|
||||
params: MParams;
|
||||
response: AResponse;
|
||||
tokens: TokenCounts;
|
||||
latency: number;
|
||||
context: Ctx;
|
||||
}) => void | Promise<void>;
|
||||
afterCacheHit?: (event: {
|
||||
timestamp: string;
|
||||
modelType: ModelType;
|
||||
modelProvider: ModelProvider;
|
||||
params: MParams;
|
||||
response: MResponse;
|
||||
context: Ctx;
|
||||
}) => void | Promise<void>;
|
||||
beforeError?: (event: {
|
||||
timestamp: string;
|
||||
modelType: ModelType;
|
||||
modelProvider: ModelProvider;
|
||||
params: MParams;
|
||||
error: unknown;
|
||||
context: Ctx;
|
||||
}) => void | Promise<void>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Cache for storing model responses
|
||||
*/
|
||||
export interface ModelCache<
|
||||
MParams extends ModelParams,
|
||||
MResponse extends ModelResponse
|
||||
> {
|
||||
get(key: MParams): Promise<MResponse | null>;
|
||||
set(key: MParams, value: MResponse): Promise<boolean>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generic interface for a model tokenizer
|
||||
*/
|
||||
export interface ITokenizer {
|
||||
/** Tokenize a string into an array of integer tokens */
|
||||
encode(text: string): Uint32Array;
|
||||
/** Decode an array of integer tokens into a string */
|
||||
decode(tokens: number[] | Uint32Array): string;
|
||||
/**
|
||||
* Count the number of tokens in a string or ChatMessage(s).
|
||||
* A single ChatMessage is counted as a completion and an array as a prompt.
|
||||
* Strings are counted as is.
|
||||
*/
|
||||
countTokens(input?: string | ChatMessage | ChatMessage[]): number;
|
||||
/** Truncate a string to a maximum number of tokens */
|
||||
truncate(args: {
|
||||
/** Text to truncate */
|
||||
text: string;
|
||||
/** Maximum number of tokens to keep (inclusive) */
|
||||
max: number;
|
||||
/** Truncate from the start or end of the text */
|
||||
from?: 'start' | 'end';
|
||||
}): string;
|
||||
}
|
||||
|
||||
export type CreateTokenizer = (model: string) => ITokenizer;
|
||||
@@ -0,0 +1,7 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
describe('should', () => {
|
||||
it('exported', () => {
|
||||
expect(1).toEqual(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"extends": "tsconfig/node-esm.json",
|
||||
"include": ["./src/**/*"],
|
||||
"exclude": ["**/node_modules", "**/.*/"],
|
||||
"compilerOptions": {
|
||||
"types": ["vitest/globals", "node"]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
import { defineConfig } from 'tsup';
|
||||
|
||||
export default defineConfig({
|
||||
entry: ['src/index.ts', 'src/openai/index.ts', 'src/custom/index.ts'],
|
||||
format: ['esm', 'cjs'],
|
||||
clean: true,
|
||||
dts: true,
|
||||
target: 'es2022',
|
||||
});
|
||||
@@ -0,0 +1,12 @@
|
||||
import { defineConfig } from 'vite';
|
||||
|
||||
export default defineConfig(() => {
|
||||
return {
|
||||
mode: 'test',
|
||||
|
||||
test: {
|
||||
environment: 'node',
|
||||
globals: true,
|
||||
},
|
||||
};
|
||||
});
|
||||
Generated
+5144
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,4 @@
|
||||
packages:
|
||||
- 'config/*'
|
||||
- 'examples/**'
|
||||
- 'packages/*'
|
||||
Executable
+13
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -o errexit # abort on nonzero exitstatus
|
||||
set -o nounset # abort on unbound variable
|
||||
set -o pipefail # don't hide errors within pipes
|
||||
|
||||
pnpm changeset version
|
||||
git add .
|
||||
git commit -m "Publish updated packages"
|
||||
pnpm changeset tag
|
||||
git push
|
||||
git push --tags
|
||||
pnpm publish -r
|
||||
@@ -0,0 +1,11 @@
|
||||
# LLM Tools
|
||||
|
||||
## Publishing Updates
|
||||
|
||||
```bash
|
||||
# Decide what to publish
|
||||
pnpm changeset
|
||||
|
||||
# Update version numbers, add tags, and publish
|
||||
./publish.sh
|
||||
```
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
{
|
||||
"$schema": "https://turborepo.org/schema.json",
|
||||
"pipeline": {
|
||||
"build": {
|
||||
"outputs": ["dist/**"],
|
||||
"dependsOn": ["^build"]
|
||||
},
|
||||
"test": {
|
||||
"dependsOn": [],
|
||||
"inputs": ["**/*.test.ts", "**/*.test.tsx"]
|
||||
},
|
||||
"lint": {
|
||||
"dependsOn": [],
|
||||
"outputs": ["node_modules/.cache/eslint/**"]
|
||||
},
|
||||
"typecheck": {
|
||||
"dependsOn": []
|
||||
},
|
||||
"dev": {
|
||||
"cache": false
|
||||
},
|
||||
"clean": {
|
||||
"cache": false
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user