@keo-ai/axiom 0.1.0 → 0.1.1

package/README.md

@@ -35,6 +35,8 @@ Axiom 是 **keo** 的 LLM 底座库。它不碰业务编排（意图路由、节
 │              └──────────┘ ┌─────────────┐ │
 │                         │  RAG基础      │ │
 │                         │  Embedding   │ │
+│                         │  pgvector    │ │
+│                         │  rerank      │ │
 │                         └─────────────┘ │
 └────────────────────────────────────────────┘
 ```
@@ -45,7 +47,7 @@ Axiom 是 **keo** 的 LLM 底座库。它不碰业务编排（意图路由、节
 - ✅ LLM 统一调用 + Provider 故障转移
 - ✅ Function Call Loop 的驱动、记录、校验、副作用传播
 - ✅ 记忆系统的抽象与基础实现
-- ✅ RAG 基础（Embedding / 向量检索）
+- ✅ RAG 基础（Embedding / pgvector 向量检索 / rerank）
 - ✅ 基础配置中心
 
 **Axiom 不做：**
@@ -60,6 +62,32 @@ Axiom 是 **keo** 的 LLM 底座库。它不碰业务编排（意图路由、节
 
 ---
 
+## 安装
+
+```bash
+npm install @keo-ai/axiom pg
+```
+
+> `pg` 是 peer dependency，用于 EmbeddingSearch 模块的 pgvector 连接。
+
+在你的项目根目录创建 `.env` 文件：
+
+```bash
+BAILIAN_API_KEY=your-api-key
+```
+
+然后在你的项目入口文件顶部引入 `dotenv` 以加载环境变量：
+
+```ts
+import 'dotenv/config';
+
+// 或指定路径
+import { config } from 'dotenv';
+config({ path: '.env.local' });
+```
+
+---
+
 ## Predict 模块（LLM 调用）
 
 Predict 是 Axiom 的 LLM 调用层。它封装了 Provider 连接、请求组装、故障转移和返回解析，让上层只需关心「用什么模型、传什么消息」。
@@ -78,7 +106,7 @@ export BAILIAN_API_KEY="your-api-key"
 import { LLM } from '@keo-ai/axiom';
 
 const res = await LLM.predict({ model: 'qwen3.7-max', prompt: '你好' });
-console.log(res.content);
+console.log(res);
 ```
 
 ### API 概览
@@ -96,9 +124,8 @@ console.log(res.content);
 
 | `responseFormat` | 返回值 | 说明 |
 |---|---|---|
-| `'text'` | `string` | 直接返回模型输出的文本 |
+| 未设置 / `'text'` | `string` | 直接返回模型输出的文本 |
 | `'json'` | `any` | 自动 `JSON.parse` |
