@retrivora-ai/rag-engine 0.2.6 → 0.2.8

package/src/providers/vectordb/MultiTablePostgresProvider.ts

@@ -1,11 +1,45 @@
-
-
-
 import { Pool, PoolClient } from 'pg';
 import { BaseVectorProvider } from './BaseVectorProvider';
 import { VectorMatch, UpsertDocument } from '../../types';
 import { VectorDBConfig } from '../../config/RagConfig';
 
+interface QueryFieldHint {
+  field?: string;
+  value: string;
+}
+
+interface TableSearchConfig {
+  availableFields: string[];
+  searchableFields: string[];
+  normalizedFieldMap: Map<string, string>;
+}
+
+function normalizeFieldName(value: string): string {
+  return value.toLowerCase().replace(/[^a-z0-9]+/g, '');
+}
+
+function sanitizeFilterHints(rawFilter?: Record<string, unknown>): QueryFieldHint[] {
+  if (!Array.isArray(rawFilter?.__fieldHints)) return [];
+
+  return (rawFilter.__fieldHints as unknown[])
+    .filter((hint): hint is QueryFieldHint => typeof hint === 'object' && hint !== null)
+    .map((hint) => ({
+      value: typeof hint.value === 'string' ? hint.value.trim().toLowerCase() : '',
+      ...(typeof hint.field === 'string' && hint.field.trim()
+        ? { field: normalizeFieldName(hint.field) }
+        : {}),
+    }))
+    .filter((hint) => Boolean(hint.value));
+}
+
+function stripInternalFilterKeys(filter?: Record<string, unknown>): Record<string, unknown> {
+  if (!filter) return {};
+
+  return Object.fromEntries(
+    Object.entries(filter).filter(([key]) => !key.startsWith('__'))
+  );
+}
+
 /**
  * MultiTablePostgresProvider — PostgreSQL implementation that searches across
  * multiple existing tables with pre-existing embeddings.
@@ -19,6 +53,7 @@ export class MultiTablePostgresProvider extends BaseVectorProvider {
   private readonly dimensions: number;
   private readonly connectionString: string;
   private tables: string[];
+  private tableSearchConfig = new Map<string, TableSearchConfig>();
 
   constructor(config: VectorDBConfig) {
     super(config);
@@ -79,6 +114,7 @@ export class MultiTablePostgresProvider extends BaseVectorProvider {
       if (this.tables.length === 0) {
         console.warn('[MultiTablePostgresProvider] No tables with "embedding" columns found in the database.');
       } else {
+        this.tableSearchConfig = await this.loadTableSearchConfig(client, this.tables);
         console.log(
           `[MultiTablePostgresProvider] Connected. Searching across ${this.tables.length} table(s): ${this.tables.join(', ')}`
         );
@@ -127,28 +163,73 @@ export class MultiTablePostgresProvider extends BaseVectorProvider {
     console.log(`[MultiTablePostgresProvider] --- Starting Multi-Table Search ---`);
 
     const queryText = _filter?.__queryText as string | undefined;
+    const fieldHints = sanitizeFilterHints(_filter);
+    const explicitFilter = stripInternalFilterKeys(_filter);
 
     const queryPromises = this.tables.map(async (table) => {
       try {
         let sqlQuery = '';
         let params: unknown[] = [];
+        const tableConfig = this.tableSearchConfig.get(table);
+        const availableFields = tableConfig?.availableFields ?? [];
+        const searchableFields = tableConfig?.searchableFields ?? [];
+        const normalizedFieldMap = tableConfig?.normalizedFieldMap ?? new Map<string, string>();
+        const whereConditions: string[] = [];
 
         if (queryText) {
+          const genericValueHints = fieldHints
+            .filter((hint) => !hint.field)
+            .map((hint) => hint.value);
+
+          params = genericValueHints.length > 0
+            ? [vectorLiteral, queryText, genericValueHints]
+            : [vectorLiteral, queryText];
+
+          const genericValueParamIndex = genericValueHints.length > 0 ? 3 : -1;
+
+          const exactValueScoreExpr = genericValueHints.length > 0 && searchableFields.length > 0
+            ? `GREATEST(${searchableFields.map((field) =>
+              `CASE WHEN LOWER(COALESCE(to_jsonb(t)->>'${field}', '')) = ANY($${genericValueParamIndex}::text[]) THEN 1 ELSE 0 END`
+            ).join(', ')})`
+            : '0';
+
+          const fieldSpecificChecks = fieldHints
+            .filter((hint) => hint.field)
+            .map((hint) => {
+              const column = hint.field ? normalizedFieldMap.get(hint.field) : undefined;
+              if (!column) return null;
+              params.push(hint.value);
+              const paramIndex = params.length;
+              return `CASE WHEN LOWER(COALESCE(to_jsonb(t)->>'${column}', '')) = $${paramIndex}::text THEN 1 ELSE 0 END`;
+            })
+            .filter((expr): expr is string => Boolean(expr));
+
+          const exactFieldScoreExpr = fieldSpecificChecks.length > 0
+            ? `GREATEST(${fieldSpecificChecks.join(', ')})`
+            : '0';
+
           // Actual Hybrid Search: Semantic Vector Score + Full-Text Keyword Rank
           // We convert plainto_tsquery's AND (&) operator to OR (|) so that natural language queries
-          // (which contain words like "price", "product" not found in the row text) still match entities like "Router".
+          // (which contain words that are not present verbatim in the row text) still match the intended records.
           sqlQuery = `
             SELECT *, 
              (1 - (embedding <=> $1::vector)) AS vector_score,
              COALESCE(ts_rank(to_tsvector('english', t.*::text), NULLIF(REPLACE(plainto_tsquery('english', $2)::text, '&', '|'), '')::tsquery), 0) AS keyword_score,
-             ((1 - (embedding <=> $1::vector)) + (COALESCE(ts_rank(to_tsvector('english', t.*::text), NULLIF(REPLACE(plainto_tsquery('english', $2)::text, '&', '|'), '')::tsquery), 0) * 2.0)) AS hybrid_score
+             ${exactValueScoreExpr} AS exact_value_score,
+             ${exactFieldScoreExpr} AS exact_field_score,
+             (
+               (1 - (embedding <=> $1::vector)) +
+               (COALESCE(ts_rank(to_tsvector('english', t.*::text), NULLIF(REPLACE(plainto_tsquery('english', $2)::text, '&', '|'), '')::tsquery), 0) * 2.0) +
+               (${exactValueScoreExpr} * 3.0) +
+               (${exactFieldScoreExpr} * 5.0)
+             ) AS hybrid_score
             FROM "${table}" t
             ORDER BY hybrid_score DESC
             LIMIT 50
           `;
-          params = [vectorLiteral, queryText];
         } else {
           // Fallback to Vector-only Search
+          params = [vectorLiteral];
           sqlQuery = `
             SELECT *, 
              (1 - (embedding <=> $1::vector)) AS hybrid_score
@@ -156,7 +237,19 @@ export class MultiTablePostgresProvider extends BaseVectorProvider {
             ORDER BY hybrid_score DESC
             LIMIT 50
           `;
-          params = [vectorLiteral];
+        }
+
+        for (const [key, value] of Object.entries(explicitFilter)) {
+          const column = availableFields.includes(key)
+            ? key
+            : normalizedFieldMap.get(normalizeFieldName(key));
+          if (!column) continue;
+          params.push(String(value));
+          whereConditions.push(`COALESCE(to_jsonb(t)->>'${column}', '') = $${params.length}::text`);
+        }
+
+        if (whereConditions.length > 0) {
+          sqlQuery = sqlQuery.replace(`FROM "${table}" t`, `FROM "${table}" t\n            WHERE ${whereConditions.join(' AND ')}`);
         }
 
         const result = await this.pool.query(sqlQuery, params);
@@ -173,6 +266,8 @@ export class MultiTablePostgresProvider extends BaseVectorProvider {
           delete rest.embedding;
           delete rest.vector_score;
           delete rest.keyword_score;
+          delete rest.exact_value_score;
+          delete rest.exact_field_score;
 
           const content = `[TYPE: ${table.replace(/s$/, '').toUpperCase()}]\n` +
             Object.entries(rest)
@@ -202,32 +297,14 @@ export class MultiTablePostgresProvider extends BaseVectorProvider {
       allResults.push(...tableResults);
     }
 
-    // 1. Group results by table to ensure we take the top ones from each
-    const resultsByTable: Record<string, VectorMatch[]> = {};
-    for (const res of allResults) {
-      const table = res.metadata?.source_table as string;
-      if (!resultsByTable[table]) resultsByTable[table] = [];
-      resultsByTable[table].push(res);
-    }
-
-    // 2. Take top 3 from EACH table guaranteed, then fill the rest with overall best
-    const balancedResults: VectorMatch[] = [];
-    const tables = Object.keys(resultsByTable);
-
-    // First, ensure every table's best match is included
-    for (const table of tables) {
-      balancedResults.push(...resultsByTable[table].slice(0, 3));
-    }
-
-    // 3. Sort the balanced list by score
-    const finalSorted = balancedResults.sort((a, b) => b.score - a.score);
+    const finalSorted = allResults.sort((a, b) => b.score - a.score);
 
     if (finalSorted.length > 0) {
       console.log(`[MultiTablePostgresProvider] --- Search Complete ---`);
       console.log(`[MultiTablePostgresProvider] Final top match from "${finalSorted[0].metadata?.source_table ?? 'unknown'}" with score ${finalSorted[0].score.toFixed(4)}`);
     }
 
-    return finalSorted.slice(0, Math.max(topK, 15));
+    return finalSorted.slice(0, topK);
   }
 
   async delete(_id: string | number, _namespace?: string): Promise<void> { // eslint-disable-line @typescript-eslint/no-unused-vars
@@ -252,4 +329,67 @@ export class MultiTablePostgresProvider extends BaseVectorProvider {
       await this.pool.end();
     }
   }
+
+  private async loadTableSearchConfig(
+    client: PoolClient,
+    tables: string[]
+  ): Promise<Map<string, TableSearchConfig>> {
+    if (tables.length === 0) return new Map();
+
+    const configuredSearchFields = this.parseConfiguredSearchFields();
+    const result = await client.query(`
+      SELECT table_name, column_name, data_type
+      FROM information_schema.columns
+      WHERE table_schema = 'public'
+        AND table_name = ANY($1::text[])
+      ORDER BY ordinal_position
+    `, [tables]);
+
+    const rowsByTable = new Map<string, Array<{ column_name: string; data_type: string }>>();
+    for (const row of result.rows as Array<{ table_name: string; column_name: string; data_type: string }>) {
+      const rows = rowsByTable.get(row.table_name) ?? [];
+      rows.push({ column_name: row.column_name, data_type: row.data_type });
+      rowsByTable.set(row.table_name, rows);
+    }
+
+    const configByTable = new Map<string, TableSearchConfig>();
+    for (const table of tables) {
+      const columns = rowsByTable.get(table) ?? [];
+      const availableFields = columns
+        .filter(({ column_name }) => column_name !== 'embedding')
+        .map(({ column_name }) => column_name);
+
+      const inferredFields = columns
+        .filter(({ column_name, data_type }) =>
+          column_name !== 'embedding' &&
+          !['ARRAY', 'json', 'jsonb', 'bytea', 'tsvector', 'USER-DEFINED'].includes(data_type)
+        )
+        .map(({ column_name }) => column_name);
+
+      const searchableFields = configuredSearchFields.length > 0
+        ? configuredSearchFields.filter((field) => columns.some((column) => column.column_name === field))
+        : inferredFields;
+
+      const normalizedFieldMap = new Map<string, string>();
+      for (const field of availableFields) {
+        normalizedFieldMap.set(normalizeFieldName(field), field);
+      }
+
+      configByTable.set(table, {
+        availableFields,
+        searchableFields,
+        normalizedFieldMap,
+      });
+    }
+
+    return configByTable;
+  }
+
+  private parseConfiguredSearchFields(): string[] {
+    const raw = this.config.options?.searchFields;
+    if (Array.isArray(raw)) {
+      return raw.filter((field): field is string => typeof field === 'string' && field.trim().length > 0);
+    }
+    return [];
+  }
 }

package/src/providers/vectordb/PostgreSQLProvider.ts

@@ -3,6 +3,14 @@ import { VectorDBConfig } from '../../config/RagConfig';
 import { BaseVectorProvider } from './BaseVectorProvider';
 import { VectorMatch, UpsertDocument } from '../../types';
 
+function stripInternalFilterKeys(filter?: Record<string, unknown>): Record<string, unknown> {
+  if (!filter) return {};
+
+  return Object.fromEntries(
+    Object.entries(filter).filter(([key]) => !key.startsWith('__'))
+  );
+}
+
 /**
  * PostgreSQLProvider — PostgreSQL implementation using the pgvector extension.
  */
@@ -103,8 +111,9 @@ export class PostgreSQLProvider extends BaseVectorProvider {
     const params: unknown[] = [vectorLiteral, topK];
     if (namespace) params.push(namespace);
 
-    if (filter && Object.keys(filter).length > 0) {
-      const filterConditions = Object.entries(filter)
+    const publicFilter = stripInternalFilterKeys(filter);
+    if (Object.keys(publicFilter).length > 0) {
+      const filterConditions = Object.entries(publicFilter)
         .map(([key, val]) => {
           const paramIdx = params.length + 1;
           params.push(JSON.stringify(val));