Filter simulation personas to agent-eligible entities

This commit is contained in:
Sami Rusani 2026-08-05 00:19:03 +02:00
parent b5b53acc57
commit eff1eba1ba
5 changed files with 291 additions and 23 deletions

View File

@ -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]

View File

@ -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 # 预分配列表保持顺序

View File

@ -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
)

View File

@ -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]

View File

@ -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"]