| 9 | const logger = getLogger(LogCategories.INFRA.VECTOR); |
| 10 | |
| 11 | export class ObVectorCtrl implements VectorControllerType { |
| 12 | private obClient: ObClass; |
| 13 | private controllerType: 'oceanbase' | 'seekdb'; |
| 14 | constructor({ type }: { type: 'oceanbase' | 'seekdb' }) { |
| 15 | this.obClient = new ObClass({ type }); |
| 16 | this.controllerType = type; |
| 17 | } |
| 18 | init: VectorControllerType['init'] = async () => { |
| 19 | try { |
| 20 | await this.obClient.query(` |
| 21 | CREATE TABLE IF NOT EXISTS ${DatasetVectorTableName} ( |
| 22 | id BIGINT AUTO_INCREMENT PRIMARY KEY, |
| 23 | vector VECTOR(1536) NOT NULL, |
| 24 | team_id VARCHAR(50) NOT NULL, |
| 25 | dataset_id VARCHAR(50) NOT NULL, |
| 26 | collection_id VARCHAR(50) NOT NULL, |
| 27 | createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP |
| 28 | ); |
| 29 | `); |
| 30 | await this.obClient.query( |
| 31 | `CREATE VECTOR INDEX IF NOT EXISTS vector_index ON ${DatasetVectorTableName}(vector) WITH (distance=${OceanBaseIndexConfig.distance}, type=${OceanBaseIndexConfig.type}, m=16, ef_construction=200);` |
| 32 | ); |
| 33 | await this.obClient.query( |
| 34 | `CREATE INDEX IF NOT EXISTS team_dataset_collection_index ON ${DatasetVectorTableName}(team_id, dataset_id, collection_id);` |
| 35 | ); |
| 36 | await this.obClient.query( |
| 37 | `CREATE INDEX IF NOT EXISTS create_time_index ON ${DatasetVectorTableName}(createtime);` |
| 38 | ); |
| 39 | |
| 40 | logger.info('Vector DB initialization completed', { |
| 41 | provider: this.controllerType |
| 42 | }); |
| 43 | } catch (error) { |
| 44 | logger.error('Vector DB initialization failed', { |
| 45 | provider: this.controllerType, |
| 46 | error |
| 47 | }); |
| 48 | } |
| 49 | }; |
| 50 | |
| 51 | insert: VectorControllerType['insert'] = async (props) => { |
| 52 | const { teamId, datasetId, collectionId, vectors } = props; |
| 53 | |
| 54 | const values = vectors.map((vector) => [ |
| 55 | { key: 'vector', value: `[${vector}]` }, |
| 56 | { key: 'team_id', value: String(teamId) }, |
| 57 | { key: 'dataset_id', value: String(datasetId) }, |
| 58 | { key: 'collection_id', value: String(collectionId) } |
| 59 | ]); |
| 60 | |
| 61 | const { rowCount, insertIds } = await this.obClient.insert(DatasetVectorTableName, { |
| 62 | values |
| 63 | }); |
| 64 | |
| 65 | if (rowCount === 0) { |
| 66 | return Promise.reject('insertDatasetData: no insert'); |
| 67 | } |
| 68 | |