kamidso-back/tests/kamidso/test_data_graph.py

import json

import pytest

from kamidso.contracts.errors import KamidsoError
from kamidso.contracts.user import CurrentUser
from kamidso.data_graph.explorer.contracts import normalize_source_map
from kamidso.data_graph.service import DataGraphService
from kamidso.data_graph.tools.graph_search import (
    format_attached_links,
    links_between_resources,
    search_maps,
)
from kamidso.data_graph.tools.source_inspection import SourceInspector
from kamidso.repositories.store import RepositoryStore
from kamidso.resources.files import FileService
from kamidso.settings.flags import load_kamidso_settings
from kamidso.storage.local import LocalStorage

USER = CurrentUser(
    id="usr_test",
    email="test@kamidso.test",
    name="Test",
    role="developer",
    workspace_id="ws_test",
)


def _user():
    return USER


@pytest.mark.asyncio
async def test_excel_inspect_returns_raw_grid_not_a_header_rule(tmp_path):
    import openpyxl

    path = tmp_path / "finance.xlsx"
    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "Sheet1"
    ws["A1"] = "گزارش مالی ۱۴۰۵"
    ws["A3"] = "پرداخت‌کننده"
    ws.merge_cells("A3:B3")
    ws["A4"] = "اطلاعات فردی"
    ws["C4"] = "حساب"
    ws["A5"] = "نام"
    ws["B5"] = "کد"
    ws["C5"] = "بانک"
    ws["D5"] = "مبلغ"
    wb.save(path)

    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(USER, "finance.xlsx", path.read_bytes())
    inspector = SourceInspector(store, storage, USER)
    overview = await inspector.inspect_resource(doc["_id"])
    assert overview["kind"] == "excel"
    assert overview["sheets"][0]["merged"]
    grid = await inspector.inspect_path(doc["_id"], "Sheet1!A1:D5")
    values = {cell["a"]: cell["v"] for cell in grid["cells"]}
    assert values["A1"] == "گزارش مالی ۱۴۰۵"
    assert values["A3"] == "پرداخت‌کننده"
    assert values["A5"] == "نام"
    assert any("A3" in item for item in grid["merged"])
    assert "header_row" not in grid


@pytest.mark.asyncio
async def test_excel_overview_lists_every_data_sheet(tmp_path):
    import openpyxl

    path = tmp_path / "employees.xlsx"
    wb = openpyxl.Workbook()
    employees = wb.active
    employees.title = "employees"
    employees["A1"] = "employee_id"
    employees["B1"] = "نام"
    departments = wb.create_sheet("departments")
    departments["A1"] = "department_id"
    departments["B1"] = "نام واحد"
    readme = wb.create_sheet("README")
    readme["A1"] = "هدف فایل"
    readme["B1"] = "دیتاست تست"
    wb.save(path)

    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(USER, "employees.xlsx", path.read_bytes())
    inspector = SourceInspector(store, storage, USER)
    overview = await inspector.inspect_resource(doc["_id"])
    assert [item["name"] for item in overview["sheets"]] == ["employees", "departments", "README"]
    assert overview["data_sheets"] == ["employees", "departments"]
    assert next(item for item in overview["sheets"] if item["name"] == "README")["role"] == "notes"
    preview = next(item for item in overview["sheets"] if item["name"] == "departments")["preview"]
    assert any(cell["v"] == "department_id" for cell in preview)


@pytest.mark.asyncio
async def test_discovery_keeps_every_excel_data_sheet(tmp_path):
    import openpyxl

    path = tmp_path / "employees.xlsx"
    wb = openpyxl.Workbook()
    employees = wb.active
    employees.title = "employees"
    employees["A1"] = "employee_id"
    employees["B1"] = "نام"
    employees["A2"] = "E001"
    departments = wb.create_sheet("departments")
    departments["A1"] = "department_id"
    departments["B1"] = "نام واحد"
    departments["A2"] = "D001"
    readme = wb.create_sheet("README")
    readme["A1"] = "هدف فایل"
    readme["B1"] = "دیتاست تست"
    wb.save(path)

    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(USER, "employees.xlsx", path.read_bytes())
    service = DataGraphService(store, storage)

    calls = []

    async def complete(prompt: str) -> str:
        calls.append(prompt)
        payload = {
            "output": {
                "kind": "excel",
                "coverage": "observed_partial",
                "datasets": [
                    {
                        "name": "employees",
                        "profile": {"domain": "hr roster", "row_grain": "one employee"},
                        "fields": [
                            {
                                "name": "employee_id",
                                "path": ["employee_id"],
                                "address": "excel://employees.xlsx/employees/employee_id",
                                "data_type": "string",
                            },
                            {
                                "name": "نام",
                                "path": ["نام"],
                                "address": "excel://employees.xlsx/employees/نام",
                                "data_type": "string",
                            },
                        ],
                    }
                ],
            }
        }
        return "Thought: employees only.\nAction: terminate\nAction Input: " + json.dumps(payload, ensure_ascii=False)

    saved = await service.discover(USER, doc["_id"], complete)
    assert calls
    names = [item["name"] for item in saved["datasets"]]
    assert names == ["employees", "departments"]
    employees = next(item for item in saved["datasets"] if item["name"] == "employees")
    assert [field["name"] for field in employees["fields"]] == ["employee_id", "نام"]
    dept = next(item for item in saved["datasets"] if item["name"] == "departments")
    assert [field["name"] for field in dept["fields"]] == ["department_id", "نام واحد"]


