kamidso-back/kamidso/data_graph/api/routes.py

from __future__ import annotations

import asyncio
import json
import time
from typing import Any, Awaitable, Callable, Dict, List, Optional

from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from kamidso.api.deps import require_user
from kamidso.contracts.errors import KamidsoError, KamidsoErrorCode
from kamidso.data_graph.discovery_jobs import DiscoveryHub, SUGGEST_JOB_ID, raise_if_stopping
from kamidso.data_graph.explorer.live_highlight import live_highlight
from kamidso.data_graph.llm import engine_complete
from kamidso.data_graph.models import public_source_map
from kamidso.data_graph.service import DataGraphService

router = APIRouter()


def _ok(data: Any) -> Dict[str, Any]:
    return {"success": True, "data": data, "err_code": None, "err_msg": None}


def _raise_http(exc: KamidsoError, status: int = 400) -> None:
    raise HTTPException(status_code=status, detail=exc.to_public_dict()) from exc


class DiscoverRequest(BaseModel):
    model: Optional[str] = None
    scope: str = "all"


class MapUpdateRequest(BaseModel):
    datasets: List[Dict[str, Any]] = Field(default_factory=list)


class EdgeCreateRequest(BaseModel):
    from_field: str
    to_field: str
    type: str = "custom"
    reason: str = ""


class EdgeStatusRequest(BaseModel):
    status: str


class SearchRequest(BaseModel):
    query: str


class FieldSemanticRequest(BaseModel):
    semantic_label: str = ""
    meaning: str = ""
    status: str = "confirmed"
    facts: dict | None = None


def _service(request: Request) -> DataGraphService:
    app = request.app.state.kamidso
    return DataGraphService(app.store, app.storage)


def _hub(request: Request) -> DiscoveryHub:
    app = request.app.state.kamidso
    hub = getattr(app, "discovery_hub", None)
    if hub is None:
        hub = DiscoveryHub()
        app.discovery_hub = hub
    return hub


@router.get("/data-graph/maps")
async def list_maps(request: Request):
    user = require_user(request)
    return _ok(await _service(request).list_maps(user))


@router.get("/data-graph/edges")
async def list_edges(request: Request):
    user = require_user(request)
    return _ok(await _service(request).list_edges(user))


_SSE_HEADERS = {
    "Cache-Control": "no-cache, no-transform",
    "X-Accel-Buffering": "no",
}


def _sse(event: Dict[str, Any]) -> str:
    return f"data: {json.dumps(event, ensure_ascii=False, default=str)}\n\n"


def _live_complete(
    model: Optional[str],
    on_step: Callable[[Dict[str, Any]], Awaitable[None]],
    stop_requested: Optional[Callable[[], None]] = None,
):
    """Stream model text and surface the field currently being written."""
    state: Dict[str, Any] = {"signature": None, "at": 0.0}

    async def publish(event: Dict[str, Any], *, force: bool = False) -> None:
        signature = (event.get("field"), event.get("meaning"), event.get("count"), event.get("thought"))
        now = time.monotonic()
        previous = state["signature"]
        same_field = previous is not None and previous[0] == signature[0]
        if signature == previous:
            return
        if not force and same_field and now - float(state["at"]) < 0.4:
            return
        state["signature"] = signature
        state["at"] = now
        await on_step(event)

    async def complete(prompt: str) -> str:
        async def on_text(text: str) -> None:
            if stop_requested:
                stop_requested()
            event = live_highlight(text)
            if event is not None:
                await publish(event)

        text = await engine_complete(prompt, model, on_text=on_text)
        final = live_highlight(text)
        if final is not None:
            await publish(final, force=True)
        return text

    return complete


