Initial commit
This commit is contained in:
@@ -0,0 +1,458 @@
|
||||
"""
|
||||
图数据库数据准备模块
|
||||
"""
|
||||
|
||||
import logging
|
||||
import json
|
||||
from typing import List, Dict, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
|
||||
from neo4j import GraphDatabase
|
||||
from langchain_core.documents import Document
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class GraphNode:
|
||||
"""图节点数据结构"""
|
||||
node_id: str
|
||||
labels: List[str]
|
||||
name: str
|
||||
properties: Dict[str, Any]
|
||||
|
||||
@dataclass
|
||||
class GraphRelation:
|
||||
"""图关系数据结构"""
|
||||
start_node_id: str
|
||||
end_node_id: str
|
||||
relation_type: str
|
||||
properties: Dict[str, Any]
|
||||
|
||||
class GraphDataPreparationModule:
|
||||
"""图数据库数据准备模块 - 从Neo4j读取数据并转换为文档"""
|
||||
|
||||
def __init__(self, uri: str, user: str, password: str, database: str = "neo4j"):
|
||||
"""
|
||||
初始化图数据库连接
|
||||
|
||||
Args:
|
||||
uri: Neo4j连接URI
|
||||
user: 用户名
|
||||
password: 密码
|
||||
database: 数据库名称
|
||||
"""
|
||||
self.uri = uri
|
||||
self.user = user
|
||||
self.password = password
|
||||
self.database = database
|
||||
self.driver = None
|
||||
self.documents: List[Document] = []
|
||||
self.chunks: List[Document] = []
|
||||
self.recipes: List[GraphNode] = []
|
||||
self.ingredients: List[GraphNode] = []
|
||||
self.cooking_steps: List[GraphNode] = []
|
||||
|
||||
self._connect()
|
||||
|
||||
def _connect(self):
|
||||
"""建立Neo4j连接"""
|
||||
try:
|
||||
self.driver = GraphDatabase.driver(
|
||||
self.uri,
|
||||
auth=(self.user, self.password),
|
||||
database=self.database
|
||||
)
|
||||
logger.info(f"已连接到Neo4j数据库: {self.uri}")
|
||||
|
||||
# 测试连接
|
||||
with self.driver.session() as session:
|
||||
result = session.run("RETURN 1 as test")
|
||||
test_result = result.single()
|
||||
if test_result:
|
||||
logger.info("Neo4j连接测试成功")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"连接Neo4j失败: {e}")
|
||||
raise
|
||||
|
||||
def close(self):
|
||||
"""关闭数据库连接"""
|
||||
if hasattr(self, 'driver') and self.driver:
|
||||
self.driver.close()
|
||||
logger.info("Neo4j连接已关闭")
|
||||
|
||||
def load_graph_data(self) -> Dict[str, Any]:
|
||||
"""
|
||||
从Neo4j加载图数据
|
||||
|
||||
Returns:
|
||||
包含节点和关系的数据字典
|
||||
"""
|
||||
logger.info("正在从Neo4j加载图数据...")
|
||||
|
||||
with self.driver.session() as session:
|
||||
# 加载所有菜谱节点，从Category关系中读取分类信息
|
||||
recipes_query = """
|
||||
MATCH (r:Recipe)
|
||||
WHERE r.nodeId >= '200000000'
|
||||
OPTIONAL MATCH (r)-[:BELONGS_TO_CATEGORY]->(c:Category)
|
||||
WITH r, collect(c.name) as categories
|
||||
RETURN r.nodeId as nodeId, labels(r) as labels, r.name as name,
|
||||
properties(r) as originalProperties,
|
||||
CASE WHEN size(categories) > 0
|
||||
THEN categories[0]
|
||||
ELSE COALESCE(r.category, '未知') END as mainCategory,
|
||||
CASE WHEN size(categories) > 0
|
||||
THEN categories
|
||||
ELSE [COALESCE(r.category, '未知')] END as allCategories
|
||||
ORDER BY r.nodeId
|
||||
"""
|
||||
|
||||
result = session.run(recipes_query)
|
||||
self.recipes = []
|
||||
for record in result:
|
||||
# 合并原始属性和新的分类信息
|
||||
properties = dict(record["originalProperties"])
|
||||
properties["category"] = record["mainCategory"]
|
||||
properties["all_categories"] = record["allCategories"]
|
||||
|
||||
node = GraphNode(
|
||||
node_id=record["nodeId"],
|
||||
labels=record["labels"],
|
||||
name=record["name"],
|
||||
properties=properties
|
||||
)
|
||||
self.recipes.append(node)
|
||||
|
||||
logger.info(f"加载了 {len(self.recipes)} 个菜谱节点")
|
||||
|
||||
# 加载所有食材节点
|
||||
ingredients_query = """
|
||||
MATCH (i:Ingredient)
|
||||
WHERE i.nodeId >= '200000000'
|
||||
RETURN i.nodeId as nodeId, labels(i) as labels, i.name as name,
|
||||
properties(i) as properties
|
||||
ORDER BY i.nodeId
|
||||
"""
|
||||
|
||||
result = session.run(ingredients_query)
|
||||
self.ingredients = []
|
||||
for record in result:
|
||||
node = GraphNode(
|
||||
node_id=record["nodeId"],
|
||||
labels=record["labels"],
|
||||
name=record["name"],
|
||||
properties=record["properties"]
|
||||
)
|
||||
self.ingredients.append(node)
|
||||
|
||||
logger.info(f"加载了 {len(self.ingredients)} 个食材节点")
|
||||
|
||||
# 加载所有烹饪步骤节点
|
||||
steps_query = """
|
||||
MATCH (s:CookingStep)
|
||||
WHERE s.nodeId >= '200000000'
|
||||
RETURN s.nodeId as nodeId, labels(s) as labels, s.name as name,
|
||||
properties(s) as properties
|
||||
ORDER BY s.nodeId
|
||||
"""
|
||||
|
||||
result = session.run(steps_query)
|
||||
self.cooking_steps = []
|
||||
for record in result:
|
||||
node = GraphNode(
|
||||
node_id=record["nodeId"],
|
||||
labels=record["labels"],
|
||||
name=record["name"],
|
||||
properties=record["properties"]
|
||||
)
|
||||
self.cooking_steps.append(node)
|
||||
|
||||
logger.info(f"加载了 {len(self.cooking_steps)} 个烹饪步骤节点")
|
||||
|
||||
return {
|
||||
'recipes': len(self.recipes),
|
||||
'ingredients': len(self.ingredients),
|
||||
'cooking_steps': len(self.cooking_steps)
|
||||
}
|
||||
|
||||
def build_recipe_documents(self) -> List[Document]:
|
||||
"""
|
||||
构建菜谱文档，集成相关的食材和步骤信息
|
||||
|
||||
Returns:
|
||||
结构化的菜谱文档列表
|
||||
"""
|
||||
logger.info("正在构建菜谱文档...")
|
||||
|
||||
documents = []
|
||||
|
||||
with self.driver.session() as session:
|
||||
for recipe in self.recipes:
|
||||
try:
|
||||
recipe_id = recipe.node_id
|
||||
recipe_name = recipe.name
|
||||
|
||||
# 获取菜谱的相关食材
|
||||
ingredients_query = """
|
||||
MATCH (r:Recipe {nodeId: $recipe_id})-[req:REQUIRES]->(i:Ingredient)
|
||||
RETURN i.name as name, i.category as category,
|
||||
req.amount as amount, req.unit as unit,
|
||||
i.description as description
|
||||
ORDER BY i.name
|
||||
"""
|
||||
|
||||
ingredients_result = session.run(ingredients_query, {"recipe_id": recipe_id})
|
||||
ingredients_info = []
|
||||
for ing_record in ingredients_result:
|
||||
amount = ing_record.get("amount", "")
|
||||
unit = ing_record.get("unit", "")
|
||||
ingredient_text = f"{ing_record['name']}"
|
||||
if amount and unit:
|
||||
ingredient_text += f"({amount}{unit})"
|
||||
if ing_record.get("description"):
|
||||
ingredient_text += f" - {ing_record['description']}"
|
||||
ingredients_info.append(ingredient_text)
|
||||
|
||||
# 获取菜谱的烹饪步骤
|
||||
steps_query = """
|
||||
MATCH (r:Recipe {nodeId: $recipe_id})-[c:CONTAINS_STEP]->(s:CookingStep)
|
||||
RETURN s.name as name, s.description as description,
|
||||
s.stepNumber as stepNumber, s.methods as methods,
|
||||
s.tools as tools, s.timeEstimate as timeEstimate,
|
||||
c.stepOrder as stepOrder
|
||||
ORDER BY COALESCE(c.stepOrder, s.stepNumber, 999)
|
||||
"""
|
||||
|
||||
steps_result = session.run(steps_query, {"recipe_id": recipe_id})
|
||||
steps_info = []
|
||||
for step_record in steps_result:
|
||||
step_text = f"步骤: {step_record['name']}"
|
||||
if step_record.get("description"):
|
||||
step_text += f"\n描述: {step_record['description']}"
|
||||
if step_record.get("methods"):
|
||||
step_text += f"\n方法: {step_record['methods']}"
|
||||
if step_record.get("tools"):
|
||||
step_text += f"\n工具: {step_record['tools']}"
|
||||
if step_record.get("timeEstimate"):
|
||||
step_text += f"\n时间: {step_record['timeEstimate']}"
|
||||
steps_info.append(step_text)
|
||||
|
||||
# 构建完整的菜谱文档内容
|
||||
content_parts = [f"# {recipe_name}"]
|
||||
|
||||
# 添加菜谱基本信息
|
||||
if recipe.properties.get("description"):
|
||||
content_parts.append(f"\n## 菜品描述\n{recipe.properties['description']}")
|
||||
|
||||
if recipe.properties.get("cuisineType"):
|
||||
content_parts.append(f"\n菜系: {recipe.properties['cuisineType']}")
|
||||
|
||||
if recipe.properties.get("difficulty"):
|
||||
content_parts.append(f"难度: {recipe.properties['difficulty']}星")
|
||||
|
||||
if recipe.properties.get("prepTime") or recipe.properties.get("cookTime"):
|
||||
time_info = []
|
||||
if recipe.properties.get("prepTime"):
|
||||
time_info.append(f"准备时间: {recipe.properties['prepTime']}")
|
||||
if recipe.properties.get("cookTime"):
|
||||
time_info.append(f"烹饪时间: {recipe.properties['cookTime']}")
|
||||
content_parts.append(f"\n时间信息: {', '.join(time_info)}")
|
||||
|
||||
if recipe.properties.get("servings"):
|
||||
content_parts.append(f"份量: {recipe.properties['servings']}")
|
||||
|
||||
# 添加食材信息
|
||||
if ingredients_info:
|
||||
content_parts.append("\n## 所需食材")
|
||||
for i, ingredient in enumerate(ingredients_info, 1):
|
||||
content_parts.append(f"{i}. {ingredient}")
|
||||
|
||||
# 添加步骤信息
|
||||
if steps_info:
|
||||
content_parts.append("\n## 制作步骤")
|
||||
for i, step in enumerate(steps_info, 1):
|
||||
content_parts.append(f"\n### 第{i}步\n{step}")
|
||||
|
||||
# 添加标签信息
|
||||
if recipe.properties.get("tags"):
|
||||
content_parts.append(f"\n## 标签\n{recipe.properties['tags']}")
|
||||
|
||||
# 组合成最终内容
|
||||
full_content = "\n".join(content_parts)
|
||||
|
||||
# 创建文档对象
|
||||
doc = Document(
|
||||
page_content=full_content,
|
||||
metadata={
|
||||
"node_id": recipe_id,
|
||||
"recipe_name": recipe_name,
|
||||
"node_type": "Recipe",
|
||||
"category": recipe.properties.get("category", "未知"),
|
||||
"cuisine_type": recipe.properties.get("cuisineType", "未知"),
|
||||
"difficulty": recipe.properties.get("difficulty", 0),
|
||||
"prep_time": recipe.properties.get("prepTime", ""),
|
||||
"cook_time": recipe.properties.get("cookTime", ""),
|
||||
"servings": recipe.properties.get("servings", ""),
|
||||
"ingredients_count": len(ingredients_info),
|
||||
"steps_count": len(steps_info),
|
||||
"doc_type": "recipe",
|
||||
"content_length": len(full_content)
|
||||
}
|
||||
)
|
||||
|
||||
documents.append(doc)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"构建菜谱文档失败 {recipe_name} (ID: {recipe_id}): {e}")
|
||||
continue
|
||||
|
||||
self.documents = documents
|
||||
logger.info(f"成功构建 {len(documents)} 个菜谱文档")
|
||||
return documents
|
||||
|
||||
def chunk_documents(self, chunk_size: int = 500, chunk_overlap: int = 50) -> List[Document]:
|
||||
"""
|
||||
对文档进行分块处理
|
||||
|
||||
Args:
|
||||
chunk_size: 分块大小
|
||||
chunk_overlap: 重叠大小
|
||||
|
||||
Returns:
|
||||
分块后的文档列表
|
||||
"""
|
||||
logger.info(f"正在进行文档分块，块大小: {chunk_size}, 重叠: {chunk_overlap}")
|
||||
|
||||
if not self.documents:
|
||||
raise ValueError("请先构建文档")
|
||||
|
||||
chunks = []
|
||||
chunk_id = 0
|
||||
|
||||
for doc in self.documents:
|
||||
content = doc.page_content
|
||||
|
||||
# 简单的按长度分块
|
||||
if len(content) <= chunk_size:
|
||||
# 内容较短，不需要分块
|
||||
chunk = Document(
|
||||
page_content=content,
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"chunk_id": f"{doc.metadata['node_id']}_chunk_{chunk_id}",
|
||||
"parent_id": doc.metadata["node_id"],
|
||||
"chunk_index": 0,
|
||||
"total_chunks": 1,
|
||||
"chunk_size": len(content),
|
||||
"doc_type": "chunk"
|
||||
}
|
||||
)
|
||||
chunks.append(chunk)
|
||||
chunk_id += 1
|
||||
else:
|
||||
# 按章节分块（基于标题）
|
||||
sections = content.split('\n## ')
|
||||
if len(sections) <= 1:
|
||||
# 没有二级标题，按长度强制分块
|
||||
total_chunks = (len(content) - 1) // (chunk_size - chunk_overlap) + 1
|
||||
|
||||
for i in range(total_chunks):
|
||||
start = i * (chunk_size - chunk_overlap)
|
||||
end = min(start + chunk_size, len(content))
|
||||
|
||||
chunk_content = content[start:end]
|
||||
|
||||
chunk = Document(
|
||||
page_content=chunk_content,
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"chunk_id": f"{doc.metadata['node_id']}_chunk_{chunk_id}",
|
||||
"parent_id": doc.metadata["node_id"],
|
||||
"chunk_index": i,
|
||||
"total_chunks": total_chunks,
|
||||
"chunk_size": len(chunk_content),
|
||||
"doc_type": "chunk"
|
||||
}
|
||||
)
|
||||
chunks.append(chunk)
|
||||
chunk_id += 1
|
||||
else:
|
||||
# 按章节分块
|
||||
total_chunks = len(sections)
|
||||
for i, section in enumerate(sections):
|
||||
if i == 0:
|
||||
# 第一个部分包含标题
|
||||
chunk_content = section
|
||||
else:
|
||||
# 其他部分添加章节标题
|
||||
chunk_content = f"## {section}"
|
||||
|
||||
chunk = Document(
|
||||
page_content=chunk_content,
|
||||
metadata={
|
||||
**doc.metadata,
|
||||
"chunk_id": f"{doc.metadata['node_id']}_chunk_{chunk_id}",
|
||||
"parent_id": doc.metadata["node_id"],
|
||||
"chunk_index": i,
|
||||
"total_chunks": total_chunks,
|
||||
"chunk_size": len(chunk_content),
|
||||
"doc_type": "chunk",
|
||||
"section_title": section.split('\n')[0] if i > 0 else "主标题"
|
||||
}
|
||||
)
|
||||
chunks.append(chunk)
|
||||
chunk_id += 1
|
||||
|
||||
self.chunks = chunks
|
||||
logger.info(f"文档分块完成，共生成 {len(chunks)} 个块")
|
||||
return chunks
|
||||
|
||||
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""
|
||||
获取数据统计信息
|
||||
|
||||
Returns:
|
||||
统计信息字典
|
||||
"""
|
||||
stats = {
|
||||
'total_recipes': len(self.recipes),
|
||||
'total_ingredients': len(self.ingredients),
|
||||
'total_cooking_steps': len(self.cooking_steps),
|
||||
'total_documents': len(self.documents),
|
||||
'total_chunks': len(self.chunks)
|
||||
}
|
||||
|
||||
if self.documents:
|
||||
# 分类统计
|
||||
categories = {}
|
||||
cuisines = {}
|
||||
difficulties = {}
|
||||
|
||||
for doc in self.documents:
|
||||
category = doc.metadata.get('category', '未知')
|
||||
categories[category] = categories.get(category, 0) + 1
|
||||
|
||||
cuisine = doc.metadata.get('cuisine_type', '未知')
|
||||
cuisines[cuisine] = cuisines.get(cuisine, 0) + 1
|
||||
|
||||
difficulty = doc.metadata.get('difficulty', 0)
|
||||
difficulties[str(difficulty)] = difficulties.get(str(difficulty), 0) + 1
|
||||
|
||||
stats.update({
|
||||
'categories': categories,
|
||||
'cuisines': cuisines,
|
||||
'difficulties': difficulties,
|
||||
'avg_content_length': sum(doc.metadata.get('content_length', 0) for doc in self.documents) / len(self.documents),
|
||||
'avg_chunk_size': sum(chunk.metadata.get('chunk_size', 0) for chunk in self.chunks) / len(self.chunks) if self.chunks else 0
|
||||
})
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
|
||||
def __del__(self):
|
||||
"""析构函数，确保关闭连接"""
|
||||
self.close()
|
||||
Reference in New Issue
Block a user