Agent Zero
511 lines
17 KiB
Python
511 lines
17 KiB
Python
"""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")
|