@pytest.mark.asyncio
async def test_excel_inspect_caps_to_header_band(tmp_path):
    import openpyxl

    path = tmp_path / "employees.xlsx"
    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "employees"
    ws["A1"] = "employee_id"
    for row in range(2, 32):
        ws.cell(row=row, column=1, value=f"E{row:03d}")
    wb.save(path)

    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(USER, "employees.xlsx", path.read_bytes())
    inspector = SourceInspector(store, storage, USER)
    grid = await inspector.inspect_path(doc["_id"], "employees!A1:A31")
    addresses = {cell["a"] for cell in grid["cells"]}
    assert "A1" in addresses
    assert "A8" in addresses
    assert "A9" not in addresses
    assert grid["truncated"] is True


@pytest.mark.asyncio
async def test_sqlite_inspect_lists_declared_columns(tmp_path):
    from kamidso.engine.ext.datasource.rdbms.conn_sqlite import SQLiteConnector

    db_path = tmp_path / "crm.sqlite"
    connector = SQLiteConnector.from_file_path(str(db_path))
    connector.run("CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, payload TEXT)")
    connector.close()
    store = RepositoryStore()
    doc = await store.create_datasource(
        {
            "workspace_id": USER.workspace_id,
            "name": "CRM",
            "type": "sqlite",
            "db_type": "sqlite",
            "db_path": str(db_path),
            "status": "active",
        }
    )
    inspector = SourceInspector(store, None, USER)
    overview = await inspector.inspect_resource(doc["_id"])
    assert "customers" in overview["tables"]
    described = await inspector.inspect_path(doc["_id"], "customers")
    names = [col["name"] for col in described["columns"]]
    assert names == ["id", "name", "payload"]
    assert described["coverage_hint"] == "declared_complete"


def test_source_map_requires_explicit_coverage():
    with pytest.raises(KamidsoError):
        normalize_source_map(
            {"datasets": [{"name": "t", "fields": [{"name": "a", "address": "sql://x"}]}]},
            resource_id="r",
            resource_name="r",
            resource_type="file",
        )


def test_search_ignores_suggested_edges():
    maps = [
        {
            "status": "ready",
            "resource_id": "hr",
            "resource_name": "HR",
            "kind": "excel",
            "datasets": [
                {
                    "name": "employees",
                    "fields": [
                        {
                            "id": "field_name",
                            "name": "full_name",
                            "path": ["full_name"],
                            "address": "excel://hr/employees/full_name",
                            "meaning": "نام کارمند",
                        }
                    ],
                }
            ],
        }
    ]
    edges = [
        {
            "from": "field_name",
            "to": "field_payer",
            "from_address": "excel://hr/employees/full_name",
            "to_address": "excel://finance/payer",
            "type": "entity_match",
            "status": "suggested",
        }
    ]
    result = search_maps(maps, edges, "نام کارمند")
    assert result["relevant_fields"][0]["address"] == "excel://hr/employees/full_name"
    assert result["confirmed_paths"] == []


def _two_source_maps():
    return [
        {
            "status": "ready",
            "resource_id": "file_hr",
            "resource_name": "employees",
            "datasets": [
                {
                    "name": "employees",
                    "fields": [
                        {
                            "id": "field_name",
                            "name": "نام",
                            "address": "excel://employees.xlsx/employees/نام",
                        }
                    ],
                }
            ],
        },
        {
            "status": "ready",
            "resource_id": "file_fin",
            "resource_name": "Book1",
            "datasets": [
                {
                    "name": "1404",
                    "fields": [
                        {
                            "id": "field_payer",
                            "name": "نام پرداخت کننده",
                            "address": "excel://Book1.xlsx/1404/نام پرداخت کننده",
                        }
                    ],
                }
            ],
        },
    ]


