From 99f7c93dac7c39b8625fb8d534989e829d6246d0 Mon Sep 17 00:00:00 2001 From: liubiren Date: Thu, 8 Oct 2026 14:52:18 +0800 Subject: [PATCH] 1 --- agent/application/workshop/research.py | 4 +- agent/document_preprocess/main.py | 291 ++++++++++++++++++++----- 2 files changed, 238 insertions(+), 57 deletions(-) diff --git a/agent/application/workshop/research.py b/agent/application/workshop/research.py index 52d86c7..cb7a24f 100644 --- a/agent/application/workshop/research.py +++ b/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={ diff --git a/agent/document_preprocess/main.py b/agent/document_preprocess/main.py index 1de4327..3c45305 100644 --- a/agent/document_preprocess/main.py +++ b/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_vector ON document_chunks USING hnsw (text_vector vector_cosine_ops); -CREATE INDEX idx_document_chunks_text_fingerprint ON document_chunks(text_fingerprint); +CREATE INDEX idx_document_chunks_text +ON document_chunks +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( - host="www.liubiren.cloud", - port=5201, - user="root", - password="PG198752", - dbname="database", - ) as connection: + with connect(**DATABASE_CONNECTION_CONFIG) 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 - 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 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") - + # 获取文本嵌入向量 + text_vectors = get_text_vectors(batch) + # 请求成功则将本次请求添加至窗口队列 + window_deque.append((time(), batch_text_tokens)) # 写入文本向量 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"]) + """