AI/tests/test_research.py

from pathlib import Path

from sqlalchemy.orm import Session

from crawler_platform.app.config.loader import load_project_config
from crawler_platform.app.core.database import models
from crawler_platform.app.core.database.repository import KnowledgeRepository
from crawler_platform.app.core.database.session import make_engine
from crawler_platform.app.core.extractor.base import ExtractedClaim, ExtractedEntity, ExtractionBundle, ExtractionPageContext
from crawler_platform.app.core.extractor.validation import attach_page_context
from crawler_platform.app.core.research.graph_query import SemanticGraphQuery
from crawler_platform.app.core.research.graph_research_loop import GraphResearchLoop
from crawler_platform.app.core.research.relevance_engine import RelevanceEngine
from crawler_platform.app.domains.perfume.extractor import PerfumeRuleBasedExtractor


def make_repo():
    config = load_project_config("configs/perfume_subscription.yaml")
    engine = make_engine("sqlite:///:memory:")
    models.Base.metadata.create_all(engine)
    session = Session(engine)
    repo = KnowledgeRepository(session)
    project = repo.upsert_project(config)
    source = repo.get_source(project.id, "official_brand_site")
    return config, session, repo, project, source


def seed_validated_claim(repo, project, source, config):
    page = repo.upsert_page(
        project.id,
        source.id,
        "https://example.com/product/neroli",
        "Neroli Summer",
        200,
        "Neroli Summer\nTop notes: Bergamot",
    )
    bundle = ExtractionBundle(
        entities=[
            ExtractedEntity("Perfume", "Neroli Summer", evidence_text="Neroli Summer", confidence=0.9),
            ExtractedEntity("FragranceNote", "Bergamot", evidence_text="Top notes: Bergamot", confidence=0.9),
        ],
        claims=[
            ExtractedClaim(
                "Neroli Summer",
                "Perfume",
                "hasTopNote",
                "Bergamot",
                "FragranceNote",
                evidence_text="Top notes: Bergamot",
                confidence=0.93,
                confidence_reason="directly stated",
            )
        ],
        extractor_name="llm_json_extractor",
        provider="lm_studio",
        raw_output={"extraction_mode": "primary"},
    )
    context = ExtractionPageContext(
        url=page.url,
        final_url=page.url,
        title=page.title,
        page_type="ProductPage",
        clean_text=page.cleaned_text_summary,
        source_zones=[
            {
                "zone_type": "product_description",
                "selector": "main",
                "text": "Top notes: Bergamot",
                "claim_allowed": True,
            }
        ],
    )
    repo.save_extraction_bundle(project.id, source, page, attach_page_context(bundle, context), config)


def test_relevance_engine_prefers_graph_related_product_urls():
    config, session, repo, project, source = make_repo()
    seed_validated_claim(repo, project, source, config)

    engine = RelevanceEngine(session)
    product = engine.score_url(
        project_id=project.id,
        url="https://example.com/product/neroli-bergamot",
        label="Neroli Summer Bergamot perfume",
        source_trust=source.trust_level,
    )
    login = engine.score_url(
        project_id=project.id,
        url="https://example.com/member/login.html",
        label="Login",
        source_trust=source.trust_level,
    )

    assert product.score > login.score
    assert product.breakdown["entity_overlap"] > 0
    session.close()


def test_semantic_graph_query_returns_trend_summary():
    config, session, repo, project, source = make_repo()
    seed_validated_claim(repo, project, source, config)

    rows = SemanticGraphQuery(session).trend_summary(project.id)

    assert rows[0]["name"] == "Bergamot"
    assert rows[0]["predicate"] == "hasTopNote" or rows[0]["support_count"] >= 1
    session.close()


def test_graph_research_loop_records_session_memory():
    config, session, repo, project, _source = make_repo()
    fixture = Path("tests/fixtures/sample_perfume.html").resolve()
    loop = GraphResearchLoop(repo, PerfumeRuleBasedExtractor())

    result = loop.run(
        project_config=config,
        source_name="official_brand_site",
        seed_url=str(fixture),
        goal="Fixture semantic exploration",
        max_depth=0,
        max_steps=1,
        min_relevance=0.0,
    )
    session.commit()

    assert result.explored_count == 1
    assert result.history
    assert result.memory["visited_targets"]
    jobs = session.query(models.CrawlJob).all()
    assert any((job.metadata_json or {}).get("kind") == "research_session" for job in jobs)
    session.close()


def test_graph_research_loop_can_start_from_knowledge_gaps():
    config, session, repo, _project, _source = make_repo()
    loop = GraphResearchLoop(repo, PerfumeRuleBasedExtractor())

    result = loop.run(
        project_config=config,
        source_name="official_brand_site",
        goal="Gap-driven semantic exploration",
        max_depth=0,
        max_steps=1,
        max_branch=3,
        min_relevance=0.0,
    )

    assert result.explored_count == 1
    assert result.history
    assert result.history[0]["item"]["target_type"] == "knowledge_gap"
    session.close()