kamidso-back/kamidso/integrations/engine/engine.py

from __future__ import annotations

import json
import logging
import os
import time
from pathlib import Path
from typing import Any, AsyncIterator, Dict, List, Optional

from kamidso.application.run_registry import ActiveRunRegistry
from kamidso.assistant.context_builder import KamidsoContextBuilder
from kamidso.contracts.chat import KamidsoChatRequest
from kamidso.contracts.resource import ResourceRef
from kamidso.contracts.run import KamidsoRun
from kamidso.contracts.user import CurrentUser
from kamidso.integrations.engine.events import (
    begin_adapter_run,
    map_engine_event,
    release_adapter_run,
)
from kamidso.integrations.engine.runtime import EngineRuntime
from kamidso.resources.resolver import MaterializedResource
from kamidso.settings.flags import KamidsoSettings

logger = logging.getLogger("kamidso.engine")


class KamidsoEngine:
    name = "kamidso"

    def __init__(
        self,
        runtime: EngineRuntime,
        settings: KamidsoSettings,
        active: ActiveRunRegistry,
    ) -> None:
        self._runtime = runtime
        self._settings = settings
        self._active = active
        self._builder = KamidsoContextBuilder()

    async def health(self) -> Dict[str, str]:
        if not self._runtime.ready:
            payload: Dict[str, str] = {"status": "unavailable"}
            if self._runtime.error:
                payload["error"] = self._runtime.error
            if self._runtime.config_path:
                payload["config"] = self._runtime.config_path
            return payload
        return {"status": "ready", "config": self._runtime.config_path}

    async def cancel(
        self, execution_id: str, *, conversation_id: str | None = None
    ) -> None:
        self._active.request_cancel(execution_id)
        if conversation_id:
            from kamidso.engine.app.openapi.api_v1.agentic_data_api import (
                _cancel_conversation_agent,
            )

            await _cancel_conversation_agent(conversation_id)

    async def execute(
        self,
        *,
        user: CurrentUser,
        run: KamidsoRun,
        request: KamidsoChatRequest,
        materialized: List[MaterializedResource],
        prior_history: Optional[List[Dict[str, str]]] = None,
        last_report: Optional[Dict[str, Any]] = None,
    ) -> AsyncIterator[Dict[str, Any]]:
        if not self._runtime.ready:
            yield {
                "type": "run.failed",
                "run_id": run.id,
                "error": {"code": "ENGINE_UNAVAILABLE", "message": "Engine runtime unavailable"},
            }
            return

        begin_adapter_run(run.id)
        try:
            async for event in self._execute_ready(
                user=user,
                run=run,
                request=request,
                materialized=materialized,
                prior_history=prior_history,
                last_report=last_report,
            ):
                yield event
        finally:
            # Failure, cancellation, and client disconnect all land here,
            # after the last event has already been yielded.
            release_adapter_run(run.id)

    async def _execute_ready(
        self,
        *,
        user: CurrentUser,
        run: KamidsoRun,
        request: KamidsoChatRequest,
        materialized: List[MaterializedResource],
        prior_history: Optional[List[Dict[str, str]]] = None,
        last_report: Optional[Dict[str, Any]] = None,
    ) -> AsyncIterator[Dict[str, Any]]:
        from kamidso.engine.app.openapi.api_v1.agentic_data_api import (
            _open_turn_attachments,
            _react_agent_stream,
        )
        from kamidso.engine.app.openapi.api_view_model import ConversationVo
        from kamidso.engine.serve.utils.auth import UserRequest

        ext_info: Dict[str, Any] = {
            "kamidso_run_id": run.id,
            "workspace_id": user.workspace_id,
        }
        if prior_history:
            ext_info["kamidso_history"] = prior_history
        if last_report:
            ext_info["kamidso_last_report"] = last_report
        staged_names, staged_paths = _stage_files_for_engine(materialized)
        if staged_paths:
            ext_info["input_files"] = staged_names
            ext_info["file_path"] = staged_paths[0]
            ext_info["kamidso_resolved_file_paths"] = staged_paths
        ds_refs = [ref for ref in request.resources if ref.type == "datasource"]
        ds_ids = [ref.id for ref in ds_refs]
        ds_labels = [ref.name or ref.id for ref in ds_refs]
        kb_ids = [ref.id for ref in request.resources if ref.type == "knowledge"]
        kb_label = next((ref.name for ref in request.resources if ref.type == "knowledge" and ref.name), None)
        if ds_ids:
            ext_info["database_name"] = ds_ids[0]
            ext_info["database_names"] = ds_ids
            ext_info["database_labels"] = ds_labels
            if ds_labels:
                ext_info["database_label"] = ds_labels[0]
        if kb_ids:
            ext_info["knowledge_space_name"] = kb_ids[0]
            if kb_label:
                ext_info["knowledge_space_label"] = kb_label
        connector_refs = [ref for ref in request.resources if ref.type == "connector" and ref.id]
        if connector_refs:
            ext_info.update(connector_selection(connector_refs))

        context = self._builder.build(
            user,
            model=request.model_id,
            skill=request.skill_id,
            files=[ref.id for ref in request.resources if ref.type == "file"],
            datasources=ds_ids,
            knowledge=kb_ids,
            ext_info=ext_info,
            skill_context=request.skill_context,
        )
        ext_info["kamidso_behavior"] = context.behavior
        ext_info["agent_language"] = "en"
        ext_info["response_language"] = context.locale
        ext_info["reasoning_mode"] = request.reasoning_mode or "fast"
        ext_info["report_mode"] = bool(getattr(request, "report_mode", False))
        has_files = any(ref.type == "file" for ref in request.resources)
        if request.skill_id:
            ext_info["skill_name"] = request.skill_id
            ext_info["skill_source"] = "user"
        elif context.skills:
            ext_info["skill_name"] = context.skills[0]
            ext_info["skill_source"] = "file" if has_files else "user"
        if has_files:
            ext_info["analysis_mode"] = context.analysis_mode or "file"
        attached_ids = [
            ref.id
            for ref in request.resources
            if ref.type in {"file", "datasource"} and ref.id
        ]
        try:
            from kamidso.data_graph.tools.graph_search import (
                confirmed_links_for_attached,
                field_facts_for_attached,
            )

            links = await confirmed_links_for_attached(user.workspace_id, attached_ids)
            facts = await field_facts_for_attached(user.workspace_id, attached_ids)
        except Exception:
            logger.exception("organization links were not attached")
            links = None
            facts = None
        if links is not None:
            ext_info["organization_links"] = links
        if facts is not None:
            ext_info["organization_field_facts"] = facts

        model_name = request.model_id or os.environ.get("LLM_MODEL_NAME")
        dialogue = ConversationVo(
            conv_uid=run.conversation_id,
            user_input=run.input,
            user_name=user.id,
            chat_mode="chat_react_agent",
            model_name=model_name,
            temperature=request.temperature,
            max_new_tokens=request.max_new_tokens,
            select_param="",
            incremental=False,
            ext_info=ext_info,
        )
        user_token = UserRequest(
            user_id=user.id,
            user_name=user.id,
            role=user.role,
            nick_name=user.name,
            email=user.email,
            real_name=user.name,
        )

        started = False
        started_at = time.perf_counter()
        agent_stream = None
        try:
            attachment_ctx = await _open_turn_attachments(dialogue, user_token)
            agent_stream = _react_agent_stream(
                dialogue, tool_mode="full", attachment_ctx=attachment_ctx
            )
            async for chunk in agent_stream:
                if self._active.is_cancelled(run.id):
                    yield {
                        "type": "run.cancelled",
                        "run_id": run.id,
                        "error": {"code": "RUN_CANCELLED"},
                    }
                    return
                payload = _parse_sse(chunk)
                if payload is None:
                    continue
                for event in map_engine_event(run.id, payload, self._settings.product.locale):
                    if not started and event.get("type") == "run.started":
                        started = True
                    if event.get("type") == "run.completed":
                        event.setdefault("usage", {})
                        event["usage"]["duration_ms"] = int((time.perf_counter() - started_at) * 1000)
                    yield event
        except Exception as exc:  # noqa: BLE001
            logger.exception("kamidso.engine.core execution failed")
            yield {
                "type": "run.failed",
                "run_id": run.id,
                "error": {
                    "code": "ENGINE_EXECUTION_FAILED",
                    "message": "اجرای تحلیل با خطا مواجه شد.",
                    "retryable": True,
                    "details": str(exc),
                },
            }
        finally:
            if agent_stream is not None:
                close = getattr(agent_stream, "aclose", None)
                if callable(close):
                    try:
                        await close()
                    except Exception:  # noqa: BLE001
                        logger.warning("Failed to close ReAct agent stream", exc_info=True)