async def _stream(work: Callable[[Callable[[Dict[str, Any]], Awaitable[None]]], Awaitable[Any]]):
    queue: asyncio.Queue = asyncio.Queue()

    async def on_step(event: Dict[str, Any]) -> None:
        await queue.put(event)

    async def run() -> None:
        try:
            result = await work(on_step)
            await queue.put({"type": "done", "result": result})
        except asyncio.CancelledError:
            await queue.put({"type": "error", "message": "کشف ساختار قطع شد"})
            raise
        except Exception as exc:
            message = exc.to_public_dict().get("message") if isinstance(exc, KamidsoError) else str(exc)
            await queue.put({"type": "error", "message": message or "خطا"})
        finally:
            await queue.put(None)

    task = asyncio.create_task(run())

    async def generate():
        finished = False
        try:
            yield _sse({"type": "status", "message": "شروع شد"})
            while True:
                try:
                    event = await asyncio.wait_for(queue.get(), timeout=5)
                except asyncio.TimeoutError:
                    yield _sse({"type": "status", "message": "در حال فکر کردن"})
                    continue
                if event is None:
                    finished = True
                    break
                yield _sse(event)
        finally:
            if not finished and not task.done():
                task.cancel()
            try:
                await asyncio.shield(task)
            except (asyncio.CancelledError, Exception):
                pass

    return StreamingResponse(generate(), media_type="text/event-stream", headers=_SSE_HEADERS)


@router.post("/data-graph/resources/{resource_id}/discover/stream")
async def discover_resource_stream(resource_id: str, body: DiscoverRequest, request: Request):
    user = require_user(request)
    service = _service(request)
    hub = _hub(request)

    def stop_requested() -> None:
        current = hub.get(user.workspace_id, resource_id)
        raise_if_stopping(current.stop_mode if current else None)

    async def work(on_step):
        complete = _live_complete(body.model, on_step, stop_requested)
        return await service.discover(
            user,
            resource_id,
            complete,
            on_step=on_step,
            stop_requested=stop_requested,
        )

    job = await hub.start_or_attach(user, resource_id, work)

    async def generate():
        async for event in hub.iter_events(job):
            yield _sse(event)

    return StreamingResponse(generate(), media_type="text/event-stream", headers=_SSE_HEADERS)


async def _seal_idle(request: Request, resource_id: str, *, summarize: bool):
    user = require_user(request)
    service = _service(request)
    resource_type, name = await service._resource_label(user, resource_id)
    previous = await service._repo.get_source_map(user.workspace_id, resource_id)

    async def complete(prompt: str) -> str:
        return await engine_complete(prompt, None)

    sealed = await service._seal_stopped(
        user,
        resource_id,
        name,
        resource_type,
        previous,
        summarize=summarize,
        complete=complete,
        on_step=None,
    )
    return _ok(public_source_map(sealed) if sealed else None)


@router.post("/data-graph/resources/{resource_id}/discover/cancel")
async def cancel_discover(resource_id: str, request: Request):
    user = require_user(request)
    job = _hub(request).get(user.workspace_id, resource_id)
    if job is not None and not job.done:
        job.request_stop("cancel")
        return _ok({"stopping": "cancel"})
    return await _seal_idle(request, resource_id, summarize=False)


@router.post("/data-graph/resources/{resource_id}/discover/wrap")
async def wrap_discover(resource_id: str, request: Request):
    user = require_user(request)
    job = _hub(request).get(user.workspace_id, resource_id)
    if job is not None and not job.done:
        job.request_stop("wrap")
        return _ok({"stopping": "wrap"})
    return await _seal_idle(request, resource_id, summarize=True)


@router.post("/data-graph/suggestions/stream")
async def suggest_edges_stream(body: DiscoverRequest, request: Request):
    user = require_user(request)
    service = _service(request)
    hub = _hub(request)

    async def work(on_step):
        complete = _live_complete(body.model, on_step)
        return await service.suggest(user, complete, on_step=on_step, scope=body.scope)

    job = await hub.start_or_attach(user, SUGGEST_JOB_ID, work)

    async def generate():
        async for event in hub.iter_events(job):
            yield _sse(event)

    return StreamingResponse(generate(), media_type="text/event-stream", headers=_SSE_HEADERS)


@router.post("/data-graph/resources/{resource_id}/discover")
async def discover_resource(resource_id: str, body: DiscoverRequest, request: Request):
    user = require_user(request)
    service = _service(request)

    async def complete(prompt: str) -> str:
        return await engine_complete(prompt, body.model)

    async def work(on_step):
        return await service.discover(user, resource_id, complete, on_step=on_step)

    try:
        job = await _hub(request).start_or_attach(user, resource_id, work)
        return _ok(await job.wait())
    except KamidsoError as exc:
        status = 404 if exc.code == KamidsoErrorCode.RESOURCE_NOT_FOUND else 400
        _raise_http(exc, status=status)


