Agent Zero
309 lines
10 KiB
Python
309 lines
10 KiB
Python
"""External Agent API — allows external systems to interact with AI agents.
|
|
|
|
Provides REST endpoints for running agents, checking status, and streaming
|
|
responses. Authenticated via Bearer API tokens with rate limiting.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Request
|
|
from fastapi.responses import StreamingResponse
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.db import get_db, set_tenant_context
|
|
from app.deps import get_current_user_bearer, require_permission
|
|
from app.plugins.builtins.ai_assistant.schemas import (
|
|
ExternalAgentRequest,
|
|
ExternalAgentResponse,
|
|
)
|
|
from app.plugins.builtins.ai_assistant.services import stream_chat_comm as stream_chat
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/api/v1/external/agent", tags=["external-agent"])
|
|
|
|
|
|
async def _check_external_rate_limit(request: Request, tenant_id: str, token_prefix: str) -> None:
|
|
"""Check rate limit: 10 requests/minute per token."""
|
|
from app.core.rate_limit import check_rate_limit
|
|
|
|
redis_key = f"rate:external_agent:{tenant_id}:{token_prefix}"
|
|
await check_rate_limit(redis_key, max_attempts=10, window_seconds=60)
|
|
|
|
|
|
@router.post(
|
|
"/{agent_id}/run",
|
|
dependencies=[Depends(require_permission("ai:write"))],
|
|
)
|
|
async def run_agent_external(
|
|
agent_id: str,
|
|
data: ExternalAgentRequest,
|
|
request: Request,
|
|
current_user: dict[str, Any] = Depends(get_current_user_bearer),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
"""Execute an AI agent with a message input.
|
|
|
|
Authenticated via Bearer API token. Rate limited to 10 requests/minute.
|
|
"""
|
|
tenant_id = uuid.UUID(current_user["tenant_id"])
|
|
await set_tenant_context(db, tenant_id)
|
|
|
|
# Rate limiting
|
|
token_prefix = current_user.get("token_prefix", "default")
|
|
await _check_external_rate_limit(request, str(tenant_id), token_prefix)
|
|
|
|
try:
|
|
aid = uuid.UUID(agent_id)
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(status_code=400, detail="Invalid agent ID") from None
|
|
|
|
# Find the agent
|
|
from app.plugins.builtins.ai_assistant.models import AIAgent
|
|
|
|
result = await db.execute(
|
|
select(AIAgent)
|
|
.where(AIAgent.id == aid)
|
|
.where(AIAgent.tenant_id == tenant_id)
|
|
.limit(1)
|
|
)
|
|
agent = result.scalar_one_or_none()
|
|
if agent is None:
|
|
raise HTTPException(status_code=404, detail="Agent not found")
|
|
if not agent.is_active:
|
|
raise HTTPException(status_code=400, detail="Agent is not active")
|
|
|
|
# Create or find a session for this external interaction
|
|
|
|
# AIChatSession/AIChatMessage removed — using comm tables
|
|
|
|
# Create a comm conversation for this external interaction
|
|
from app.plugins.builtins.kommunikation.models import CommConversation
|
|
session = CommConversation(
|
|
tenant_id=tenant_id,
|
|
title=f"External: {data.message[:50]}" if data.message else "External Agent Run",
|
|
owner_id=uuid.UUID(current_user["user_id"]),
|
|
created_by=uuid.UUID(current_user["user_id"]),
|
|
created_by_type="user",
|
|
metadata_={"conversation_type": "ai", "agent_id": str(agent.id)},
|
|
)
|
|
db.add(session)
|
|
await db.flush()
|
|
|
|
# Store the user message
|
|
from app.plugins.builtins.kommunikation.models import CommMessage
|
|
user_msg = CommMessage(
|
|
conversation_id=session.id,
|
|
sender_id=uuid.UUID(current_user["user_id"]),
|
|
sender_type="user",
|
|
content=data.message,
|
|
content_format="text",
|
|
tenant_id=tenant_id,
|
|
)
|
|
db.add(user_msg)
|
|
await db.flush()
|
|
|
|
# Build user context for RBAC
|
|
user_context = {
|
|
"user_id": current_user["user_id"],
|
|
"tenant_id": current_user["tenant_id"],
|
|
"role": current_user.get("role", ""),
|
|
"permissions": current_user.get("permissions", []),
|
|
"denied_permissions": current_user.get("denied_permissions", []),
|
|
"is_system_admin": current_user.get("is_system_admin", False),
|
|
"field_permissions": current_user.get("field_permissions", {}),
|
|
}
|
|
|
|
# Run the agent via streaming chat (non-streaming mode)
|
|
|
|
full_response = ""
|
|
async with get_db() as stream_db:
|
|
await set_tenant_context(stream_db, tenant_id)
|
|
async for chunk in stream_chat(
|
|
stream_db,
|
|
session.id,
|
|
agent,
|
|
data.message,
|
|
user_context,
|
|
tenant_id,
|
|
uuid.UUID(current_user["user_id"]),
|
|
):
|
|
if chunk.startswith("data: ") and chunk != "data: [DONE]\n\n":
|
|
try:
|
|
payload = json.loads(chunk[6:])
|
|
if "content" in payload:
|
|
full_response += payload["content"]
|
|
except (json.JSONDecodeError, IndexError):
|
|
pass
|
|
|
|
# Count tokens (approximate)
|
|
tokens_used = len(full_response.split()) + len(data.message.split())
|
|
|
|
return ExternalAgentResponse(
|
|
response=full_response,
|
|
agent_id=str(agent.id),
|
|
session_id=str(session.id),
|
|
tokens_used=tokens_used,
|
|
)
|
|
|
|
|
|
@router.get(
|
|
"/{agent_id}/status",
|
|
dependencies=[Depends(require_permission("ai:read"))],
|
|
)
|
|
async def get_agent_status_external(
|
|
agent_id: str,
|
|
request: Request,
|
|
current_user: dict[str, Any] = Depends(get_current_user_bearer),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
"""Get the status of an AI agent.
|
|
|
|
Returns agent configuration, active status, and recent run stats.
|
|
"""
|
|
tenant_id = uuid.UUID(current_user["tenant_id"])
|
|
await set_tenant_context(db, tenant_id)
|
|
|
|
# Rate limiting
|
|
token_prefix = current_user.get("token_prefix", "default")
|
|
await _check_external_rate_limit(request, str(tenant_id), token_prefix)
|
|
|
|
try:
|
|
aid = uuid.UUID(agent_id)
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(status_code=400, detail="Invalid agent ID") from None
|
|
|
|
from app.plugins.builtins.ai_assistant.models import AIAgent
|
|
|
|
result = await db.execute(
|
|
select(AIAgent)
|
|
.where(AIAgent.id == aid)
|
|
.where(AIAgent.tenant_id == tenant_id)
|
|
.limit(1)
|
|
)
|
|
agent = result.scalar_one_or_none()
|
|
if agent is None:
|
|
raise HTTPException(status_code=404, detail="Agent not found")
|
|
|
|
# Get recent run stats
|
|
from sqlalchemy import func
|
|
|
|
from app.plugins.builtins.automation.contracts import AutomationContract
|
|
|
|
recent_runs = await db.execute(
|
|
select(func.count())
|
|
.select_from(AutomationContract.AgentRun)
|
|
.where(AutomationContract.AgentRun.agent_id == aid)
|
|
.where(AutomationContract.AgentRun.tenant_id == tenant_id)
|
|
)
|
|
total_runs = recent_runs.scalar() or 0
|
|
|
|
return {
|
|
"agent_id": str(agent.id),
|
|
"name": agent.name,
|
|
"description": agent.description,
|
|
"is_active": agent.is_active,
|
|
"tool_count": len(agent.tool_ids or []),
|
|
"total_runs": total_runs,
|
|
"created_at": agent.created_at.isoformat() if agent.created_at else None,
|
|
"updated_at": agent.updated_at.isoformat() if agent.updated_at else None,
|
|
}
|
|
|
|
|
|
@router.post(
|
|
"/{agent_id}/stream",
|
|
dependencies=[Depends(require_permission("ai:write"))],
|
|
)
|
|
async def stream_agent_external(
|
|
agent_id: str,
|
|
data: ExternalAgentRequest,
|
|
request: Request,
|
|
current_user: dict[str, Any] = Depends(get_current_user_bearer),
|
|
db: AsyncSession = Depends(get_db),
|
|
):
|
|
"""Stream an AI agent response via SSE.
|
|
|
|
Authenticated via Bearer API token. Rate limited to 10 requests/minute.
|
|
Returns Server-Sent Events stream.
|
|
"""
|
|
tenant_id = uuid.UUID(current_user["tenant_id"])
|
|
await set_tenant_context(db, tenant_id)
|
|
|
|
# Rate limiting
|
|
token_prefix = current_user.get("token_prefix", "default")
|
|
await _check_external_rate_limit(request, str(tenant_id), token_prefix)
|
|
|
|
try:
|
|
aid = uuid.UUID(agent_id)
|
|
except (ValueError, TypeError):
|
|
raise HTTPException(status_code=400, detail="Invalid agent ID") from None
|
|
|
|
from app.plugins.builtins.ai_assistant.models import AIAgent
|
|
|
|
result = await db.execute(
|
|
select(AIAgent)
|
|
.where(AIAgent.id == aid)
|
|
.where(AIAgent.tenant_id == tenant_id)
|
|
.limit(1)
|
|
)
|
|
agent = result.scalar_one_or_none()
|
|
if agent is None:
|
|
raise HTTPException(status_code=404, detail="Agent not found")
|
|
if not agent.is_active:
|
|
raise HTTPException(status_code=400, detail="Agent is not active")
|
|
|
|
# Create session
|
|
# AIChatSession removed — using comm tables
|
|
|
|
from app.plugins.builtins.kommunikation.models import CommConversation
|
|
session = CommConversation(
|
|
tenant_id=tenant_id,
|
|
title=f"External Stream: {data.message[:50]}" if data.message else "External Agent Stream",
|
|
owner_id=uuid.UUID(current_user["user_id"]),
|
|
created_by=uuid.UUID(current_user["user_id"]),
|
|
created_by_type="user",
|
|
metadata_={"conversation_type": "ai", "agent_id": str(agent.id)},
|
|
)
|
|
db.add(session)
|
|
await db.commit()
|
|
|
|
# Build user context
|
|
user_context = {
|
|
"user_id": current_user["user_id"],
|
|
"tenant_id": current_user["tenant_id"],
|
|
"role": current_user.get("role", ""),
|
|
"permissions": current_user.get("permissions", []),
|
|
"denied_permissions": current_user.get("denied_permissions", []),
|
|
"is_system_admin": current_user.get("is_system_admin", False),
|
|
"field_permissions": current_user.get("field_permissions", {}),
|
|
}
|
|
|
|
from app.plugins.builtins.ai_assistant.services import stream_chat_comm
|
|
|
|
async def event_stream():
|
|
from app.core.db import get_session_factory
|
|
factory = get_session_factory()
|
|
async with factory() as stream_db:
|
|
await set_tenant_context(stream_db, tenant_id)
|
|
async for chunk in stream_chat_comm(
|
|
stream_db, session.id, agent, data.message, user_context, tenant_id, uuid.UUID(current_user["user_id"])
|
|
):
|
|
yield chunk
|
|
yield "data: [DONE]\n\n"
|
|
|
|
return StreamingResponse(
|
|
event_stream(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|