-| 未设置 | `LLMResponse` | 完整响应（含 `content`、`usage`、`model`） |
 
 ```ts
 // text → string
@@ -170,6 +197,212 @@ try {
 
 ---
 
+## EmbeddingSearch 模块（向量检索）
+
+EmbeddingSearch 是 Axiom 的 RAG 底座。它封装了 `query → embedding → pgvector 检索 → [可选 rerank]` 的完整链路，只需一行代码即可实现语义检索。
+
+### 需要的列
+
+使用 `EmbeddingSearch` 前，先创建 pgvector 表。**只有 `embedding` 列是必需的**（默认列名 `embedding`，可通过 `embeddingColumn` 改），其余列完全自由：
+
+| 列名 | 类型 | 说明 |
+|---|---|---|
+| `embedding` | `vector(1024)` | 向量字段，必需，默认列名 `embedding`，可通过 `embeddingColumn` 自定义 |
+| *(任意)* | *(任意)* | 你自己的业务列，可直接用于过滤和返回 |
+
+**开 rerank 时的额外要求**：需指定 `rerankKey` 告诉系统哪一列是文本内容（默认 `'content'`）。
+
+```sql
+CREATE EXTENSION IF NOT EXISTS vector;
+
+-- 最简表：不开 rerank，只有向量和业务列
+CREATE TABLE products (
+  product_id SERIAL PRIMARY KEY,
+  name VARCHAR(100),
+  category VARCHAR(50),
+  price NUMERIC,
+  embedding VECTOR(1024)
+);
+
+-- 开 rerank 的表：有文本列供 rerank 使用
+CREATE TABLE documents (
+  doc_id SERIAL PRIMARY KEY,
+  title VARCHAR(100),
+  body TEXT,           -- rerankKey: 'body'
+  embedding VECTOR(1024)
+);
+
+CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops);
+```
+
+### 快速开始
+
+不传 `select` 时默认返回所有列（`SELECT *`）：
+
+```ts
+import { EmbeddingSearch } from '@keo-ai/axiom';
+import { Pool } from 'pg';
+
+const pool = new Pool({ connectionString: 'postgresql://...' });
+
+// 默认 SELECT *，返回所有列
+const results = await EmbeddingSearch.query('衣服质量怎么样', {
+  tableName: 'documents',
+}, pool);
+
+for (const r of results) {
+  console.log(r.id, r.content, r.embeddingScore, r.rerankScore);
+}
+```
+
+指定 `select` 只返回需要的列：
+
+```ts
+const results = await EmbeddingSearch.query('衣服质量怎么样', {
+  tableName: 'documents',
+  select: ['id', 'category', 'price'],
+}, pool);
+
+// 只返回 id、category、price 三列，以及 embeddingScore 和 rerankScore
+for (const r of results) {
+  console.log(r.category, r.price, r.embeddingScore);
+}
+```
+
+文本列通过 `rerankKey` 指定：
+
+```ts
+const results = await EmbeddingSearch.query('衣服质量怎么样', {
+  tableName: 'documents',
+  rerankKey: 'body',
+}, pool);
+```
+
+### 返回结构
+
+`SearchResult` 不固定字段，取决于你查询了哪些列，但始终包含两个 score：
+
+```ts
+interface SearchResult {
+  [key: string]: unknown;  // 你 select 的列
+  embeddingScore: number;   // cosine similarity，范围 -1 ~ 1
+  rerankScore?: number;     // 仅 enableRerank 时存在
+}
+```
+
+### EmbeddingSearch.query 配置项
+
+| 配置项 | 类型 | 默认值 | 说明 |
+|---|---|---|---|
+| `tableName` | `string` | — | 检索表名（必填） |
+| `enableRerank` | `boolean` | `true` | 是否启用 rerank |
+| `embeddingTopK` | `number` | `10` | pgvector 检索返回条数（rerank 的候选集） |
+| `finalTopN` | `number` | `5` | 最终结果条数 |
+| `embeddingThreshold` | `number` | `0.8` | embedding 相似度阈值过滤（cosine similarity） |
+| `rerankThreshold` | `number` | `0.1` | rerank 分数阈值过滤 |
+| `embeddingColumn` | `string` | `'embedding'` | 向量列名 |
+| `rerankKey` | `string` | `'content'` | 启用 rerank 时，用于重排序的文本列名 |
+| `dimensions` | `number` | `1024` | embedding 输出维度（1 ~ 1024） |
+| `select` | `string[]` | `undefined` | 指定返回哪些列，不传则 `SELECT *` |
+| `filter` | `Record<string, unknown>` | — | 对表独立列做等值/范围过滤 |
+
+### 过滤
+
+支持对表的任意独立列做等值和范围过滤：
+
+```ts
+// 等值过滤
+const results = await EmbeddingSearch.query('query', {
+  filter: { category: '衣服', shop: '旗舰店' },
+}, pool);
+
+// 混合：等值 + 数值范围
+const ranked = await EmbeddingSearch.query('query', {
+  filter: {
+    category: '衣服',
+    price: { $gt: 100, $lt: 500 },
+    rating: { $gte: 4 },
+  },
+}, pool);
+
+// 数组：IN 查询
+const multi = await EmbeddingSearch.query('query', {
+  filter: {
+    category: ['衣服', '裤子', '鞋子'],
+  },
+}, pool);
+```
+
+### 生成 Embedding
+
+如果你只需要把文本转成向量，直接用 `embed`：
+
+```ts
+import { embed } from '@keo-ai/axiom';
+
+const vector = await embed('衣服质量怎么样');
+// vector: number[]，默认 1024 维
+
+// 指定维度（1 ~ 1024）
+const vector256 = await embed('衣服质量怎么样', 256);
+```
+
+| 参数 | 类型 | 默认值 | 说明 |
+|---|---|---|---|
+| `text` | `string` | — | 要转成向量的文本（必填） |
+| `dimensions` | `number` | `1024` | 输出维度（1 ~ 1024） |
+
+### 纯向量检索（已有向量）
+
+`EmbeddingSearch.query` 会自动把文本转成向量并做 rerank。如果你**已经有了向量**（比如自己生成的 embedding），只想做最原始的 pgvector 近邻搜索，用 `EmbeddingSearch.vectorSearch`：
+
+```ts
+import { EmbeddingSearch } from '@keo-ai/axiom';
+
+const results = await EmbeddingSearch.vectorSearch({
+  pool,
+  tableName: 'documents',
+  embeddingColumn: 'vec',
+  vector: queryVector,  // 你自己生成的向量
+  topK: 10,
+  threshold: 0.7,
+  filter: { category: '衣服' },
+});
+```
+
+> `EmbeddingSearch.vectorSearch` **不做 embedding，也不做 rerank**，只负责拿你给的向量去 pgvector 里搜近邻。
+
+| 参数 | 类型 | 默认值 | 说明 |
+|---|---|---|---|
+| `pool` | `Pool` | — | pg 连接池（必填） |
+| `tableName` | `string` | — | 检索表名（必填） |
+| `vector` | `number[]` | — | 查询向量（必填） |
+| `topK` | `number` | — | 返回条数（必填） |
+| `threshold` | `number` | — | 相似度阈值 |
+| `embeddingColumn` | `string` | `'embedding'` | 向量列名 |
+| `filter` | `Record<string, unknown>` | — | 对表独立列做过滤 |
+| `select` | `string[]` | `undefined` | 指定返回哪些列 |
+
+### 依赖
+
+EmbeddingSearch 模块内部调用百炼的 embedding 和 rerank 服务，无需额外配置：
+
+- **Embedding**：`text-embedding-v4`，1024 维，兼容 OpenAI 协议
+- **Rerank**：`qwen3-rerank`，百炼原生 API
+
+只需确保 `BAILIAN_API_KEY` 已设置。
+
+### 错误处理
+
+同 Predict 模块：**直接抛异常**。常见错误：
+
+- `BAILIAN_API_KEY environment variable is not set` — 未配置 API Key
+- `[embedding] HTTP 4xx/5xx` — embedding 接口异常
+- `[rerank] HTTP 4xx/5xx` — rerank 接口异常
+- `Invalid table name` — 表名包含非法字符
+
+---
+
 ## 设计理念
 
 ```

