leocrm/app/plugins/builtins/ai_assistant/services.py

"""Business logic for the AI Assistant plugin.

Uses LiteLLM for multi-provider LLM calls and PydanticAI for the agent loop.
Tools from the global ToolRegistry are wrapped as PydanticAI tools with
RBAC permission checks.
"""

from __future__ import annotations

import json
import logging
import uuid
from collections.abc import AsyncGenerator
from typing import Any

import litellm
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.ai.llm_client import llm_complete
from app.core.permissions import check_permission
from app.plugins.builtins.ai_assistant.models import (
    AIAgent,
    AIChatFolder,
    AIModel,
    AIPreset,
    AIProvider,
)
from app.plugins.builtins.ai_assistant.tool_registry import (
    AITool,
    get_tool_registry,
)

logger = logging.getLogger(__name__)

# Suppress litellm verbose logging
litellm.suppress_debug_info = True


# ─── Helpers ───

def mask_api_key(key: str) -> str:
    """Mask API key for display: show first 8 and last 4 chars."""
    if not key or len(key) <= 12:
        return "***"
    return f"{key[:8]}...{key[-4:]}"


def provider_to_response(provider: AIProvider) -> dict[str, Any]:
    return {
        "id": str(provider.id),
        "name": provider.name,
        "provider_type": provider.provider_type,
        "api_key": mask_api_key(provider.api_key),
        "base_url": provider.base_url,
        "is_active": provider.is_active,
        "is_default": provider.is_default,
        "config": provider.config or {},
        "created_at": provider.created_at.isoformat() if provider.created_at else None,
        "updated_at": provider.updated_at.isoformat() if provider.updated_at else None,
    }


def model_to_response(model: AIModel) -> dict[str, Any]:
    return {
        "id": str(model.id),
        "provider_id": str(model.provider_id),
        "model_id": model.model_id,
        "display_name": model.display_name,
        "context_window": model.context_window,
        "supports_tools": model.supports_tools,
        "supports_streaming": model.supports_streaming,
        "is_active": model.is_active,
        "config": model.config or {},
    }


def preset_to_response(preset: AIPreset) -> dict[str, Any]:
    return {
        "id": str(preset.id),
        "name": preset.name,
        "model_id": preset.model_id,
        "provider_id": str(preset.provider_id) if preset.provider_id else None,
        "temperature": preset.temperature,
        "max_tokens": preset.max_tokens,
        "top_p": preset.top_p,
        "system_prompt": preset.system_prompt,
        "config": preset.config or {},
        "is_active": preset.is_active,
    }


def agent_to_response(agent: AIAgent) -> dict[str, Any]:
    return {
        "id": str(agent.id),
        "name": agent.name,
        "description": agent.description,
        "system_prompt": agent.system_prompt,
        "preset_id": str(agent.preset_id) if agent.preset_id else None,
        "tool_ids": agent.tool_ids or [],
        "is_default": agent.is_default,
        "is_active": agent.is_active,
        "config": agent.config or {},
        "created_at": agent.created_at.isoformat() if agent.created_at else None,
        "updated_at": agent.updated_at.isoformat() if agent.updated_at else None,
    }


def folder_to_response(folder: AIChatFolder) -> dict[str, Any]:
    return {
        "id": str(folder.id),
        "name": folder.name,
        "parent_id": str(folder.parent_id) if folder.parent_id else None,
        "user_id": str(folder.user_id),
        "sort_order": folder.sort_order,
        "created_at": folder.created_at.isoformat() if folder.created_at else None,
    }


# ─── Provider/Model/Preset/Agent CRUD ───

async def get_default_provider(db: AsyncSession, tenant_id: uuid.UUID) -> AIProvider | None:
    """Get the default provider for a tenant."""
    result = await db.execute(
        select(AIProvider)
        .where(AIProvider.tenant_id == tenant_id)
        .where(AIProvider.is_default.is_(True))
        .limit(1)
    )
    return result.scalar_one_or_none()


