@keo-ai/axiom 0.1.1 → 0.1.2

package/dist/embedding_search/embed.d.ts

@@ -0,0 +1 @@
+export declare function embed(text: string, dimensions?: number): Promise<number[]>;

package/dist/embedding_search/embed.js

@@ -0,0 +1,52 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.embed = embed;
+const EMBEDDING_URL = 'https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings';
+const EMBEDDING_MODEL = 'text-embedding-v4';
+function embed(text_1) {
+    return __awaiter(this, arguments, void 0, function* (text, dimensions = 1024) {
+        var _a;
+        const apiKey = process.env.BAILIAN_API_KEY;
+        if (!apiKey) {
+            throw new Error('BAILIAN_API_KEY environment variable is not set');
+        }
+        const body = {
+            model: EMBEDDING_MODEL,
+            input: text,
+            dimensions,
+        };
+        let response;
+        try {
+            response = yield fetch(EMBEDDING_URL, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                    Authorization: `Bearer ${apiKey}`,
+                },
+                body: JSON.stringify(body),
+            });
+        }
+        catch (cause) {
+            throw new Error(`[embedding] ${cause instanceof Error ? cause.message : String(cause)}`);
+        }
+        if (!response.ok) {
+            const text = yield response.text();
+            throw new Error(`[embedding] HTTP ${response.status}: ${text}`);
+        }
+        const data = (yield response.json());
+        const vector = (_a = data.data[0]) === null || _a === void 0 ? void 0 : _a.embedding;
+        if (!vector) {
+            throw new Error('[embedding] No embedding in response');
+        }
+        return vector;
+    });
+}

package/dist/embedding_search/index.d.ts

@@ -0,0 +1,27 @@
+import type { Pool } from 'pg';
+import type { EmbeddingSearchConfig, SearchResult } from './types';
+/**
+ * 向量检索静态入口。
+ *
+ * 封装 embedding → pgvector 检索 → [可选 rerank] 的完整链路。
+ * PG 连接由外部传入，Axiom 不管理连接池生命周期。
+ *
+ * @example
+ * ```ts
+ * const results = await EmbeddingSearch.query('query文本', {
+ *   enableRerank: true,
+ *   embeddingTopK: 20,
+ *   finalTopN: 5,
+ *   embeddingThreshold: 0.7,
+ *   rerankThreshold: 0.5,
+ * }, pgPool);
+ * ```
+ */
+export declare class EmbeddingSearch {
+    static query(query: string, config: EmbeddingSearchConfig, pool: Pool): Promise<SearchResult[]>;
+    static vectorSearch(options: Omit<import('./pgvector').PgVectorSearchOptions, 'pool'> & {
+        pool: Pool;
+    }): Promise<SearchResult[]>;
+}
+export type { EmbeddingSearchConfig, SearchResult } from './types';
+export { embed } from './embed';

package/dist/embedding_search/index.js

