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 typing import Any, AsyncGenerator

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

from app.core.permissions import check_permission
from app.plugins.builtins.ai_assistant.models import (
    AIAgent,
    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,
        "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,
    }


# ─── 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 == 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 == 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_model = model_id
    if provider and provider.provider_type != "openai":
        litellm_model = f"{provider.provider_type}/{model_id}"

    # Build messages with system prompt
    litellm_messages = []
    if 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 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.

    Yields SSE-formatted strings.
    """
    # Load conversation history
    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})

    # Add user message
    messages.append({"role": "user", "content": user_message})
    await save_message(db, session.id, "user", user_message, tenant_id)

    # Get agent tools
    registry = get_tool_registry()
    tools = registry.get_by_names(agent.tool_ids or [])
    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
        if tool_schemas:
            params["tools"] = tool_schemas

        # Stream LLM response
        collected_content = ""
        collected_tool_calls: list[dict[str, Any]] = []
        try:
            response = await litellm.acompletion(**params)
            async for chunk in response:
                delta = chunk.choices[0].delta
                if delta.content:
                    collected_content += delta.content
                    yield f"data: {json.dumps({'type': 'token', 'content': delta.content})}\n\n"
                if delta.tool_calls:
                    for tc in delta.tool_calls:
                        idx = tc.index
                        while len(collected_tool_calls) <= idx:
                            collected_tool_calls.append({"id": "", "function": {"name": "", "arguments": ""}})
                        if tc.id:
                            collected_tool_calls[idx]["id"] = tc.id
                        if tc.function:
                            if tc.function.name:
                                collected_tool_calls[idx]["function"]["name"] += tc.function.name
                            if tc.function.arguments:
                                collected_tool_calls[idx]["function"]["arguments"] += tc.function.arguments
        except Exception as exc:
            logger.error("LLM streaming 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")