liubiren
/
Python
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Activity

1

This commit is contained in:
liubiren 2026-10-08 14:52:18 +08:00
parent 18d8489cdc
commit 99f7c93dac
2 changed files with 238 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
agent/application/workshop/research.py
View File

@ -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={

291
agent/document_preprocess/main.py
View File

@ -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"])
"""