@@ -0,0 +1,162 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __rest = (this && this.__rest) || function (s, e) {
+    var t = {};
+    for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
+        t[p] = s[p];
+    if (s != null && typeof Object.getOwnPropertySymbols === "function")
+        for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
+            if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
+                t[p[i]] = s[p[i]];
+        }
+    return t;
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.embed = exports.EmbeddingSearch = void 0;
+const pgvector_1 = require("./pgvector");
+const embed_1 = require("./embed");
+const RERANK_URL = 'https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank';
+const RERANK_MODEL = 'qwen3-rerank';
+function rerank(query, documents, topN, apiKey) {
+    return __awaiter(this, void 0, void 0, function* () {
+        const body = {
+            model: RERANK_MODEL,
+            input: {
+                query,
+                documents: documents.map((d) => d.content),
+            },
+            parameters: {
+                return_documents: true,
+                top_n: topN,
+            },
+        };
+        let response;
+        try {
+            response = yield fetch(RERANK_URL, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                    Authorization: `Bearer ${apiKey}`,
+                },
+                body: JSON.stringify(body),
+            });
+        }
+        catch (cause) {
+            throw new Error(`[rerank] ${cause instanceof Error ? cause.message : String(cause)}`);
+        }
+        if (!response.ok) {
+            const text = yield response.text();
+            throw new Error(`[rerank] HTTP ${response.status}: ${text}`);
+        }
+        const data = (yield response.json());
+        return data.output.results.map((r) => {
+            var _a, _b, _c, _d;
+            const doc = documents[r.index];
+            return {
+                index: r.index,
+                id: (_a = doc === null || doc === void 0 ? void 0 : doc.id) !== null && _a !== void 0 ? _a : r.index,
+                content: (_d = (_c = (_b = r.document) === null || _b === void 0 ? void 0 : _b.text) !== null && _c !== void 0 ? _c : doc === null || doc === void 0 ? void 0 : doc.content) !== null && _d !== void 0 ? _d : '',
+                score: r.relevance_score,
+            };
+        });
+    });
+}
+/**
+ * 向量检索静态入口。
+ *
+ * 封装 embedding → pgvector 检索 → [可选 rerank] 的完整链路。
+ * PG 连接由外部传入，Axiom 不管理连接池生命周期。
+ *
+ * @example
+ * ```ts
+ * const results = await EmbeddingSearch.query('query文本', {
+ *   enableRerank: true,
+ *   embeddingTopK: 20,
+ *   finalTopN: 5,
+ *   embeddingThreshold: 0.7,
+ *   rerankThreshold: 0.5,
+ * }, pgPool);
+ * ```
+ */
+class EmbeddingSearch {
+    static query(query, config, pool) {
+        return __awaiter(this, void 0, void 0, function* () {
+            const { enableRerank = true, embeddingTopK = 10, finalTopN = 5, embeddingThreshold = 0.8, rerankThreshold = 0.1, tableName, embeddingColumn, rerankKey = 'content', dimensions, filter, select, } = config;
+            const idColumn = 'id';
+            const actualTopK = enableRerank
+                ? Math.max(embeddingTopK, finalTopN)
+                : finalTopN;
+            const apiKey = process.env.BAILIAN_API_KEY;
+            if (!apiKey) {
+                throw new Error('BAILIAN_API_KEY environment variable is not set');
+            }
+            const queryVector = yield (0, embed_1.embed)(query, dimensions);
+            // 启用 rerank 时，确保能拿到 rerankKey 对应列的内容
+            let actualSelect = select;
+            if (enableRerank && actualSelect && !actualSelect.includes(rerankKey)) {
+                actualSelect = [...actualSelect, rerankKey];
+            }
+            const searchResults = yield (0, pgvector_1.vectorSearch)({
+                pool,
+                tableName,
+                vector: queryVector,
+                topK: actualTopK,
+                threshold: embeddingThreshold,
+                embeddingColumn,
+                filter,
+                select: actualSelect,
+            });
+            if (searchResults.length === 0) {
+                return [];
+            }
+            if (enableRerank) {
+                const missingRerank = searchResults.find((r) => r[rerankKey] === undefined);
+                if (missingRerank) {
+                    throw new Error(`Rerank requires column "${rerankKey}" but it was not found. ` +
+                        `Either create this column in your table, or set rerankKey to the correct column name.`);
+                }
+            }
+            if (!enableRerank) {
+                return searchResults.slice(0, finalTopN);
+            }
+            const rerankResults = yield rerank(query, searchResults.map((r) => ({
+                id: r[idColumn],
+                content: r[rerankKey],
+            })), finalTopN, apiKey);
+            let reranked = [];
+            for (const rr of rerankResults) {
+                if (rerankThreshold !== undefined && rr.score < rerankThreshold) {
+                    continue;
+                }
+                const original = searchResults[rr.index];
+                if (original) {
+                    reranked.push(Object.assign(Object.assign({}, original), { rerankScore: rr.score }));
+                }
+            }
+            // 如果用户没选 rerankKey 对应列，从返回结果中剔除
+            if (select && !select.includes(rerankKey)) {
+                reranked = reranked.map((r) => {
+                    const _a = r, _b = rerankKey, _ = _a[_b], rest = __rest(_a, [typeof _b === "symbol" ? _b : _b + ""]);
+                    return rest;
+                });
+            }
+            return reranked;
+        });
+    }
+    static vectorSearch(options) {
+        return __awaiter(this, void 0, void 0, function* () {
+            return (0, pgvector_1.vectorSearch)(options);
+        });
+    }
+}
+exports.EmbeddingSearch = EmbeddingSearch;
+var embed_2 = require("./embed");
+Object.defineProperty(exports, "embed", { enumerable: true, get: function () { return embed_2.embed; } });

