fix: Koloxarto/fix ragknowledge for postgres (#2153)

* fix formatting out of the way

* fix postgress chunk uuid handling for ragKnowledge

---------

Co-authored-by: Odilitime <janesmith@airmail.cc>

Author: web3gh

packages/adapter-postgres/src/index.ts

@@ -4,31 +4,31 @@ import { v4 } from "uuid";
 import pg from "pg";
 type Pool = pg.Pool;
 
-import {
-    QueryConfig,
-    QueryConfigValues,
-    QueryResult,
-    QueryResultRow,
-} from "pg";
 import {
     Account,
     Actor,
+    DatabaseAdapter,
+    EmbeddingProvider,
     GoalStatus,
+    Participant,
+    RAGKnowledgeItem,
+    elizaLogger,
+    getEmbeddingConfig,
     type Goal,
+    type IDatabaseCacheAdapter,
     type Memory,
     type Relationship,
     type UUID,
-    type IDatabaseCacheAdapter,
-    Participant,
-    elizaLogger,
-    getEmbeddingConfig,
-    DatabaseAdapter,
-    EmbeddingProvider,
-    RAGKnowledgeItem
 } from "@elizaos/core";
 import fs from "fs";
-import { fileURLToPath } from "url";
 import path from "path";
+import {
+    QueryConfig,
+    QueryConfigValues,
+    QueryResult,
+    QueryResultRow,
+} from "pg";
+import { fileURLToPath } from "url";
 
 const __filename = fileURLToPath(import.meta.url); // get the resolved path to the file
 const __dirname = path.dirname(__filename); // get the name of the directory
@@ -199,7 +199,7 @@ export class PostgresDatabaseAdapter
             return true;
         } catch (error) {
             elizaLogger.error("Failed to validate vector extension:", {
-                error: error instanceof Error ? error.message : String(error)
+                error: error instanceof Error ? error.message : String(error),
             });
             return false;
         }
