522 lines
17 KiB
Python
522 lines
17 KiB
Python
"""工具专项测试 - 文档提取和知识图谱工具
|
|
|
|
本模块测试基于数据基底的工具功能，包括：
|
|
1. 文档读取和解析工具
|
|
2. 实体提取工具
|
|
3. 知识图谱查询工具
|
|
4. 推理工具
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
from agentlite import Agent, Message, TextPart, tool
|
|
|
|
|
|
def tool_output(result: Any) -> Any:
|
|
"""兼容旧式返回值和 ToolResult 返回值."""
|
|
return getattr(result, "output", result)
|
|
|
|
|
|
# =============================================================================
|
|
# 数据加载 fixtures
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.fixture
|
|
def data_dir() -> Path:
|
|
"""返回测试数据目录路径."""
|
|
return Path(__file__).parent.parent / "data"
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_article(data_dir: Path) -> str:
|
|
"""加载样例文章."""
|
|
return (data_dir / "documents" / "sample_article.md").read_text(encoding="utf-8")
|
|
|
|
|
|
@pytest.fixture
|
|
def technical_spec(data_dir: Path) -> str:
|
|
"""加载技术规范文档."""
|
|
return (data_dir / "documents" / "technical_spec.md").read_text(encoding="utf-8")
|
|
|
|
|
|
@pytest.fixture
|
|
def meeting_notes(data_dir: Path) -> str:
|
|
"""加载会议记录."""
|
|
return (data_dir / "documents" / "meeting_notes.txt").read_text(encoding="utf-8")
|
|
|
|
|
|
@pytest.fixture
|
|
def knowledge_graph_entities(data_dir: Path) -> dict[str, Any]:
|
|
"""加载知识图谱实体数据."""
|
|
with open(data_dir / "knowledge_base" / "entities.json") as f:
|
|
return json.load(f)
|
|
|
|
|
|
@pytest.fixture
|
|
def knowledge_graph_relations(data_dir: Path) -> dict[str, Any]:
|
|
"""加载知识图谱关系数据."""
|
|
with open(data_dir / "knowledge_base" / "relations.json") as f:
|
|
return json.load(f)
|
|
|
|
|
|
@pytest.fixture
|
|
def graph_queries(data_dir: Path) -> list[dict[str, Any]]:
|
|
"""加载图谱查询测试用例."""
|
|
with open(data_dir / "knowledge_base" / "graph_queries.yaml") as f:
|
|
data = yaml.safe_load(f)
|
|
return data.get("queries", [])
|
|
|
|
|
|
# =============================================================================
|
|
# 知识图谱工具实现
|
|
# =============================================================================
|
|
|
|
|
|
class KnowledgeGraph:
|
|
"""知识图谱内存存储."""
|
|
|
|
def __init__(self, entities: list[dict], relations: list[dict]):
|
|
self._entities = {e["id"]: e for e in entities}
|
|
self._relations = relations
|
|
self._index_by_type: dict[str, list[str]] = {}
|
|
self._index_by_name: dict[str, str] = {}
|
|
|
|
# 构建索引
|
|
for entity_id, entity in self._entities.items():
|
|
entity_type = entity.get("type", "Unknown")
|
|
if entity_type not in self._index_by_type:
|
|
self._index_by_type[entity_type] = []
|
|
self._index_by_type[entity_type].append(entity_id)
|
|
|
|
name = entity.get("name", "")
|
|
if name:
|
|
self._index_by_name[name] = entity_id
|
|
|
|
def get_entity(self, entity_id: str) -> dict | None:
|
|
"""获取实体."""
|
|
return self._entities.get(entity_id)
|
|
|
|
def get_entity_by_name(self, name: str) -> dict | None:
|
|
"""通过名称获取实体."""
|
|
entity_id = self._index_by_name.get(name)
|
|
if entity_id:
|
|
return self._entities.get(entity_id)
|
|
return None
|
|
|
|
def get_entities_by_type(self, entity_type: str) -> list[dict]:
|
|
"""获取特定类型的所有实体."""
|
|
entity_ids = self._index_by_type.get(entity_type, [])
|
|
return [self._entities[eid] for eid in entity_ids if eid in self._entities]
|
|
|
|
def get_relations(
|
|
self, from_id: str | None = None, to_id: str | None = None, relation_type: str | None = None
|
|
) -> list[dict]:
|
|
"""获取关系."""
|
|
results = []
|
|
for rel in self._relations:
|
|
if from_id and rel.get("from") != from_id:
|
|
continue
|
|
if to_id and rel.get("to") != to_id:
|
|
continue
|
|
if relation_type and rel.get("type") != relation_type:
|
|
continue
|
|
results.append(rel)
|
|
return results
|
|
|
|
def get_neighbors(self, entity_id: str, relation_type: str | None = None) -> list[dict]:
|
|
"""获取邻居实体."""
|
|
relations = self.get_relations(from_id=entity_id, relation_type=relation_type)
|
|
neighbors = []
|
|
for rel in relations:
|
|
target_id = rel.get("to")
|
|
if target_id and target_id in self._entities:
|
|
neighbors.append({"entity": self._entities[target_id], "relation": rel})
|
|
return neighbors
|
|
|
|
def find_path(self, start_id: str, end_id: str, max_depth: int = 3) -> list[list[str]] | None:
|
|
"""查找两个实体之间的路径."""
|
|
if start_id == end_id:
|
|
return [[start_id]]
|
|
|
|
if max_depth <= 0:
|
|
return None
|
|
|
|
# BFS
|
|
from collections import deque
|
|
|
|
queue = deque([(start_id, [start_id])])
|
|
visited = {start_id}
|
|
all_paths = []
|
|
|
|
while queue:
|
|
current_id, path = queue.popleft()
|
|
|
|
if len(path) > max_depth + 1:
|
|
continue
|
|
|
|
relations = self.get_relations(from_id=current_id)
|
|
for rel in relations:
|
|
next_id = rel.get("to")
|
|
if not next_id:
|
|
continue
|
|
|
|
new_path = path + [next_id]
|
|
|
|
if next_id == end_id:
|
|
all_paths.append(new_path)
|
|
elif next_id not in visited and len(new_path) <= max_depth:
|
|
visited.add(next_id)
|
|
queue.append((next_id, new_path))
|
|
|
|
return all_paths if all_paths else None
|
|
|
|
|
|
@pytest.fixture
|
|
def knowledge_graph(knowledge_graph_entities, knowledge_graph_relations) -> KnowledgeGraph:
|
|
"""创建知识图谱实例."""
|
|
return KnowledgeGraph(
|
|
entities=knowledge_graph_entities.get("entities", []),
|
|
relations=knowledge_graph_relations.get("relations", []),
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# 工具定义
|
|
# =============================================================================
|
|
|
|
|
|
@tool()
|
|
async def read_document(file_path: str) -> str:
|
|
"""读取文档内容.
|
|
|
|
Args:
|
|
file_path: 文档路径
|
|
|
|
Returns:
|
|
文档内容
|
|
"""
|
|
path = Path(file_path)
|
|
if not path.exists():
|
|
return f"Error: File not found: {file_path}"
|
|
try:
|
|
return path.read_text(encoding="utf-8")
|
|
except Exception as e:
|
|
return f"Error reading file: {e}"
|
|
|
|
|
|
@tool()
|
|
async def extract_entities(text: str) -> str:
|
|
"""从文本中提取实体.
|
|
|
|
Args:
|
|
text: 输入文本
|
|
|
|
Returns:
|
|
JSON 格式的实体列表
|
|
"""
|
|
# 简化的实体提取 - 实际应使用 NLP 模型
|
|
import re
|
|
|
|
entities = []
|
|
|
|
# 提取人名（简单的中文姓名匹配）
|
|
person_pattern = r"[\u4e00-\u9fa5]{2,4}"
|
|
potential_names = re.findall(person_pattern, text)
|
|
common_names = ["张三", "李四", "王五", "赵六", "李飞飞", "吴恩达", "Yann LeCun"]
|
|
|
|
for name in potential_names:
|
|
if name in common_names or len(name) == 3:
|
|
entities.append({"type": "Person", "name": name})
|
|
|
|
# 提取公司/组织名
|
|
org_pattern = r"(TechCorp|OpenAI|GitHub|Google)"
|
|
orgs = re.findall(org_pattern, text)
|
|
for org in set(orgs):
|
|
entities.append({"type": "Organization", "name": org})
|
|
|
|
# 提取技术术语
|
|
tech_pattern = r"(Python|TensorFlow|PyTorch|GPT-4|AI|LLM)"
|
|
techs = re.findall(tech_pattern, text)
|
|
for tech in set(techs):
|
|
entities.append({"type": "Technology", "name": tech})
|
|
|
|
return json.dumps(entities, ensure_ascii=False)
|
|
|
|
|
|
@tool()
|
|
async def query_knowledge_graph(query_type: str, params: str) -> str:
|
|
"""查询知识图谱.
|
|
|
|
Args:
|
|
query_type: 查询类型 (person_relations, company_employees, technology_users, etc.)
|
|
params: JSON 格式的查询参数
|
|
|
|
Returns:
|
|
查询结果
|
|
"""
|
|
# 这里使用全局的 kg 实例，实际应在 Agent 初始化时注入
|
|
try:
|
|
params_dict = json.loads(params)
|
|
except json.JSONDecodeError:
|
|
return json.dumps({"error": "Invalid JSON params"})
|
|
|
|
# 简化实现 - 实际应基于知识图谱查询
|
|
result = {"query_type": query_type, "params": params_dict, "results": []}
|
|
|
|
return json.dumps(result, ensure_ascii=False)
|
|
|
|
|
|
@tool()
|
|
async def reason_about_path(start_entity: str, end_entity: str) -> str:
|
|
"""推理两个实体之间的关系路径.
|
|
|
|
Args:
|
|
start_entity: 起始实体名称
|
|
end_entity: 目标实体名称
|
|
|
|
Returns:
|
|
推理结果
|
|
"""
|
|
return json.dumps(
|
|
{
|
|
"start": start_entity,
|
|
"end": end_entity,
|
|
"reasoning": f"分析 {start_entity} 到 {end_entity} 的关系链...",
|
|
"path": [],
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# 测试用例
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.mark.tools
|
|
class TestDocumentTools:
|
|
"""文档工具测试."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_document(self, data_dir: Path, sample_article: str):
|
|
"""测试文档读取工具."""
|
|
result = tool_output(await read_document(str(data_dir / "documents" / "sample_article.md")))
|
|
|
|
assert "人工智能" in result
|
|
assert "GitHub Copilot" in result
|
|
assert "张三" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_document_not_found(self):
|
|
"""测试读取不存在的文档."""
|
|
result = tool_output(await read_document("/nonexistent/file.md"))
|
|
|
|
assert "Error" in result
|
|
assert "not found" in result.lower()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_entities_from_article(self, sample_article: str):
|
|
"""测试从文章中提取实体."""
|
|
result = tool_output(await extract_entities(sample_article))
|
|
entities = json.loads(result)
|
|
|
|
# 验证提取到实体
|
|
assert len(entities) > 0
|
|
|
|
# 验证实体类型
|
|
entity_names = [e["name"] for e in entities]
|
|
assert "张三" in entity_names
|
|
assert "TechCorp" in entity_names or "OpenAI" in entity_names
|
|
|
|
|
|
@pytest.mark.tools
|
|
class TestKnowledgeGraphTools:
|
|
"""知识图谱工具测试."""
|
|
|
|
def test_knowledge_graph_initialization(self, knowledge_graph: KnowledgeGraph):
|
|
"""测试知识图谱初始化."""
|
|
# 验证实体数量
|
|
entity = knowledge_graph.get_entity_by_name("张三")
|
|
assert entity is not None
|
|
assert entity["type"] == "Person"
|
|
|
|
# 验证公司实体
|
|
company = knowledge_graph.get_entity_by_name("TechCorp")
|
|
assert company is not None
|
|
assert company["type"] == "Company"
|
|
|
|
def test_get_entities_by_type(self, knowledge_graph: KnowledgeGraph):
|
|
"""测试按类型获取实体."""
|
|
persons = knowledge_graph.get_entities_by_type("Person")
|
|
assert len(persons) >= 3 # 张三、李四、李飞飞
|
|
|
|
technologies = knowledge_graph.get_entities_by_type("Technology")
|
|
assert len(technologies) >= 2 # Python、OpenAI API
|
|
|
|
def test_get_relations(self, knowledge_graph: KnowledgeGraph):
|
|
"""测试获取关系."""
|
|
# 获取张三的所有关系
|
|
zhangsan = knowledge_graph.get_entity_by_name("张三")
|
|
assert zhangsan is not None
|
|
|
|
relations = knowledge_graph.get_relations(from_id=zhangsan["id"])
|
|
assert len(relations) >= 2 # works_for, uses
|
|
|
|
# 验证关系类型
|
|
relation_types = [r["type"] for r in relations]
|
|
assert "works_for" in relation_types
|
|
assert "uses" in relation_types
|
|
|
|
def test_get_neighbors(self, knowledge_graph: KnowledgeGraph):
|
|
"""测试获取邻居节点."""
|
|
zhangsan = knowledge_graph.get_entity_by_name("张三")
|
|
assert zhangsan is not None
|
|
|
|
neighbors = knowledge_graph.get_neighbors(zhangsan["id"])
|
|
assert len(neighbors) >= 2
|
|
|
|
# 验证邻居包含 TechCorp
|
|
neighbor_names = [n["entity"]["name"] for n in neighbors]
|
|
assert "TechCorp" in neighbor_names
|
|
|
|
def test_find_path(self, knowledge_graph: KnowledgeGraph):
|
|
"""测试查找路径."""
|
|
zhangsan = knowledge_graph.get_entity_by_name("张三")
|
|
techcorp = knowledge_graph.get_entity_by_name("TechCorp")
|
|
|
|
assert zhangsan is not None
|
|
assert techcorp is not None
|
|
|
|
paths = knowledge_graph.find_path(zhangsan["id"], techcorp["id"])
|
|
assert paths is not None
|
|
assert len(paths) > 0
|
|
|
|
# 验证路径长度
|
|
first_path = paths[0]
|
|
assert len(first_path) == 2 # 张三 -> TechCorp
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_knowledge_graph(self):
|
|
"""测试知识图谱查询工具."""
|
|
params = json.dumps({"entity_name": "张三"})
|
|
result = tool_output(await query_knowledge_graph("person_relations", params))
|
|
|
|
data = json.loads(result)
|
|
assert data["query_type"] == "person_relations"
|
|
assert "params" in data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reason_about_path(self):
|
|
"""测试路径推理工具."""
|
|
result = tool_output(await reason_about_path("张三", "OpenAI"))
|
|
|
|
data = json.loads(result)
|
|
assert data["start"] == "张三"
|
|
assert data["end"] == "OpenAI"
|
|
assert "reasoning" in data
|
|
|
|
|
|
@pytest.mark.tools
|
|
class TestDataIntegrity:
|
|
"""数据完整性测试."""
|
|
|
|
def test_entities_json_valid(self, knowledge_graph_entities: dict):
|
|
"""验证实体 JSON 格式正确."""
|
|
assert "entities" in knowledge_graph_entities
|
|
entities = knowledge_graph_entities["entities"]
|
|
assert len(entities) > 0
|
|
|
|
# 验证每个实体都有必需的字段
|
|
for entity in entities:
|
|
assert "id" in entity
|
|
assert "type" in entity
|
|
assert "name" in entity
|
|
|
|
def test_relations_json_valid(
|
|
self, knowledge_graph_relations: dict, knowledge_graph_entities: dict
|
|
):
|
|
"""验证关系 JSON 格式正确且引用的实体存在."""
|
|
assert "relations" in knowledge_graph_relations
|
|
relations = knowledge_graph_relations["relations"]
|
|
|
|
entity_ids = {e["id"] for e in knowledge_graph_entities["entities"]}
|
|
|
|
for relation in relations:
|
|
assert "from" in relation
|
|
assert "to" in relation
|
|
assert "type" in relation
|
|
|
|
# 验证引用的实体存在
|
|
assert relation["from"] in entity_ids, f"Entity {relation['from']} not found"
|
|
assert relation["to"] in entity_ids, f"Entity {relation['to']} not found"
|
|
|
|
def test_graph_queries_yaml_valid(self, graph_queries: list):
|
|
"""验证查询 YAML 格式正确."""
|
|
assert len(graph_queries) > 0
|
|
|
|
for query in graph_queries:
|
|
assert "id" in query
|
|
assert "description" in query
|
|
assert "query" in query
|
|
assert "expected_results" in query
|
|
|
|
def test_documents_exist(self, data_dir: Path):
|
|
"""验证测试文档存在且非空."""
|
|
docs_dir = data_dir / "documents"
|
|
|
|
sample_article = docs_dir / "sample_article.md"
|
|
assert sample_article.exists()
|
|
assert sample_article.stat().st_size > 0
|
|
|
|
tech_spec = docs_dir / "technical_spec.md"
|
|
assert tech_spec.exists()
|
|
assert tech_spec.stat().st_size > 0
|
|
|
|
meeting_notes = docs_dir / "meeting_notes.txt"
|
|
assert meeting_notes.exists()
|
|
assert meeting_notes.stat().st_size > 0
|
|
|
|
|
|
@pytest.mark.tools
|
|
class TestAgentWithTools:
|
|
"""Agent 集成工具测试."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_with_document_tools(self, mock_provider, data_dir: Path):
|
|
"""测试带有文档工具的 Agent."""
|
|
mock_provider.add_text_response("我已经读取了文档")
|
|
|
|
agent = Agent(
|
|
provider=mock_provider,
|
|
tools=[read_document],
|
|
system_prompt="你是一个文档助手，可以读取和分析文档。",
|
|
)
|
|
|
|
response = await agent.run(f"请读取文档 {data_dir / 'documents' / 'sample_article.md'}")
|
|
|
|
assert "文档" in response or "读取" in response
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_with_kg_tools(self, mock_provider):
|
|
"""测试带有知识图谱工具的 Agent."""
|
|
mock_provider.add_text_response("张三在 TechCorp 工作")
|
|
|
|
agent = Agent(
|
|
provider=mock_provider,
|
|
tools=[query_knowledge_graph, reason_about_path],
|
|
system_prompt="你是一个知识图谱助手，可以查询实体关系。",
|
|
)
|
|
|
|
response = await agent.run("张三在哪里工作？")
|
|
|
|
assert response is not None
|
|
assert len(response) > 0
|