def test_one_attached_source_gets_no_graph_context():
    assert links_between_resources(_two_source_maps(), [], ["file_hr"]) is None


def test_two_attached_sources_get_only_confirmed_cross_links():
    edges = [
        {
            "from": "field_name",
            "to": "field_payer",
            "type": "same_concept",
            "status": "confirmed",
        },
        {
            "from": "field_name",
            "to": "field_payer",
            "type": "same_concept",
            "status": "suggested",
        },
    ]
    attached = links_between_resources(_two_source_maps(), edges, ["file_hr", "file_fin"])
    assert attached is not None
    assert attached["confirmed_links"] == [
        {
            "from_resource": "employees",
            "from_field": "نام",
            "from_address": "excel://employees.xlsx/employees/نام",
            "to_resource": "Book1",
            "to_field": "نام پرداخت کننده",
            "to_address": "excel://Book1.xlsx/1404/نام پرداخت کننده",
            "type": "same_concept",
        }
    ]
    text = format_attached_links(attached)
    assert "Confirmed links between attached sources" in text
    assert "excel://employees.xlsx/employees/نام" in text

    missing = links_between_resources(_two_source_maps(), [], ["file_hr", "file_fin"])
    assert missing is not None
    assert missing["confirmed_links"] == []
    assert "Do not invent a join" in format_attached_links(missing)


@pytest.mark.asyncio
async def test_structure_agent_saves_contract(tmp_path):
    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(
        USER, "note.json", json.dumps({"customer": {"id": "C1"}}).encode("utf-8"), mime_type="application/json"
    )
    service = DataGraphService(store, storage)

    payload = {
        "output": {
            "kind": "json",
            "coverage": "observed_partial",
            "coverage_detail": {"documents_sampled": 1},
            "datasets": [
                {
                    "name": "customer",
                    "fields": [
                        {
                            "name": "id",
                            "path": ["customer", "id"],
                            "address": "json://note.json/customer/id",
                            "data_type": "string",
                            "confidence": 0.8,
                        }
                    ],
                }
            ],
        }
    }

    async def complete(prompt: str) -> str:
        if "Observation" not in prompt:
            return (
                "Thought: look at the file.\n"
                "Action: inspect_resource\n"
                f'Action Input: {{"resource_id": "{doc["_id"]}"}}\n'
            )
        return (
            "Thought: the sample is one object.\n"
            "Action: terminate\n"
            f"Action Input: {json.dumps(payload, ensure_ascii=False)}\n"
        )

    saved = await service.discover(USER, doc["_id"], complete)
    assert saved["coverage"] == "observed_partial"
    assert saved["datasets"][0]["fields"][0]["address"].endswith("/customer/id")


def _json_map_payload(semantic=None):
    field = {
        "name": "id",
        "path": ["customer", "id"],
        "address": "json://note.json/customer/id",
        "data_type": "string",
        "confidence": 0.8,
    }
    if semantic:
        field.update(semantic)
    return {
        "output": {
            "kind": "json",
            "coverage": "observed_partial",
            "coverage_detail": {"documents_sampled": 1},
            "datasets": [{"name": "customer", "fields": [field]}],
        }
    }


def _inspect_then(payload):
    async def complete(prompt: str) -> str:
        if "Observation" not in prompt:
            return (
                "Thought: look at the file.\n"
                "Action: inspect_resource\n"
                'Action Input: {"resource_id": "ignored"}\n'
            )
        return (
            "Thought: structure is ready.\n"
            "Action: terminate\n"
            f"Action Input: {json.dumps(payload, ensure_ascii=False)}\n"
        )

    return complete


@pytest.mark.asyncio
async def test_discovery_failure_keeps_the_last_good_map(tmp_path):
    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(
        USER, "note.json", json.dumps({"customer": {"id": "C1"}}).encode("utf-8"), mime_type="application/json"
    )
    service = DataGraphService(store, storage)
    saved = await service.discover(USER, doc["_id"], _inspect_then(_json_map_payload()))
    assert saved["status"] == "ready"

    async def boom(_prompt: str) -> str:
        raise RuntimeError("timeout")

    with pytest.raises(Exception):
        await service.discover(USER, doc["_id"], boom)
    kept = await service.list_maps(USER)
    assert kept[0]["status"] == "ready"
    assert kept[0]["datasets"][0]["fields"][0]["address"].endswith("/customer/id")
    assert kept[0]["last_discovery_error"]