async def get_provider_by_id(db: AsyncSession, provider_id: uuid.UUID, tenant_id: uuid.UUID) -> AIProvider | None:
    result = await db.execute(
        select(AIProvider)
        .where(AIProvider.id == provider_id)
        .where(AIProvider.tenant_id == tenant_id)
        .limit(1)
    )
    return result.scalar_one_or_none()


async def get_preset_by_id(db: AsyncSession, preset_id: uuid.UUID, tenant_id: uuid.UUID) -> AIPreset | None:
    result = await db.execute(
        select(AIPreset)
        .where(AIPreset.id == preset_id)
        .where(AIPreset.tenant_id == tenant_id)
        .limit(1)
    )
    return result.scalar_one_or_none()


async def get_agent_by_id(db: AsyncSession, agent_id: uuid.UUID, tenant_id: uuid.UUID) -> AIAgent | None:
    result = await db.execute(
        select(AIAgent)
        .where(AIAgent.id == agent_id)
        .where(AIAgent.tenant_id == tenant_id)
        .limit(1)
    )
    return result.scalar_one_or_none()


async def get_default_agent(db: AsyncSession, tenant_id: uuid.UUID) -> AIAgent | None:
    result = await db.execute(
        select(AIAgent)
        .where(AIAgent.tenant_id == tenant_id)
        .where(AIAgent.is_default.is_(True))
        .limit(1)
    )
    return result.scalar_one_or_none()


# ─── Comm-based Chat Helpers (replaces AIChatSession/AIChatMessage) ───

async def get_comm_messages(
    db: AsyncSession, conversation_id: uuid.UUID, tenant_id: uuid.UUID
) -> list[dict[str, Any]]:
    """Get message history from comm_messages for an AI conversation."""
    from app.plugins.builtins.kommunikation.models import CommMessage
    result = await db.execute(
        select(CommMessage)
        .where(CommMessage.conversation_id == conversation_id)
        .where(CommMessage.tenant_id == tenant_id)
        .order_by(CommMessage.created_at.asc())
    )
    msgs = list(result.scalars().all())
    return [{"role": m.sender_type if m.sender_type != "ai" else "assistant", "content": m.content} for m in msgs]


async def save_comm_message(
    db: AsyncSession,
    conversation_id: uuid.UUID,
    role: str,
    content: str,
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
) -> None:
    """Save a message to comm_messages for an AI conversation."""
    from app.plugins.builtins.kommunikation.models import CommMessage
    sender_type = "user" if role == "user" else "ai"
    msg = CommMessage(
        conversation_id=conversation_id,
        sender_id=user_id if role == "user" else None,
        sender_type=sender_type,
        content=content,
        content_format="text",
        tenant_id=tenant_id,
    )
    db.add(msg)
    await db.flush()