package/dist/index.d.ts

@@ -1,5 +1,31 @@
+/**
+ * Axiom — LLM 调用底座库。
+ *
+ * 使用前需确保环境变量已加载（推荐 dotenv）：
+ * ```ts
+ * import 'dotenv/config';
+ * import { LLM } from '@keo-ai/axiom';
+ *
+ * const res = await LLM.predict({ model: 'qwen3.7-max', prompt: '你好' });
+ * ```
+ *
+ * 必需环境变量：
+ * - `BAILIAN_API_KEY` — 百炼 API Key
+ *
+ * 可选环境变量：
+ * - `BAILIAN_BASE_URL` — 自定义百炼 API 地址
+ * - `BAILIAN_DEFAULT_MODEL` — 默认模型
+ *
+ * @packageDocumentation
+ */
+/** LLM 静态入口类，封装 Provider 连接、故障转移和返回解析 */
 export { LLM } from './predict';
 export type { LLMRequest, LLMResponse, Message, ProviderConfig, StreamChunk, } from './predict';
 export type { LLMProvider, PredictorOptions } from './predict';
-export { Predictor, BailianProvider, MODEL_REGISTRY } from './predict';
+/** 预测器核心类，负责模型路由与故障转移 */
+export { Predictor } from './predict';
+/** 百炼 Provider 实现 */
+export { BailianProvider } from './predict';
+/** 模型注册表，定义模型与 Provider 的映射关系 */
+export { MODEL_REGISTRY } from './predict';
 export type { Model, PredictConfig, PredictWithMessagesConfig } from './predict';

package/dist/index.js

@@ -1,9 +1,35 @@
 "use strict";
