microfish/backend/tests/test_local_golden_flow.py

from types import SimpleNamespace
from typing import Any, cast

from sqlalchemy import create_engine

from app.config import Config
from app.db import Base, create_session_factory
from app.services.local_graph_builder import LocalGraphBuilderService
from app.services.memory_entity_reader import LocalEntityReader
from app.services.memory_extraction import (
    ExtractedEdge,
    ExtractedEntity,
    MemoryExtractionResult,
)
from app.services.memory_service import MemoryExtractionService
from app.services.memory_tools import LocalMemoryTools
from app.services.oasis_profile_generator import OasisProfileGenerator
from app.services.report_agent import ReportAgent
from app.services.report_agent import ReportManager, ReportOutline, ReportSection, ReportStatus
from app.services.simulation_config_generator import (
    AgentActivityConfig,
    SimulationConfigGenerator,
    SimulationParameters,
)
from app.services.simulation_manager import SimulationManager, SimulationStatus
from app.services.memory_entity_reader import make_local_entity_reader_factory


class DeterministicExtraction:
    def extract(self, *, language, ontology, episode_text, context=""):
        return MemoryExtractionResult(
            entities=[
                ExtractedEntity(
                    mention="Alice",
                    canonical_name="Alice",
                    labels=["Entity", "Person"],
                    summary="A founder building Orbit.",
                    confidence=0.99,
                ),
                ExtractedEntity(
                    mention="Orbit",
                    canonical_name="Orbit",
                    labels=["Entity", "Organization"],
                    summary="A local project.",
                    confidence=0.99,
                ),
            ],
            edges=[
                ExtractedEdge(
                    source_entity_ref="Alice",
                    target_entity_ref="Orbit",
                    relation="FOUNDED",
                    fact="Alice founded Orbit.",
                    confidence=0.98,
                )
            ],
            episode_summary="Alice founded Orbit.",
        )

    def persist(self, repository, result, *, source_type, source_ref, episode_text):
        return MemoryExtractionService(cast(Any, None)).persist(
            repository,
            result,
            source_type=source_type,
            source_ref=source_ref,
            episode_text=episode_text,
        )


def test_local_graph_profile_report_golden_flow(monkeypatch):
    monkeypatch.setattr(Config, "MEMORY_BACKEND", "local")
    monkeypatch.setattr(Config, "LLM_API_KEY", "test-key")
    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(engine)
    session_factory = create_session_factory(engine)
    try:
        builder = LocalGraphBuilderService(
            session_factory,
            organization_id="org-a",
            project_id="project-a",
            extraction_service=cast(Any, DeterministicExtraction()),
            language="en",
        )
        graph_id = builder.create_graph("golden")
        builder.add_text_batches(graph_id, ["Alice founded Orbit."])

        reader = LocalEntityReader(
            session_factory(),
            organization_id="org-a",
            graph_id=graph_id,
            owns_session=True,
        )
        try:
            filtered = reader.filter_defined_entities(defined_entity_types=["Person"])
            assert [entity.name for entity in filtered.entities] == ["Alice"]

            local_tools = LocalMemoryTools(reader.repository)
            profile_generator = OasisProfileGenerator(
                graph_id=graph_id,
                use_zep_context=True,
                local_memory_tools=local_tools,
            )
            profile = profile_generator.generate_profile_from_entity(
                cast(Any, filtered.entities[0]),
                user_id=1,
                use_llm=False,
            )
            assert profile.name == "Alice"
            assert profile.source_entity_uuid == filtered.entities[0].uuid
            assert profile_generator.zep_client is None

            report_agent = ReportAgent(
                graph_id=graph_id,
                simulation_id="simulation-a",
                simulation_requirement="Understand the project origin.",
                llm_client=cast(Any, SimpleNamespace()),
                memory_tools=local_tools,
            )
            report_context = report_agent._execute_tool(
                "quick_search",
                {"query": "Alice", "limit": 10},
            )
            assert "Alice founded Orbit." in report_context
        finally:
            reader.close()
    finally:
        engine.dispose()