@@ -239,8 +239,10 @@ export class PostgresDatabaseAdapter
                 );
             `);
 
-            if (!rows[0].exists || !await this.validateVectorSetup()) {
-                elizaLogger.info("Applying database schema - tables or vector extension missing");
+            if (!rows[0].exists || !(await this.validateVectorSetup())) {
+                elizaLogger.info(
+                    "Applying database schema - tables or vector extension missing"
+                );
                 const schema = fs.readFileSync(
                     path.resolve(__dirname, "../schema.sql"),
                     "utf8"
@@ -1523,12 +1525,17 @@ export class PostgresDatabaseAdapter
 
             const { rows } = await this.pool.query(sql, queryParams);
 
-            return rows.map(row => ({
+            return rows.map((row) => ({
                 id: row.id,
                 agentId: row.agentId,
-                content: typeof row.content === 'string' ? JSON.parse(row.content) : row.content,
-                embedding: row.embedding ? new Float32Array(row.embedding) : undefined,
-                createdAt: row.createdAt.getTime()
+                content:
+                    typeof row.content === "string"
+                        ? JSON.parse(row.content)
+                        : row.content,
+                embedding: row.embedding
+                    ? new Float32Array(row.embedding)
+                    : undefined,
+                createdAt: row.createdAt.getTime(),
             }));
         }, "getKnowledge");
     }
@@ -1544,7 +1551,7 @@ export class PostgresDatabaseAdapter
             const cacheKey = `embedding_${params.agentId}_${params.searchText}`;
             const cachedResult = await this.getCache({
                 key: cacheKey,
-                agentId: params.agentId
+                agentId: params.agentId,
             });
 
             if (cachedResult) {
@@ -1594,24 +1601,29 @@ export class PostgresDatabaseAdapter
             const { rows } = await this.pool.query(sql, [
                 vectorStr,
                 params.agentId,
-                `%${params.searchText || ''}%`,
+                `%${params.searchText || ""}%`,
                 params.match_threshold,
-                params.match_count
+                params.match_count,
             ]);
 
-            const results = rows.map(row => ({
+            const results = rows.map((row) => ({
                 id: row.id,
                 agentId: row.agentId,
-                content: typeof row.content === 'string' ? JSON.parse(row.content) : row.content,
-                embedding: row.embedding ? new Float32Array(row.embedding) : undefined,
+                content:
+                    typeof row.content === "string"
+                        ? JSON.parse(row.content)
+                        : row.content,
+                embedding: row.embedding
+                    ? new Float32Array(row.embedding)
+                    : undefined,
                 createdAt: row.createdAt.getTime(),
-                similarity: row.combined_score
+                similarity: row.combined_score,
             }));
 
             await this.setCache({
                 key: cacheKey,
                 agentId: params.agentId,
-                value: JSON.stringify(results)
+                value: JSON.stringify(results),
             });
 
             return results;
@@ -1622,35 +1634,52 @@ export class PostgresDatabaseAdapter
         return this.withDatabase(async () => {
             const client = await this.pool.connect();
             try {
-                await client.query('BEGIN');
-
-                const sql = `
-                    INSERT INTO knowledge (
-                        id, "agentId", content, embedding, "createdAt",
-                        "isMain", "originalId", "chunkIndex", "isShared"
-                    ) VALUES ($1, $2, $3, $4, to_timestamp($5/1000.0), $6, $7, $8, $9)
-                    ON CONFLICT (id) DO NOTHING
-                `;
+                await client.query("BEGIN");
 
                 const metadata = knowledge.content.metadata || {};
-                const vectorStr = knowledge.embedding ?
-                `[${Array.from(knowledge.embedding).join(",")}]` : null;
-
-                await client.query(sql, [
-                    knowledge.id,
-                    metadata.isShared ? null : knowledge.agentId,
-                    knowledge.content,
-                    vectorStr,
-                    knowledge.createdAt || Date.now(),
-                    metadata.isMain || false,
-                    metadata.originalId || null,
-                    metadata.chunkIndex || null,
-                    metadata.isShared || false
-                ]);
+                const vectorStr = knowledge.embedding
+                    ? `[${Array.from(knowledge.embedding).join(",")}]`
+                    : null;
+
+                // If this is a chunk, use createKnowledgeChunk
+                if (metadata.isChunk && metadata.originalId) {
+                    await this.createKnowledgeChunk({
+                        id: knowledge.id,
+                        originalId: metadata.originalId,
+                        agentId: metadata.isShared ? null : knowledge.agentId,
+                        content: knowledge.content,
+                        embedding: knowledge.embedding,
+                        chunkIndex: metadata.chunkIndex || 0,
+                        isShared: metadata.isShared || false,
+                        createdAt: knowledge.createdAt || Date.now(),
+                    });
+                } else {
+                    // This is a main knowledge item
+                    await client.query(
+                        `
+                        INSERT INTO knowledge (
+                            id, "agentId", content, embedding, "createdAt",
+                            "isMain", "originalId", "chunkIndex", "isShared"
+                        ) VALUES ($1, $2, $3, $4, to_timestamp($5/1000.0), $6, $7, $8, $9)
+                        ON CONFLICT (id) DO NOTHING
+                    `,
+                        [
+                            knowledge.id,
+                            metadata.isShared ? null : knowledge.agentId,
+                            knowledge.content,
+                            vectorStr,
+                            knowledge.createdAt || Date.now(),
+                            true,
+                            null,
+                            null,
+                            metadata.isShared || false,
+                        ]
+                    );
+                }
 
-                await client.query('COMMIT');
+                await client.query("COMMIT");
             } catch (error) {
-                await client.query('ROLLBACK');
+                await client.query("ROLLBACK");
                 throw error;
             } finally {
                 client.release();
@@ -1660,19 +1689,100 @@ export class PostgresDatabaseAdapter
 
     async removeKnowledge(id: UUID): Promise<void> {
         return this.withDatabase(async () => {
-            await this.pool.query('DELETE FROM knowledge WHERE id = $1', [id]);
+            const client = await this.pool.connect();
+            try {
+                await client.query("BEGIN");
+
+                // Check if this is a pattern-based chunk deletion (e.g., "id-chunk-*")
+                if (typeof id === "string" && id.includes("-chunk-*")) {
+                    const mainId = id.split("-chunk-")[0];
+                    // Delete chunks for this main ID
+                    await client.query(
+                        'DELETE FROM knowledge WHERE "originalId" = $1',
+                        [mainId]
+                    );
+                } else {
+                    // First delete all chunks associated with this knowledge item
+                    await client.query(
+                        'DELETE FROM knowledge WHERE "originalId" = $1',
+                        [id]
+                    );
+                    // Then delete the main knowledge item
+                    await client.query("DELETE FROM knowledge WHERE id = $1", [
+                        id,
+                    ]);
+                }
+
+                await client.query("COMMIT");
+            } catch (error) {
+                await client.query("ROLLBACK");
+                elizaLogger.error("Error removing knowledge", {
+                    error:
+                        error instanceof Error ? error.message : String(error),
+                    id,
+                });
+                throw error;
+            } finally {
+                client.release();
+            }
         }, "removeKnowledge");
     }
 
     async clearKnowledge(agentId: UUID, shared?: boolean): Promise<void> {
         return this.withDatabase(async () => {
-            const sql = shared ?
-                'DELETE FROM knowledge WHERE ("agentId" = $1 OR "isShared" = true)' :
-                'DELETE FROM knowledge WHERE "agentId" = $1';
+            const sql = shared
+                ? 'DELETE FROM knowledge WHERE ("agentId" = $1 OR "isShared" = true)'
+                : 'DELETE FROM knowledge WHERE "agentId" = $1';
 
             await this.pool.query(sql, [agentId]);
         }, "clearKnowledge");
     }
+
+    private async createKnowledgeChunk(params: {
+        id: UUID;
+        originalId: UUID;
+        agentId: UUID | null;
+        content: any;
+        embedding: Float32Array | undefined | null;
+        chunkIndex: number;
+        isShared: boolean;
+        createdAt: number;
+    }): Promise<void> {
+        const vectorStr = params.embedding
+            ? `[${Array.from(params.embedding).join(",")}]`
+            : null;
+
+        // Store the pattern-based ID in the content metadata for compatibility
+        const patternId = `${params.originalId}-chunk-${params.chunkIndex}`;
+        const contentWithPatternId = {
+            ...params.content,
+            metadata: {
+                ...params.content.metadata,
+                patternId,
+            },
+        };
+
+        await this.pool.query(
+            `
+            INSERT INTO knowledge (
+                id, "agentId", content, embedding, "createdAt",
+                "isMain", "originalId", "chunkIndex", "isShared"
+            ) VALUES ($1, $2, $3, $4, to_timestamp($5/1000.0), $6, $7, $8, $9)
+            ON CONFLICT (id) DO NOTHING
+        `,
+            [
+                v4(), // Generate a proper UUID for PostgreSQL
+                params.agentId,
+                contentWithPatternId, // Store the pattern ID in metadata
+                vectorStr,
+                params.createdAt,
+                false,
+                params.originalId,
+                params.chunkIndex,
+                params.isShared,
+            ]
+        );
+    }
 }
 
 export default PostgresDatabaseAdapter;