Filter simulation personas to agent-eligible entities
This commit is contained in:
parent
b5b53acc57
commit
eff1eba1ba
|
|
@ -509,10 +509,17 @@ def prepare_simulation():
|
|||
defined_entity_types=entity_types_list,
|
||||
enrich_with_edges=False # 不获取边信息,加快速度
|
||||
)
|
||||
agent_entities_preview = OasisProfileGenerator.filter_agent_persona_entities(
|
||||
filtered_preview.entities,
|
||||
allow_group_accounts=False
|
||||
)
|
||||
# 保存实体数量到状态(供前端立即获取)
|
||||
state.entities_count = filtered_preview.filtered_count
|
||||
state.entity_types = list(filtered_preview.entity_types)
|
||||
logger.info(f"预期实体数量: {filtered_preview.filtered_count}, 类型: {filtered_preview.entity_types}")
|
||||
state.entities_count = len(agent_entities_preview)
|
||||
state.entity_types = sorted({
|
||||
entity.get_entity_type() or "Unknown"
|
||||
for entity in agent_entities_preview
|
||||
})
|
||||
logger.info(f"预期Agent数量: {len(agent_entities_preview)}, 类型: {state.entity_types}")
|
||||
except Exception as e:
|
||||
logger.warning(f"同步获取实体数量失败(将在后台任务中重试): {e}")
|
||||
# 失败不影响后续流程,后台任务会重新获取
|
||||
|
|
@ -1434,7 +1441,8 @@ def generate_profiles():
|
|||
"graph_id": "mirofish_xxxx", // 必填
|
||||
"entity_types": ["Student"], // 可选
|
||||
"use_llm": true, // 可选
|
||||
"platform": "reddit" // 可选
|
||||
"platform": "reddit", // 可选
|
||||
"allow_group_accounts": false // 可选,是否允许机构/群体账号
|
||||
}
|
||||
"""
|
||||
try:
|
||||
|
|
@ -1465,10 +1473,18 @@ def generate_profiles():
|
|||
}), 400
|
||||
|
||||
generator = OasisProfileGenerator()
|
||||
allow_group_accounts = data.get('allow_group_accounts', False)
|
||||
profiles = generator.generate_profiles_from_entities(
|
||||
entities=filtered.entities,
|
||||
use_llm=use_llm
|
||||
use_llm=use_llm,
|
||||
allow_group_accounts=allow_group_accounts
|
||||
)
|
||||
|
||||
if not profiles:
|
||||
return jsonify({
|
||||
"success": False,
|
||||
"error": t('api.noMatchingEntities')
|
||||
}), 400
|
||||
|
||||
if platform == "reddit":
|
||||
profiles_data = [p.to_reddit_format() for p in profiles]
|
||||
|
|
|
|||
|
|
@ -239,6 +239,57 @@ class OasisProfileGenerator:
|
|||
"university", "governmentagency", "organization", "ngo",
|
||||
"mediaoutlet", "company", "institution", "group", "community"
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def _normalize_entity_type(cls, entity_type: Optional[str]) -> str:
|
||||
if not entity_type:
|
||||
return ""
|
||||
return "".join(ch for ch in entity_type.lower() if ch.isalnum())
|
||||
|
||||
@classmethod
|
||||
def is_agent_persona_entity(
|
||||
cls,
|
||||
entity: EntityNode,
|
||||
allow_group_accounts: bool = False
|
||||
) -> bool:
|
||||
"""Return whether a graph entity should be promoted into an agent."""
|
||||
allowed_types = {
|
||||
cls._normalize_entity_type(entity_type)
|
||||
for entity_type in cls.INDIVIDUAL_ENTITY_TYPES
|
||||
}
|
||||
if allow_group_accounts:
|
||||
allowed_types.update(
|
||||
cls._normalize_entity_type(entity_type)
|
||||
for entity_type in cls.GROUP_ENTITY_TYPES
|
||||
)
|
||||
|
||||
labels = [
|
||||
label for label in entity.labels
|
||||
if label not in ["Entity", "Node"]
|
||||
]
|
||||
entity_type = entity.get_entity_type()
|
||||
if entity_type:
|
||||
labels.insert(0, entity_type)
|
||||
|
||||
return any(
|
||||
cls._normalize_entity_type(label) in allowed_types
|
||||
for label in labels
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def filter_agent_persona_entities(
|
||||
cls,
|
||||
entities: List[EntityNode],
|
||||
allow_group_accounts: bool = False
|
||||
) -> List[EntityNode]:
|
||||
"""Filter graph entities down to nodes that can act as simulation agents."""
|
||||
return [
|
||||
entity for entity in entities
|
||||
if cls.is_agent_persona_entity(
|
||||
entity,
|
||||
allow_group_accounts=allow_group_accounts
|
||||
)
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -900,7 +951,8 @@ class OasisProfileGenerator:
|
|||
graph_id: Optional[str] = None,
|
||||
parallel_count: int = 5,
|
||||
realtime_output_path: Optional[str] = None,
|
||||
output_platform: str = "reddit"
|
||||
output_platform: str = "reddit",
|
||||
allow_group_accounts: bool = True
|
||||
) -> List[OasisAgentProfile]:
|
||||
"""
|
||||
批量从实体生成Agent Profile(支持并行生成)
|
||||
|
|
@ -913,6 +965,7 @@ class OasisProfileGenerator:
|
|||
parallel_count: 并行生成数量,默认5
|
||||
realtime_output_path: 实时写入的文件路径(如果提供,每生成一个就写入一次)
|
||||
output_platform: 输出平台格式 ("reddit" 或 "twitter")
|
||||
allow_group_accounts: 是否允许机构/群体账号成为Agent
|
||||
|
||||
Returns:
|
||||
Agent Profile列表
|
||||
|
|
@ -923,6 +976,19 @@ class OasisProfileGenerator:
|
|||
# 设置graph_id用于Zep检索
|
||||
if graph_id:
|
||||
self.graph_id = graph_id
|
||||
|
||||
original_count = len(entities)
|
||||
entities = self.filter_agent_persona_entities(
|
||||
entities,
|
||||
allow_group_accounts=allow_group_accounts
|
||||
)
|
||||
if len(entities) != original_count:
|
||||
logger.info(
|
||||
"Agent人设实体过滤: 原始 %s, 可用 %s, 跳过 %s",
|
||||
original_count,
|
||||
len(entities),
|
||||
original_count - len(entities),
|
||||
)
|
||||
|
||||
total = len(entities)
|
||||
profiles = [None] * total # 预分配列表保持顺序
|
||||
|
|
|
|||
|
|
@ -301,26 +301,40 @@ class SimulationManager:
|
|||
defined_entity_types=defined_entity_types,
|
||||
enrich_with_edges=True
|
||||
)
|
||||
|
||||
state.entities_count = filtered.filtered_count
|
||||
state.entity_types = list(filtered.entity_types)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(
|
||||
"reading", 100,
|
||||
t('progress.readingComplete', count=filtered.filtered_count),
|
||||
current=filtered.filtered_count,
|
||||
total=filtered.filtered_count
|
||||
)
|
||||
|
||||
|
||||
if filtered.filtered_count == 0:
|
||||
state.status = SimulationStatus.FAILED
|
||||
state.error = "没有找到符合条件的实体,请检查图谱是否正确构建"
|
||||
self._save_simulation_state(state)
|
||||
raise ValueError(state.error)
|
||||
|
||||
agent_entities = OasisProfileGenerator.filter_agent_persona_entities(
|
||||
filtered.entities,
|
||||
allow_group_accounts=False
|
||||
)
|
||||
|
||||
state.entities_count = len(agent_entities)
|
||||
state.entity_types = sorted({
|
||||
entity.get_entity_type() or "Unknown"
|
||||
for entity in agent_entities
|
||||
})
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(
|
||||
"reading", 100,
|
||||
t('progress.readingComplete', count=len(agent_entities)),
|
||||
current=len(agent_entities),
|
||||
total=len(agent_entities)
|
||||
)
|
||||
|
||||
if not agent_entities:
|
||||
state.status = SimulationStatus.FAILED
|
||||
state.error = "没有找到可用于生成人设的个人实体,请检查图谱是否包含Person类实体"
|
||||
self._save_simulation_state(state)
|
||||
raise ValueError(state.error)
|
||||
|
||||
# ========== 阶段2: 生成Agent Profile ==========
|
||||
total_entities = len(filtered.entities)
|
||||
total_entities = len(agent_entities)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(
|
||||
|
|
@ -355,13 +369,14 @@ class SimulationManager:
|
|||
realtime_platform = "twitter"
|
||||
|
||||
profiles = generator.generate_profiles_from_entities(
|
||||
entities=filtered.entities,
|
||||
entities=agent_entities,
|
||||
use_llm=use_llm_for_profiles,
|
||||
progress_callback=profile_progress,
|
||||
graph_id=state.graph_id, # 传入graph_id用于Zep检索
|
||||
parallel_count=parallel_profile_count, # 并行生成数量
|
||||
realtime_output_path=realtime_output_path, # 实时保存路径
|
||||
output_platform=realtime_platform # 输出格式
|
||||
output_platform=realtime_platform, # 输出格式
|
||||
allow_group_accounts=False
|
||||
)
|
||||
|
||||
state.profiles_count = len(profiles)
|
||||
|
|
@ -426,7 +441,7 @@ class SimulationManager:
|
|||
graph_id=state.graph_id,
|
||||
simulation_requirement=simulation_requirement,
|
||||
document_text=document_text,
|
||||
entities=filtered.entities,
|
||||
entities=agent_entities,
|
||||
enable_twitter=state.enable_twitter,
|
||||
enable_reddit=state.enable_reddit
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,77 @@
|
|||
from app.services.oasis_profile_generator import OasisProfileGenerator
|
||||
from app.services.zep_entity_reader import EntityNode
|
||||
|
||||
|
||||
def _entity(name, labels):
|
||||
return EntityNode(
|
||||
uuid=f"uuid-{name}",
|
||||
name=name,
|
||||
labels=labels,
|
||||
summary=f"{name} summary",
|
||||
attributes={},
|
||||
)
|
||||
|
||||
|
||||
def test_agent_persona_filter_keeps_only_person_entities_by_default():
|
||||
entities = [
|
||||
_entity("Alice", ["Entity", "Person"]),
|
||||
_entity("MiroFish", ["Entity", "Company"]),
|
||||
_entity("market event", ["Entity", "Topic"]),
|
||||
_entity("Bob", ["Entity", "PublicFigure"]),
|
||||
_entity("raw fragment", ["Entity"]),
|
||||
]
|
||||
|
||||
filtered = OasisProfileGenerator.filter_agent_persona_entities(entities)
|
||||
|
||||
assert [entity.name for entity in filtered] == ["Alice", "Bob"]
|
||||
|
||||
|
||||
def test_agent_persona_filter_can_keep_group_accounts_when_requested():
|
||||
entities = [
|
||||
_entity("Alice", ["Entity", "Person"]),
|
||||
_entity("MiroFish", ["Entity", "Company"]),
|
||||
_entity("market event", ["Entity", "Topic"]),
|
||||
]
|
||||
|
||||
filtered = OasisProfileGenerator.filter_agent_persona_entities(
|
||||
entities,
|
||||
allow_group_accounts=True,
|
||||
)
|
||||
|
||||
assert [entity.name for entity in filtered] == ["Alice", "MiroFish"]
|
||||
|
||||
|
||||
def test_generate_profiles_from_entities_skips_non_person_entities(monkeypatch):
|
||||
entities = [
|
||||
_entity("Alice", ["Entity", "Person"]),
|
||||
_entity("MiroFish", ["Entity", "Company"]),
|
||||
_entity("market event", ["Entity", "Topic"]),
|
||||
_entity("Bob", ["Entity", "Person"]),
|
||||
]
|
||||
generator = object.__new__(OasisProfileGenerator)
|
||||
generator.graph_id = None
|
||||
|
||||
monkeypatch.setattr(generator, "_print_generated_profile", lambda *args: None)
|
||||
|
||||
def generate_profile(entity, user_id, use_llm):
|
||||
return type(
|
||||
"Profile",
|
||||
(),
|
||||
{
|
||||
"name": entity.name,
|
||||
"user_id": user_id,
|
||||
"to_reddit_format": lambda self: {"name": self.name},
|
||||
},
|
||||
)()
|
||||
|
||||
monkeypatch.setattr(generator, "generate_profile_from_entity", generate_profile)
|
||||
|
||||
profiles = generator.generate_profiles_from_entities(
|
||||
entities,
|
||||
use_llm=False,
|
||||
parallel_count=1,
|
||||
allow_group_accounts=False,
|
||||
)
|
||||
|
||||
assert [profile.name for profile in profiles] == ["Alice", "Bob"]
|
||||
assert [profile.user_id for profile in profiles] == [0, 1]
|
||||
|
|
@ -5,12 +5,13 @@ import pytest
|
|||
from app import create_app
|
||||
from app.config import Config
|
||||
from app.services import simulation_manager as simulation_manager_module
|
||||
from app.services.oasis_profile_generator import OasisProfileGenerator
|
||||
from app.services.simulation_manager import (
|
||||
SimulationManager,
|
||||
SimulationState,
|
||||
SimulationStatus,
|
||||
)
|
||||
from app.services.zep_entity_reader import FilteredEntities
|
||||
from app.services.zep_entity_reader import EntityNode, FilteredEntities
|
||||
|
||||
|
||||
def _write_failed_state(root, simulation_id="sim_failed"):
|
||||
|
|
@ -100,3 +101,96 @@ def test_zero_entities_persists_failed_state_and_raises(tmp_path, monkeypatch):
|
|||
assert persisted["config_generated"] is False
|
||||
assert persisted["config_reasoning"] == ""
|
||||
assert "没有找到符合条件的实体" in persisted["error"]
|
||||
|
||||
|
||||
def test_prepare_uses_persona_filtered_entities_for_profiles_and_config(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
entities = [
|
||||
EntityNode(
|
||||
uuid="person-1",
|
||||
name="Alice",
|
||||
labels=["Entity", "Person"],
|
||||
summary="A person",
|
||||
attributes={},
|
||||
),
|
||||
EntityNode(
|
||||
uuid="company-1",
|
||||
name="MiroFish",
|
||||
labels=["Entity", "Company"],
|
||||
summary="The project",
|
||||
attributes={},
|
||||
),
|
||||
EntityNode(
|
||||
uuid="topic-1",
|
||||
name="market event",
|
||||
labels=["Entity", "Topic"],
|
||||
summary="A discussion topic",
|
||||
attributes={},
|
||||
),
|
||||
]
|
||||
captured = {}
|
||||
|
||||
class Reader:
|
||||
def filter_defined_entities(self, **kwargs):
|
||||
return FilteredEntities(
|
||||
entities=entities,
|
||||
entity_types={"Person", "Company", "Topic"},
|
||||
total_count=3,
|
||||
filtered_count=3,
|
||||
)
|
||||
|
||||
class ConfigGenerator:
|
||||
def generate_config(self, **kwargs):
|
||||
captured["config_entities"] = kwargs["entities"]
|
||||
|
||||
class Params:
|
||||
generation_reasoning = "ok"
|
||||
|
||||
def to_json(self):
|
||||
return "{}"
|
||||
|
||||
return Params()
|
||||
|
||||
def generate_profiles(self, **kwargs):
|
||||
captured["profile_entities"] = kwargs["entities"]
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(SimulationManager, "SIMULATION_DATA_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(simulation_manager_module, "ZepEntityReader", Reader)
|
||||
monkeypatch.setattr(
|
||||
simulation_manager_module,
|
||||
"SimulationConfigGenerator",
|
||||
lambda: ConfigGenerator(),
|
||||
)
|
||||
monkeypatch.setattr(OasisProfileGenerator, "__init__", lambda self, **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
OasisProfileGenerator,
|
||||
"generate_profiles_from_entities",
|
||||
generate_profiles,
|
||||
)
|
||||
monkeypatch.setattr(OasisProfileGenerator, "save_profiles", lambda *args, **kwargs: None)
|
||||
|
||||
manager = SimulationManager()
|
||||
state = SimulationState(
|
||||
simulation_id="sim_filtered",
|
||||
project_id="project",
|
||||
graph_id="graph",
|
||||
status=SimulationStatus.CREATED,
|
||||
enable_reddit=False,
|
||||
enable_twitter=False,
|
||||
)
|
||||
manager._save_simulation_state(state)
|
||||
|
||||
result = manager.prepare_simulation(
|
||||
simulation_id=state.simulation_id,
|
||||
simulation_requirement="Agent personas MUST be individual people.",
|
||||
document_text="document",
|
||||
use_llm_for_profiles=False,
|
||||
)
|
||||
|
||||
assert result.entities_count == 1
|
||||
assert result.entity_types == ["Person"]
|
||||
assert [entity.name for entity in captured["profile_entities"]] == ["Alice"]
|
||||
assert [entity.name for entity in captured["config_entities"]] == ["Alice"]
|
||||
|
|
|
|||
Loading…
Reference in New Issue