ForcePilot/backend/package/yuxi/knowledge/implementations/notion.py
Wenjie Zhang 99c590d003 refactor: 知识库系统重构 - 统一命名、优化图检索与向量存储逻辑
- base/manager: db_id → kb_id 命名统一,重构元数据加载与文件记录逻辑
- implementations: dify/milvus/notion/read_only 适配新命名与接口
- graphs: milvus_graph_service/vector_store/graph_utils 优化图构建与检索逻辑
- eval: benchmark_generation/evaluator/service 适配新命名并优化评估流程
- utils/kb_utils: 对应更新工具函数
2026-05-26 17:40:17 +08:00

745 lines
30 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import asyncio
import os
import time
import traceback
from typing import Any
import httpx
from yuxi.knowledge.implementations.read_only_connectors import ReadOnlyConnectors
from yuxi.utils import logger
NOTION_API_BASE = "https://api.notion.com/v1"
NOTION_DEFAULT_VERSION = "2026-03-11"
NOTION_RETRY_STATUSES = {429, 500, 502, 503, 504}
NOTION_MAX_BLOCKS = 2000
NOTION_MAX_DEPTH = 8
NOTION_DEFAULT_MAX_HYDRATE_PAGES = 20
NOTION_PAGE_CACHE_TTL_SECONDS = 300
NOTION_PAGE_CACHE_MAX_SIZE = 64
class NotionAPIError(RuntimeError):
pass
class _NotionClient:
def __init__(self, token: str, notion_version: str) -> None:
self.token = token
self.notion_version = notion_version or NOTION_DEFAULT_VERSION
self._client: httpx.AsyncClient | None = None
async def __aenter__(self):
self._client = httpx.AsyncClient(timeout=45.0)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
if self._client:
await self._client.aclose()
self._client = None
return False
async def request(
self,
method: str,
path: str,
*,
json: dict[str, Any] | None = None,
params: dict[str, Any] | None = None,
retries: int = 4,
) -> dict[str, Any]:
headers = {
"Authorization": f"Bearer {self.token}",
"Notion-Version": self.notion_version,
"Content-Type": "application/json",
"User-Agent": "yuxi-notion-kb/1.0",
}
url = f"{NOTION_API_BASE}{path}"
for attempt in range(retries + 1):
try:
response = await self._request(method, url, json=json, params=params, headers=headers)
if response.status_code in NOTION_RETRY_STATUSES and attempt < retries:
await asyncio.sleep(self._retry_delay(response, attempt))
continue
try:
response.raise_for_status()
except httpx.HTTPStatusError as exc:
body_preview = response.text[:1000] if response.text else ""
logger.error(f"Notion HTTP error: status={response.status_code}, path={path}, body={body_preview}")
raise NotionAPIError(str(exc)) from exc
return response.json()
except (httpx.HTTPError, NotionAPIError) as exc:
if isinstance(exc, NotionAPIError) or attempt >= retries:
raise
await asyncio.sleep(min(2**attempt, 8))
raise NotionAPIError(f"Notion request failed: {method} {path}")
async def _request(self, method: str, url: str, **kwargs) -> httpx.Response:
if self._client:
return await self._client.request(method, url, **kwargs)
async with httpx.AsyncClient(timeout=45.0) as client:
return await client.request(method, url, **kwargs)
@staticmethod
def _retry_delay(response: httpx.Response, attempt: int) -> float:
retry_after = response.headers.get("Retry-After")
if retry_after:
try:
return float(retry_after)
except ValueError:
pass
return float(min(2**attempt, 8))
async def paginate(
self,
method: str,
path: str,
*,
body: dict[str, Any] | None = None,
params: dict[str, Any] | None = None,
limit: int | None = None,
) -> list[dict[str, Any]]:
results: list[dict[str, Any]] = []
cursor: str | None = None
method_upper = method.upper()
while True:
request_body = dict(body or {})
request_params = dict(params or {})
if cursor:
if method_upper == "GET":
request_params["start_cursor"] = cursor
else:
request_body["start_cursor"] = cursor
payload = await self.request(method_upper, path, json=request_body or None, params=request_params or None)
page_results = payload.get("results", []) if isinstance(payload, dict) else []
if isinstance(page_results, list):
results.extend(item for item in page_results if isinstance(item, dict))
if limit is not None and len(results) >= limit:
return results[:limit]
if not payload.get("has_more"):
return results
cursor = payload.get("next_cursor")
if not cursor:
return results
async def search_pages(self, query_text: str, limit: int) -> list[dict[str, Any]]:
return await self.paginate(
"POST",
"/search",
body={
"query": query_text,
"filter": {"property": "object", "value": "page"},
"sort": {"direction": "descending", "timestamp": "last_edited_time"},
"page_size": min(max(limit, 1), 100),
},
limit=limit,
)
async def query_data_source(self, data_source_id: str, limit: int) -> list[dict[str, Any]]:
return await self.paginate(
"POST",
f"/data_sources/{data_source_id}/query",
body={
"page_size": min(max(limit, 1), 100),
"result_type": "page",
"sorts": [{"timestamp": "last_edited_time", "direction": "descending"}],
},
limit=limit,
)
async def retrieve_page(self, page_id: str) -> dict[str, Any]:
return await self.request("GET", f"/pages/{page_id}")
async def retrieve_block_children(self, block_id: str, limit: int) -> list[dict[str, Any]]:
return await self.paginate(
"GET",
f"/blocks/{block_id}/children",
params={"page_size": min(max(limit, 1), 100)},
limit=limit,
)
class NotionKB(ReadOnlyConnectors):
"""连接 Notion Data Source 的只读知识库实现"""
kb_type = "notion"
name = "Notion"
description = "连接 Notion Data Source 的只读知识库,支持检索、打开页面和页内查找"
def __init__(self, work_dir: str, **kwargs):
del kwargs
super().__init__(work_dir)
self._page_markdown_cache: dict[tuple[str, str, str, str], tuple[float, str]] = {}
@classmethod
def get_create_params_config(cls) -> dict[str, Any]:
return {
"options": [
{
"key": "notion_token",
"label": "Notion Token",
"type": "password",
"required": False,
"placeholder": "留空则使用 NOTION_TOKEN / NOTION_API_KEY",
"description": "Notion integration token需要 read content 权限",
},
{
"key": "notion_data_source_id",
"label": "Data Source ID",
"type": "text",
"required": True,
"placeholder": "请输入 Notion data_source_id",
},
{
"key": "notion_version",
"label": "Notion API Version",
"type": "text",
"required": False,
"default": NOTION_DEFAULT_VERSION,
"placeholder": NOTION_DEFAULT_VERSION,
},
]
}
@classmethod
def validate_additional_params(cls, additional_params: dict | None) -> dict:
params = dict(additional_params or {})
token = str(params.get("notion_token") or "").strip()
data_source_id = str(params.get("notion_data_source_id") or "").strip()
notion_version = str(params.get("notion_version") or NOTION_DEFAULT_VERSION).strip() or NOTION_DEFAULT_VERSION
if not data_source_id:
raise ValueError("Notion 参数缺失: notion_data_source_id")
if not token and not (os.getenv("NOTION_TOKEN") or os.getenv("NOTION_API_KEY")):
raise ValueError("Notion 参数缺失: notion_token 或环境变量 NOTION_TOKEN/NOTION_API_KEY")
return {
"notion_token": token,
"notion_data_source_id": data_source_id,
"notion_version": notion_version,
}
async def aquery(self, query_text: str, kb_id: str, agent_call: bool = False, **kwargs) -> list[dict]:
del agent_call
try:
token, data_source_id, notion_version = self._get_connection_config(kb_id)
query_params = self._get_query_params(kb_id)
merged = {**query_params, **kwargs}
final_top_k = min(max(int(merged.get("final_top_k", 10)), 1), 50)
max_scan_pages = min(max(int(merged.get("max_scan_pages", 100)), 10), 1000)
max_hydrate_pages = min(
max(int(merged.get("max_hydrate_pages", NOTION_DEFAULT_MAX_HYDRATE_PAGES)), final_top_k),
100,
)
snippet_window_lines = min(max(int(merged.get("snippet_window_lines", 12)), 3), 40)
search_mode = str(merged.get("search_mode", "hybrid")).lower()
if search_mode not in {"notion_search", "data_source_scan", "hybrid"}:
search_mode = "hybrid"
async with _NotionClient(token, notion_version) as client:
candidate_pages = await self._search_candidate_pages(
client,
query_text,
data_source_id,
search_mode=search_mode,
final_top_k=final_top_k,
max_scan_pages=max_scan_pages,
)
candidate_pages = self._rank_candidate_pages(candidate_pages, query_text)[:max_hydrate_pages]
result_tasks = [
self._build_search_result(
client,
page,
query_text,
rank=rank,
source=source,
snippet_window_lines=snippet_window_lines,
)
for rank, (page, source) in enumerate(candidate_pages, start=1)
if page.get("id")
]
results = [result for result in await asyncio.gather(*result_tasks) if result]
return sorted(results, key=lambda item: item.get("score", 0.0), reverse=True)[:final_top_k]
except (NotionAPIError, httpx.HTTPError, ValueError) as exc:
logger.error(f"Notion query failed for kb_id={kb_id}: {exc}, {traceback.format_exc()}")
return []
async def _search_candidate_pages(
self,
client: _NotionClient,
query_text: str,
data_source_id: str,
*,
search_mode: str,
final_top_k: int,
max_scan_pages: int,
) -> list[tuple[dict[str, Any], str]]:
candidates: list[tuple[dict[str, Any], str]] = []
seen: set[str] = set()
if search_mode in {"notion_search", "hybrid"}:
search_limit = min(max(final_top_k * 5, 20), 100)
for page in await client.search_pages(query_text, search_limit):
if not self._page_belongs_to_data_source(page, data_source_id):
continue
page_id = str(page.get("id") or "")
if page_id and page_id not in seen:
seen.add(page_id)
candidates.append((page, "notion_search"))
if search_mode in {"data_source_scan", "hybrid"} and len(candidates) < final_top_k:
for page in await client.query_data_source(data_source_id, max_scan_pages):
page_id = str(page.get("id") or "")
if page_id and page_id not in seen:
seen.add(page_id)
candidates.append((page, "data_source_scan"))
return candidates
@classmethod
def _rank_candidate_pages(
cls,
candidate_pages: list[tuple[dict[str, Any], str]],
query_text: str,
) -> list[tuple[dict[str, Any], str]]:
terms = cls._query_terms(query_text)
def score_candidate(candidate: tuple[dict[str, Any], str]) -> float:
page, source = candidate
title = cls._extract_page_title(page).lower()
properties = "\n".join(cls._extract_property_texts(page).values()).lower()
score = 2.0 if source == "notion_search" else 0.0
score += sum(5.0 for term in terms if term in title)
score += sum(2.0 for term in terms if term in properties)
return score
return sorted(candidate_pages, key=score_candidate, reverse=True)
async def _build_search_result(
self,
client: _NotionClient,
page: dict[str, Any],
query_text: str,
*,
rank: int,
source: str,
snippet_window_lines: int,
) -> dict[str, Any] | None:
page_id = str(page.get("id") or "")
if not page_id:
return None
markdown = await self._page_to_markdown(client, page_id, page)
score, snippet, match_source, line_start, line_end = self._score_page(
query_text,
markdown["title"],
markdown["properties"],
markdown["content"],
snippet_window_lines=snippet_window_lines,
)
if score <= 0 and source != "notion_search":
return None
if score <= 0:
score = 0.1
snippet = self._preview(markdown["content"], snippet_window_lines)
match_source = "notion_search"
line_start = 1
line_end = min(len(markdown["content"].splitlines()), snippet_window_lines)
return {
"content": snippet,
"score": score + 1.0 / (rank + 1000),
"metadata": {
"source": markdown["title"],
"file_id": page_id,
"chunk_id": f"{page_id}:{line_start}",
"chunk_index": rank,
"notion_url": page.get("url"),
"created_time": page.get("created_time"),
"last_edited_time": page.get("last_edited_time"),
"line_start": line_start,
"line_end": line_end,
"match_source": match_source,
},
}
async def open_file_content(self, kb_id: str, file_id: str, offset: int = 0, limit: int = 800) -> dict:
content = await self._read_page_markdown(kb_id, file_id)
return self._build_open_file_window(content, offset=offset, limit=limit)
async def find_file_content(
self,
kb_id: str,
file_id: str,
patterns: list[str],
*,
use_regex: bool = False,
case_sensitive: bool = False,
max_windows: int = 5,
window_size: int = 80,
) -> dict:
content = await self._read_page_markdown(kb_id, file_id)
return self._build_find_file_windows(
content,
patterns=patterns,
use_regex=use_regex,
case_sensitive=case_sensitive,
max_windows=max_windows,
window_size=window_size,
)
async def _read_page_markdown(self, kb_id: str, page_id: str) -> str:
token, data_source_id, notion_version = self._get_connection_config(kb_id)
cache_key = (kb_id, page_id, data_source_id, notion_version)
cached = self._get_cached_page_markdown(cache_key)
if cached is not None:
return cached
async with _NotionClient(token, notion_version) as client:
page = await client.retrieve_page(page_id)
if not self._page_belongs_to_data_source(page, data_source_id):
raise ValueError(f"Notion 页面 {page_id} 不属于当前 Data Source")
markdown = await self._page_to_markdown(client, page_id, page)
content = markdown["content"]
self._set_cached_page_markdown(cache_key, content)
return content
def _get_cached_page_markdown(self, cache_key: tuple[str, str, str, str]) -> str | None:
cached = self._page_markdown_cache.get(cache_key)
if not cached:
return None
cached_at, content = cached
if time.monotonic() - cached_at > NOTION_PAGE_CACHE_TTL_SECONDS:
self._page_markdown_cache.pop(cache_key, None)
return None
return content
def _set_cached_page_markdown(self, cache_key: tuple[str, str, str, str], content: str) -> None:
if len(self._page_markdown_cache) >= NOTION_PAGE_CACHE_MAX_SIZE:
oldest_key = min(self._page_markdown_cache, key=lambda key: self._page_markdown_cache[key][0])
self._page_markdown_cache.pop(oldest_key, None)
self._page_markdown_cache[cache_key] = (time.monotonic(), content)
def _get_connection_config(self, kb_id: str) -> tuple[str, str, str]:
metadata = self.databases_meta.get(kb_id, {}).get("metadata", {}) or {}
token = str(
metadata.get("notion_token") or os.getenv("NOTION_TOKEN") or os.getenv("NOTION_API_KEY") or ""
).strip()
data_source_id = str(metadata.get("notion_data_source_id") or "").strip()
notion_version = str(metadata.get("notion_version") or NOTION_DEFAULT_VERSION).strip() or NOTION_DEFAULT_VERSION
if not token or not data_source_id:
raise ValueError(f"Notion config incomplete for kb_id={kb_id}")
return token, data_source_id, notion_version
async def _page_to_markdown(
self,
client: _NotionClient,
page_id: str,
page: dict[str, Any] | None = None,
) -> dict[str, Any]:
page = page or await client.retrieve_page(page_id)
title = self._extract_page_title(page)
properties = self._extract_property_texts(page)
lines = [f"# {title}", ""]
for name, value in properties.items():
if value and name != title:
lines.append(f"- {name}: {value}")
if len(lines) > 2:
lines.append("")
block_state = {"count": 0, "truncated": False}
lines.extend(await self._blocks_to_markdown(client, page_id, depth=0, state=block_state))
if block_state["truncated"]:
lines.append("")
lines.append("[内容过长,已截断]")
return {"title": title, "properties": properties, "content": "\n".join(lines).strip()}
async def _blocks_to_markdown(
self,
client: _NotionClient,
block_id: str,
*,
depth: int,
state: dict[str, Any],
) -> list[str]:
if depth > NOTION_MAX_DEPTH or state["count"] >= NOTION_MAX_BLOCKS:
state["truncated"] = True
return []
remaining = NOTION_MAX_BLOCKS - state["count"]
blocks = await client.retrieve_block_children(block_id, remaining)
lines: list[str] = []
for block in blocks:
if state["count"] >= NOTION_MAX_BLOCKS:
state["truncated"] = True
break
state["count"] += 1
lines.extend(self._block_to_lines(block, depth))
if block.get("has_children"):
lines.extend(await self._blocks_to_markdown(client, block["id"], depth=depth + 1, state=state))
return lines
def _block_to_lines(self, block: dict[str, Any], depth: int) -> list[str]:
block_type = str(block.get("type") or "")
value = block.get(block_type) if isinstance(block.get(block_type), dict) else {}
text = self._rich_text_plain_text(value.get("rich_text") or [])
indent = " " * depth
if block_type == "paragraph":
return [f"{indent}{text}".rstrip()] if text else []
if block_type in {"heading_1", "heading_2", "heading_3"}:
level = {"heading_1": "#", "heading_2": "##", "heading_3": "###"}[block_type]
return [f"{level} {text}".rstrip()] if text else []
if block_type == "bulleted_list_item":
return [f"{indent}- {text}".rstrip()] if text else []
if block_type == "numbered_list_item":
return [f"{indent}1. {text}".rstrip()] if text else []
if block_type == "to_do":
marker = "x" if value.get("checked") else " "
return [f"{indent}- [{marker}] {text}".rstrip()] if text else []
if block_type in {"toggle", "quote", "callout"}:
return [f"{indent}> {text}".rstrip()] if text else []
if block_type == "code":
language = value.get("language") or ""
return [f"```{language}", text, "```"] if text else []
if block_type == "child_page":
title = value.get("title") or "Untitled"
return [f"{indent}## {title}"]
if block_type == "table_row":
cells = [self._rich_text_plain_text(cell) for cell in value.get("cells") or []]
return [" | ".join(cells)] if cells else []
if block_type in {"bookmark", "embed", "image", "file", "pdf", "video"}:
url = self._file_or_url(value)
caption = self._rich_text_plain_text(value.get("caption") or [])
label = caption or url
return [f"{indent}[{block_type}] {label}".rstrip()] if label else []
if text:
return [f"{indent}{text}".rstrip()]
return [f"{indent}[unsupported block: {block_type}]"] if block_type else []
@staticmethod
def _file_or_url(value: dict[str, Any]) -> str:
if value.get("url"):
return str(value["url"])
for key in ("external", "file"):
nested = value.get(key)
if isinstance(nested, dict) and nested.get("url"):
return str(nested["url"])
return ""
@classmethod
def _extract_page_title(cls, page: dict[str, Any]) -> str:
for prop in (page.get("properties") or {}).values():
if isinstance(prop, dict) and prop.get("type") == "title":
title = cls._rich_text_plain_text(prop.get("title") or [])
if title:
return title
return "Untitled"
@classmethod
def _extract_property_texts(cls, page: dict[str, Any]) -> dict[str, str]:
properties: dict[str, str] = {}
for name, prop in (page.get("properties") or {}).items():
if not isinstance(prop, dict):
continue
text = cls._property_to_text(prop)
if text:
properties[str(name)] = text
return properties
@classmethod
def _property_to_text(cls, prop: dict[str, Any]) -> str:
prop_type = prop.get("type")
if prop_type in {"title", "rich_text"}:
return cls._rich_text_plain_text(prop.get(prop_type) or [])
if prop_type == "select":
value = prop.get("select") or {}
return str(value.get("name") or "")
if prop_type == "multi_select":
return ", ".join(str(item.get("name") or "") for item in prop.get("multi_select") or [] if item.get("name"))
if prop_type == "status":
value = prop.get("status") or {}
return str(value.get("name") or "")
if prop_type == "date":
value = prop.get("date") or {}
return " - ".join(str(value.get(key) or "") for key in ("start", "end") if value.get(key))
if prop_type in {"url", "email", "phone_number", "number", "checkbox", "created_time", "last_edited_time"}:
value = prop.get(prop_type)
return "" if value is None else str(value)
if prop_type == "people":
return ", ".join(str(item.get("name") or item.get("id") or "") for item in prop.get("people") or [])
if prop_type == "files":
return ", ".join(str(item.get("name") or "") for item in prop.get("files") or [] if item.get("name"))
if prop_type == "formula":
formula = prop.get("formula") or {}
return cls._property_to_text(
{"type": formula.get("type"), formula.get("type"): formula.get(formula.get("type"))}
)
return ""
@staticmethod
def _rich_text_plain_text(items: list[dict[str, Any]]) -> str:
return "".join(str(item.get("plain_text") or item.get("text", {}).get("content") or "") for item in items)
@classmethod
def _page_belongs_to_data_source(
cls,
page: dict[str, Any],
data_source_id: str,
*,
allow_unknown_parent: bool = False,
) -> bool:
parent = page.get("parent") or {}
normalized_data_source_id = cls._normalize_notion_id(data_source_id)
data_source_parent_ids = [parent.get("data_source_id")]
parent_type = parent.get("type")
if parent_type == "data_source_id":
data_source_parent_ids.append(parent.get(parent_type))
known_data_source_ids = [candidate_id for candidate_id in data_source_parent_ids if candidate_id]
if known_data_source_ids:
return any(
cls._normalize_notion_id(str(candidate_id)) == normalized_data_source_id
for candidate_id in known_data_source_ids
)
database_id = parent.get("database_id") or (parent.get(parent_type) if parent_type == "database_id" else None)
if database_id and cls._normalize_notion_id(str(database_id)) == normalized_data_source_id:
return True
return allow_unknown_parent
@staticmethod
def _normalize_notion_id(value: str) -> str:
return value.replace("-", "").lower()
@classmethod
def _score_page(
cls,
query_text: str,
title: str,
properties: dict[str, str],
content: str,
*,
snippet_window_lines: int,
) -> tuple[float, str, str, int, int]:
terms = cls._query_terms(query_text)
title_lower = title.lower()
property_text = "\n".join(properties.values())
property_lower = property_text.lower()
lines = content.splitlines()
score = 0.0
match_source = "content"
if any(term in title_lower for term in terms):
score += 5.0
match_source = "title"
if any(term in property_lower for term in terms):
score += 3.0
if match_source == "content":
match_source = "properties"
matched_indexes = []
for index, line in enumerate(lines):
lower_line = line.lower()
matches = sum(1 for term in terms if term in lower_line)
if matches:
matched_indexes.append(index)
score += float(matches)
if not matched_indexes:
return score, "", match_source, 1, 0
first_match = matched_indexes[0]
half_window = snippet_window_lines // 2
start = max(first_match - half_window, 0)
end = min(start + snippet_window_lines, len(lines))
start = max(end - snippet_window_lines, 0)
snippet = "\n".join(lines[start:end])
return score, snippet, match_source, start + 1, end
@staticmethod
def _query_terms(query_text: str) -> list[str]:
full_query = query_text.strip().lower()
terms = [full_query] if full_query else []
terms.extend(part.lower() for part in full_query.split() if part)
return list(dict.fromkeys(terms))
@staticmethod
def _preview(content: str, line_count: int) -> str:
return "\n".join(content.splitlines()[:line_count])
def get_query_params_config(self, kb_id: str, **kwargs) -> dict:
del kb_id, kwargs
return {
"type": "notion",
"options": [
{
"key": "search_mode",
"label": "检索模式",
"type": "select",
"default": "hybrid",
"options": [
{
"value": "hybrid",
"label": "混合检索",
"description": "先 Notion 搜索,不足时扫描 Data Source",
},
{"value": "notion_search", "label": "Notion 搜索", "description": "使用 Notion search API"},
{
"value": "data_source_scan",
"label": "Data Source 扫描",
"description": "扫描 Data Source 并本地匹配",
},
],
"description": "选择 Notion 检索方式",
},
{
"key": "final_top_k",
"label": "最终返回数量",
"type": "number",
"default": 10,
"min": 1,
"max": 50,
"description": "返回给前端和智能体的页面数量",
},
{
"key": "max_scan_pages",
"label": "最大扫描页面数",
"type": "number",
"default": 100,
"min": 10,
"max": 1000,
"description": "Data Source 扫描模式最多读取的页面数量",
},
{
"key": "max_hydrate_pages",
"label": "最大读取全文页数",
"type": "number",
"default": NOTION_DEFAULT_MAX_HYDRATE_PAGES,
"min": 1,
"max": 100,
"description": "候选排序后最多读取 block 全文的页面数量",
},
{
"key": "snippet_window_lines",
"label": "片段窗口行数",
"type": "number",
"default": 12,
"min": 3,
"max": 40,
"description": "搜索结果片段包含的上下文行数",
},
],
}