1

agent/application/workshop/research.py

@@ -63,7 +63,7 @@ async def web_search(args: Args) -> str:
                 response = await client.post(
                     "https://api.bocha.cn/v1/web-search",
                     headers={
-                        "Authorization": f"Bearer sk-a9ed73d8b26745799cf970f84d70f512",
+                        "Authorization": "Bearer sk-a9ed73d8b26745799cf970f84d70f512",
                         "Content-Type": "application/json",
                     },
                     json={
@@ -98,7 +98,7 @@ async def web_search(args: Args) -> str:
                 response = await client.post(
                     "https://api.siliconflow.cn/v1/rerank",
                     headers={
-                        "Authorization": f"Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
+                        "Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
                         "Content-Type": "application/json",
                     },
                     json={

agent/document_preprocess/main.py

@@ -21,11 +21,17 @@ from requests import Response, exceptions, post
 from transformers import AutoTokenizer
 
 """
+文档同步记录表建表SQL
+CREATE TABLE IF NOT EXISTS document_sync_records (
+    id TEXT PRIMARY KEY,
+    modified_at FLOAT NOT NULL
+);
+
 文档分片表建表SQL
 CREATE EXTENSION IF NOT EXISTS vector;
-CREATE EXTENSION IF NOT EXISTS pg_trgm;
+CREATE EXTENSION IF NOT EXISTS pg_search;
 
-CREATE TABLE document_chunks (
+CREATE TABLE IF NOT EXISTS document_chunks (
     id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
     document_id TEXT NOT NULL,
     index INT NOT NULL,
@@ -36,9 +42,18 @@ CREATE TABLE document_chunks (
     return_text TEXT
 );
 
-CREATE INDEX idx_document_chunks_text ON document_chunks USING bm25 (text);
+CREATE INDEX idx_document_chunks_text
-CREATE INDEX idx_document_chunks_text_vector ON document_chunks USING hnsw (text_vector vector_cosine_ops);
+ON document_chunks
-CREATE INDEX idx_document_chunks_text_fingerprint ON document_chunks(text_fingerprint);
+USING paradedb (
+    id,
+    (text::pdb.jieba)
+);
+
+CREATE INDEX idx_document_chunks_text_vector
+ON document_chunks USING hnsw (text_vector vector_cosine_ops);
+
+CREATE INDEX idx_document_chunks_text_fingerprint
+ON document_chunks(text_fingerprint);
 """
 
 
@@ -61,13 +76,72 @@ REMOTE_DIRECTORY_PATH = (
 )
 LOCAL_DIRECTORY_PATH = Path(__file__).parent / "vault"  # 本地库文件夹路径
 
+DATABASE_CONNECTION_CONFIG = {
+    "host": "www.liubiren.cloud",
+    "port": 5201,
+    "user": "root",
+    "password": "PG198752",
+    "dbname": "database",
+}
+
 
 CHUNK_SIZE = 8192  # 分片大小，BAAI/BGE-M3标值
 
-TPM = 500_000 * 0.8  # 每分钟词元数限额，硅基流动标值，安全冗余20%
+TPM = 500_000 * 0.8  # 每分钟词元数限额，硅基流动标值，安全冗余
 WINDOW_PERIOD = 60  # 窗口期
 
 
+def get_text_vectors(input: list[Chunk] | str) -> list[list[float]]:
+    """
+    获取文本嵌入向量
+    :param input: 输入
+    :return: 文本嵌入向量列表
+    """
+    request_count = 0  # 请求次数
+    response_payload: dict | None = None
+    while request_count < 3:
+        request_count += 1
+        response: Response | None = None
+        try:
+            response = post(
+                "https://api.siliconflow.cn/v1/embeddings",
+                headers={
+                    "Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
+                    "Content-Type": "application/json",
+                },
+                json={
+                    "model": "BAAI/bge-m3",
+                    "input": (
+                        [chunk.text for chunk in input]
+                        if isinstance(input, list)
+                        else [input]
+                    ),
+                },
+                timeout=60,
+            )
+            response.raise_for_status()
+            response_payload = response.json()
+            break
+
+        except exceptions.HTTPError as exception:
+            # 触发限速
+            if response is not None and response.status_code == 429:
+                sleep(2**request_count)
+                continue
+
+            raise RuntimeError(
+                f"HTTPError: {exception}, resp={response.text if response is not None else 'response is None'}"
+            )
+
+        except Exception as exception:
+            print(f"Exception: {exception}")
+            sleep(2**request_count)
+
+    if response_payload is None:
+        raise RuntimeError("response_payload is None")
+    return [item["embedding"] for item in response_payload["data"]]
+
+
 def sync_to_local_vault() -> None:
     """
     同步Obsidian云端库至本地库
@@ -212,13 +286,7 @@ def process_documents() -> None:
             window_deque.popleft()
         return sum(text_tokens for _, text_tokens in window_deque)
 
-    with connect(
+    with connect(**DATABASE_CONNECTION_CONFIG) as connection:
-        host="www.liubiren.cloud",
-        port=5201,
-        user="root",
-        password="PG198752",
-        dbname="database",
-    ) as connection:
         print("📥 正在加载文档同步记录...")
         document_sync_records: dict[str, float] = {}  # 文档同步记录字典
         with connection.cursor() as cursor:
@@ -383,51 +451,14 @@ def process_documents() -> None:
                                 max((window_deque[0][0] + WINDOW_PERIOD) - time(), 0.2)
                             )
 
-                        request_count = 0  # 请求次数
+                        # 获取文本嵌入向量
-                        response_payload: dict | None = None
+                        text_vectors = get_text_vectors(batch)
-                        while request_count < 3:
+                        # 请求成功则将本次请求添加至窗口队列
-                            request_count += 1
+                        window_deque.append((time(), batch_text_tokens))
-                            response: Response | None = None
-                            try:
-                                response = post(
-                                    "https://api.siliconflow.cn/v1/embeddings",
-                                    headers={
-                                        "Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
-                                        "Content-Type": "application/json",
-                                    },
-                                    json={
-                                        "model": "BAAI/bge-m3",
-                                        "input": [chunk.text for chunk in batch],
-                                    },
-                                    timeout=60,
-                                )
-                                response.raise_for_status()
-                                response_payload = response.json()
-                                # 请求成功则将本次请求添加至窗口队列
-                                window_deque.append((time(), batch_text_tokens))
-                                break
-
-                            except exceptions.HTTPError as exception:
-                                # 触发限速
-                                if response is not None and response.status_code == 429:
-                                    sleep(2**request_count)
-                                    continue
-
-                                raise RuntimeError(
-                                    f"HTTPError: {exception}, resp={response.text if response is not None else 'response is None'}"
-                                )
-
-                            except Exception as exception:
-                                print(f"Exception: {exception}")
-                                sleep(2**request_count)
-
-                        if response_payload is None:
-                            raise RuntimeError("response_payload is None")
-
                         # 写入文本向量
                         for chunk, text_vector in zip(
                             batch,
-                            [item["embedding"] for item in response_payload["data"]],
+                            text_vectors,
                         ):
                             chunk.text_vector = text_vector
 
@@ -491,8 +522,158 @@ def process_documents() -> None:
     print("===== 处理文档完成 =====\n")
 
 
+def hybrid_retrieve(
+    text: str,
+    keywords: list[str] | None = None,
+) -> list[dict]:
+    """
+    混合检索
+    :param text: 检索文本
+    :param keywords: 检索关键词列表
+    :return: 检索结果
+    """
+    candidates: list[dict] = []
+    with connect(**DATABASE_CONNECTION_CONFIG) as connection:
+        with connection.cursor() as cursor:
+            cursor.execute(
+                """
+                WITH bm25_retrieve_results AS (
+                    SELECT
+                        id,
+                        document_id,
+                        index,
+                        text,
+                        return_text,
+                        text_fingerprint,
+                        ROW_NUMBER() OVER (ORDER BY paradedb.score(id) DESC) AS bm25_retrieve_rank
+                    FROM document_chunks
+                    WHERE text @@@ %s
+                    ORDER BY paradedb.score(id) DESC
+                    LIMIT 50
+                ),
+                vector_retrieve_results AS (
+                    SELECT
+                        id,
+                        document_id,
+                        index,
+                        text,
+                        return_text,
+                        text_fingerprint,
+                        ROW_NUMBER() OVER (ORDER BY text_vector <-> %s::vector(1024)) AS vector_retrieve_rank
+                    FROM document_chunks
+                    ORDER BY text_vector <-> %s::vector(1024)
+                    LIMIT 50
+                ),
+                merged_results AS (
+                    SELECT
+                        COALESCE(bm25_retrieve_results.id, vector_retrieve_results.id) AS id,
+                        COALESCE(bm25_retrieve_results.document_id, vector_retrieve_results.document_id) AS document_id,
+                        COALESCE(bm25_retrieve_results.index, vector_retrieve_results.index) AS index,
+                        COALESCE(bm25_retrieve_results.text, vector_retrieve_results.text) AS text,
+                        COALESCE(bm25_retrieve_results.return_text, vector_retrieve_results.return_text) AS return_text,
+                        COALESCE(bm25_retrieve_results.text_fingerprint, vector_retrieve_results.text_fingerprint) AS text_fingerprint,
+                        COALESCE(1.0/(bm25_retrieve_rank + 60), 0) + COALESCE(1.0/(vector_retrieve_rank + 60), 0) AS rank_score
+                    FROM bm25_retrieve_results
+                    FULL OUTER JOIN vector_retrieve_results USING(id)
+                ),
+                deduplication AS (
+                    SELECT
+                        *,
+                        ROW_NUMBER() OVER (PARTITION BY text_fingerprint ORDER BY rank_score DESC) AS duplicate_index
+                    FROM merged_results
+                )
+                SELECT
+                    id,
+                    document_id,
+                    index,
+                    text,
+                    return_text,
+                    rank_score
+                FROM deduplication
+                WHERE duplicate_index = 1
+                ORDER BY rank_score DESC
+                LIMIT 100;
+                """,
+                (
+                    " ".join(keywords) if keywords is not None else text,
+                    text_vector := get_text_vectors(text)[0],  # 获取文本嵌入向量
+                    text_vector,
+                ),
+            )
+            rows = cursor.fetchall()
+            if cursor.description is None:
+                return candidates
+            candidates = [
+                dict(
+                    zip(
+                        ([metadata[0] for metadata in cursor.description]),
+                        row,
+                    )
+                )
+                for row in rows
+            ]
+
+    try:
+        request_count = 0  # 请求次数
+        response_payload: dict | None = None
+        while request_count < 3:
+            request_count += 1
+            response: Response | None = None
+            try:
+                response = post(
+                    "https://api.siliconflow.cn/v1/rerank",
+                    headers={
+                        "Authorization": "Bearer sk-xsnuwirjjphhfdbvznfdfjqlinfdlrnlxuhkbbqynfnbhiqz",
+                        "Content-Type": "application/json",
+                    },
+                    json={
+                        "model": "BAAI/bge-reranker-v2-m3",
+                        "query": text,
+                        "documents": [item["return_text"] for item in candidates],
+                        "return_documents": False,  # 不返回文档内容
+                        "top_n": 10,  # 返回结果数
+                    },
+                    timeout=60,
+                )
+                response.raise_for_status()
+                response_payload = response.json()
+                break
+
+            except exceptions.HTTPError as exception:
+                # 触发限速
+                if response is not None and response.status_code == 429:
+                    sleep(2**request_count)
+                    continue
+
+                raise RuntimeError(
+                    f"HTTPError: {exception}, resp={response.text if response is not None else 'response is None'}"
+                )
+
+            except Exception as exception:
+                print(f"Exception: {exception}")
+                sleep(2**request_count)
+
+        if response_payload is None:
+            raise RuntimeError("response_payload is None")
+
+    # 若发生异常则返回候选前10个分片
+    except Exception:
+        return candidates[:10]
+
+    return [
+        candidates[result["index"]]
+        for result in response_payload["results"]
+        if result["relevance_score"] >= 0.85
+    ]
+
+
 if __name__ == "__main__":
     # 同步Obsidian云端库至本地库
     sync_to_local_vault()
     # 处理文档
     process_documents()
+    """
+    results = hybrid_retrieve("如何实现事实校验")
+    for result in results:
+        print(result["return_text"])
+    """