def connector_selection(resources: List[Any]) -> Dict[str, List[str]]:
    """Map attached connector resources onto the existing connector_ids field."""
    refs = [ref for ref in resources if getattr(ref, "type", None) == "connector" and getattr(ref, "id", None)]
    return {
        "connector_ids": [ref.id for ref in refs],
        "connector_labels": [ref.name or ref.id for ref in refs],
    }


def _stage_files_for_engine(
    materialized: List[MaterializedResource],
) -> tuple[List[str], List[str]]:
    """Pass run files from the unified data tree without copying them.

    Returns relative names and resolved absolute paths under ``data/uploads``
    or ``data/runtime``.
    """
    relative: List[str] = []
    resolved: List[str] = []
    for item in materialized:
        if not item.runtime_path:
            continue
        source = Path(item.runtime_path)
        if not source.is_file():
            logger.warning("Materialized file missing for kamidso.engine.core: %s", source)
            continue
        relative.append(source.name)
        resolved.append(str(source.resolve()))
    return relative, resolved


def _parse_sse(chunk: str) -> Optional[Dict[str, Any]]:
    if not isinstance(chunk, str):
        return None
    line = chunk.splitlines()[0] if chunk else ""
    if not line.startswith("data:"):
        return None
    raw = line[5:].strip()
    if not raw:
        return None
    try:
        payload = json.loads(raw)
    except (json.JSONDecodeError, TypeError, ValueError):
        return None
    return payload if isinstance(payload, dict) else None