def test_local_graph_profile_simulation_report_is_persisted(monkeypatch, tmp_path):
    monkeypatch.setattr(Config, "MEMORY_BACKEND", "local")
    monkeypatch.setattr(Config, "LLM_API_KEY", "test-key")
    monkeypatch.setattr(Config, "UPLOAD_FOLDER", str(tmp_path / "uploads"))
    monkeypatch.setattr(ReportManager, "REPORTS_DIR", str(tmp_path / "reports"))

    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(engine)
    session_factory = create_session_factory(engine)
    simulation_dir = tmp_path / "simulations"
    try:
        builder = LocalGraphBuilderService(
            session_factory,
            organization_id="org-a",
            project_id="project-a",
            extraction_service=cast(Any, DeterministicExtraction()),
            language="en",
        )
        graph_id = builder.create_graph("e2e")
        builder.add_text_batches(graph_id, ["Alice founded Orbit."])

        monkeypatch.setattr(
            SimulationConfigGenerator,
            "generate_config",
            lambda self, **kwargs: SimulationParameters(
                simulation_id=kwargs["simulation_id"],
                project_id=kwargs["project_id"],
                graph_id=kwargs["graph_id"],
                simulation_requirement=kwargs["simulation_requirement"],
                agent_configs=[
                    AgentActivityConfig(
                        agent_id=1,
                        entity_uuid="node-alice",
                        entity_name="Alice",
                        entity_type="Person",
                    )
                ],
                generation_reasoning="deterministic-test",
            ),
        )

        simulation_manager = SimulationManager(
            entity_reader_factory=make_local_entity_reader_factory(
                session_factory,
                organization_id="org-a",
            )
        )
        simulation_manager.SIMULATION_DATA_DIR = str(simulation_dir)
        simulation = simulation_manager.create_simulation(
            project_id="project-a",
            graph_id=graph_id,
            enable_twitter=False,
            enable_reddit=True,
        )
        prepared = simulation_manager.prepare_simulation(
            simulation.simulation_id,
            simulation_requirement="Understand the project origin.",
            document_text="Alice founded Orbit.",
            defined_entity_types=["Person"],
            use_llm_for_profiles=False,
            parallel_profile_count=1,
        )
        assert prepared.status is SimulationStatus.READY
        assert prepared.entities_count == 1
        assert prepared.profiles_count == 1
        assert simulation_manager.get_profiles(prepared.simulation_id) [0]["name"] == "Alice"

        report_session = session_factory()
        try:
            report_tools = LocalMemoryTools(
                report_session,
                organization_id="org-a",
                graph_id=graph_id,
            )
            report_agent = ReportAgent(
                graph_id=graph_id,
                simulation_id=prepared.simulation_id,
                simulation_requirement="Understand the project origin.",
                llm_client=cast(Any, SimpleNamespace()),
                memory_tools=report_tools,
            )
            monkeypatch.setattr(
                ReportAgent,
                "plan_outline",
                lambda self, progress_callback=None: ReportOutline(
                    title="Local E2E Report",
                    summary="Evidence-backed local report.",
                    sections=[ReportSection(title="Evidence")],
                ),
            )
            monkeypatch.setattr(
                ReportAgent,
                "_generate_section_react",
                lambda self, section, outline, previous_sections, progress_callback=None, section_index=0: self._execute_tool(
                    "quick_search",
                    {"query": "Alice", "limit": 10},
                ),
            )

            report = report_agent.generate_report(report_id="report-local-e2e")
            assert report.status is ReportStatus.COMPLETED
            assert "Alice founded Orbit." in report.markdown_content
            persisted = ReportManager.get_report("report-local-e2e")
            assert persisted is not None
            assert persisted.status is ReportStatus.COMPLETED
            assert "Alice founded Orbit." in persisted.markdown_content
        finally:
            report_session.close()
    finally:
        engine.dispose()