545 lines
23 KiB
Python
545 lines
23 KiB
Python
"""
|
|
Dịch vụ xây dựng Đồ thị Tri thức (Knowledge Graph)
|
|
API 2: Sử dụng Zep API để xây dựng một Standalone Graph (Đồ thị độc lập)
|
|
"""
|
|
|
|
import os
|
|
import re # changed
|
|
import uuid
|
|
import time
|
|
import threading
|
|
from typing import Dict, Any, List, Optional, Callable
|
|
from dataclasses import dataclass
|
|
|
|
from zep_cloud.client import Zep
|
|
from zep_cloud import EpisodeData, EntityEdgeSourceTarget
|
|
|
|
from ..config import Config
|
|
from ..models.task import TaskManager, TaskStatus
|
|
from ..utils.zep_paging import fetch_all_nodes, fetch_all_edges
|
|
from .text_processor import TextProcessor
|
|
|
|
|
|
@dataclass
|
|
class GraphInfo:
|
|
"""Các trường thông tin cơ bản của Graph"""
|
|
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,
|
|
}
|
|
|
|
|
|
class GraphBuilderService:
|
|
"""
|
|
Dịch vụ tạo lập Graph
|
|
Đảm nhiệm logic gọi request lên Zep API để thiết lập Graph
|
|
"""
|
|
|
|
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 has not been configured.")
|
|
|
|
self.client = Zep(api_key=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 = 3
|
|
) -> str:
|
|
"""
|
|
Khởi chạy tiến trình bất đồng bộ xây dựng Graph
|
|
|
|
Args:
|
|
text: Văn bản toàn văn làm nguồn vào
|
|
ontology: Từ điển chuẩn cấu trúc Ontology (Đầu ra từ API số 1)
|
|
graph_name: Tên đặt cho Graph
|
|
chunk_size: Kích thước từng khối text (chunk)
|
|
chunk_overlap: Giới hạn những từ đè lên nhau giữa các chunk (bảo toàn flow hội thoại / ngữ cảnh)
|
|
batch_size: Chuyển dữ liệu theo mảng batch để tiết kiệm số lần Request
|
|
|
|
Returns:
|
|
Trạng thái Task ID vừa khởi tạo
|
|
"""
|
|
# Đưa Task vào danh sách quản lý
|
|
task_id = self.task_manager.create_task(
|
|
task_type="graph_build",
|
|
metadata={
|
|
"graph_name": graph_name,
|
|
"chunk_size": chunk_size,
|
|
"text_length": len(text),
|
|
}
|
|
)
|
|
|
|
# Bắt đầu gọi Workder ở luồng ảo (Back ground Thread) để người dùng không tắc giao diện đợi xử lý
|
|
# Asynchronous Execution: Graph building runs in background threads with progress tracking through a task management system
|
|
thread = threading.Thread(
|
|
target=self._build_graph_worker,
|
|
args=(task_id, text, ontology, graph_name, chunk_size, chunk_overlap, batch_size)
|
|
)
|
|
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
|
|
):
|
|
"""Tiến trình cài đặt ngầm tạo Graph với các bước tuần tự"""
|
|
try:
|
|
self.task_manager.update_task(
|
|
task_id,
|
|
status=TaskStatus.PROCESSING,
|
|
progress=5,
|
|
message="Building Knowledge Graph..."
|
|
)
|
|
|
|
# Bước 1. Init tạo khung xương Graph trên Zep
|
|
graph_id = self.create_graph(graph_name)
|
|
self.task_manager.update_task(
|
|
task_id,
|
|
progress=10,
|
|
message=f"Created empty graph: {graph_id}"
|
|
)
|
|
|
|
# 2. Thiết lập ontology
|
|
self.set_ontology(graph_id, ontology)
|
|
self.task_manager.update_task(
|
|
task_id,
|
|
progress=15,
|
|
message="Ontology scheme applied successfully"
|
|
)
|
|
|
|
# Bước 3. Chia nhỏ văn bản gốc
|
|
# Text Processing: Chunks documents and sends them to Zep for entity extraction
|
|
chunks = TextProcessor.split_text(text, chunk_size, chunk_overlap) # Split into chunks
|
|
total_chunks = len(chunks)
|
|
self.task_manager.update_task(
|
|
task_id,
|
|
progress=20,
|
|
message=f"Split text into {total_chunks} chunk(s)"
|
|
)
|
|
|
|
# Bước 4. Gửi các đợt chunk tới Zep dưới dạng batch
|
|
# Send in batches (size=3)
|
|
episode_uuids = 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), # Thể hiện từ 20-60%
|
|
message=msg
|
|
)
|
|
)
|
|
|
|
# Bước 5. Đợi hàm Backend của Cloud Zep xử lý đồng bộ xong các episode
|
|
self.task_manager.update_task(
|
|
task_id,
|
|
progress=60,
|
|
message="Waiting for Zep to process data..."
|
|
)
|
|
|
|
# Wait for Zep processing
|
|
self._wait_for_episodes(
|
|
episode_uuids,
|
|
lambda msg, prog: self.task_manager.update_task(
|
|
task_id,
|
|
progress=60 + int(prog * 0.3), # Thể hiện từ 60-90%
|
|
message=msg
|
|
)
|
|
)
|
|
|
|
# Bước 6. Thống kê lại Graph hoàn thiện
|
|
self.task_manager.update_task(
|
|
task_id,
|
|
progress=90,
|
|
message="Fetching finalized graph info..."
|
|
)
|
|
|
|
# Retrieve graph statistics
|
|
graph_info = self._get_graph_info(graph_id)
|
|
|
|
# Thông báo hoàn tất
|
|
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) -> str:
|
|
"""Graph Creation: Initializes a new graph in Zep Cloud with a unique ID"""
|
|
|
|
# Generate unique graph_id
|
|
graph_id = f"mirofish_{uuid.uuid4().hex[:16]}"
|
|
|
|
# Call Zep API create()
|
|
self.client.graph.create(
|
|
graph_id=graph_id,
|
|
name=name,
|
|
description="MiroFish Social Simulation Graph"
|
|
)
|
|
|
|
return graph_id
|
|
|
|
def set_ontology(self, graph_id: str, ontology: Dict[str, Any]):
|
|
"""
|
|
Cấu hình dữ liệu Ontology (Bản thể học) cho Graph trên server Zep (Public access)
|
|
Ontology Setup: Defines the schema for entities and relationships using dynamic class creation.
|
|
"""
|
|
import warnings
|
|
from typing import Optional
|
|
from pydantic import Field
|
|
from zep_cloud.external_clients.ontology import EntityModel, EntityText, EdgeModel
|
|
|
|
# Ẩn bỏ đi các Warning (Cảnh báo) của thư viện Pydantic v2 liên quan đến Field(default=None)
|
|
# Vì đây là format bắt buộc phải có từ Zep SDK, các cảnh báo này phát sinh do tự động khởi tạo lớp ảo, hoàn toàn có thể bỏ qua được.
|
|
warnings.filterwarnings('ignore', category=UserWarning, module='pydantic')
|
|
|
|
# Danh sách các tên định danh (variable/name) trùng với từ khoá bảo lưu của Zep, không được dùng làm tên thuộc tính
|
|
RESERVED_NAMES = {'uuid', 'name', 'group_id', 'name_embedding', 'summary', 'created_at'}
|
|
|
|
def safe_attr_name(attr_name: str) -> str:
|
|
"""Hàm thay đổi các tên thuộc tính bị trùng với keyword của hệ thống để an toàn hơn"""
|
|
if attr_name.lower() in RESERVED_NAMES:
|
|
return f"entity_{attr_name}"
|
|
return attr_name
|
|
|
|
# Khởi tạo động (Dynamic Class Creation) các Model Loại Thực thể từ JSON đầu vào
|
|
entity_types = {}
|
|
|
|
# Processes each entity type from ontology definition
|
|
for entity_def in ontology.get("entity_types", []):
|
|
name = entity_def["name"]
|
|
description = entity_def.get("description", f"A {name} entity.")
|
|
|
|
# Chuẩn bị file từ điển cho Attribute và kiểu chú thích (Theo chuẩn Pydantic v2)
|
|
attrs = {"__doc__": description}
|
|
annotations = {}
|
|
|
|
for attr_def in entity_def.get("attributes", []):
|
|
attr_name = safe_attr_name(attr_def["name"]) # Áp dụng hàm chống bị trùng từ khoá
|
|
attr_desc = attr_def.get("description", attr_name)
|
|
# Zep API bắt buộc phải nhận vào field description
|
|
attrs[attr_name] = Field(description=attr_desc, default=None)
|
|
annotations[attr_name] = Optional[EntityText] # Chú thích kiểu dữ liệu
|
|
|
|
attrs["__annotations__"] = annotations
|
|
|
|
# Dynamic class generation - Create Pydnatic model class for entity type at runtime
|
|
entity_class = type(name, (EntityModel,), attrs)
|
|
entity_class.__doc__ = description
|
|
entity_types[name] = entity_class # Store in entity_types dict
|
|
|
|
# Tương tự, dựa vào JSON để khởi tạo động khai báo các Model Loại Quan Hệ
|
|
edge_definitions = {}
|
|
for edge_def in ontology.get("edge_types", []):
|
|
name = edge_def["name"]
|
|
description = edge_def.get("description", f"A {name} relationship.")
|
|
|
|
# Dọn các attribute dictionary và typing tương tự
|
|
attrs = {"__doc__": description}
|
|
annotations = {}
|
|
|
|
for attr_def in edge_def.get("attributes", []):
|
|
attr_name = safe_attr_name(attr_def["name"]) # Filter an toàn
|
|
attr_desc = attr_def.get("description", attr_name)
|
|
# Đảm bảo giữ format Zep API
|
|
attrs[attr_name] = Field(description=attr_desc, default=None)
|
|
annotations[attr_name] = Optional[str] # Định dạng Data cho thuộc tình của loại Quan Hệ là chuỗi String
|
|
|
|
attrs["__annotations__"] = annotations
|
|
|
|
# Khởi tạo Class động với Tên chuẩn format (PascalCase)
|
|
class_name = ''.join(word.capitalize() for word in name.split('_'))
|
|
# Creates Pydantic model class for relationship type
|
|
edge_class = type(class_name, (EdgeModel,), attrs)
|
|
edge_class.__doc__ = description
|
|
|
|
# Mapping thông số luồng thực thể gắn kết với Quan Hệ (Source/Targets config)
|
|
source_targets = []
|
|
for st in 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) # Store in edge_definitions dict
|
|
|
|
# Action Gọi lệnh thay đổi Ontology cho môi trường GraphID của Zep
|
|
if entity_types or edge_definitions:
|
|
# Call Zep set_ontology()
|
|
self.client.graph.set_ontology(
|
|
graph_ids=[graph_id],
|
|
entities=entity_types if entity_types else None,
|
|
edges=edge_definitions if edge_definitions else None,
|
|
)
|
|
|
|
def add_text_batches(
|
|
self,
|
|
graph_id: str,
|
|
chunks: List[str],
|
|
batch_size: int = 3,
|
|
progress_callback: Optional[Callable] = None
|
|
) -> List[str]:
|
|
"""Tải các đoạn văn bản (text chunks) lên Graph theo từng gói nhỏ (batch) và trả về id (Episode UUID) của mọi phân đoạn dữ liệu gửi đi."""
|
|
episode_uuids = []
|
|
total_chunks = len(chunks)
|
|
|
|
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(
|
|
f"Sending data batch {batch_num}/{total_batches} ({len(batch_chunks)} chunks)...",
|
|
progress
|
|
)
|
|
|
|
# Create EpisodeData objects
|
|
episodes = [
|
|
EpisodeData(data=chunk, type="text")
|
|
for chunk in batch_chunks
|
|
]
|
|
|
|
# Khởi chạy gửi cho Zep Server
|
|
max_retries = 10 # changed
|
|
for attempt in range(max_retries): # changed
|
|
try:
|
|
# UPLOAD BATCH TO ZEP - Sends batch of episodes for entity extraction
|
|
batch_result = self.client.graph.add_batch(
|
|
graph_id=graph_id,
|
|
episodes=episodes
|
|
)
|
|
|
|
# Cập nhật và thu thập lại UUID của các Episode được trả về sau khi tạo mới
|
|
if batch_result and isinstance(batch_result, list):
|
|
for ep in batch_result:
|
|
ep_uuid = getattr(ep, 'uuid_', None) or getattr(ep, 'uuid', None)
|
|
if ep_uuid:
|
|
episode_uuids.append(ep_uuid)
|
|
|
|
# Cài thời gian chờ (delay) nhỏ để tránh rate-limit bị quá tải số lượng requests
|
|
time.sleep(3)
|
|
break # changed: thoát retry loop nếu thành công
|
|
|
|
except Exception as e: # changed
|
|
err_str = str(e)
|
|
if "episode usage limit" in err_str or ("status_code: 429" in err_str): # changed: bắt lỗi rate-limit
|
|
# Đọc thời điểm reset từ error message để biết cần chờ bao lâu
|
|
reset_match = re.search(r"x-ratelimit-reset['\"]:\s*['\"]?(\d+)", err_str) # changed
|
|
if reset_match: # changed
|
|
wait_seconds = max(int(reset_match.group(1)) - int(time.time()) + 2, 5) # changed
|
|
else: # changed
|
|
wait_seconds = 20 # changed: fallback 12s (5 calls/phút → cách nhau 12s)
|
|
if progress_callback: # changed
|
|
progress_callback( # changed
|
|
f"Rate limited by Zep. Waiting {wait_seconds}s then retry (attempt {attempt + 1}/{max_retries})...", # changed
|
|
(i + len(batch_chunks)) / total_chunks # changed
|
|
) # changed
|
|
time.sleep(wait_seconds) # changed
|
|
else: # changed: lỗi khác thì raise luôn, không retry
|
|
if progress_callback:
|
|
progress_callback(f"Failed to send batch {batch_num}: {err_str}", 0)
|
|
raise # changed
|
|
else: # changed: for...else — chạy khi hết max_retries mà vẫn chưa break
|
|
raise Exception(f"Batch {batch_num} failed after {max_retries} retries due to rate limiting") # changed
|
|
|
|
return episode_uuids
|
|
|
|
def _wait_for_episodes(
|
|
self,
|
|
episode_uuids: List[str],
|
|
progress_callback: Optional[Callable] = None,
|
|
timeout: int = 1000
|
|
):
|
|
"""Chạy vòng lặp để kiểm tra và chờ cho tới khi mọi Episode (các khối Text) đều hoàn tất quá trình process từ hệ thống"""
|
|
if not episode_uuids:
|
|
if progress_callback:
|
|
progress_callback("No episodes to scan (Progress 100%)", 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(f"Waiting for analysis of {total_episodes} text chunks to begin...", 0)
|
|
|
|
while pending_episodes:
|
|
# Ngắt thoát và trả về lỗi nếu bị Timeout (Chạy quá thời gian cho phép)
|
|
if time.time() - start_time > timeout:
|
|
if progress_callback:
|
|
progress_callback(
|
|
f"Some text segments have timed out, but {completed_count}/{total_episodes} have completed successfully",
|
|
completed_count / total_episodes
|
|
)
|
|
break
|
|
|
|
# Duyệt vòng lặp mỗi episode uuid để lấy cập nhật tiến trình check của từng episode một
|
|
for ep_uuid in list(pending_episodes):
|
|
try:
|
|
# CHECK EPISODE STATUS - Queries individual episode processing status
|
|
episode = self.client.graph.episode.get(uuid_=ep_uuid)
|
|
# CHECK PROCESSED FLAG - Determines if Zep
|
|
is_processed = getattr(episode, 'processed', False)
|
|
|
|
if is_processed:
|
|
# if processed: remove from set
|
|
pending_episodes.remove(ep_uuid)
|
|
completed_count += 1
|
|
|
|
except Exception as e:
|
|
# Tạm thời bỏ qua nếu request lỗi, vòng lặp kế theo sẽ tự động call tiếp để get status
|
|
pass
|
|
|
|
elapsed = int(time.time() - start_time)
|
|
# update progress callback
|
|
if progress_callback:
|
|
progress_callback(
|
|
f"Zep is processing in the background... {completed_count}/{total_episodes} done, {len(pending_episodes)} tasks remaining ({elapsed}s elapsed)",
|
|
completed_count / total_episodes if total_episodes > 0 else 0
|
|
)
|
|
|
|
if pending_episodes:
|
|
time.sleep(3) # Lặp chu kỳ check mỗi 3 giây
|
|
|
|
if progress_callback:
|
|
# Final completion message
|
|
progress_callback(f"Data upload process completed: {completed_count}/{total_episodes}", 1.0)
|
|
|
|
def _get_graph_info(self, graph_id: str) -> GraphInfo:
|
|
"""Lấy/Get dữ liệu Graph Info hiện tại"""
|
|
# Load các điểm nút/Entity đang có (Qua trình duyệt web/paging)
|
|
nodes = fetch_all_nodes(self.client, graph_id)
|
|
|
|
# Lấy theo mảng phân trang thông tin các Edges/Mối Liên Kết
|
|
edges = fetch_all_edges(self.client, graph_id)
|
|
|
|
# Nối lại và thống kê những Entity Types
|
|
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]:
|
|
"""
|
|
Gói gọn toàn bộ dữ liệu cấu trúc (Bao gồm dữ liệu Graph chi tiết)
|
|
|
|
Args:
|
|
graph_id: ID của đồ thị
|
|
|
|
Returns:
|
|
Một object Dictionary bao hàm thông tin dữ liệu về Mạng lưới Cụm (nodes) và Cạnh (edges),
|
|
và toàn bộ chi tiết đi kèm khác (Time khởi tạo, Property).
|
|
"""
|
|
nodes = fetch_all_nodes(self.client, graph_id)
|
|
edges = fetch_all_edges(self.client, graph_id)
|
|
|
|
# Giữ một map tra cứu để phục vụ lấy 'Tên' nhanh theo ID UUID
|
|
node_map = {}
|
|
for node in nodes:
|
|
node_map[node.uuid_] = node.name or ""
|
|
|
|
nodes_data = []
|
|
for node in nodes:
|
|
# Lấy thông số về Thời gian được ghi nhận/khởi tạo
|
|
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:
|
|
# Thu thập các timestamp gắn với cạnh
|
|
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)
|
|
|