package/dist/embedding_search/pgvector.d.ts

@@ -0,0 +1,19 @@
+import type { Pool } from 'pg';
+import type { SearchResult } from './types';
+export interface PgVectorSearchOptions {
+    readonly pool: Pool;
+    readonly tableName: string;
+    readonly vector: number[];
+    readonly topK: number;
+    readonly threshold?: number;
+    readonly embeddingColumn?: string;
+    readonly filter?: Record<string, unknown>;
+    readonly select?: string[];
+}
+/**
+ * 使用 pgvector 执行 cosine similarity 检索。
+ *
+ * - `embedding <=> query` 为 cosine distance（范围 0~2，越小越相似）
+ * - `1 - cosine_distance` 为 cosine similarity（范围 -1~1，越大越相似）
+ */
+export declare function vectorSearch(options: PgVectorSearchOptions): Promise<SearchResult[]>;

package/dist/embedding_search/pgvector.js

@@ -0,0 +1,116 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __rest = (this && this.__rest) || function (s, e) {
+    var t = {};
+    for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
+        t[p] = s[p];
+    if (s != null && typeof Object.getOwnPropertySymbols === "function")
+        for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
+            if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
+                t[p[i]] = s[p[i]];
+        }
+    return t;
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.vectorSearch = vectorSearch;
+function validateIdentifier(name) {
+    if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(name)) {
+        throw new Error(`Invalid identifier: ${name}`);
+    }
+}
+/**
+ * 使用 pgvector 执行 cosine similarity 检索。
+ *
+ * - `embedding <=> query` 为 cosine distance（范围 0~2，越小越相似）
+ * - `1 - cosine_distance` 为 cosine similarity（范围 -1~1，越大越相似）
+ */
+function vectorSearch(options) {
+    return __awaiter(this, void 0, void 0, function* () {
+        const { pool, tableName, vector, topK, threshold, embeddingColumn = 'embedding', filter, select } = options;
+        validateIdentifier(tableName);
+        validateIdentifier(embeddingColumn);
+        const vectorStr = JSON.stringify(vector);
+        const params = [vectorStr, topK];
+        const conditions = [];
+        let paramIndex = 3;
+        if (threshold !== undefined) {
+            conditions.push(`1 - ("${embeddingColumn}" <=> $1::vector) >= $${paramIndex}`);
+            params.push(threshold);
+            paramIndex++;
+        }
+        if (filter && Object.keys(filter).length > 0) {
+            for (const [key, value] of Object.entries(filter)) {
+                validateIdentifier(key);
+                if (Array.isArray(value)) {
+                    const placeholders = value.map((_, i) => `$${paramIndex + i}`).join(', ');
+                    conditions.push(`"${key}" IN (${placeholders})`);
+                    params.push(...value);
+                    paramIndex += value.length;
+                }
+                else if (value && typeof value === 'object') {
+                    const ops = value;
+                    for (const [op, opValue] of Object.entries(ops)) {
+                        switch (op) {
+                            case '$eq':
+                                conditions.push(`"${key}" = $${paramIndex}`);
+                                params.push(opValue);
+                                paramIndex++;
+                                break;
+                            case '$gt':
+                                conditions.push(`"${key}" > $${paramIndex}`);
+                                params.push(opValue);
+                                paramIndex++;
+                                break;
+                            case '$lt':
+                                conditions.push(`"${key}" < $${paramIndex}`);
+                                params.push(opValue);
+                                paramIndex++;
+                                break;
+                            case '$gte':
+                                conditions.push(`"${key}" >= $${paramIndex}`);
+                                params.push(opValue);
+                                paramIndex++;
+                                break;
+                            case '$lte':
+                                conditions.push(`"${key}" <= $${paramIndex}`);
+                                params.push(opValue);
+                                paramIndex++;
+                                break;
+                            default:
+                                throw new Error(`Unsupported filter operator: ${op}`);
+                        }
+                    }
+                }
+                else {
+                    conditions.push(`"${key}" = $${paramIndex}`);
+                    params.push(value);
+                    paramIndex++;
+                }
+            }
+        }
+        const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '';
+        const selectClause = select && select.length > 0
+            ? select.map((c) => { validateIdentifier(c); return `"${c}"`; }).join(', ')
+            : '*';
+        const sql = `
+    SELECT ${selectClause}, 1 - ("${embeddingColumn}" <=> $1::vector) as score
+    FROM "${tableName}"
+    ${whereClause}
+    ORDER BY "${embeddingColumn}" <=> $1::vector
+    LIMIT $2
+  `;
+        const result = yield pool.query(sql, params);
+        return result.rows.map((row) => {
+            const { score } = row, rest = __rest(row, ["score"]);
+            return Object.assign(Object.assign({}, rest), { embeddingScore: score });
+        });
+    });
+}

