879 lines
33 KiB
Python
879 lines
33 KiB
Python
"""
|
||
图谱构建服务
|
||
接口2：使用Zep API构建Standalone Graph
|
||
"""
|
||
|
||
import hashlib
|
||
import uuid
|
||
import time
|
||
import threading
|
||
from typing import Dict, Any, List, Optional, Callable
|
||
from dataclasses import dataclass
|
||
|
||
from zep_cloud import BatchAddItem, EntityEdgeSourceTarget, NotFoundError
|
||
|
||
from ..config import Config
|
||
from ..models.task import TaskManager, TaskStatus
|
||
from ..utils.zep_paging import fetch_all_nodes, fetch_all_edges
|
||
from ..utils.ontology import (
|
||
MAX_ONTOLOGY_TYPES,
|
||
RESERVED_ONTOLOGY_ATTRIBUTE_NAMES,
|
||
normalize_ontology_attributes,
|
||
normalize_ontology_source_targets,
|
||
)
|
||
from ..utils.zep import (
|
||
call_zep_read_with_retry,
|
||
get_zep_client,
|
||
is_retryable_zep_error,
|
||
)
|
||
from .text_processor import TextProcessor
|
||
from ..utils.locale import t, get_locale, set_locale
|
||
|
||
|
||
@dataclass
|
||
class GraphInfo:
|
||
"""图谱信息"""
|
||
graph_id: str
|
||
node_count: int
|
||
edge_count: int
|
||
entity_types: List[str]
|
||
|
||
def to_dict(self) -> Dict[str, Any]:
|
||
return {
|
||
"graph_id": self.graph_id,
|
||
"node_count": self.node_count,
|
||
"edge_count": self.edge_count,
|
||
"entity_types": self.entity_types,
|
||
}
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class BatchSubmission:
|
||
"""Durable identity for one Zep Batch API ingestion operation."""
|
||
|
||
batch_id: str
|
||
operation_id: str
|
||
episode_uuids: List[str]
|
||
item_count: int
|
||
|
||
|
||
class GraphBuilderService:
|
||
"""
|
||
图谱构建服务
|
||
负责调用Zep API构建知识图谱
|
||
"""
|
||
|
||
def __init__(self, api_key: Optional[str] = None):
|
||
self.api_key = api_key or Config.ZEP_API_KEY
|
||
if not self.api_key:
|
||
raise ValueError("ZEP_API_KEY 未配置")
|
||
|
||
self.client = get_zep_client(self.api_key)
|
||
self.task_manager = TaskManager()
|
||
|
||
def build_graph_async(
|
||
self,
|
||
text: str,
|
||
ontology: Dict[str, Any],
|
||
graph_name: str = "MiroFish Graph",
|
||
chunk_size: int = 500,
|
||
chunk_overlap: int = 50,
|
||
batch_size: int = 350
|
||
) -> str:
|
||
"""
|
||
异步构建图谱
|
||
|
||
Args:
|
||
text: 输入文本
|
||
ontology: 本体定义（来自接口1的输出）
|
||
graph_name: 图谱名称
|
||
chunk_size: 文本块大小
|
||
chunk_overlap: 块重叠大小
|
||
batch_size: 每批发送的块数量
|
||
|
||
Returns:
|
||
任务ID
|
||
"""
|
||
# 创建任务
|
||
task_id = self.task_manager.create_task(
|
||
task_type="graph_build",
|
||
metadata={
|
||
"graph_name": graph_name,
|
||
"chunk_size": chunk_size,
|
||
"text_length": len(text),
|
||
}
|
||
)
|
||
|
||
# Capture locale before spawning background thread
|
||
current_locale = get_locale()
|
||
|
||
# 在后台线程中执行构建
|
||
thread = threading.Thread(
|
||
target=self._build_graph_worker,
|
||
args=(task_id, text, ontology, graph_name, chunk_size, chunk_overlap, batch_size, current_locale)
|
||
)
|
||
thread.daemon = True
|
||
thread.start()
|
||
|
||
return task_id
|
||
|
||
def _build_graph_worker(
|
||
self,
|
||
task_id: str,
|
||
text: str,
|
||
ontology: Dict[str, Any],
|
||
graph_name: str,
|
||
chunk_size: int,
|
||
chunk_overlap: int,
|
||
batch_size: int,
|
||
locale: str = 'zh'
|
||
):
|
||
"""图谱构建工作线程"""
|
||
set_locale(locale)
|
||
try:
|
||
self.task_manager.update_task(
|
||
task_id,
|
||
status=TaskStatus.PROCESSING,
|
||
progress=5,
|
||
message=t('progress.startBuildingGraph')
|
||
)
|
||
|
||
# Validate the complete ingestion payload before the first Cloud
|
||
# mutation, including this legacy service entry point.
|
||
chunks = TextProcessor.split_text(text, chunk_size, chunk_overlap)
|
||
self.validate_batch_chunks(chunks, batch_size=batch_size)
|
||
total_chunks = len(chunks)
|
||
|
||
# 1. 创建图谱
|
||
graph_id = self.create_graph(graph_name)
|
||
self.task_manager.update_task(
|
||
task_id,
|
||
progress=10,
|
||
message=t('progress.graphCreated', graphId=graph_id)
|
||
)
|
||
|
||
# 2. 设置本体
|
||
self.set_ontology(graph_id, ontology)
|
||
self.task_manager.update_task(
|
||
task_id,
|
||
progress=15,
|
||
message=t('progress.ontologySet')
|
||
)
|
||
|
||
# 3. 文本分块已在 Cloud mutation 前完成并验证
|
||
self.task_manager.update_task(
|
||
task_id,
|
||
progress=20,
|
||
message=t('progress.textSplit', count=total_chunks)
|
||
)
|
||
|
||
# 4. 分批发送数据
|
||
submission = self.add_text_batches(
|
||
graph_id, chunks, batch_size,
|
||
lambda msg, prog: self.task_manager.update_task(
|
||
task_id,
|
||
progress=20 + int(prog * 0.4), # 20-60%
|
||
message=msg
|
||
)
|
||
)
|
||
|
||
# 5. 等待Zep处理完成
|
||
self.task_manager.update_task(
|
||
task_id,
|
||
progress=60,
|
||
message=t('progress.waitingZepProcess')
|
||
)
|
||
|
||
self._wait_for_batch(
|
||
submission,
|
||
lambda msg, prog: self.task_manager.update_task(
|
||
task_id,
|
||
progress=60 + int(prog * 0.3), # 60-90%
|
||
message=msg
|
||
)
|
||
)
|
||
|
||
# 6. 获取图谱信息
|
||
self.task_manager.update_task(
|
||
task_id,
|
||
progress=90,
|
||
message=t('progress.fetchingGraphInfo')
|
||
)
|
||
|
||
graph_info = self._get_graph_info(graph_id)
|
||
|
||
# 完成
|
||
self.task_manager.complete_task(task_id, {
|
||
"graph_id": graph_id,
|
||
"graph_info": graph_info.to_dict(),
|
||
"chunks_processed": total_chunks,
|
||
})
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
error_msg = f"{str(e)}\n{traceback.format_exc()}"
|
||
self.task_manager.fail_task(task_id, error_msg)
|
||
|
||
def create_graph(
|
||
self,
|
||
name: str,
|
||
*,
|
||
graph_id: str | None = None,
|
||
graph_id_callback: Optional[Callable[[str], None]] = None,
|
||
) -> str:
|
||
"""Create a graph with a caller-durable ID and reconcile lost replies."""
|
||
|
||
graph_id = graph_id or f"mirofish_{uuid.uuid4().hex[:16]}"
|
||
# Persist the client-generated ID before the non-idempotent POST so a
|
||
# later reset can clean up a graph whose successful response was lost.
|
||
if graph_id_callback:
|
||
graph_id_callback(graph_id)
|
||
|
||
try:
|
||
self.client.graph.create(
|
||
graph_id=graph_id,
|
||
name=name,
|
||
description="MiroFish Social Simulation Graph"
|
||
)
|
||
except Exception as error:
|
||
if not is_retryable_zep_error(error):
|
||
raise
|
||
reconciliation_error = None
|
||
for attempt in range(3):
|
||
try:
|
||
call_zep_read_with_retry(
|
||
lambda: self.client.graph.get(graph_id),
|
||
operation_name=f"reconcile graph create {graph_id}",
|
||
)
|
||
reconciliation_error = None
|
||
break
|
||
except NotFoundError as not_found:
|
||
reconciliation_error = not_found
|
||
if attempt < 2:
|
||
time.sleep(attempt + 1)
|
||
except Exception as read_error:
|
||
reconciliation_error = read_error
|
||
break
|
||
if reconciliation_error is not None:
|
||
raise error from reconciliation_error
|
||
|
||
return graph_id
|
||
|
||
@staticmethod
|
||
def build_operation_id(graph_id: str, chunks: List[str]) -> str:
|
||
payload_hash = hashlib.sha256("\0".join(chunks).encode("utf-8")).hexdigest()
|
||
return hashlib.sha256(
|
||
f"{graph_id}:{payload_hash}".encode("utf-8")
|
||
).hexdigest()
|
||
|
||
def _find_batch_by_operation_id(
|
||
self,
|
||
graph_id: str,
|
||
operation_id: str,
|
||
*,
|
||
max_attempts: int = 3,
|
||
) -> Any | None:
|
||
"""Find one server-created batch after an ambiguous create reply."""
|
||
|
||
for attempt in range(1, max_attempts + 1):
|
||
matches: List[Any] = []
|
||
cursor: int | None = None
|
||
seen_cursors: set[int] = set()
|
||
while True:
|
||
page = call_zep_read_with_retry(
|
||
lambda: self.client.batch.list(limit=100, cursor=cursor),
|
||
operation_name=f"reconcile batch create {operation_id}",
|
||
)
|
||
for batch in getattr(page, "batches", None) or []:
|
||
metadata = getattr(batch, "metadata", None) or {}
|
||
if (
|
||
metadata.get("mirofish_operation_id") == operation_id
|
||
and metadata.get("graph_id") == graph_id
|
||
):
|
||
matches.append(batch)
|
||
next_cursor = getattr(page, "next_cursor", None)
|
||
if next_cursor is None:
|
||
break
|
||
if next_cursor == cursor or next_cursor in seen_cursors:
|
||
raise RuntimeError("Zep batch list cursor did not advance")
|
||
seen_cursors.add(next_cursor)
|
||
cursor = next_cursor
|
||
|
||
if len(matches) > 1:
|
||
raise RuntimeError(
|
||
f"Multiple Zep batches match operation {operation_id}; refusing ambiguity"
|
||
)
|
||
if matches:
|
||
return matches[0]
|
||
if attempt < max_attempts:
|
||
time.sleep(attempt)
|
||
return None
|
||
|
||
def set_ontology(self, graph_id: str, ontology: Dict[str, Any]):
|
||
"""设置图谱本体（公开方法）"""
|
||
import warnings
|
||
from typing import Optional
|
||
from pydantic import Field
|
||
from zep_cloud.external_clients.ontology import EntityModel, EntityText, EdgeModel
|
||
|
||
# 抑制 Pydantic v2 关于 Field(default=None) 的警告
|
||
# 这是 Zep SDK 要求的用法，警告来自动态类创建，可以安全忽略
|
||
warnings.filterwarnings('ignore', category=UserWarning, module='pydantic')
|
||
|
||
def safe_attr_name(attr_name: str) -> str:
|
||
"""将保留名称转换为安全名称"""
|
||
if attr_name.lower() in RESERVED_ONTOLOGY_ATTRIBUTE_NAMES:
|
||
return f"entity_{attr_name}"
|
||
return attr_name
|
||
|
||
# 动态创建实体类型
|
||
entity_types = {}
|
||
for entity_def in ontology.get("entity_types", [])[:MAX_ONTOLOGY_TYPES]:
|
||
name = entity_def["name"]
|
||
description = entity_def.get("description", f"A {name} entity.")
|
||
|
||
# 创建属性字典和类型注解（Pydantic v2 需要）
|
||
attrs = {"__doc__": description}
|
||
annotations = {}
|
||
|
||
for normalized in normalize_ontology_attributes(
|
||
entity_def.get("attributes", [])
|
||
):
|
||
attr_name = safe_attr_name(normalized["name"]) # 使用安全名称
|
||
attr_desc = normalized["description"]
|
||
# Zep API 需要 Field 的 description，这是必需的
|
||
attrs[attr_name] = Field(description=attr_desc, default=None)
|
||
annotations[attr_name] = Optional[EntityText] # 类型注解
|
||
|
||
attrs["__annotations__"] = annotations
|
||
|
||
# 动态创建类
|
||
entity_class = type(name, (EntityModel,), attrs)
|
||
entity_class.__doc__ = description
|
||
entity_types[name] = entity_class
|
||
|
||
# 动态创建边类型
|
||
edge_definitions = {}
|
||
for edge_def in ontology.get("edge_types", [])[:MAX_ONTOLOGY_TYPES]:
|
||
name = edge_def["name"]
|
||
description = edge_def.get("description", f"A {name} relationship.")
|
||
|
||
# 创建属性字典和类型注解
|
||
attrs = {"__doc__": description}
|
||
annotations = {}
|
||
|
||
for normalized in normalize_ontology_attributes(
|
||
edge_def.get("attributes", [])
|
||
):
|
||
attr_name = safe_attr_name(normalized["name"]) # 使用安全名称
|
||
attr_desc = normalized["description"]
|
||
# Zep API 需要 Field 的 description，这是必需的
|
||
attrs[attr_name] = Field(description=attr_desc, default=None)
|
||
annotations[attr_name] = Optional[str] # 边属性用str类型
|
||
|
||
attrs["__annotations__"] = annotations
|
||
|
||
# 动态创建类
|
||
class_name = ''.join(word.capitalize() for word in name.split('_'))
|
||
edge_class = type(class_name, (EdgeModel,), attrs)
|
||
edge_class.__doc__ = description
|
||
|
||
# 构建source_targets
|
||
source_targets = []
|
||
for st in normalize_ontology_source_targets(
|
||
edge_def.get("source_targets", [])
|
||
):
|
||
source_targets.append(
|
||
EntityEdgeSourceTarget(
|
||
source=st.get("source", "Entity"),
|
||
target=st.get("target", "Entity")
|
||
)
|
||
)
|
||
|
||
if source_targets:
|
||
edge_definitions[name] = (edge_class, source_targets)
|
||
|
||
# 调用Zep API设置本体
|
||
if entity_types or edge_definitions:
|
||
self.client.graph.set_ontology(
|
||
graph_ids=[graph_id],
|
||
# Zep iterates entities.items(), so edge-only ontologies must
|
||
# pass an empty dictionary rather than None.
|
||
entities=entity_types,
|
||
edges=edge_definitions if edge_definitions else None,
|
||
)
|
||
|
||
def add_text_batches(
|
||
self,
|
||
graph_id: str,
|
||
chunks: List[str],
|
||
batch_size: int = 350,
|
||
progress_callback: Optional[Callable] = None,
|
||
batch_created_callback: Optional[Callable[[str | None, str], None]] = None,
|
||
) -> BatchSubmission:
|
||
"""Submit document chunks through Zep's current Batch API.
|
||
|
||
Mutating calls are deliberately not retried: create/add are not
|
||
documented as idempotent, and an ambiguous replay can duplicate graph
|
||
episodes. The returned batch identity allows callers to persist and
|
||
reconcile the operation instead.
|
||
"""
|
||
|
||
if not graph_id:
|
||
raise ValueError("graph_id is required")
|
||
self.validate_batch_chunks(chunks, batch_size=batch_size)
|
||
|
||
total_chunks = len(chunks)
|
||
operation_id = self.build_operation_id(graph_id, chunks)
|
||
if batch_created_callback:
|
||
# Journal the deterministic operation before the server-generated
|
||
# batch ID POST. This leaves enough identity for later diagnosis
|
||
# even if both the response and immediate list reconciliation fail.
|
||
batch_created_callback(None, operation_id)
|
||
|
||
try:
|
||
batch = self.client.batch.create(
|
||
metadata={
|
||
"mirofish_operation_id": operation_id,
|
||
"graph_id": graph_id,
|
||
"chunk_count": total_chunks,
|
||
}
|
||
)
|
||
except Exception as error:
|
||
if not is_retryable_zep_error(error):
|
||
raise
|
||
batch = self._find_batch_by_operation_id(graph_id, operation_id)
|
||
if batch is None:
|
||
raise RuntimeError(
|
||
"Zep batch creation is unconfirmed and no matching operation was found"
|
||
) from error
|
||
batch_id = getattr(batch, "batch_id", None)
|
||
if not batch_id:
|
||
raise RuntimeError("Zep Batch API returned no batch_id")
|
||
if batch_created_callback:
|
||
batch_created_callback(batch_id, operation_id)
|
||
|
||
episode_uuids: List[str] = []
|
||
for i in range(0, total_chunks, batch_size):
|
||
batch_chunks = chunks[i:i + batch_size]
|
||
batch_num = i // batch_size + 1
|
||
total_batches = (total_chunks + batch_size - 1) // batch_size
|
||
|
||
if progress_callback:
|
||
progress = (i + len(batch_chunks)) / total_chunks
|
||
progress_callback(
|
||
t('progress.sendingBatch', current=batch_num, total=total_batches, chunks=len(batch_chunks)),
|
||
progress
|
||
)
|
||
|
||
items = [
|
||
BatchAddItem(
|
||
type="graph_episode",
|
||
graph_id=graph_id,
|
||
data=chunk,
|
||
data_type="text",
|
||
source_description="MiroFish source document chunk",
|
||
metadata={
|
||
"mirofish_operation_id": operation_id,
|
||
"chunk_index": i + offset,
|
||
"chunk_sha256": hashlib.sha256(
|
||
chunk.encode("utf-8")
|
||
).hexdigest(),
|
||
},
|
||
)
|
||
for offset, chunk in enumerate(batch_chunks)
|
||
]
|
||
|
||
expected_item_count = i + len(items)
|
||
try:
|
||
item_details = self.client.batch.add(
|
||
batch_id=batch_id,
|
||
items=items,
|
||
)
|
||
except Exception as e:
|
||
if progress_callback:
|
||
progress_callback(t('progress.batchFailed', batch=batch_num, error=str(e)), 0)
|
||
if is_retryable_zep_error(e):
|
||
recovered_items = self._reconcile_batch_item_count(
|
||
batch_id,
|
||
expected_item_count,
|
||
)
|
||
recovered_indexes = {
|
||
getattr(item, "sequence_index", None)
|
||
for item in recovered_items
|
||
}
|
||
if (
|
||
len(recovered_items) == expected_item_count
|
||
and recovered_indexes == set(range(expected_item_count))
|
||
):
|
||
item_details = recovered_items[i:expected_item_count]
|
||
else:
|
||
raise RuntimeError(
|
||
f"Zep batch {batch_id} item submission is unconfirmed; "
|
||
"the draft was not processed or replayed"
|
||
) from e
|
||
else:
|
||
raise RuntimeError(
|
||
f"Zep batch {batch_id} item submission failed"
|
||
) from e
|
||
|
||
if len(item_details or []) != len(items):
|
||
recovered_items = self._reconcile_batch_item_count(
|
||
batch_id,
|
||
expected_item_count,
|
||
)
|
||
recovered_indexes = {
|
||
getattr(item, "sequence_index", None)
|
||
for item in recovered_items
|
||
}
|
||
if (
|
||
len(recovered_items) == expected_item_count
|
||
and recovered_indexes == set(range(expected_item_count))
|
||
):
|
||
item_details = recovered_items[i:expected_item_count]
|
||
else:
|
||
raise RuntimeError(
|
||
f"Zep batch {batch_id} acknowledged {len(item_details or [])} "
|
||
f"of {len(items)} items"
|
||
)
|
||
for item in item_details:
|
||
episode_uuid = getattr(item, "episode_uuid", None)
|
||
if episode_uuid:
|
||
episode_uuids.append(episode_uuid)
|
||
|
||
try:
|
||
self.client.batch.process(batch_id=batch_id)
|
||
except Exception as error:
|
||
# A process response can be lost after the server accepted it.
|
||
# Reconcile with a safe GET instead of issuing a second POST.
|
||
summary = call_zep_read_with_retry(
|
||
lambda: self.client.batch.get(batch_id=batch_id),
|
||
operation_name=f"reconcile batch {batch_id}",
|
||
)
|
||
if getattr(summary, "status", None) in {None, "draft"}:
|
||
raise RuntimeError(
|
||
f"Zep batch {batch_id} processing is unconfirmed"
|
||
) from error
|
||
|
||
return BatchSubmission(
|
||
batch_id=batch_id,
|
||
operation_id=operation_id,
|
||
episode_uuids=episode_uuids,
|
||
item_count=total_chunks,
|
||
)
|
||
|
||
@staticmethod
|
||
def validate_batch_chunks(chunks: List[str], *, batch_size: int = 350) -> None:
|
||
"""Validate every Batch API limit before the first Cloud mutation."""
|
||
|
||
if not chunks:
|
||
raise ValueError("At least one text chunk is required")
|
||
if not 1 <= batch_size <= 350:
|
||
raise ValueError("batch_size must be between 1 and 350")
|
||
if len(chunks) > 50_000:
|
||
raise ValueError("A Zep batch cannot contain more than 50,000 items")
|
||
oversized = [index for index, chunk in enumerate(chunks) if len(chunk) > 10_000]
|
||
if oversized:
|
||
raise ValueError(
|
||
f"Zep batch item exceeds 10,000 characters at chunk {oversized[0]}"
|
||
)
|
||
|
||
def _list_batch_items(self, batch_id: str) -> List[Any]:
|
||
items: List[Any] = []
|
||
cursor: int | None = None
|
||
seen_cursors: set[int] = set()
|
||
while True:
|
||
page = call_zep_read_with_retry(
|
||
lambda: self.client.batch.list_items(
|
||
batch_id=batch_id,
|
||
limit=100,
|
||
cursor=cursor,
|
||
),
|
||
operation_name=f"list batch items {batch_id}",
|
||
)
|
||
items.extend(getattr(page, "items", None) or [])
|
||
next_cursor = getattr(page, "next_cursor", None)
|
||
if next_cursor is None:
|
||
break
|
||
if next_cursor == cursor or next_cursor in seen_cursors:
|
||
raise RuntimeError(f"Zep batch {batch_id} item cursor did not advance")
|
||
seen_cursors.add(next_cursor)
|
||
cursor = next_cursor
|
||
return items
|
||
|
||
def _reconcile_batch_item_count(
|
||
self,
|
||
batch_id: str,
|
||
expected_item_count: int,
|
||
*,
|
||
max_attempts: int = 3,
|
||
) -> List[Any]:
|
||
"""Allow a short propagation window after an ambiguous add reply."""
|
||
|
||
items: List[Any] = []
|
||
for attempt in range(1, max_attempts + 1):
|
||
items = self._list_batch_items(batch_id)
|
||
if len(items) >= expected_item_count:
|
||
return items
|
||
if attempt < max_attempts:
|
||
time.sleep(attempt)
|
||
return items
|
||
|
||
def get_batch_summary(self, batch_id: str) -> Any:
|
||
"""Read a persisted batch identity for restart reconciliation."""
|
||
|
||
return call_zep_read_with_retry(
|
||
lambda: self.client.batch.get(batch_id=batch_id),
|
||
operation_name=f"get batch {batch_id}",
|
||
)
|
||
|
||
def _wait_for_batch(
|
||
self,
|
||
submission: BatchSubmission,
|
||
progress_callback: Optional[Callable] = None,
|
||
timeout: int | None = None,
|
||
) -> List[str]:
|
||
"""Wait for a Batch API terminal state and validate every item."""
|
||
|
||
timeout = timeout or Config.ZEP_INGESTION_TIMEOUT_SECONDS
|
||
start_time = time.time()
|
||
terminal_states = {"succeeded", "partial", "failed", "invalid", "canceled"}
|
||
|
||
while True:
|
||
if time.time() - start_time > timeout:
|
||
raise TimeoutError(
|
||
f"Zep batch {submission.batch_id} did not finish within {timeout}s"
|
||
)
|
||
|
||
summary = call_zep_read_with_retry(
|
||
lambda: self.client.batch.get(batch_id=submission.batch_id),
|
||
operation_name=f"poll batch {submission.batch_id}",
|
||
)
|
||
status = getattr(summary, "status", None)
|
||
progress = getattr(summary, "progress", None)
|
||
percent = float(getattr(progress, "percent_complete", 0) or 0) / 100
|
||
if progress_callback:
|
||
completed = int(getattr(progress, "succeeded_items", 0) or 0)
|
||
progress_callback(
|
||
t(
|
||
'progress.zepProcessing',
|
||
completed=completed,
|
||
total=submission.item_count,
|
||
pending=max(submission.item_count - completed, 0),
|
||
elapsed=int(time.time() - start_time),
|
||
),
|
||
min(max(percent, 0.0), 1.0),
|
||
)
|
||
|
||
if status in terminal_states:
|
||
break
|
||
time.sleep(3)
|
||
|
||
items = self._list_batch_items(submission.batch_id)
|
||
if status != "succeeded":
|
||
failed_items = [
|
||
item for item in items
|
||
if getattr(item, "status", None) not in {"succeeded", "skipped"}
|
||
]
|
||
first_error = getattr(failed_items[0], "error", None) if failed_items else None
|
||
raise RuntimeError(
|
||
f"Zep batch {submission.batch_id} ended as {status}; "
|
||
f"failed_items={len(failed_items)}; first_error={first_error}"
|
||
)
|
||
if len(items) != submission.item_count:
|
||
raise RuntimeError(
|
||
f"Zep batch {submission.batch_id} contains {len(items)} items, "
|
||
f"expected {submission.item_count}"
|
||
)
|
||
|
||
ordered_items = sorted(
|
||
items,
|
||
key=lambda item: getattr(item, "sequence_index", 0) or 0,
|
||
)
|
||
episode_uuids: List[str] = []
|
||
for item in ordered_items:
|
||
item_status = getattr(item, "status", None)
|
||
episode_uuid = getattr(item, "episode_uuid", None)
|
||
source_uuid = getattr(item, "source_uuid", None)
|
||
if item_status != "succeeded" or not episode_uuid:
|
||
raise RuntimeError(
|
||
f"Zep batch {submission.batch_id} returned an incomplete item"
|
||
)
|
||
if source_uuid and source_uuid != episode_uuid:
|
||
raise RuntimeError(
|
||
f"Zep batch {submission.batch_id} returned mismatched episode UUIDs"
|
||
)
|
||
episode_uuids.append(episode_uuid)
|
||
|
||
if progress_callback:
|
||
progress_callback(
|
||
t(
|
||
'progress.processingComplete',
|
||
completed=len(episode_uuids),
|
||
total=submission.item_count,
|
||
),
|
||
1.0,
|
||
)
|
||
return episode_uuids
|
||
|
||
def _wait_for_episodes(
|
||
self,
|
||
episode_uuids: List[str],
|
||
progress_callback: Optional[Callable] = None,
|
||
timeout: int = 600
|
||
):
|
||
"""等待所有 episode 处理完成（通过查询每个 episode 的 processed 状态）"""
|
||
if not episode_uuids:
|
||
if progress_callback:
|
||
progress_callback(t('progress.noEpisodesWait'), 1.0)
|
||
return
|
||
|
||
start_time = time.time()
|
||
pending_episodes = set(episode_uuids)
|
||
completed_count = 0
|
||
total_episodes = len(episode_uuids)
|
||
|
||
if progress_callback:
|
||
progress_callback(t('progress.waitingEpisodes', count=total_episodes), 0)
|
||
|
||
while pending_episodes:
|
||
if time.time() - start_time > timeout:
|
||
if progress_callback:
|
||
progress_callback(
|
||
t('progress.episodesTimeout', completed=completed_count, total=total_episodes),
|
||
completed_count / total_episodes
|
||
)
|
||
raise TimeoutError(
|
||
f"Zep episode processing timed out with "
|
||
f"{len(pending_episodes)} episode(s) still pending"
|
||
)
|
||
|
||
# 检查每个 episode 的处理状态
|
||
for ep_uuid in list(pending_episodes):
|
||
episode = call_zep_read_with_retry(
|
||
lambda: self.client.graph.episode.get(uuid_=ep_uuid),
|
||
operation_name=f"poll episode {ep_uuid}",
|
||
)
|
||
is_processed = getattr(episode, 'processed', False)
|
||
|
||
if is_processed:
|
||
pending_episodes.remove(ep_uuid)
|
||
completed_count += 1
|
||
|
||
elapsed = int(time.time() - start_time)
|
||
if progress_callback:
|
||
progress_callback(
|
||
t('progress.zepProcessing', completed=completed_count, total=total_episodes, pending=len(pending_episodes), elapsed=elapsed),
|
||
completed_count / total_episodes if total_episodes > 0 else 0
|
||
)
|
||
|
||
if pending_episodes:
|
||
time.sleep(3) # 每3秒检查一次
|
||
|
||
if progress_callback:
|
||
progress_callback(t('progress.processingComplete', completed=completed_count, total=total_episodes), 1.0)
|
||
|
||
def _get_graph_info(self, graph_id: str) -> GraphInfo:
|
||
"""获取图谱信息"""
|
||
# 获取节点（分页）
|
||
nodes = fetch_all_nodes(self.client, graph_id)
|
||
|
||
# 获取边（分页）
|
||
edges = fetch_all_edges(self.client, graph_id)
|
||
|
||
# 统计实体类型
|
||
entity_types = set()
|
||
for node in nodes:
|
||
if node.labels:
|
||
for label in node.labels:
|
||
if label not in ["Entity", "Node"]:
|
||
entity_types.add(label)
|
||
|
||
return GraphInfo(
|
||
graph_id=graph_id,
|
||
node_count=len(nodes),
|
||
edge_count=len(edges),
|
||
entity_types=list(entity_types)
|
||
)
|
||
|
||
def get_graph_data(self, graph_id: str) -> Dict[str, Any]:
|
||
"""
|
||
获取完整图谱数据（包含详细信息）
|
||
|
||
Args:
|
||
graph_id: 图谱ID
|
||
|
||
Returns:
|
||
包含nodes和edges的字典，包括时间信息、属性等详细数据
|
||
"""
|
||
nodes = fetch_all_nodes(self.client, graph_id)
|
||
edges = fetch_all_edges(self.client, graph_id)
|
||
|
||
# 创建节点映射用于获取节点名称
|
||
node_map = {}
|
||
for node in nodes:
|
||
node_map[node.uuid_] = node.name or ""
|
||
|
||
nodes_data = []
|
||
for node in nodes:
|
||
# 获取创建时间
|
||
created_at = getattr(node, 'created_at', None)
|
||
if created_at:
|
||
created_at = str(created_at)
|
||
|
||
nodes_data.append({
|
||
"uuid": node.uuid_,
|
||
"name": node.name,
|
||
"labels": node.labels or [],
|
||
"summary": node.summary or "",
|
||
"attributes": node.attributes or {},
|
||
"created_at": created_at,
|
||
})
|
||
|
||
edges_data = []
|
||
for edge in edges:
|
||
# 获取时间信息
|
||
created_at = getattr(edge, 'created_at', None)
|
||
valid_at = getattr(edge, 'valid_at', None)
|
||
invalid_at = getattr(edge, 'invalid_at', None)
|
||
expired_at = getattr(edge, 'expired_at', None)
|
||
|
||
# 获取 episodes
|
||
episodes = getattr(edge, 'episodes', None) or getattr(edge, 'episode_ids', None)
|
||
if episodes and not isinstance(episodes, list):
|
||
episodes = [str(episodes)]
|
||
elif episodes:
|
||
episodes = [str(e) for e in episodes]
|
||
|
||
# 获取 fact_type
|
||
fact_type = getattr(edge, 'fact_type', None) or edge.name or ""
|
||
|
||
edges_data.append({
|
||
"uuid": edge.uuid_,
|
||
"name": edge.name or "",
|
||
"fact": edge.fact or "",
|
||
"fact_type": fact_type,
|
||
"source_node_uuid": edge.source_node_uuid,
|
||
"target_node_uuid": edge.target_node_uuid,
|
||
"source_node_name": node_map.get(edge.source_node_uuid, ""),
|
||
"target_node_name": node_map.get(edge.target_node_uuid, ""),
|
||
"attributes": edge.attributes or {},
|
||
"created_at": str(created_at) if created_at else None,
|
||
"valid_at": str(valid_at) if valid_at else None,
|
||
"invalid_at": str(invalid_at) if invalid_at else None,
|
||
"expired_at": str(expired_at) if expired_at else None,
|
||
"episodes": episodes or [],
|
||
})
|
||
|
||
return {
|
||
"graph_id": graph_id,
|
||
"nodes": nodes_data,
|
||
"edges": edges_data,
|
||
"node_count": len(nodes_data),
|
||
"edge_count": len(edges_data),
|
||
}
|
||
|
||
def delete_graph(self, graph_id: str):
|
||
"""删除图谱"""
|
||
self.client.graph.delete(graph_id=graph_id)
|