Skip to content

Commit

Permalink
Add plugin qps
Browse files Browse the repository at this point in the history
  • Loading branch information
allevo committed Oct 8, 2024
1 parent e424632 commit 0bb6169
Show file tree
Hide file tree
Showing 11 changed files with 837 additions and 167 deletions.
2 changes: 1 addition & 1 deletion packages/orama/src/components/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ export function insert(
}
}

function insertVector(index: Index, prop: string, value: number[] | VectorType, id: DocumentID): void {
export function insertVector(index: AnyIndexStore, prop: string, value: number[] | VectorType, id: DocumentID): void {
if (!(value instanceof Float32Array)) {
value = new Float32Array(value)
}
Expand Down
1 change: 1 addition & 0 deletions packages/orama/src/trees.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,4 @@ export * as avl from './trees/avl.js'
export * as zip from './trees/zip.js'
export * as bkd from './trees/bkd.js'
export * as flat from './trees/flat.js'
export * as bool from './trees/bool.js'
13 changes: 13 additions & 0 deletions packages/plugin-qps/LICENSE.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
Copyright 2024 OramaSearch Inc

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
96 changes: 96 additions & 0 deletions packages/plugin-qps/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
# Orama Plugin Embeddings

**Orama Plugin Embeddings** allows you to generate fast text embeddings at insert and search time offline, directly on your machine - no OpenAI needed!

## Installation

To get started with **Orama Plugin Embeddings**, just install it with npm:

```sh
npm i @orama/plugin-embeddings
```

**Important note**: to use this plugin, you'll also need to install one of the following TensorflowJS backend:

- `@tensorflow/tfjs`
- `@tensorflow/tfjs-node`
- `@tensorflow/tfjs-backend-webgl`
- `@tensorflow/tfjs-backend-cpu`
- `@tensorflow/tfjs-node-gpu`
- `@tensorflow/tfjs-backend-wasm`

For example, if you're running Orama on the browser, we highly recommend using `@tensorflow/tfjs-backend-webgl`:

```sh
npm i @tensorflow/tfjs-backend-webgl
```

If you're using Orama in Node.js, we recommend using `@tensorflow/tfjs-node`:

```sh
npm i @tensorflow/tfjs-node
```

## Usage

```js
import { create } from '@orama/orama'
import { pluginEmbeddings } from '@orama/plugin-embeddings'
import '@tensorflow/tfjs-node' // Or any other appropriate TensorflowJS backend

const plugin = await pluginEmbeddings({
embeddings: {
defaultProperty: 'embeddings', // Property used to store generated embeddings
onInsert: {
generate: true, // Generate embeddings at insert-time
properties: ['description'], // properties to use for generating embeddings at insert time
verbose: true,
}
}
})

const db = await create({
schema: {
description: 'string',
embeddings: 'vector[512]' // Orama generates 512-dimensions vectors
},
plugins: [plugin]
})
```

Example usage at insert time:

```js
await insert(db, {
description: 'Classroom Headphones Bulk 5 Pack, Student On Ear Color Varieties'
})

await insert(db, {
description: 'Kids Wired Headphones for School Students K-12'
})

await insert(db, {
description: 'Kids Headphones Bulk 5-Pack for K-12 School'
})

await insert(db, {
description: 'Bose QuietComfort Bluetooth Headphones'
})
```

Orama will automatically generate text embeddings and store them into the `embeddings` property.

Then, you can use the `vector` or `hybrid` setting to perform hybrid or vector search at runtime:

```js
await search(db, {
term: 'Headphones for 12th grade students',
mode: 'vector'
})
```

Orama will generate embeddings at search time and perform vector or hybrid search for you.

# License

[Apache 2.0](/LICENSE.md)
53 changes: 53 additions & 0 deletions packages/plugin-qps/package.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
{
"name": "@orama/plugin-qps",
"version": "3.0.0-rc-2",
"description": "Performant search algorithm optimized for descriptive texts",
"keywords": [
"orama",
"embeddings",
"secure proxy",
"vector search"
],
"license": "Apache-2.0",
"main": "./dist/index.js",
"type": "module",
"exports": {
".": {
"require": "./dist/index.cjs",
"import": "./dist/index.js",
"types": "./dist/index.d.ts",
"browser": "./dist/index.global.js"
}
},
"bugs": {
"url": "https://github.com/askorama/orama/issues"
},
"homepage": "https://github.com/askorama/orama#readme",
"repository": {
"type": "git",
"url": "git+https://github.com/askorama/orama.git"
},
"sideEffects": false,
"types": "./dist/index.d.ts",
"files": [
"dist"
],
"scripts": {
"build": "tsup --config tsup.lib.js",
"lint": "exit 0",
"test": "exit 0"
},
"publishConfig": {
"access": "public"
},
"devDependencies": {
"@types/node": "^20.9.0",
"tap": "^21.0.1",
"tsup": "^7.2.0",
"tsx": "^4.19.1",
"typescript": "^5.0.0"
},
"dependencies": {
"@orama/orama": "workspace:*"
}
}
140 changes: 140 additions & 0 deletions packages/plugin-qps/src/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
import type { create, AnyOrama, SearchableType, IIndex, AnyIndexStore, VectorIndex, AnySchema, SearchableValue, Tokenizer } from '@orama/orama'
import {
getVectorSize, index as Index, internalDocumentIDStore, isVectorType } from '@orama/orama/components'
import { avl, bkd, flat, radix, bool } from '@orama/orama/trees'

type InternalDocumentID = internalDocumentIDStore.InternalDocumentID;

type CreateParams = Parameters<typeof create<AnyOrama, IIndex<QPSIndex>>>[0]
type Component = NonNullable<CreateParams['components']>
type IndexParameter = NonNullable<Component['index']>


export type Tree =
| Index.TTree<'Radix', radix.RadixNode>
| Index.TTree<'AVL', avl.AVLTree<number, InternalDocumentID[]>>
| Index.TTree<'Bool', bool.BoolNode>
| Index.TTree<'Flat', flat.FlatTree>
| Index.TTree<'BKD', bkd.BKDTree>

export interface QPSIndex extends AnyIndexStore {
indexes: Record<string, Tree>
vectorIndexes: Record<string, VectorIndex>
searchableProperties: string[]
searchablePropertiesWithTypes: Record<string, SearchableType>
}

function recursiveCreate<T extends AnyOrama>(indexDatastore: QPSIndex, schema: T['schema'], prefix: string) {
for (const [prop, type] of Object.entries<SearchableType>(schema)) {
const path = `${prefix}${prefix ? '.' : ''}${prop}`

if (typeof type === 'object' && !Array.isArray(type)) {
// Nested
recursiveCreate(indexDatastore, type, path)
continue
}

if (isVectorType(type)) {
indexDatastore.searchableProperties.push(path)
indexDatastore.searchablePropertiesWithTypes[path] = type
indexDatastore.vectorIndexes[path] = {
size: getVectorSize(type),
vectors: {}
}
} else {
const isArray = /\[/.test(type as string)
switch (type) {
case 'boolean':
case 'boolean[]':
indexDatastore.indexes[path] = { type: 'Bool', node: new bool.BoolNode(), isArray }
break
case 'number':
case 'number[]':
indexDatastore.indexes[path] = { type: 'AVL', node: new avl.AVLTree<number, InternalDocumentID[]>(0, []), isArray }
break
case 'string':
case 'string[]':
indexDatastore.indexes[path] = { type: 'Radix', node: new radix.RadixTree(), isArray }
// indexDatastore.avgFieldLength[path] = 0
// indexDatastore.frequencies[path] = {}
// indexDatastore.tokenOccurrences[path] = {}
// indexDatastore.fieldLengths[path] = {}
break
case 'enum':
case 'enum[]':
indexDatastore.indexes[path] = { type: 'Flat', node: new flat.FlatTree(), isArray }
break
case 'geopoint':
indexDatastore.indexes[path] = { type: 'BKD', node: new bkd.BKDTree(), isArray }
break
default:
throw new Error('INVALID_SCHEMA_TYPE: ' + path)
}

indexDatastore.searchableProperties.push(path)
indexDatastore.searchablePropertiesWithTypes[path] = type
}
}
}

export function qpsComponents(): {
index: IndexParameter,
} {

return {
index: {
create: function create<T extends AnyOrama>(orama: T, sharedInternalDocumentStore: T['internalDocumentIDStore'], schema: T['schema']) {
const indexDatastore: QPSIndex = {
indexes: {},
vectorIndexes: {},
searchableProperties: [],
searchablePropertiesWithTypes: {},
}

recursiveCreate(indexDatastore, schema, '')

return indexDatastore
},
insert: function insert(
implementation: IIndex<QPSIndex>,
indexDatastorage: QPSIndex,
prop: string,
id: DocumentID,
internalId: InternalDocumentID,
value: SearchableValue,
schemaType: SearchableType,
language: string | undefined,
tokenizer: Tokenizer,
docsCount: number
) {
if (isVectorType(schemaType)) {
return Index.insertVector(indexDatastorage, prop, value as number[] | Float32Array, id)
}

const insertScalar = insertScalarBuilder(implementation, index, prop, internalId, language, tokenizer, docsCount, options)

if (!isArrayType(schemaType)) {
return insertScalar(value)
}

const elements = value as Array<string | number | boolean>
const elementsLength = elements.length
for (let i = 0; i < elementsLength; i++) {
insertScalar(elements[i])
}
},
remove: Index.remove,
insertDocumentScoreParameters: Index.insertDocumentScoreParameters,
insertTokenScoreParameters: Index.insertTokenScoreParameters,
removeDocumentScoreParameters: Index.removeDocumentScoreParameters,
removeTokenScoreParameters: Index.removeTokenScoreParameters,
calculateResultScores: Index.calculateResultScores,
search: Index.search,
searchByWhereClause: Index.searchByWhereClause,
getSearchableProperties: Index.getSearchableProperties,
getSearchablePropertiesWithTypes: Index.getSearchablePropertiesWithTypes,
load: Index.load,
save: Index.save
}
}
}
25 changes: 25 additions & 0 deletions packages/plugin-qps/tests/index.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
import t from 'tap'
import { create, insertMultiple, search } from '@orama/orama'

t.test('plugin-qps', async t => {
const db = create({
schema: {
name: 'string',
age: 'number',
isCool: 'boolean',
algo: 'string[]',
preferredNumbers: 'number[]',
} as const
})

insertMultiple(db, [
{ name: 'foo', age: 33, isCool: true, algo: ['algo1', 'algo2'], preferredNumbers: [20] },
{ name: 'bar', age: 32, isCool: true, algo: ['algo3'], preferredNumbers: [55] },
{ name: 'baz', age: 22, isCool: false, algo: ['algo4'], preferredNumbers: [22] }
])

const result = search(db, {
term: 'b'
})
console.log(result)
})
Loading

0 comments on commit 0bb6169

Please sign in to comment.