@meaningfully/core 0.1.1 → 0.1.3

package/src/__tests__/Meaningfully.test.ts

@@ -0,0 +1,327 @@
+import { describe, it, beforeEach, expect, vi } from 'vitest';
+import { MetadataManager } from '../MetadataManager';
+import fs from 'fs';
+import path from 'path';
+import { sanitizeProjectName } from '../utils.js';
+import { createVectorStore, createDocumentStore, createIndexStore} from '../services/embeddings.js';
+import { IndexStruct } from 'llamaindex';
+import { Client } from 'pg'; // Import the real Postgres client
+
+
+vi.mock('../MetadataManager');
+vi.mock('fs');
+vi.mock('path');
+
+// Mock the embedding module before importing MeaningfullyAPI
+vi.doMock('../api/embedding.js', () => ({
+  getIndex: vi.fn(),
+  search: vi.fn().mockResolvedValue([{ id: 1, text: 'result' }]),
+  createEmbeddings: vi.fn().mockResolvedValue({ success: true, error: null }),
+}));
+vi.doMock('../services/csvLoader.js', () => ({
+  loadDocumentsFromCsv: vi.fn().mockResolvedValue([]),
+}));
+import { BaseNode } from 'llamaindex';
+
+// Mock BaseNode so that getEmbeddings returns made up numbers
+vi.mock('llamaindex', async (importOriginal) => {
+  const actual = await importOriginal();
+  return {
+    ...actual,
+    BaseNode: class extends actual.BaseNode {
+      async getEmbeddings() {
+        return [0.1, 0.2, 0.3, 0.4];
+      }
+      generateHash() {
+        return 'hash';
+      }
+      getContent(){
+        return "content";
+      }
+    }
+  };
+});
+
+// Import MeaningfullyAPI after mocking
+const { MeaningfullyAPI } = await import('../Meaningfully');
+
+const FAKE_SETTINGS = {
+        openAIKey: 'sk-proj-testtesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttest',
+        azureOpenAIKey: 'testtesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttesttest',
+        mistralApiKey: 'testtesttesttesttesttesttesttest',
+        geminiApiKey:  'testtesttesttesttesttesttesttesttesttes',
+        azureOpenAIApiVersion: "2024-02-01",
+        azureOpenAIEndpoint: "https://test.openai.azure.com",
+        oLlamaBaseURL: "http://localhost:11434",
+      }
+
+describe('MeaningfullyAPI', () => {
+  let api: MeaningfullyAPI;
+  let mockMetadataManager: MetadataManager;
+
+  beforeEach(() => {
+    // @ts-ignore
+    mockMetadataManager = new MetadataManager() as MetadataManager;
+    vi.spyOn(mockMetadataManager, 'addDocumentSet').mockResolvedValue(1);
+    vi.spyOn(mockMetadataManager, 'getSettings').mockResolvedValue(FAKE_SETTINGS);
+    vi.spyOn(mockMetadataManager, 'deleteDocumentSet').mockResolvedValue();
+    api = new MeaningfullyAPI({
+      storagePath: 'mock_storage_path',
+      metadataManager: mockMetadataManager,
+    });
+  });
+
+  describe('uploadCsv', () => {
+    it('should upload a CSV and create embeddings successfully', async () => {
+      const mockData = {
+        filePath: '/mock/file.csv',
+        datasetName: 'testDataset',
+        textColumns: ['text'],
+        metadataColumns: [],
+        splitIntoSentences: true,
+        combineSentencesIntoChunks: false,
+        sploderMaxSize: 100,
+        chunkSize: 512,
+        chunkOverlap: 0,
+        modelName: 'testModel',
+        modelProvider: 'openai',
+        description: 'Test dataset',
+      };
+
+      // Mock createEmbeddings for this test
+      const createEmbeddingsMock = vi.spyOn(await import('../api/embedding.js'), 'createEmbeddings');
+      createEmbeddingsMock.mockResolvedValue({ success: true });
+
+      const result = await api.uploadCsv(mockData);
+
+      expect(createEmbeddingsMock).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.any(String),
+        expect.objectContaining({ modelName: 'testModel' }),
+        expect.any(Object),
+        expect.any(Object)
+      );
+      expect(result).toEqual({ success: true, documentSetId: 1 });
+
+      createEmbeddingsMock.mockRestore(); // Restore the original implementation after the test
+    });
+
+    it('should handle errors during embeddings creation', async () => {
+      const mockData = {
+        filePath: '/mock/file.csv',
+        datasetName: 'testDataset',
+        textColumns: ['text'],
+        metadataColumns: [],
+        splitIntoSentences: true,
+        combineSentencesIntoChunks: false,
+        sploderMaxSize: 100,
+        chunkSize: 512,
+        chunkOverlap: 0,
+        modelName: 'testModel',
+        modelProvider: 'openai',
+        description: 'Test dataset',
+      };
+
+      // Mock createEmbeddings to simulate an error
+      const createEmbeddingsMock = vi.spyOn(await import('../api/embedding.js'), 'createEmbeddings');
+      createEmbeddingsMock.mockResolvedValue({ success: false, error: 'Embedding error' });
+
+      await expect(api.uploadCsv(mockData)).rejects.toThrow('Embedding error');
+      expect(mockMetadataManager.deleteDocumentSet).toHaveBeenCalledWith(1);
+
+      createEmbeddingsMock.mockRestore(); // Restore the original implementation after the test
+    });
+  });
+
+  describe('searchDocumentSet', () => {
+    it('should search a document set and return results', async () => {
+      vi.spyOn(mockMetadataManager, 'getDocumentSet').mockResolvedValue({
+        parameters: { modelName: 'testModel', modelProvider: 'openai', vectorStoreType: 'simple' },
+        name: 'testDataset',
+        documentSetId: 5,
+        uploadDate: new Date(),
+        totalDocuments: 420
+      });
+
+      const results = await api.searchDocumentSet(1, 'query', 10);
+
+      expect(results).toEqual([{ id: 1, text: 'result' }]);
+      expect(mockMetadataManager.getDocumentSet).toHaveBeenCalledWith(1);
+    });
+  });
+
+  describe('deleteDocumentSet', () => {
+    it('should delete a document set and associated files', async () => {
+      vi.spyOn(mockMetadataManager, 'getDocumentSet').mockResolvedValue({
+        parameters: { vectorStoreType: 'simple' },
+        name: 'testDataset',
+        documentSetId: 1,
+        uploadDate: new Date(),
+        totalDocuments: 100
+      });
+
+      vi.spyOn(fs, 'rmSync').mockImplementation(() => {});
+
+      const result = await api.deleteDocumentSet(1);
+
+      expect(mockMetadataManager.deleteDocumentSet).toHaveBeenCalledWith(1);
+      expect(fs.rmSync).toHaveBeenCalledWith(
+        path.join('mock_storage_path', 'testDataset'),
+        { recursive: true, force: true }
+      );
+      expect(result).toEqual({ success: true });
+    });
+  });
+
+  describe('getMaskedSettings', () => {
+    it('should return masked settings', async () => {
+      vi.spyOn(mockMetadataManager, 'getSettings').mockResolvedValue(FAKE_SETTINGS);
+
+      const settings = await api.getMaskedSettings();
+
+      expect(settings).toEqual({
+        openAIKey: 'sk-proj-*******testtest',
+        azureOpenAIKey: 'testtest*******testtest',
+        mistralApiKey: 'testtest*******testtest',
+        geminiApiKey: 'testtest*******ttesttes',
+        azureOpenAIApiVersion: "2024-02-01",
+        azureOpenAIEndpoint: "https://test.openai.azure.com",
+        oLlamaBaseURL:  "http://localhost:11434",
+      });
+    });
+  });
+});
+
+
+describe('MeaningfullyAPI - Store Deletion with Real Implementation', () => {
+  let api: MeaningfullyAPI;
+  let mockMetadataManager: MetadataManager;
+  let realPostgresClient: Client;
+
+  beforeEach(async () => {
+    // @ts-ignore
+    mockMetadataManager = new MetadataManager() as MetadataManager;
+
+    // Initialize a real Postgres client
+    realPostgresClient = new Client({
+      connectionString: process.env.POSTGRES_CONNECTION_STRING,
+    });
+    await realPostgresClient.connect();
+
+    api = new MeaningfullyAPI({
+      storagePath: 'mock_storage_path',
+      metadataManager: mockMetadataManager,
+      postgresClient: realPostgresClient, // Use the real client
+    });
+    vi.unmock('fs')
+    if (!fs.existsSync("mock_storage_path")){
+      fs.mkdirSync("mock_storage_path");
+    }
+  });
+
+  it('should create and delete a Postgres vector store using real implementation', async () => {
+    const projectName = 'test_project';
+    const sanitizedProjectName = sanitizeProjectName(projectName);
+    const tableName = `vecs_${sanitizedProjectName}`;
+
+    // Create the vector store using the real implementation
+    const vectorStore = await createVectorStore(
+      { vectorStoreType: 'postgres', projectName, storagePath: 'mock_storage_path', modelProvider: 'openai', modelName: 'text-embedding-ada-002' },
+      FAKE_SETTINGS,
+      api.getClients()
+    );
+    await vectorStore.add([new BaseNode({ id: '1', text: 'test document', embedding: Array(1536).fill(0.01) })]);
+
+    // Verify the table exists
+    const tableExistsQuery = `
+      SELECT EXISTS (
+        SELECT FROM information_schema.tables 
+        WHERE table_name = $1
+      );
+    `;
+    const tableExistsResult = await realPostgresClient.query(tableExistsQuery, [tableName]);
+    expect(tableExistsResult.rows[0].exists).toBe(true);
+
+    // Call the delete method
+    await api.deletePostgresVectorStore(projectName);
+
+    // Verify the table no longer exists
+    const tableDeletedResult = await realPostgresClient.query(tableExistsQuery, [tableName]);
+    expect(tableDeletedResult.rows[0].exists).toBe(false);
+  });
+
+  // it('should delete Simple vector store using real implementation', async () => {
+  //   const projectName = 'test_project';
+  //   const sanitizedProjectName = sanitizeProjectName(projectName);
+  //   const storagePath = 'mock_storage_path';
+  //   const persistDir = path.join(storagePath, sanitizedProjectName);
+
+  //   // Create the vector store using the real implementation
+  //   const vectorStore = await createVectorStore(
+  //     { vectorStoreType: 'simple', projectName, storagePath, modelProvider: "openai", modelName: 'text-embedding-3-small' },
+  //     FAKE_SETTINGS,
+  //     api.getClients()
+  //   );
+  //   await vectorStore.add([new BaseNode({ id: '1', text: 'test document', embedding: [1,2,3] })]);
+  //   await vectorStore.persist(path.join(persistDir, 'vector_store.json'));
+
+  //   // Verify the vector store exists
+  //   expect(fs.existsSync(path.join(persistDir, 'vector_store.json'))).toBe(true);
+
+  //   // Call the delete method
+  //   await api.deleteSimpleVectorStore(projectName);
+
+  //   // Verify the vector store no longer exists
+  //   expect(fs.existsSync(path.join(persistDir, 'vector_store.json'))).toBe(false);
+  // });
+
+  // it('should delete Simple document store using real implementation', async () => {
+  //   const projectName = 'test_project';
+  //   const sanitizedProjectName = sanitizeProjectName(projectName);
+  //   const storagePath = 'mock_storage_path';
+  //   const persistDir = path.join(storagePath, sanitizedProjectName);
+
+  //   // Create the document store using the real implementation
+  //   const docStore = await createDocumentStore(
+  //     { vectorStoreType: 'simple', projectName, storagePath },
+  //     FAKE_SETTINGS,
+  //     api.getClients()
+  //   );
+  //   await docStore.addDocuments([new BaseNode({ id: '1', text: 'test document' })], true);
+  //   await docStore.persist(path.join(persistDir, 'doc_store.json'));
+
+  //   // Verify the document store exists
+  //   expect(fs.existsSync(path.join(persistDir, 'doc_store.json'))).toBe(true);
+
+  //   // Call the delete method
+  //   await api.deleteSimpleDocStore(projectName);
+
+  //   // Verify the document store no longer exists
+  //   expect(fs.existsSync(path.join(persistDir, 'doc_store.json'))).toBe(false);
+  // });
+
+  // it('should delete Simple index store using real implementation', async () => {
+  //   const projectName = 'test_project';
+  //   const sanitizedProjectName = sanitizeProjectName(projectName);
+  //   const storagePath = 'mock_storage_path';
+  //   const persistDir = path.join(storagePath, sanitizedProjectName);
+
+  //   // Create the index store using the real implementation
+  //   const indexStore = await createIndexStore(
+  //     { vectorStoreType: 'simple', projectName, storagePath },
+  //     FAKE_SETTINGS,
+  //     api.getClients()
+  //   );
+  //   indexStore.addIndexStruct(new IndexStruct({ summary: 'test document' }));
+  //   await indexStore.persist(path.join(persistDir, 'index_store.json'));
+
+  //   // Verify the index store exists
+  //   expect(fs.existsSync(path.join(persistDir, 'index_store.json'))).toBe(true);
+
+  //   // Call the delete method
+  //   await api.deleteSimpleIndexStore(projectName);
+
+  //   // Verify the index store no longer exists
+  //   expect(fs.existsSync(path.join(persistDir, 'index_store.json'))).toBe(false);
+  // });
+});

