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 aiofiles
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,
    AIChatAttachment,
    AIChatFolder,
    AIChatMessage,
    AIChatSession,
    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 session_to_response(session: AIChatSession) -> dict[str, Any]:
    return {
        "id": str(session.id),
        "user_id": str(session.user_id),
        "agent_id": str(session.agent_id) if session.agent_id else None,
        "title": session.title,
        "is_pinned": session.is_pinned,
        "is_sidebar": session.is_sidebar,
        "folder_id": str(session.folder_id) if session.folder_id else None,
        "sort_order": session.sort_order,
        "created_at": session.created_at.isoformat() if session.created_at else None,
        "updated_at": session.updated_at.isoformat() if session.updated_at else None,
    }


def message_to_response(msg: AIChatMessage) -> dict[str, Any]:
    return {
        "id": str(msg.id),
        "session_id": str(msg.session_id),
        "role": msg.role,
        "content": msg.content,
        "tool_calls": msg.tool_calls if msg.tool_calls else None,
        "tool_results": msg.tool_results if msg.tool_results else None,
        "tokens": msg.tokens,
        "model_used": msg.model_used,
        "created_at": msg.created_at.isoformat() if msg.created_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,
    }


def attachment_to_response(att: AIChatAttachment) -> dict[str, Any]:
    return {
        "id": str(att.id),
        "message_id": str(att.message_id) if att.message_id else None,
        "session_id": str(att.session_id),
        "filename": att.filename,
        "mime_type": att.mime_type,
        "size_bytes": att.size_bytes,
    }


# ─── 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()


# ─── Session/Message helpers ───

async def get_session_by_id(
    db: AsyncSession, session_id: uuid.UUID, user_id: uuid.UUID, tenant_id: uuid.UUID
) -> AIChatSession | None:
    result = await db.execute(
        select(AIChatSession)
        .where(AIChatSession.id == session_id)
        .where(AIChatSession.user_id == user_id)
        .where(AIChatSession.tenant_id == tenant_id)
        .limit(1)
    )
    return result.scalar_one_or_none()


async def get_session_messages(
    db: AsyncSession, session_id: uuid.UUID, tenant_id: uuid.UUID
) -> list[AIChatMessage]:
    result = await db.execute(
        select(AIChatMessage)
        .where(AIChatMessage.session_id == session_id)
        .where(AIChatMessage.tenant_id == tenant_id)
        .order_by(AIChatMessage.created_at.asc())
    )
    return list(result.scalars().all())


async def save_message(
    db: AsyncSession,
    session_id: uuid.UUID,
    role: str,
    content: str,
    tenant_id: uuid.UUID,
    tool_calls: list | None = None,
    tool_results: list | None = None,
    tokens: int = 0,
    model_used: str = "",
) -> AIChatMessage:
    msg = AIChatMessage(
        session_id=session_id,
        role=role,
        content=content,
        tool_calls=tool_calls,
        tool_results=tool_results,
        tokens=tokens,
        model_used=model_used,
        tenant_id=tenant_id,
    )
    db.add(msg)
    await db.flush()
    return msg


# ─── LLM Chat with Tool Loop ───

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}"


async def _extract_attachment_content(
    db: AsyncSession,
    session_id: uuid.UUID,
    tenant_id: uuid.UUID,
) -> str:
    """Extract text content from session attachments for LLM context."""
    result = await db.execute(
        select(AIChatAttachment)
        .where(AIChatAttachment.session_id == session_id)
        .where(AIChatAttachment.tenant_id == tenant_id)
        .order_by(AIChatAttachment.created_at.asc())
    )
    attachments = list(result.scalars().all())
    if not attachments:
        return ""

    parts: list[str] = []
    for att in attachments:
        try:
            async with aiofiles.open(att.storage_path, "rb") as f:
                content = await f.read()

            text_content = ""
            mime = att.mime_type.lower()

            if mime.startswith("text/") or att.filename.endswith((".txt", ".md", ".csv", ".json", ".yaml", ".yml", ".py", ".js", ".ts", ".html", ".xml")):
                text_content = content.decode("utf-8", errors="replace")
            elif mime == "application/pdf" or att.filename.endswith(".pdf"):
                try:
                    from io import BytesIO

                    from pypdf import PdfReader
                    reader = PdfReader(BytesIO(content))
                    text_content = "\n".join(page.extract_text() or "" for page in reader.pages)
                except ImportError:
                    text_content = f"[PDF file: {att.filename} - extraction not available]"
            elif mime.startswith("image/"):
                text_content = f"[Image file: {att.filename} ({att.mime_type}, {att.size_bytes} bytes)]"
            else:
                text_content = f"[Binary file: {att.filename} ({att.mime_type}, {att.size_bytes} bytes)]"

            if len(text_content) > 10000:
                text_content = text_content[:10000] + "\n... [truncated]"

            parts.append(f"--- Attachment: {att.filename} ---\n{text_content}")
        except Exception as exc:
            logger.warning("Failed to extract attachment %s: %s", att.filename, exc)
            parts.append(f"--- Attachment: {att.filename} (extraction failed) ---")

    return "\n\n".join(parts)


async def stream_chat(
    db: AsyncSession,
    session: AIChatSession,
    agent: AIAgent,
    user_message: str,
    user_context: dict[str, Any],
    tenant_id: uuid.UUID,
) -> AsyncGenerator[str, None]:
    """Stream chat response via SSE with tool-calling loop."""
    history = await get_session_messages(db, session.id, tenant_id)
    messages: list[dict[str, Any]] = []
    for msg in history:
        messages.append({"role": msg.role, "content": msg.content})

    attachment_content = await _extract_attachment_content(db, session.id, tenant_id)
    full_message = user_message
    if attachment_content:
        full_message = f"{user_message}\n\n--- Attached Files ---\n{attachment_content}"

    messages.append({"role": "user", "content": full_message})
    await save_message(db, session.id, "user", user_message, tenant_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 [])
    # Ensure call_crm_api is always available
    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

    # 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):
        # Add tools to params if available, but NOT on the last iteration
        # to force the LLM to give a final answer instead of looping
        if tool_schemas and iteration < max_iterations - 1:
            params["tools"] = tool_schemas
        elif "tools" in params:
            del params["tools"]

        # LLM response via llm_complete (non-streaming)
        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"
            # Extract tool calls from raw response
            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_message(db, session.id, "assistant", f"Error: {exc}", tenant_id, model_used=model_id)
            await db.commit()
            return

        # If tool calls, execute them and continue loop
        if collected_tool_calls:
            # Save assistant message with tool calls
            await save_message(
                db, session.id, "assistant", collected_content, tenant_id,
                tool_calls=collected_tool_calls, model_used=model_id,
            )
            yield f"data: {json.dumps({'type': 'tool_calls', 'tools': [tc['function']['name'] for tc in collected_tool_calls]})}\n\n"

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

                tool = registry.get(tool_name)
                if 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"

                # Add tool result to messages for next iteration
                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,
                })

            # Update params with new messages for next iteration
            params, model_id = await build_litellm_params(db, agent, messages, tenant_id)
            continue

        # No tool calls — final response
        await save_message(db, session.id, "assistant", collected_content, tenant_id, model_used=model_id)
        await db.commit()
        yield f"data: {json.dumps({'type': 'done', 'content': collected_content})}\n\n"
        return

    # Max iterations reached
    await save_message(db, session.id, "assistant", collected_content, tenant_id, model_used=model_id)
    await db.commit()
    yield f"data: {json.dumps({'type': 'done', 'content': collected_content})}\n\n"


# ─── 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")