Merge pull request #64 from zackify/feature/context-routes

feat: Add context storage and retrieval routes with MCP tools

Author: zackify

src/database/database.ts

@@ -6,43 +6,9 @@ if (process.env.SQLITE_PATH) {
   Database.setCustomSQLite(process.env.SQLITE_PATH);
 }
 
-/**
- * Get the database instance
- * This function always returns the current database, 
- * checking for a test database first, then falling back to the singleton
- */
-export function getDb(): Database {
-  // For tests, we support injecting a test database instance
-  const testDb = (globalThis as any).testDb;
-  if (testDb) {
-    return testDb;
-  }
-  
-  // For normal operation, create a singleton
-  if (!(globalThis as any)._dbSingleton) {
-    const dbPath = process.env.DATABASE_PATH || "./data/db.sqlite";
-    console.log(`Creating database connection to ${dbPath}`);
-    
-    (globalThis as any)._dbSingleton = new Database(dbPath);
-    (globalThis as any)._dbSingleton.exec("PRAGMA journal_mode = WAL;");
-    sqliteVec.load((globalThis as any)._dbSingleton);
-  }
-  
-  return (globalThis as any)._dbSingleton;
-}
+const dbPath = process.env.DATABASE_PATH || "./data/db.sqlite";
+console.log(`Creating database connection to ${dbPath}`);
 
-// Create and export a db proxy that always returns the current database instance
-// This ensures that even code using the imported db directly will get the test db when appropriate
-export const db = new Proxy({} as Database, {
-  get: function(target, prop) {
-    const currentDb = getDb();
-    const value = currentDb[prop as keyof Database];
-    
-    // If it's a function, bind it to the correct database instance
-    if (typeof value === 'function') {
-      return value.bind(currentDb);
-    }
-    
-    return value;
-  }
-});
+export const db = new Database(dbPath);
+db.exec("PRAGMA journal_mode = WAL;");
+sqliteVec.load(db);

src/database/metadata.ts

@@ -1,12 +1,32 @@
 import { db } from "./database";
 