package/dist/embedding_search/types.d.ts

@@ -0,0 +1,26 @@
+export interface ComparisonOp {
+    readonly $eq?: unknown;
+    readonly $gt?: number;
+    readonly $lt?: number;
+    readonly $gte?: number;
+    readonly $lte?: number;
+}
+export type FilterValue = string | number | boolean | null | unknown[] | ComparisonOp;
+export interface EmbeddingSearchConfig {
+    readonly enableRerank?: boolean;
+    readonly embeddingTopK?: number;
+    readonly finalTopN?: number;
+    readonly embeddingThreshold?: number;
+    readonly rerankThreshold?: number;
+    readonly tableName: string;
+    readonly embeddingColumn?: string;
+    readonly rerankKey?: string;
+    readonly dimensions?: number;
+    readonly filter?: Record<string, FilterValue>;
+    readonly select?: string[];
+}
+export interface SearchResult {
+    readonly [key: string]: unknown;
+    readonly embeddingScore: number;
+    readonly rerankScore?: number;
+}

package/dist/embedding_search/types.js

@@ -0,0 +1,2 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });

package/dist/index.d.ts

@@ -29,3 +29,6 @@ export { BailianProvider } from './predict';
 /** 模型注册表，定义模型与 Provider 的映射关系 */
 export { MODEL_REGISTRY } from './predict';
 export type { Model, PredictConfig, PredictWithMessagesConfig } from './predict';
+/** 向量检索（Embedding + pgvector + Rerank） */
+export { EmbeddingSearch, embed } from './embedding_search';
+export type { EmbeddingSearchConfig, SearchResult } from './embedding_search';

package/dist/index.js

@@ -20,7 +20,7 @@
  * @packageDocumentation
  */
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.MODEL_REGISTRY = exports.BailianProvider = exports.Predictor = exports.LLM = void 0;
+exports.embed = exports.EmbeddingSearch = exports.MODEL_REGISTRY = exports.BailianProvider = exports.Predictor = exports.LLM = void 0;
 /** LLM 静态入口类，封装 Provider 连接、故障转移和返回解析 */
 var predict_1 = require("./predict");
 Object.defineProperty(exports, "LLM", { enumerable: true, get: function () { return predict_1.LLM; } });