package/src/api/__tests__/embedding.test.ts

@@ -2,7 +2,7 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { createEmbeddings, previewResults, getDocStore, getIndex, search } from '../embedding.js';
 import { loadDocumentsFromCsv } from '../../services/csvLoader.js';
-import { transformDocumentsToNodes, estimateCost, searchDocuments, getExistingVectorStoreIndex, persistNodes, getExistingDocStore } from '../../services/embeddings.js';
+import { transformDocumentsToNodes, estimateCost, searchDocuments, getExistingVectorStoreIndex, persistNodes, getStorageContext } from '../../services/embeddings.js';
 import { MetadataMode } from 'llamaindex';
 
 // filepath: /Users/jeremybmerrill/code/meaningfully/src/main/api/embedding.test.ts
@@ -82,9 +82,9 @@ describe('embedding.ts', () => {
     describe('getDocStore', () => {
         it('should return existing doc store', async () => {
             const mockDocStore = 'docStore';
-            getExistingDocStore.mockResolvedValue(mockDocStore);
+            getStorageContext.mockResolvedValue({ docStore: mockDocStore });
 
-            const result = await getDocStore({});
+            const result = await getDocStore({}, {}, {});
 
             expect(result).toBe(mockDocStore);
         });
@@ -173,7 +173,7 @@ describe('embedding.ts', () => {
   describe('getDocStore', () => {
     it('should return existing doc store', async () => {
       const mockDocStore = 'docStore';
-      getExistingDocStore.mockResolvedValue(mockDocStore);
+      getStorageContext.mockResolvedValue({ docStore: mockDocStore });
 
       const result = await getDocStore({});
 

package/src/api/embedding.ts

@@ -1,4 +1,4 @@
-import { transformDocumentsToNodes, estimateCost, searchDocuments, getExistingVectorStoreIndex, persistNodes, persistDocuments, getExistingDocStore } from "../services/embeddings.js";
+import { transformDocumentsToNodes, estimateCost, searchDocuments, getExistingVectorStoreIndex, persistNodes, persistDocuments, getStorageContext } from "../services/embeddings.js";
 import type { EmbeddingConfig, EmbeddingResult, SearchResult, PreviewResult, Settings, MetadataFilter, Clients } from "../types/index.js";
 import { loadDocumentsFromCsv } from "../services/csvLoader.js";
 import { MetadataMode } from "llamaindex";
@@ -97,8 +97,8 @@ export async function previewResults(
   }
 } 
 
-export async function getDocStore(config: EmbeddingConfig) {
-  return await getExistingDocStore(config);
+export async function getDocStore(config: EmbeddingConfig, settings: Settings, clients: Clients) {
+  return (await getStorageContext(config, settings, clients)).docStore;
 }
 
 export async function getIndex(config: EmbeddingConfig, settings: Settings, clients: Clients) {

package/src/services/embeddings.d.ts

@@ -12,7 +12,6 @@ export declare function estimateCost(nodes: TextNode[], modelName: string): {
     pricePer1M: number;
 };
 export declare function getExistingVectorStoreIndex(config: EmbeddingConfig, settings: Settings, clients: Clients): Promise<VectorStoreIndex>;
-export declare function getExistingDocStore(config: EmbeddingConfig): Promise<import("llamaindex").BaseDocumentStore>;
 export declare function transformDocumentsToNodes(documents: Document[], config: EmbeddingConfig): Promise<TextNode<import("llamaindex").Metadata>[]>;
 export declare function getEmbedModel(config: EmbeddingConfig, settings: Settings): MockEmbedding | OpenAIEmbedding | OllamaEmbedding | MistralAIEmbedding | GeminiEmbedding;
 export declare function getStorageContext(config: EmbeddingConfig, settings: Settings, clients: Clients): Promise<StorageContext>;

package/src/services/embeddings.ts

@@ -11,12 +11,15 @@ import {
   SimpleVectorStore,
   type StorageContext,
   Settings as LlamaindexSettings,
-  SimpleDocumentStore
+  SimpleDocumentStore,
+  BaseDocumentStore,
+  BaseIndexStore,
+  SimpleIndexStore
 } from "llamaindex";
 import { OllamaEmbedding} from '@llamaindex/ollama'
 import { MistralAIEmbedding, MistralAIEmbeddingModelType } from '@llamaindex/mistral'
 import { GeminiEmbedding } from '@llamaindex/google'
-import { PGVectorStore } from '@llamaindex/postgres';
+import { PGVectorStore, PostgresDocumentStore, PostgresIndexStore } from '@llamaindex/postgres';
 import { AzureOpenAIEmbedding } from "@llamaindex/azure";
 import { Sploder } from "./sploder.js";
 import { CustomSentenceSplitter } from "./sentenceSplitter.js";
@@ -97,79 +100,14 @@ export function estimateCost(nodes: TextNode[], modelName: string): {
 }
 
 export async function getExistingVectorStoreIndex(config: EmbeddingConfig, settings: Settings, clients: Clients) {
-  const embedModel = getEmbedModel(config, settings);
-  switch (config.vectorStoreType) {
-    case "simple":
-      const persistDir = join(config.storagePath, sanitizeProjectName(config.projectName));
-      const storageContext = await storageContextFromDefaults({
-        persistDir: persistDir,
-      });
-      let vsi = await VectorStoreIndex.init({
-        storageContext: storageContext,
-      });
-      vsi.embedModel = embedModel;
-      return vsi;
-
-    case "postgres":
-      if (!clients.postgresClient) {
-        throw new Error("Postgres client required but not provided");
-      }
-      const pgStore = new PGVectorStore({
-        clientConfig: { connectionString: process.env.POSTGRES_CONNECTION_STRING }, 
-        tableName: sanitizeProjectName(config.projectName),
-        dimensions: MODEL_DIMENSIONS[config.modelName] || 1536, // default to 1536 if model not found
-        embeddingModel: embedModel
-      });
-      const pgStorageContext = await storageContextFromDefaults({
-        vectorStores: { [ModalityType.TEXT]: pgStore },
-      });
-      return await VectorStoreIndex.init({
-        storageContext: pgStorageContext,
-      });
-    case "weaviate":
-      if (!clients.weaviateClient) {
-        throw new Error("Weaviate client required but not provided");
-      }
-      const weaviateStore = new BatchingWeaviateVectorStore({
-        indexName: capitalizeFirstLetter(sanitizeProjectName(config.projectName)),
-        weaviateClient: clients.weaviateClient,
-        embeddingModel: embedModel
-      });
-
-      // WeaviateVectorStore's getNodeSimilarity method looks for distance, but current weaviate provides score
-      // (WeaviateVectorStore would get `score` if we were doing hybrid search)
-      // Overwrite the private getNodeSimilarity method to use 'score' from metadata
-      // @ts-ignore
-      weaviateStore.getNodeSimilarity = (entry, _similarityKey = "score") => {
-        return  entry.metadata.score;
-      }
-
-      return await VectorStoreIndex.fromVectorStore(weaviateStore)
-
-    default:
-      throw new Error(`Unsupported vector store type: ${config.vectorStoreType}`);
+  const storageContext = await getStorageContext(config, settings, clients);
+  const vectorStore = storageContext.vectorStores[ModalityType.TEXT];
+  if (!vectorStore) {
+    throw new Error("Vector store for ModalityType.TEXT is undefined");
   }
+  return await VectorStoreIndex.fromVectorStore(vectorStore);
 }
 
-export async function getExistingDocStore(config: EmbeddingConfig) {
-  // switch (config.vectorStoreType) {
-  //   case "simple":
-      const persistDir = join(config.storagePath, sanitizeProjectName(config.projectName) );
-      const storageContext = await storageContextFromDefaults({
-        persistDir: persistDir,
-      });
-      return storageContext.docStore;
-
-  //   case "postgres":
-  //     throw new Error(`Not yet implemented vector store type: ${config.vectorStoreType}`);
-  //     // return await createVectorStore(config);
-  //   default:
-  //     throw new Error(`Unsupported vector store type: ${config.vectorStoreType}`);
-  // }
-}
-
-
-
 export async function transformDocumentsToNodes(
   documents: Document[],
   config: EmbeddingConfig,
@@ -249,12 +187,15 @@ export function getEmbedModel(
 
 export async function getStorageContext(config: EmbeddingConfig, settings: Settings, clients: Clients): Promise<StorageContext> {
   const vectorStore = await createVectorStore(config, settings, clients);
+  const docStore = await createDocumentStore(config, settings, clients); // new SimpleDocumentStore()
+  const indexStore = await createIndexStore(config, settings, clients);
   fs.mkdirSync(config.storagePath, { recursive: true }); 
   const persistDir = join(config.storagePath, sanitizeProjectName(config.projectName) );
   return await storageContextFromDefaults({
     persistDir: persistDir,
     vectorStores: {[ModalityType.TEXT]: vectorStore},
-    docStore: new SimpleDocumentStore()
+    docStore: docStore,
+    indexStore: indexStore
       /*
         if docStore is created with a persist path (as it is by default in storageContextFromDefaults)
         then it will write to disk after every put(), which happens 2+ times per document.
@@ -273,8 +214,13 @@ export async function persistDocuments(documents: Document[], config: EmbeddingC
 
   // see comments in getStorageContext
   const persistDir = join(config.storagePath, sanitizeProjectName(config.projectName) );
-  // @ts-ignore
-  await (storageContext.docStore as SimpleDocumentStore).kvStore.persist(join(persistDir, "doc_store.json"));
+  if (storageContext.docStore instanceof SimpleDocumentStore) {
+    // @ts-ignore
+    await (storageContext.docStore as SimpleDocumentStore).kvStore.persist(join(persistDir, "doc_store.json"));
+  }else if (storageContext.docStore instanceof PostgresDocumentStore) {
+    // PostgresDocumentStore does not need to be explicitly persisted, so we don't include it in the OR conditional here..
+    console.log("Pretending to persist Postgres document store, but it actually persists automatically.");
+  }
 
   console.timeEnd("persistDocuments Run Time");
 }
@@ -303,11 +249,11 @@ export async function persistNodes(nodes: TextNode[], config: EmbeddingConfig, s
   // all the if statements are just type-checking boilerplate.
   // N.B. WeaviateVectorStore does not need to be explicitly persisted, so we don't include it in the OR conditional here..
   if (vectorStore) {
-    if (vectorStore instanceof PGVectorStore || vectorStore instanceof SimpleVectorStore) {
+    if (vectorStore instanceof SimpleVectorStore) {
       await vectorStore.persist(join(config.storagePath, sanitizeProjectName(config.projectName), "vector_store.json"));
-    } else if (vectorStore instanceof BatchingWeaviateVectorStore) {
+    } else if (vectorStore instanceof PGVectorStore || vectorStore instanceof BatchingWeaviateVectorStore) {
       // WeaviateVectorStore does not have a persist method, it persists automatically
-      console.log("Pretending to persist Weaviate vector store, but it actually persists automatically.");
+      console.log("Pretending to persist Weaviate or Postgres vector store, but it actually persists automatically.");
     } else {
       throw new Error("Vector store does not support persist method");
     }
@@ -318,7 +264,8 @@ export async function persistNodes(nodes: TextNode[], config: EmbeddingConfig, s
   return index;
 }
 
-async function createVectorStore(config: EmbeddingConfig, settings: Settings, clients: Clients): Promise<PGVectorStore | SimpleVectorStore | BatchingWeaviateVectorStore> {
+// exported only for tests
+export async function createVectorStore(config: EmbeddingConfig, settings: Settings, clients: Clients): Promise<PGVectorStore | SimpleVectorStore | BatchingWeaviateVectorStore> {
   const embeddingModel = getEmbedModel(config, settings);
   switch (config.vectorStoreType) {
 
@@ -326,8 +273,8 @@ async function createVectorStore(config: EmbeddingConfig, settings: Settings, cl
     // otherwise it defaults to Ada.
     case "postgres":
       return new PGVectorStore({
-        clientConfig: {connectionString: process.env.POSTGRES_CONNECTION_STRING},
-        tableName: sanitizeProjectName(config.projectName),
+        client: clients.postgresClient,
+        tableName: "vecs_" + sanitizeProjectName(config.projectName),
         dimensions: MODEL_DIMENSIONS[config.modelName] || 1536, // default to 1536 if model not found
         embeddingModel: embeddingModel
       });
@@ -357,6 +304,40 @@ async function createVectorStore(config: EmbeddingConfig, settings: Settings, cl
   }
 }
 
+// exported only for tests
+export async function createDocumentStore(config: EmbeddingConfig, settings: Settings, clients: Clients): Promise<BaseDocumentStore> {
+  // we create the doc store without a persist path, so it doesn't write to disk after every put()
+  switch (config.documentStoreType || config.vectorStoreType) {
+    case "postgres":
+      return new PostgresDocumentStore({
+        client: clients.postgresClient,
+        tableName: "docs_" + sanitizeProjectName(config.projectName),
+      });
+    case "simple":
+    case "weaviate": 
+      return new SimpleDocumentStore();
+    default:
+      throw new Error(`Unsupported vector store type: ${config.vectorStoreType}`);
+  }
+}
+
+// exported only for tests
+export async function createIndexStore(config: EmbeddingConfig, settings: Settings, clients: Clients): Promise<BaseIndexStore> {
+  switch (config.documentStoreType || config.vectorStoreType) {
+    case "postgres":
+      return new PostgresIndexStore({
+        client: clients.postgresClient,
+        tableName: "idx_" + sanitizeProjectName(config.projectName),
+      });
+    case "simple":
+    case "weaviate": 
+      return new SimpleIndexStore();
+    default:
+      throw new Error(`Unsupported vector store type: ${config.vectorStoreType}`);
+  }
+
+}
+
 export async function searchDocuments(
   index: VectorStoreIndex,
   query: string,

package/src/types/index.ts

@@ -57,6 +57,8 @@ export interface EmbeddingConfig {
   modelName: string;
   modelProvider: string
   vectorStoreType: "simple" | "postgres" | "weaviate";
+  documentStoreType?: "simple" | "postgres";
+  indexStoreType?: "simple" | "postgres";
   projectName: string;
   storagePath: string;
   splitIntoSentences: boolean;