async def stream_chat_comm(
    db: AsyncSession,
    conversation_id: uuid.UUID,
    agent: AIAgent,
    user_message: str,
    user_context: dict[str, Any],
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
) -> AsyncGenerator[str, None]:
    """Stream chat response via SSE with tool-calling loop, using comm_messages."""
    history = await get_comm_messages(db, conversation_id, tenant_id)
    messages: list[dict[str, Any]] = list(history)

    messages.append({"role": "user", "content": user_message})
    await save_comm_message(db, conversation_id, "user", user_message, tenant_id, user_id)

    # Get agent tools — always include call_crm_api for full system access
    registry = get_tool_registry()
    tools = registry.get_by_names(agent.tool_ids or [])
    crm_api_tool = registry.get("call_crm_api")
    if crm_api_tool and crm_api_tool not in tools:
        tools.append(crm_api_tool)
    tool_schemas = [t.to_openai_schema() for t in tools] if tools else None

    # F01 (Astra P0): allowlist — only the tools offered above may execute.
    # A hallucinated/injected tool name must never reach a handler.
    allowed_tool_names = {t.name for t in tools}

    # Build LLM params
    params, model_id = await build_litellm_params(db, agent, messages, tenant_id)

    # Agent loop: LLM → tool calls → execute → feed back → repeat
    max_iterations = 5
    for iteration in range(max_iterations):
        if tool_schemas and iteration < max_iterations - 1:
            params["tools"] = tool_schemas
        elif "tools" in params:
            del params["tools"]

        collected_content = ""
        collected_tool_calls: list[dict[str, Any]] = []
        try:
            result = await llm_complete(
                model=params.get("model", "gpt-4o-mini"),
                messages=params.get("messages", []),
                temperature=params.get("temperature", 0.7),
                max_tokens=params.get("max_tokens", 2048),
                api_key=params.get("api_key"),
                api_base=params.get("api_base"),
                tools=params.get("tools"),
            )
            collected_content = result["content"]
            if collected_content:
                yield f"data: {json.dumps({'type': 'token', 'content': collected_content})}\n\n"
            raw_response = result["raw_response"]
            if hasattr(raw_response.choices[0].message, "tool_calls") and raw_response.choices[0].message.tool_calls:
                for tc in raw_response.choices[0].message.tool_calls:
                    collected_tool_calls.append({
                        "id": tc.id or "",
                        "function": {
                            "name": tc.function.name if tc.function else "",
                            "arguments": tc.function.arguments if tc.function and tc.function.arguments else "",
                        },
                    })
        except Exception as exc:
            logger.error("LLM error: %s", exc)
            yield f"data: {json.dumps({'type': 'error', 'content': str(exc)})}\n\n"
            await save_comm_message(db, conversation_id, "assistant", f"Error: {exc}", tenant_id, user_id)
            await db.commit()
            return

        if collected_tool_calls:
            await save_comm_message(
                db, conversation_id, "assistant", collected_content, tenant_id, user_id,
            )
            yield f"data: {json.dumps({'type': 'tool_calls', 'tools': [tc['function']['name'] for tc in collected_tool_calls]})}\n\n"

            for tc in collected_tool_calls:
                tool_name = tc["function"]["name"]
                try:
                    tool_args = json.loads(tc["function"]["arguments"])
                except json.JSONDecodeError:
                    tool_args = {}

                # F01 (Astra P0): allowlist enforcement — reject any tool
                # name that was not offered to the LLM before it reaches a
                # handler. registered ≠ permitted.
                tool = registry.get(tool_name)
                if tool_name not in allowed_tool_names:
                    logger.warning(
                        "stream_chat_comm F01 guard: tool '%s' is registered but NOT offered to this agent — rejected",
                        tool_name,
                    )
                    result = f"Error: Tool '{tool_name}' is not available to this agent"
                elif tool is None:
                    result = f"Tool '{tool_name}' not found"
                else:
                    result = await execute_tool_call(tool, tool_args, user_context)

                yield f"data: {json.dumps({'type': 'tool_result', 'tool': tool_name, 'result': result[:500]})}\n\n"

                messages.append({
                    "role": "assistant",
                    "content": collected_content,
                    "tool_calls": collected_tool_calls,
                })
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc["id"],
                    "name": tool_name,
                    "content": result,
                })

            params, model_id = await build_litellm_params(db, agent, messages, tenant_id)
            continue

        # No tool calls — final response
        await save_comm_message(db, conversation_id, "assistant", collected_content, tenant_id, user_id)
        await db.commit()
        yield f"data: {json.dumps({'type': 'done', 'content': collected_content})}\n\n"
        return

    # Max iterations reached
    await save_comm_message(db, conversation_id, "assistant", collected_content, tenant_id, user_id)
    await db.commit()
    yield f"data: {json.dumps({'type': 'done', 'content': collected_content})}\n\n"