@@ -33,3 +33,7 @@ Object.defineProperty(exports, "BailianProvider", { enumerable: true, get: funct
 /** 模型注册表，定义模型与 Provider 的映射关系 */
 var predict_4 = require("./predict");
 Object.defineProperty(exports, "MODEL_REGISTRY", { enumerable: true, get: function () { return predict_4.MODEL_REGISTRY; } });
+/** 向量检索（Embedding + pgvector + Rerank） */
+var embedding_search_1 = require("./embedding_search");
+Object.defineProperty(exports, "EmbeddingSearch", { enumerable: true, get: function () { return embedding_search_1.EmbeddingSearch; } });
+Object.defineProperty(exports, "embed", { enumerable: true, get: function () { return embedding_search_1.embed; } });

package/dist/predict/predict.d.ts

@@ -1,4 +1,4 @@
-import type { Message, LLMResponse, StreamChunk } from './types';
+import type { Message, StreamChunk } from './types';
 import type { PredictConfig, PredictWithMessagesConfig } from './config';
 /**
  * LLM 静态类。封装 Provider 连接、请求组装、故障转移和返回解析。
@@ -9,7 +9,7 @@ import type { PredictConfig, PredictWithMessagesConfig } from './config';
  * import { LLM } from 'axiom';
  *
  * const res = await LLM.predict({ model: 'qwen-max', prompt: '你好' });
- * console.log(res.content);
+ * console.log(res);
  * ```
  */
 export declare class LLM {
@@ -20,41 +20,33 @@ export declare class LLM {
      * 内部自动构造单条 user message，支持可选的 system prompt。
      *
      * 根据 `responseFormat` 返回不同类型：
-     * - `'text'` → `string`
+     * - 未设置 / `'text'` → `string`
      * - `'json'` → 解析后的对象（`any`）
-     * - 未设置 → `LLMResponse`
      *
      * @param config - 预测配置（模型、prompt、温度等）
      * @returns 模型生成的完整响应
      * @throws Error - Provider 未配置、模型不存在、或所有 Provider 均失败时抛出
      */
-    static predict(config: PredictConfig & {
-        responseFormat: 'text';
-    }): Promise<string>;
     static predict(config: PredictConfig & {
         responseFormat: 'json';
     }): Promise<any>;
-    static predict(config: PredictConfig): Promise<LLMResponse>;
+    static predict(config: PredictConfig): Promise<string>;
     /**
      * 使用自定义消息列表调用 LLM。适用于多轮对话等需要精细控制 message 结构的场景。
      *
      * 根据 `responseFormat` 返回不同类型：
-     * - `'text'` → `string`
+     * - 未设置 / `'text'` → `string`
      * - `'json'` → 解析后的对象（可用泛型指定类型）
-     * - 未设置 → `LLMResponse`
      *
      * @param messages - 消息列表（user/assistant 角色）
      * @param config - 预测配置（不含 prompt，因为由 messages 提供）
      * @returns 模型生成的完整响应
      * @throws Error - Provider 未配置、模型不存在、或所有 Provider 均失败时抛出
      */
-    static predictWithMessages(messages: ReadonlyArray<Message>, config: PredictWithMessagesConfig & {
-        responseFormat: 'text';
-    }): Promise<string>;
     static predictWithMessages(messages: ReadonlyArray<Message>, config: PredictWithMessagesConfig & {
         responseFormat: 'json';
     }): Promise<any>;
-    static predictWithMessages(messages: ReadonlyArray<Message>, config: PredictWithMessagesConfig): Promise<LLMResponse>;
+    static predictWithMessages(messages: ReadonlyArray<Message>, config: PredictWithMessagesConfig): Promise<string>;
     /**
      * 流式调用 LLM，逐块返回模型输出。
      *

package/dist/predict/predict.js

@@ -69,16 +69,14 @@ function toLLMRequest(config, messages) {
 }
 function unwrapResponse(response, responseFormat) {
     var _a;
-    if (responseFormat === 'text') {
-        return (_a = response.content) !== null && _a !== void 0 ? _a : '';
-    }
     if (responseFormat === 'json') {
         if (!response.content) {
             throw new Error('Empty response content when responseFormat is json');
         }
         return JSON.parse(response.content);
     }
-    return response;
+    // 默认按 text 返回（包括未设置 responseFormat 的情况）
+    return (_a = response.content) !== null && _a !== void 0 ? _a : '';
 }
 /**
  * LLM 静态类。封装 Provider 连接、请求组装、故障转移和返回解析。
@@ -89,7 +87,7 @@ function unwrapResponse(response, responseFormat) {
  * import { LLM } from 'axiom';
  *
  * const res = await LLM.predict({ model: 'qwen-max', prompt: '你好' });
- * console.log(res.content);
+ * console.log(res);
  * ```
  */
 class LLM {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@keo-ai/axiom",
-  "version": "0.1.1",
+  "version": "0.1.2",
   "description": "基于 LLM 的预测与推理库，支持多 Provider 切换",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",