@meaningfully/core 0.1.0 → 0.1.2

package/{tests → src/api/__tests__}/embedding.test.ts

@@ -1,15 +1,15 @@
 //@ts-nocheck
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { createEmbeddings, previewResults, getDocStore, getIndex, search } from '../src/api/embedding.js';
-import { loadDocumentsFromCsv } from '../src/services/csvLoader.js';
-import { transformDocumentsToNodes, estimateCost, searchDocuments, getExistingVectorStoreIndex, persistNodes, getExistingDocStore } from '../src/services/embeddings.js';
+import { createEmbeddings, previewResults, getDocStore, getIndex, search } from '../embedding.js';
+import { loadDocumentsFromCsv } from '../../services/csvLoader.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
 
 
-vi.mock('../src/services/csvLoader');
-vi.mock('../src/services/embeddings');
+vi.mock('../../services/csvLoader');
+vi.mock('../../services/embeddings');
 
 describe('embedding.ts', () => {
     describe('createEmbeddings', () => {
@@ -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({});
 
@@ -262,7 +262,7 @@ describe('embedding.ts', () => {
 
       it('should correctly track progress through ProgressManager', async () => {
         // Setup
-        vi.mock('../src/services/progressManager', () => {
+        vi.mock('../../services/progressManager', () => {
           const mockInstance = {
             startOperation: vi.fn(),
             updateProgress: vi.fn(),
@@ -278,8 +278,8 @@ describe('embedding.ts', () => {
         });
         
         // Re-import to use mocked version
-        const { createEmbeddings } = await import('../src/api/embedding.js');
-        const { ProgressManager } = await import('../src/services/progressManager.js');
+        const { createEmbeddings } = await import('../embedding.js');
+        const { ProgressManager } = await import('../../services/progressManager.js');
         
         const mockDocuments = [{ text: 'doc1' }, { text: 'doc2' }];
         const mockNodes = [{ text: 'node1', metadata: {} }];
@@ -305,7 +305,7 @@ describe('embedding.ts', () => {
 
       it('should properly calculate percentage in progress callback', async () => {
         // Setup mocks with spy on updateProgress
-        vi.mock('../src/services/progressManager', () => {
+        vi.mock('../../services/progressManager', () => {
           const mockInstance = {
             startOperation: vi.fn(),
             updateProgress: vi.fn(),
@@ -321,8 +321,8 @@ describe('embedding.ts', () => {
         });
         
         // Re-import to use mocked version
-        const { createEmbeddings } = await import('../src/api/embedding.js');
-        const { ProgressManager } = await import('../src/services/progressManager.js');
+        const { createEmbeddings } = await import('../embedding.js');
+        const { ProgressManager } = await import('../../services/progressManager.js');
 
         const mockDocuments = [{ text: 'doc1' }];
         const mockNodes = [{ text: 'node1', metadata: {} }];
@@ -355,7 +355,7 @@ describe('embedding.ts', () => {
 
       it('should clear operation on empty documents', async () => {
         // Setup
-        vi.mock('../src/services/progressManager', () => {
+        vi.mock('../../services/progressManager', () => {
           const mockInstance = {
             startOperation: vi.fn(),
             updateProgress: vi.fn(),
@@ -371,8 +371,8 @@ describe('embedding.ts', () => {
         });
         
         // Re-import to use mocked version
-        const { createEmbeddings } = await import('../src/api/embedding.js');
-        const { ProgressManager } = await import('../src/services/progressManager.js');
+        const { createEmbeddings } = await import('../embedding.js');
+        const { ProgressManager } = await import('../../services/progressManager.js');
 
         loadDocumentsFromCsv.mockResolvedValue([]);
         
@@ -390,7 +390,7 @@ describe('embedding.ts', () => {
 
       it('shoulde complete operation on successful embedding', async () => {
         // Setup
-        vi.mock('../src/services/progressManager', () => {
+        vi.mock('../../services/progressManager', () => {
           const mockInstance = {
             startOperation: vi.fn(),
             updateProgress: vi.fn(),
@@ -406,8 +406,8 @@ describe('embedding.ts', () => {
         });
         
         // Re-import to use mocked version
-        const { createEmbeddings } = await import('../src/api/embedding.js');
-        const { ProgressManager } = await import('../src/services/progressManager.js');
+        const { createEmbeddings } = await import('../embedding.js');
+        const { ProgressManager } = await import('../../services/progressManager.js');
 
         const mockDocuments = [{ text: 'doc1' }];
         const mockNodes = [{ text: 'node1', metadata: {} }];

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/{tests → src/services/__tests__}/csvLoader.test.ts

@@ -1,7 +1,7 @@
 //@ts-nocheck
 import { describe, it, expect, vi } from 'vitest';
 import { readFileSync } from 'fs';
-import { loadDocumentsFromCsv } from '../src/services/csvLoader.js';
+import { loadDocumentsFromCsv } from '../csvLoader.js';
 import { Document } from 'llamaindex';
 import Papa from 'papaparse';
 

package/{tests → src/services/__tests__}/embeddings.test.ts

@@ -4,7 +4,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { Document, TextNode } from 'llamaindex';
 
 // First, set up the mock before importing the module
-vi.mock(import("../src/services/embeddings.js"), async (importOriginal) => {
+vi.mock(import("../embeddings.js"), async (importOriginal) => {
   const actual = await importOriginal()
   return {
     ...actual,
@@ -19,7 +19,7 @@ vi.mock(import("../src/services/embeddings.js"), async (importOriginal) => {
 })
 
 // Now import the mocked functions
-import { transformDocumentsToNodes, getEmbedModel } from '../src/services/embeddings.js';
+import { transformDocumentsToNodes, getEmbedModel } from '../embeddings.js';
 
 describe('transformDocumentsToNodes', () => {
   beforeEach(() => {

package/{tests → src/services/__tests__}/sentenceSplitter.test.ts

@@ -1,6 +1,6 @@
 //@ts-nocheck
 import { expect, test } from 'vitest'
-import { CustomSentenceSplitter } from '../src/services/sentenceSplitter.js'
+import { CustomSentenceSplitter } from '../sentenceSplitter.js'
 import { SentenceSplitter, IngestionPipeline, Document } from "llamaindex";
 
 // do these tests just to make sure that we can factor out my hacky fixes when llamaindex is fixed.

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,12 @@ 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}`);
-  }
-}
-
-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}`);
-  // }
+  const storageContext = await getStorageContext(config, settings, clients);
+  return await VectorStoreIndex.init({
+    storageContext: storageContext,
+  });
 }
 
-
-
 export async function transformDocumentsToNodes(
   documents: Document[],
   config: EmbeddingConfig,
@@ -249,12 +185,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 +212,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 +247,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");
     }
@@ -326,8 +270,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 +301,38 @@ async function createVectorStore(config: EmbeddingConfig, settings: Settings, cl
   }
 }
 
+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}`);
+  }
+}
+
+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;

package/tsconfig.json

@@ -27,5 +27,5 @@
     /* AND if you're building for a library in a monorepo: */
     "declarationMap": true
   },
-  "exclude": ["node_modules/natural", "dist/**"]
+  "exclude": ["node_modules/natural", "dist/**", "src/**/__tests__/**"]
 }

package/dist/api/embedding.test.d.ts

@@ -1,2 +0,0 @@
-export {};
-//# sourceMappingURL=embedding.test.d.ts.map

package/dist/api/embedding.test.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"embedding.test.d.ts","sourceRoot":"","sources":["../../src/api/embedding.test.ts"],"names":[],"mappings":""}