async def build_litellm_params(
    db: AsyncSession,
    agent: AIAgent,
    messages: list[dict[str, Any]],
    tenant_id: uuid.UUID,
) -> dict[str, Any]:
    """Build LiteLLM completion parameters from agent + preset."""
    # Get preset
    preset: AIPreset | None = None
    if agent.preset_id:
        preset = await get_preset_by_id(db, agent.preset_id, tenant_id)

    model_id = "gpt-4o-mini"  # fallback
    temperature = 0.7
    max_tokens = 2048
    top_p = 1.0
    system_prompt = agent.system_prompt or ""

    if preset:
        model_id = preset.model_id
        temperature = preset.temperature
        max_tokens = preset.max_tokens
        top_p = preset.top_p
        if preset.system_prompt and not system_prompt:
            system_prompt = preset.system_prompt

    # Get provider for API key
    provider: AIProvider | None = None
    if preset and preset.provider_id:
        provider = await get_provider_by_id(db, preset.provider_id, tenant_id)
    if not provider:
        provider = await get_default_provider(db, tenant_id)

    # Build litellm model string with provider prefix
    # LiteLLM always needs a provider prefix, even for OpenAI-compatible APIs
    litellm_model = model_id
    if provider:
        litellm_model = f"{provider.provider_type}/{model_id}"

    # Build messages with system prompt
    litellm_messages = []
    if system_prompt:
        # Inject CRM API context into system prompt
        try:
            from app.plugins.builtins.ai_assistant.crm_api_tool import get_api_context_for_prompt
            api_context = await get_api_context_for_prompt()
            system_prompt = system_prompt + api_context
        except Exception:
            logger.warning("Failed to inject API context into system prompt")
        litellm_messages.append({"role": "system", "content": system_prompt})
    litellm_messages.extend(messages)

    params: dict[str, Any] = {
        "model": litellm_model,
        "messages": litellm_messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "top_p": top_p,
        "stream": True,
    }

    # Set API key from provider
    if provider and provider.api_key:
        if provider.provider_type == "openai":
            params["api_key"] = provider.api_key
        elif provider.provider_type == "anthropic":
            params["api_key"] = provider.api_key
        else:
            params["api_key"] = provider.api_key

    if provider and provider.base_url:
        params["api_base"] = provider.base_url

    return params, model_id


async def execute_tool_call(
    tool: AITool,
    arguments: dict[str, Any],
    user_context: dict[str, Any],
) -> str:
    """Execute a tool call with RBAC permission check."""
    # Check RBAC permission if tool requires one
    if tool.required_permission:
        if not check_permission(user_context, tool.required_permission):
            return f"Error: Permission '{tool.required_permission}' required for tool '{tool.name}'"

    try:
        result = await tool.handler(arguments=arguments, context=user_context)
        return result
    except Exception as exc:
        logger.warning("Tool '%s' execution error: %s", tool.name, exc)
        return f"Error executing tool '{tool.name}': {exc}"


# ─── Seed Defaults ───

async def seed_defaults(db: AsyncSession) -> None:
    """Seed default provider, preset, and agent for existing tenants."""
    from app.models import Tenant

    result = await db.execute(select(Tenant))
    tenants = list(result.scalars().all())

    for tenant in tenants:
        # Check if default provider already exists
        existing = await get_default_provider(db, tenant.id)
        if existing:
            continue

        # Create default OpenAI provider (no key, user must configure)
        provider = AIProvider(
            name="OpenAI",
            provider_type="openai",
            api_key="",
            base_url="",
            is_active=True,
            is_default=True,
            config={},
            tenant_id=tenant.id,
        )
        db.add(provider)
        await db.flush()

        # Create default preset
        preset = AIPreset(
            name="Standard",
            model_id="gpt-4o-mini",
            provider_id=provider.id,
            temperature=0.7,
            max_tokens=2048,
            top_p=1.0,
            system_prompt="Du bist ein hilfreicher KI-Assistent für ein CRM-System.",
            is_active=True,
            tenant_id=tenant.id,
        )
        db.add(preset)
        await db.flush()

        # Create default agent
        agent = AIAgent(
            name="Standard Assistent",
            description="Allgemeiner KI-Assistent",
            system_prompt="Du bist ein hilfreicher KI-Assistent. Antworte präzise und hilfreich.",
            preset_id=preset.id,
            tool_ids=[],
            is_default=True,
            is_active=True,
            config={},
            tenant_id=tenant.id,
        )
        db.add(agent)

    await db.commit()
    logger.info("AI Assistant: default providers, presets, and agents seeded")