118 lines
4.6 KiB
Python
118 lines
4.6 KiB
Python
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
from sqlalchemy import and_, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from app.models import Base, KGEntity, KGRelation
|
|
from app.services.project_knowledge_graph import DDL_GUARDRAIL_POLICY_KEY, ProjectKnowledgeGraphService
|
|
|
|
|
|
class ProjectKnowledgeGraphGuardrailSyncTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
self._old_rules_path = os.environ.get("OPENCLAW_ROUTING_RULES_PATH")
|
|
|
|
tmp_db = tempfile.NamedTemporaryFile(prefix="pmflow-kg-guardrail-", suffix=".db", delete=False)
|
|
tmp_db.close()
|
|
self.db_path = tmp_db.name
|
|
self.engine = create_async_engine(f"sqlite+aiosqlite:///{self.db_path}", future=True)
|
|
self.session_factory = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)
|
|
async with self.engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
tmp_rules = tempfile.NamedTemporaryFile(prefix="pmflow-routing-rules-", suffix=".json", delete=False)
|
|
tmp_rules.write(
|
|
json.dumps(
|
|
{
|
|
"rules": [
|
|
{
|
|
"name": "low-intent-rule",
|
|
"enabled": True,
|
|
"priority": 300,
|
|
"tier": "low",
|
|
"provider": "glm",
|
|
"model": "GLM-4.7",
|
|
"api_env": "GLM_API_KEY",
|
|
},
|
|
{
|
|
"name": "high-intent-rule",
|
|
"enabled": True,
|
|
"priority": 280,
|
|
"tier": "high",
|
|
"provider": "openai",
|
|
"model": "GLM-5",
|
|
"api_env": "OPENAI_API_KEY",
|
|
},
|
|
]
|
|
}
|
|
).encode("utf-8")
|
|
)
|
|
tmp_rules.flush()
|
|
tmp_rules.close()
|
|
self.rules_path = tmp_rules.name
|
|
os.environ["OPENCLAW_ROUTING_RULES_PATH"] = self.rules_path
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
await self.engine.dispose()
|
|
for path in (getattr(self, "db_path", None), getattr(self, "rules_path", None)):
|
|
if path:
|
|
try:
|
|
os.remove(path)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
if self._old_rules_path is None:
|
|
os.environ.pop("OPENCLAW_ROUTING_RULES_PATH", None)
|
|
else:
|
|
os.environ["OPENCLAW_ROUTING_RULES_PATH"] = self._old_rules_path
|
|
|
|
async def test_sync_model_guardrails_now_builds_dual_model_policy_graph(self) -> None:
|
|
async with self.session_factory() as db:
|
|
service = ProjectKnowledgeGraphService(db, tenant_id=1)
|
|
payload = await service.sync_model_guardrails_now()
|
|
|
|
self.assertEqual(payload["models"], 2)
|
|
self.assertGreaterEqual(payload["relations"], 4)
|
|
|
|
policy_entity = (
|
|
await db.execute(
|
|
select(KGEntity).where(
|
|
and_(
|
|
KGEntity.tenant_id == 1,
|
|
KGEntity.entity_type == "governance_policy",
|
|
KGEntity.entity_key == DDL_GUARDRAIL_POLICY_KEY,
|
|
)
|
|
)
|
|
)
|
|
).scalar_one_or_none()
|
|
self.assertIsNotNone(policy_entity)
|
|
|
|
model_entities = (
|
|
await db.execute(
|
|
select(KGEntity).where(
|
|
and_(
|
|
KGEntity.tenant_id == 1,
|
|
KGEntity.entity_type == "ai_model",
|
|
)
|
|
)
|
|
)
|
|
).scalars().all()
|
|
self.assertEqual(len(model_entities), 2)
|
|
model_names = {str(row.display_name or "") for row in model_entities}
|
|
self.assertEqual(model_names, {"GLM-4.7", "GLM-5"})
|
|
|
|
policy_relations = (
|
|
await db.execute(
|
|
select(KGRelation).where(
|
|
and_(
|
|
KGRelation.tenant_id == 1,
|
|
KGRelation.relation_type == "model_must_follow_policy",
|
|
KGRelation.object_entity_id == int(policy_entity.id),
|
|
)
|
|
)
|
|
)
|
|
).scalars().all()
|
|
self.assertEqual(len(policy_relations), 2)
|