"""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 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: 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")