ChromaDb.ts 1009 B

1234567891011121314151617181920212223242526272829303132333435363738
  1. import type { Collection } from 'chromadb';
  2. import { ChromaClient, OpenAIEmbeddingFunction } from 'chromadb';
  3. import { BaseVectorDB } from './BaseVectorDb';
  4. const embedder = new OpenAIEmbeddingFunction({
  5. openai_api_key: process.env.OPENAI_API_KEY ?? '',
  6. });
  7. class ChromaDB extends BaseVectorDB {
  8. client: ChromaClient | undefined;
  9. collection: Collection | null = null;
  10. // eslint-disable-next-line @typescript-eslint/no-useless-constructor
  11. constructor() {
  12. super();
  13. }
  14. protected async getClientAndCollection(): Promise<void> {
  15. this.client = new ChromaClient({ path: 'http://localhost:8000' });
  16. try {
  17. this.collection = await this.client.getCollection({
  18. name: 'embedchain_store',
  19. embeddingFunction: embedder,
  20. });
  21. } catch (err) {
  22. if (!this.collection) {
  23. this.collection = await this.client.createCollection({
  24. name: 'embedchain_store',
  25. embeddingFunction: embedder,
  26. });
  27. }
  28. }
  29. }
  30. }
  31. export { ChromaDB };