+/**
+ * Checks if the metadata table exists
+ * @returns true if metadata table exists, false otherwise
+ */
+export const metadataTableExists = (): boolean => {
+  try {
+    const result = db
+      .query("SELECT name FROM sqlite_master WHERE type='table' AND name='metadata'")
+      .get();
+    return !!result;
+  } catch (error) {
+    return false;
+  }
+};
+
 /**
  * Gets a metadata value from the database
  * @param key The metadata key to retrieve
  * @returns The value as string or null if not found
  */
 export const getMetadataValue = (key: string): string | null => {
   try {
+    // Return null if metadata table doesn't exist (brand new database)
+    if (!metadataTableExists()) {
+      return null;
+    }
+
     const result = db
       .query("SELECT value FROM metadata WHERE key = ?")
       .get(key);
@@ -26,7 +46,12 @@ export const getMetadataValue = (key: string): string | null => {
  */
 export const setMetadataValue = (key: string, value: string): boolean => {
   try {
-    // Ensure the metadata table exists before writing to it
+    // If metadata table doesn't exist, we can't set metadata yet
+    // This will be handled after migrations run
+    if (!metadataTableExists()) {
+      console.warn(`Cannot set metadata for key ${key}: metadata table does not exist yet`);
+      return false;
+    }
 
     // Check if the key already exists
     const exists = db.query("SELECT 1 FROM metadata WHERE key = ?").get(key);

src/database/migrations.ts

@@ -98,6 +98,18 @@ export const migrations: Migration[] = [
       return db.exec(`DROP TABLE metadata;`);
     },
   },
+  {
+    name: "create_external_id_source_index",
+    up: () => {
+      return db.exec(`
+        CREATE INDEX IF NOT EXISTS idx_documents_external_id_source 
+        ON documents(external_id, source);
+      `);
+    },
+    down: () => {
+      return db.exec(`DROP INDEX IF EXISTS idx_documents_external_id_source;`);
+    },
+  },
 
   // Add more migrations here
 ];

src/database/reembedding.ts

@@ -24,7 +24,7 @@ export const reembedAllDocuments = async (): Promise<void> => {
     console.log("Updating database schema to accommodate new embedding size...");
 
     // Alter documents table to handle potential changes in embedding dimensions
-    await db.exec(`
+    db.exec(`
       PRAGMA foreign_keys=off;
       BEGIN TRANSACTION;
       
@@ -33,10 +33,12 @@ export const reembedAllDocuments = async (): Promise<void> => {
       
       -- Drop original table
       DROP TABLE documents;
-      
-      -- Recreate table with updated schema and embedding size
-      ${createDocumentsTableSQL(embeddingSize)}
-      
+    `);
+    
+    // Recreate table with updated schema and embedding size
+    db.exec(createDocumentsTableSQL(embeddingSize));
+    
+    db.exec(`
       -- Copy data back
       INSERT INTO documents (id, external_id, text, metadata, source, created_at)
       SELECT id, external_id, text, metadata, source, created_at FROM temp_documents;
@@ -49,7 +51,7 @@ export const reembedAllDocuments = async (): Promise<void> => {
     `);
 
     // Similar process for document_chunks table
-    await db.exec(`
+    db.exec(`
       PRAGMA foreign_keys=off;
       BEGIN TRANSACTION;
       
@@ -58,10 +60,12 @@ export const reembedAllDocuments = async (): Promise<void> => {
       
       -- Drop original table
       DROP TABLE document_chunks;
-      
-      -- Recreate table with updated schema and embedding size
-      ${createDocumentChunksTableSQL(embeddingSize)}
-      
+    `);
+    
+    // Recreate table with updated schema and embedding size
+    db.exec(createDocumentChunksTableSQL(embeddingSize));
+    
+    db.exec(`
       -- Copy data back
       INSERT INTO document_chunks (id, document_id, text, created_at)
       SELECT id, document_id, text, created_at FROM temp_document_chunks;
@@ -76,7 +80,7 @@ export const reembedAllDocuments = async (): Promise<void> => {
     console.log("Database schema updated successfully");
 
     // Prepare database for batch operations
-    await db.run("BEGIN TRANSACTION");
+    db.exec("BEGIN TRANSACTION");
 
     // Get all documents
     const documents = db.query("SELECT id, text FROM documents").all() as { id: number; text: string }[];
@@ -105,8 +109,8 @@ export const reembedAllDocuments = async (): Promise<void> => {
 
             // Commit in batches to avoid holding transaction too long
             if (processedDocs % 100 === 0 && processedDocs !== documents.length) {
-              await db.run("COMMIT");
-              await db.run("BEGIN TRANSACTION");
+              db.exec("COMMIT");
+              db.exec("BEGIN TRANSACTION");
             }
           }
         }
@@ -116,10 +120,10 @@ export const reembedAllDocuments = async (): Promise<void> => {
     }
 
     // Commit document changes
-    await db.run("COMMIT");
+    db.exec("COMMIT");
 
     // Start new transaction for chunks
-    await db.run("BEGIN TRANSACTION");
+    db.exec("BEGIN TRANSACTION");
 
     // Get all document chunks
     const chunks = db.query("SELECT id, document_id, text FROM document_chunks").all() as { 
@@ -153,8 +157,8 @@ export const reembedAllDocuments = async (): Promise<void> => {
 
             // Commit in batches to avoid holding transaction too long
             if (processedChunks % 200 === 0 && processedChunks !== chunks.length) {
-              await db.run("COMMIT");
-              await db.run("BEGIN TRANSACTION");
+              db.exec("COMMIT");
+              db.exec("BEGIN TRANSACTION");
             }
           }
         }
@@ -168,13 +172,13 @@ export const reembedAllDocuments = async (): Promise<void> => {
     setMetadataValue("AI_EMBEDDING_SIZE", embeddingSize || "");
 
     // Commit chunk changes
-    await db.run("COMMIT");
+    db.exec("COMMIT");
 
     console.log("Re-embedding process completed successfully");
   } catch (error) {
     // Ensure transaction is rolled back if an error occurs
     try {
-      await db.run("ROLLBACK");
+      db.exec("ROLLBACK");
     } catch (rollbackError) {
       console.error("Error during rollback:", rollbackError);
     }

src/routes/context/getContext.ts

@@ -0,0 +1,57 @@
+import { z } from "zod";
+import { corsHeaders as headers } from "../../shared/corsHeaders";
+import { db } from "../../database/database";
+
+const schema = z.object({
+  key: z.string({ required_error: "Key field is required" }),
+});
+
+export type GetContextResult = { error: string } | { context: { key: string; message: string; metadata: Record<string, any> } };
+export type GetContextProps = z.infer<typeof schema>;
+
+export const getContext = async (body: GetContextProps): Promise<GetContextResult> => {
+  const { error, data, success } = schema.safeParse(body);
+
+  if (!success) {
+    return {
+      error: error.issues?.[0]?.message || "Validation failed",
+    };
+  }
+
+  try {
+    const context = db
+      .query("SELECT text, metadata FROM documents WHERE external_id = ? AND source = 'context'")
+      .get(data.key) as
+      | { text: string; metadata: string }
+      | undefined;
+
+    if (!context) {
+      return { error: "Context not found" };
+    }
+
+    return {
+      context: {
+        key: data.key,
+        message: context.text,
+        metadata: JSON.parse(context.metadata),
+      },
+    };
+  } catch (e) {
+    console.error("Error getting context:", e);
+    return { error: "Failed to get context" };
+  }
+};
+
+export const getContextRoute = async (request: Request) => {
+  const body = await request.json();
+  const result = await getContext(body);
+
+  if ("error" in result) {
+    return Response.json(result, {
+      status: 400,
+      headers,
+    });
+  }
+
+  return Response.json(result, { headers });
+};

src/routes/context/setContext.ts

@@ -0,0 +1,78 @@
+import { z } from "zod";
+import { generateEmbeddings } from "../../shared/generateEmbeddings";
+import { corsHeaders as headers } from "../../shared/corsHeaders";
+import { db } from "../../database/database";
+
+const schema = z.object({
+  key: z.string({ required_error: "Key field is required" }),
+  message: z.string({ required_error: "Message field is required" }),
+});
+
+export type SetContextResult = { error: string } | { message: string };
+export type SetContextProps = z.infer<typeof schema>;
+
+export const setContext = async (body: SetContextProps): Promise<SetContextResult> => {
+  const { error, data, success } = schema.safeParse(body);
+
+  if (!success) {
+    return {
+      error: error.issues?.[0]?.message || "Validation failed",
+    };
+  }
+
+  const embeddings = await generateEmbeddings(data.message, {
+    apiKey: process.env.AI_API_KEY as string,
+    baseURL: process.env.AI_BASE_URL,
+  });
+
+  if (!embeddings) {
+    return { error: "Failed to generate embeddings" };
+  }
+
+  try {
+    // Use a transaction for atomic operation
+    db.run("BEGIN TRANSACTION");
+    
+    try {
+      // Use JSON.stringify for embeddings array (more efficient than join)
+      const embeddingsStr = JSON.stringify(embeddings);
+      const metadata = JSON.stringify({ type: "context" });
+      
+      // Use INSERT OR REPLACE for simpler logic
+      db.query(
+        `
+          INSERT OR REPLACE INTO documents (external_id, text, embeddings, source, metadata)
+          VALUES (?, ?, ?, 'context', ?)
+        `
+      ).run(
+        data.key,
+        data.message,
+        embeddingsStr,
+        metadata
+      );
+      
+      db.run("COMMIT");
+      return { message: "Context successfully set" };
+    } catch (e) {
+      db.run("ROLLBACK");
+      throw e;
+    }
+  } catch (e) {
+    console.error("Error setting context:", e);
+    return { error: "Failed to set context" };
+  }
+};
+
+export const setContextRoute = async (request: Request) => {
+  const body = await request.json();
+  const result = await setContext(body);
+
+  if ("error" in result) {
+    return Response.json(result, {
+      status: 400,
+      headers,
+    });
+  }
+
+  return Response.json(result, { headers });
+};