@pytest.mark.asyncio
async def test_user_confirm_survives_rediscovery_and_chat(tmp_path):
    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(
        USER, "note.json", json.dumps({"customer": {"id": "C1"}}).encode("utf-8"), mime_type="application/json"
    )
    service = DataGraphService(store, storage)
    first = await service.discover(
        USER,
        doc["_id"],
        _inspect_then(
            _json_map_payload(
                {
                    "semantic_label": "customer_id",
                    "meaning": "شناسه مشتری",
                    "semantic_status": "inferred",
                    "semantic_confidence": 0.7,
                }
            )
        ),
    )
    field_id = first["datasets"][0]["fields"][0]["id"]
    confirmed = await service.update_field_semantic(
        USER,
        field_id,
        semantic_label="customer_id",
        meaning="شناسه یکتای مشتری",
        status="confirmed",
    )
    assert confirmed["semantic_status"] == "confirmed"
    assert confirmed["semantic_source"] == "user"

    second = await service.discover(
        USER,
        doc["_id"],
        _inspect_then(
            _json_map_payload(
                {
                    "semantic_label": "invoice_id",
                    "meaning": "شماره فاکتور",
                    "semantic_status": "inferred",
                    "semantic_confidence": 0.99,
                }
            )
        ),
    )
    kept = second["datasets"][0]["fields"][0]
    assert kept["semantic_status"] == "confirmed"
    assert kept["semantic_label"] == "customer_id"
    assert kept["meaning"] == "شناسه یکتای مشتری"

    async def semantic_complete(prompt: str) -> str:
        if "Observation" not in prompt:
            return "Thought: context\nAction: get_field_context\nAction Input: {}\n"
        return (
            "Thought: different guess\n"
            "Action: terminate\n"
            "Action Input: "
            + json.dumps(
                {
                    "output": {
                        "fields": [
                            {
                                "field_id": field_id,
                                "semantic_label": "invoice_number",
                                "meaning": "شماره فاکتور",
                                "semantic_status": "inferred",
                                "semantic_confidence": 0.99,
                            }
                        ]
                    }
                },
                ensure_ascii=False,
            )
        )

    inspected = await service.inspect_field_semantics(
        USER, [field_id], "مبلغ فاکتور", semantic_complete
    )
    assert inspected["fields"][0]["semantic_status"] == "confirmed"
    assert inspected["fields"][0]["needs_user_input"] is True
    stored = (await service.list_maps(USER))[0]["datasets"][0]["fields"][0]
    assert stored["semantic_label"] == "customer_id"
    assert stored["semantic_status"] == "confirmed"


@pytest.mark.asyncio
async def test_semantic_unresolved_asks_for_user_input(tmp_path):
    settings = load_kamidso_settings()
    store = RepositoryStore()
    storage = LocalStorage(str(tmp_path / "storage"))
    doc = await FileService(store, storage, settings).upload(
        USER, "note.json", json.dumps({"customer": {"id": "C1"}}).encode("utf-8"), mime_type="application/json"
    )
    service = DataGraphService(store, storage)
    saved = await service.discover(
        USER,
        doc["_id"],
        _inspect_then(_json_map_payload({"name": "x17", "path": ["customer", "x17"], "address": "json://note.json/customer/x17"})),
    )
    field_id = saved["datasets"][0]["fields"][0]["id"]

    async def semantic_complete(prompt: str) -> str:
        if "Observation" not in prompt:
            return "Thought: context\nAction: get_field_context\nAction Input: {}\n"
        return (
            "Thought: still unclear\n"
            "Action: terminate\n"
            "Action Input: "
            + json.dumps(
                {
                    "output": {
                        "fields": [
                            {
                                "field_id": field_id,
                                "semantic_label": "",
                                "meaning": "",
                                "semantic_status": "unknown",
                                "semantic_confidence": 0.2,
                                "needs_user_input": True,
                            }
                        ]
                    }
                }
            )
        )

    inspected = await service.inspect_field_semantics(USER, [field_id], "مطالبات مشتری", semantic_complete)
    assert inspected["fields"][0]["semantic_status"] == "unknown"
    assert inspected["fields"][0]["needs_user_input"] is True
    review = await service.list_semantic_review(USER)
    assert all(item["field_id"] != field_id for item in review)