+/**
+ * Axiom — LLM 调用底座库。
+ *
+ * 使用前需确保环境变量已加载（推荐 dotenv）：
+ * ```ts
+ * import 'dotenv/config';
+ * import { LLM } from '@keo-ai/axiom';
+ *
+ * const res = await LLM.predict({ model: 'qwen3.7-max', prompt: '你好' });
+ * ```
+ *
+ * 必需环境变量：
+ * - `BAILIAN_API_KEY` — 百炼 API Key
+ *
+ * 可选环境变量：
+ * - `BAILIAN_BASE_URL` — 自定义百炼 API 地址
+ * - `BAILIAN_DEFAULT_MODEL` — 默认模型
+ *
+ * @packageDocumentation
+ */
 Object.defineProperty(exports, "__esModule", { value: true });
 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; } });
+/** 预测器核心类，负责模型路由与故障转移 */
 var predict_2 = require("./predict");
 Object.defineProperty(exports, "Predictor", { enumerable: true, get: function () { return predict_2.Predictor; } });
-Object.defineProperty(exports, "BailianProvider", { enumerable: true, get: function () { return predict_2.BailianProvider; } });
-Object.defineProperty(exports, "MODEL_REGISTRY", { enumerable: true, get: function () { return predict_2.MODEL_REGISTRY; } });
+/** 百炼 Provider 实现 */
+var predict_3 = require("./predict");
+Object.defineProperty(exports, "BailianProvider", { enumerable: true, get: function () { return predict_3.BailianProvider; } });
+/** 模型注册表，定义模型与 Provider 的映射关系 */
+var predict_4 = require("./predict");
+Object.defineProperty(exports, "MODEL_REGISTRY", { enumerable: true, get: function () { return predict_4.MODEL_REGISTRY; } });

package/dist/predict/config.d.ts

@@ -8,8 +8,6 @@ export interface PredictConfig {
     readonly model: Model;
     /** 用户输入的 prompt */
     readonly prompt: string;
-    /** 系统提示（可选） */
-    readonly systemPrompt?: string;
     readonly temperature?: number;
     readonly maxTokens?: number;
     readonly topP?: number;
@@ -21,7 +19,6 @@ export interface PredictConfig {
  */
 export interface PredictWithMessagesConfig {
     readonly model: Model;
-    readonly systemPrompt?: string;
     readonly temperature?: number;
     readonly maxTokens?: number;
     readonly topP?: number;
@@ -29,9 +26,9 @@ export interface PredictWithMessagesConfig {
     readonly responseFormat?: 'text' | 'json';
 }
 /**
- * 将 system prompt 前置到消息列表中。
- * @param config - 含可选 systemPrompt 字段的配置
+ * 构建消息列表。
+ * @param _config - 预测配置（保留以兼容接口，暂不消费）
  * @param userMessages - 用户提供的消息列表
- * @returns 带 system message 前缀的完整列表
+ * @returns 完整消息列表
  */
-export declare function buildMessages(config: PredictConfig | PredictWithMessagesConfig, userMessages: ReadonlyArray<Message>): Message[];
+export declare function buildMessages(_config: PredictConfig | PredictWithMessagesConfig, userMessages: ReadonlyArray<Message>): Message[];

package/dist/predict/config.js

@@ -2,16 +2,11 @@
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.buildMessages = buildMessages;
 /**
- * 将 system prompt 前置到消息列表中。
- * @param config - 含可选 systemPrompt 字段的配置
+ * 构建消息列表。
+ * @param _config - 预测配置（保留以兼容接口，暂不消费）
  * @param userMessages - 用户提供的消息列表
- * @returns 带 system message 前缀的完整列表
+ * @returns 完整消息列表
  */
-function buildMessages(config, userMessages) {
-    const messages = [];
-    if (config.systemPrompt) {
-        messages.push({ role: 'system', content: config.systemPrompt });
-    }
-    messages.push(...userMessages);
-    return messages;
+function buildMessages(_config, userMessages) {
+    return [...userMessages];
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@keo-ai/axiom",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "基于 LLM 的预测与推理库，支持多 Provider 切换",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -10,7 +10,7 @@
   "license": "MIT",
   "scripts": {
     "build": "tsc",
-    "test": "vitest run"
+    "test": "vitest run --exclude '**/*.integration.test.ts'"
   },
   "dependencies": {
     "dotenv": "^17.4.2"