from __future__ import annotations import os import uuid import pytest from sqlalchemy import create_engine, text from sqlalchemy.orm import Session from app.core.knowledge.access import KnowledgeAccessContext from app.core.knowledge.point_builder import build_knowledge_snapshot from app.core.knowledge.repository import SqlKnowledgeRepository from app.core.knowledge.retrieval.sql import SqlLexicalRetriever, SqlVectorRetriever from app.core.knowledge.sync import KnowledgeSyncService pytestmark = pytest.mark.integration class FixedEmbeddingProvider: profile_key = "qwen:retrieval-test:1024" def embed(self, texts: list[str]) -> list[list[float]]: return [[1.0] * 1024 for _text in texts] @pytest.fixture() def database_url(): value = os.environ.get("TEST_DATABASE_URL") if not value: pytest.skip("TEST_DATABASE_URL is not configured") return value def test_sql_retrievers_filter_business_domain_before_recall(database_url): engine = create_engine(database_url) profile_id = str(uuid.uuid4()) domain_a, domain_b = str(uuid.uuid4()), str(uuid.uuid4()) source_a, source_b = str(uuid.uuid4()), str(uuid.uuid4()) sources = ((source_a, domain_a, "customer_a"), (source_b, domain_b, "customer_b")) try: with engine.begin() as connection: connection.execute( text( """ INSERT INTO public.knowledge_embedding_profiles ( id, provider, model, dimension, status, config_hash, activated_at ) VALUES ( CAST(:id AS uuid), 'qwen', 'retrieval-test', 1024, 'active', :config_hash, CURRENT_TIMESTAMP ) """ ), {"id": profile_id, "config_hash": "b" * 64}, ) for source_uid, domain_uid, name in sources: with Session(engine) as session, session.begin(): embedder = FixedEmbeddingProvider() repository = SqlKnowledgeRepository( session, embedding_profile_id=profile_id, embedding_profile_key=embedder.profile_key, generation=1, workspace=f"dataops-global-{domain_uid}-g1", ) snapshot = build_knowledge_snapshot( "DataFlow", { "uid": source_uid, "version": 1, "name": name, "purpose": "sync customer data", "business_domain_uid": domain_uid, }, ) KnowledgeSyncService(repository=repository, embedder=embedder).sync( snapshot ) context = KnowledgeAccessContext( subject_id="user-a", roles=frozenset({"viewer"}), permissions=frozenset({"governance:read"}), business_domain_uids=frozenset({domain_a}), correlation_id="retrieval-test", ) with Session(engine) as session: lexical = SqlLexicalRetriever(session).retrieve("customer_a", context, 20) vector = SqlVectorRetriever(session, FixedEmbeddingProvider()).retrieve( "customer", context, 20 ) assert lexical assert {item.object_uid for item in lexical} == {source_a} assert vector assert {item.business_domain_uid for item in vector} == {domain_a} assert all(item.point_revisions for item in vector) finally: with engine.begin() as connection: connection.execute( text( "DELETE FROM public.governance_documents " "WHERE object_uid = ANY(CAST(:uids AS uuid[]))" ), {"uids": [source_a, source_b]}, ) connection.execute( text( "DELETE FROM public.knowledge_points " "WHERE source_uid = ANY(CAST(:uids AS uuid[]))" ), {"uids": [source_a, source_b]}, ) connection.execute( text( "DELETE FROM public.knowledge_change_sets " "WHERE source_uid = ANY(CAST(:uids AS uuid[]))" ), {"uids": [source_a, source_b]}, ) connection.execute( text( "DELETE FROM public.knowledge_embedding_profiles " "WHERE id = CAST(:id AS uuid)" ), {"id": profile_id}, ) engine.dispose()