@router.put("/data-graph/maps/{resource_id}")
async def update_map(resource_id: str, body: MapUpdateRequest, request: Request):
    user = require_user(request)
    try:
        return _ok(await _service(request).update_map(user, resource_id, body.datasets))
    except KamidsoError as exc:
        status = 404 if exc.code == KamidsoErrorCode.RESOURCE_NOT_FOUND else 400
        _raise_http(exc, status=status)


@router.post("/data-graph/edges")
async def create_edge(body: EdgeCreateRequest, request: Request):
    user = require_user(request)
    try:
        edge = await _service(request).create_edge(
            user,
            left=body.from_field,
            right=body.to_field,
            edge_type=body.type,
            reason=body.reason,
        )
        return _ok(edge)
    except KamidsoError as exc:
        _raise_http(exc)


@router.patch("/data-graph/edges/{edge_id}")
async def update_edge(edge_id: str, body: EdgeStatusRequest, request: Request):
    user = require_user(request)
    try:
        return _ok(await _service(request).set_edge_status(user, edge_id, body.status))
    except KamidsoError as exc:
        status = 404 if exc.code == KamidsoErrorCode.RESOURCE_NOT_FOUND else 400
        _raise_http(exc, status=status)


@router.delete("/data-graph/edges/{edge_id}")
async def delete_edge(edge_id: str, request: Request):
    user = require_user(request)
    try:
        await _service(request).delete_edge(user, edge_id)
        return _ok({"id": edge_id})
    except KamidsoError as exc:
        status = 404 if exc.code == KamidsoErrorCode.RESOURCE_NOT_FOUND else 400
        _raise_http(exc, status=status)


@router.post("/data-graph/suggestions")
async def suggest_edges(body: DiscoverRequest, request: Request):
    user = require_user(request)

    async def complete(prompt: str) -> str:
        return await engine_complete(prompt, body.model)

    try:
        return _ok(await _service(request).suggest(user, complete, scope=body.scope))
    except KamidsoError as exc:
        _raise_http(exc)


@router.post("/data-graph/search")
async def search_graph(body: SearchRequest, request: Request):
    user = require_user(request)
    return _ok(await _service(request).search(user, body.query))


@router.get("/data-graph/semantics/review")
async def list_semantic_review(request: Request):
    user = require_user(request)
    return _ok(await _service(request).list_semantic_review(user))


@router.get("/data-graph/fields/{field_id}/sample")
async def sample_field(field_id: str, request: Request):
    user = require_user(request)
    try:
        return _ok(await _service(request).sample_field(user, field_id))
    except KamidsoError as exc:
        status = 404 if exc.code == KamidsoErrorCode.RESOURCE_NOT_FOUND else 400
        _raise_http(exc, status=status)


@router.patch("/data-graph/fields/{field_id}/semantic")
async def update_field_semantic(
    field_id: str, body: FieldSemanticRequest, request: Request
):
    user = require_user(request)
    try:
        return _ok(
            await _service(request).update_field_semantic(
                user,
                field_id,
                semantic_label=body.semantic_label,
                meaning=body.meaning,
                status=body.status,
                facts=body.facts,
            )
        )
    except KamidsoError as exc:
        status = 404 if exc.code == KamidsoErrorCode.RESOURCE_NOT_FOUND else 400
        _raise_http(exc, status=status)


@router.post("/data-graph/fields/{field_id}/semantic/refresh")
async def refresh_field_semantic(
    field_id: str, body: DiscoverRequest, request: Request
):
    user = require_user(request)

    async def complete(prompt: str) -> str:
        return await engine_complete(prompt, body.model)

    try:
        return _ok(
            await _service(request).refresh_field_semantic(user, field_id, complete)
        )
    except KamidsoError as exc:
        status = 404 if exc.code == KamidsoErrorCode.RESOURCE_NOT_FOUND else 400
        _raise_http(exc, status=status)