DataOps-platform/tests/integration/test_knowledge_retrieval.py

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