kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/app/policies/tool_enforcement/production_safeguards.rego

@@ -0,0 +1,80 @@
+# Production Safeguards Policy
+# Block high-risk and critical operations in production environment
+
+package kubiya.tool_enforcement
+
+import future.keywords.if
+
+# Default deny
+default allow = false
+
+# Block high-risk tools in production
+deny if {
+    input.execution.environment == "production"
+    input.tool.risk_level == "critical"
+}
+
+deny if {
+    input.execution.environment == "production"
+    input.tool.risk_level == "high"
+    input.tool.category == "command_execution"
+    contains(lower(input.tool.arguments.command), "delete")
+}
+
+deny if {
+    input.execution.environment == "production"
+    input.tool.category == "file_operation"
+    input.tool.name == "Write"
+    contains(input.tool.arguments.file_path, "/etc/")
+}
+
+# Allow everything in non-production
+allow if {
+    input.execution.environment != "production"
+    input.user.email != ""
+}
+
+# Allow safe operations in production
+allow if {
+    input.execution.environment == "production"
+    input.tool.risk_level == "low"
+}
+
+allow if {
+    input.execution.environment == "production"
+    input.tool.risk_level == "medium"
+    not deny
+}
+
+# Allow with admin override
+allow if {
+    input.user.roles[_] == "admin"
+}
+
+# Violation messages
+violation[msg] if {
+    deny
+    input.tool.risk_level == "critical"
+    msg := sprintf("Critical tool '%s' is blocked in production. Use dev/staging for testing.", [input.tool.name])
+}
+
+violation[msg] if {
+    deny
+    input.tool.risk_level == "high"
+    not input.tool.risk_level == "critical"
+    msg := sprintf("High-risk tool '%s' is restricted in production. Contact admin for approval.", [input.tool.name])
+}
+
+# Helper function
+lower(s) := lower_s {
+    lower_s := lower(s)
+}
+
+# Metadata
+metadata := {
+    "name": "production_safeguards",
+    "description": "Block high-risk operations in production environment",
+    "version": "1.0.0",
+    "author": "Kubiya",
+    "tags": ["production", "security", "risk_management"]
+}

control_plane_api/app/policies/tool_enforcement/role_based_tool_access.rego

@@ -0,0 +1,44 @@
+# Role-Based Tool Access Policy
+# Only admin and devops roles can execute command execution tools
+
+package kubiya.tool_enforcement
+
+import future.keywords.if
+import future.keywords.in
+
+# Default deny
+default allow = false
+
+# Allow if user has required role for command execution
+allow if {
+    input.tool.category == "command_execution"
+    input.user.roles[_] == "admin"
+}
+
+allow if {
+    input.tool.category == "command_execution"
+    input.user.roles[_] == "devops"
+}
+
+# Allow all other tool categories for any authenticated user
+allow if {
+    input.tool.category != "command_execution"
+    input.user.email != ""
+}
+
+# Violation message
+violation[msg] if {
+    input.tool.category == "command_execution"
+    not input.user.roles[_] == "admin"
+    not input.user.roles[_] == "devops"
+    msg := sprintf("Command execution tools require 'admin' or 'devops' role. User %s has roles: %v", [input.user.email, input.user.roles])
+}
+
+# Metadata
+metadata := {
+    "name": "role_based_tool_access",
+    "description": "Only admin and devops roles can execute command tools",
+    "version": "1.0.0",
+    "author": "Kubiya",
+    "tags": ["rbac", "security", "command_execution"]
+}

control_plane_api/app/policies/tool_restrictions.rego

@@ -0,0 +1,86 @@
+# Tool Restrictions Policy
+# Controls which tools agents can use, with different restrictions per environment
+
+package tool_restrictions
+
+import future.keywords.if
+import future.keywords.in
+
+# Default deny
+default allow = false
+
+# Dangerous tools that require explicit approval
+dangerous_tools := {
+    "docker",
+    "shell",
+    "file_system"
+}
+
+# Production-safe tools
+production_safe_tools := {
+    "python",
+    "sleep",
+    "file_generation"
+}
+
+# Development environment - allow all tools
+allow if {
+    input.environment == "development"
+}
+
+# Staging environment - allow non-dangerous tools
+allow if {
+    input.environment == "staging"
+    input.tool
+    not input.tool in dangerous_tools
+}
+
+# Production environment - only allow production-safe tools
+allow if {
+    input.environment == "production"
+    input.tool in production_safe_tools
+}
+
+# Allow dangerous tools in production with explicit approval
+allow if {
+    input.environment == "production"
+    input.tool in dangerous_tools
+    input.approved_by
+    input.approval_id
+}
+
+# Violations
+violations[msg] if {
+    input.environment == "production"
+    input.tool in dangerous_tools
+    not input.approved_by
+    msg := sprintf("Tool '%v' requires approval in production environment", [input.tool])
+}
+
+violations[msg] if {
+    input.environment == "staging"
+    input.tool in dangerous_tools
+    msg := sprintf("Dangerous tool '%v' not allowed in staging environment", [input.tool])
+}
+
+violations[msg] if {
+    input.environment == "production"
+    input.tool
+    not input.tool in production_safe_tools
+    not input.approved_by
+    msg := sprintf("Tool '%v' not allowed in production without approval", [input.tool])
+}
+
+violations[msg] if {
+    not input.tool
+    msg := "No tool specified in the request"
+}
+
+# Metadata
+metadata := {
+    "name": "tool-restrictions",
+    "description": "Control which tools agents can use based on environment and approval status",
+    "version": "1.0.0",
+    "author": "Kubiya",
+    "tags": ["security", "tools", "environment"]
+}

control_plane_api/app/routers/__init__.py

@@ -0,0 +1,4 @@
+# API Routers
+from . import agents, teams, workflows, health, executions, presence
+
+__all__ = ["agents", "teams", "workflows", "health", "executions", "presence"]

control_plane_api/app/routers/agents.py

@@ -0,0 +1,382 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi.responses import StreamingResponse
+from sqlalchemy.orm import Session
+from typing import List, Optional
+from datetime import datetime
+
+from control_plane_api.app.database import get_db
+from control_plane_api.app.middleware.auth import get_current_organization
+from control_plane_api.app.models.agent import Agent, AgentStatus
+from control_plane_api.app.models.execution import Execution, ExecutionStatus, ExecutionType
+from control_plane_api.app.services.litellm_service import litellm_service
+from control_plane_api.app.observability import (
+    create_span_with_context,
+    add_span_event,
+    add_span_error,
+    instrument_endpoint,
+)
+from pydantic import BaseModel, Field
+
+router = APIRouter()
+
+
+def agent_to_response(agent: Agent) -> dict:
+    """Convert Agent model to response dict, mapping model_config to llm_config"""
+    return {
+        "id": agent.id,
+        "name": agent.name,
+        "description": agent.description,
+        "status": agent.status,
+        "capabilities": agent.capabilities,
+        "configuration": agent.configuration,
+        "model_id": agent.model_id,
+        "llm_config": agent.model_config,
+        "team_id": agent.team_id,
+        "created_at": agent.created_at,
+        "updated_at": agent.updated_at,
+        "last_active_at": agent.last_active_at,
+        "state": agent.state,
+        "error_message": agent.error_message,
+    }
+
+
+# Pydantic schemas
+class AgentCreate(BaseModel):
+    name: str = Field(..., description="Agent name")
+    description: str | None = Field(None, description="Agent description")
+    capabilities: list = Field(default_factory=list, description="Agent capabilities")
+    configuration: dict = Field(default_factory=dict, description="Agent configuration")
+    model_id: str | None = Field(None, description="LiteLLM model identifier")
+    llm_config: dict = Field(default_factory=dict, description="Model-specific configuration")
+    team_id: str | None = Field(None, description="Team ID to assign this agent to")
+
+
+class AgentUpdate(BaseModel):
+    name: str | None = None
+    description: str | None = None
+    status: AgentStatus | None = None
+    capabilities: list | None = None
+    configuration: dict | None = None
+    state: dict | None = None
+    model_id: str | None = None
+    llm_config: dict | None = None
+    team_id: str | None = None
+
+
+class AgentResponse(BaseModel):
+    id: str
+    name: str
+    description: str | None
+    status: AgentStatus
+    capabilities: list
+    configuration: dict
+    model_id: str | None
+    llm_config: dict
+    team_id: str | None
+    created_at: datetime
+    updated_at: datetime
+    last_active_at: datetime | None
+    state: dict
+    error_message: str | None
+
+    class Config:
+        from_attributes = True
+
+
+class AgentExecutionRequest(BaseModel):
+    prompt: str = Field(..., description="The prompt to execute")
+    system_prompt: str | None = Field(None, description="Optional system prompt")
+    stream: bool = Field(False, description="Whether to stream the response")
+    config: dict | None = Field(None, description="Optional configuration including user metadata")
+
+
+class AgentExecutionResponse(BaseModel):
+    execution_id: str
+    success: bool
+    response: str | None = None
+    error: str | None = None
+    model: str
+    usage: dict | None = None
+    finish_reason: str | None = None
+
+
+@router.post("", response_model=AgentResponse, status_code=status.HTTP_201_CREATED)
+@instrument_endpoint("agents.create")
+def create_agent(agent_data: AgentCreate, organization: dict = Depends(get_current_organization), db: Session = Depends(get_db)):
+    """Create a new agent"""
+    agent = Agent(
+        name=agent_data.name,
+        description=agent_data.description,
+        capabilities=agent_data.capabilities,
+        configuration=agent_data.configuration,
+        model_id=agent_data.model_id,
+        model_config=agent_data.llm_config,
+        team_id=agent_data.team_id,
+    )
+    db.add(agent)
+    db.commit()
+    db.refresh(agent)
+    return agent_to_response(agent)
+
+
+@router.get("", response_model=List[AgentResponse])
+@instrument_endpoint("agents.list_agents")
+def list_agents(
+    skip: int = 0,
+    limit: int = 100,
+    status_filter: AgentStatus | None = None,
+    db: Session = Depends(get_db),
+):
+    """List all agents"""
+    query = db.query(Agent)
+    if status_filter:
+        query = query.filter(Agent.status == status_filter)
+    agents = query.offset(skip).limit(limit).all()
+    return [agent_to_response(agent) for agent in agents]
+
+
+@router.get("/{agent_id}", response_model=AgentResponse)
+@instrument_endpoint("agents.get_agent")
+def get_agent(agent_id: str, db: Session = Depends(get_db)):
+    """Get a specific agent by ID"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+    return agent_to_response(agent)
+
+
+@router.patch("/{agent_id}", response_model=AgentResponse)
+@instrument_endpoint("agents.update_agent")
+def update_agent(agent_id: str, agent_data: AgentUpdate, db: Session = Depends(get_db)):
+    """Update an agent"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    update_data = agent_data.model_dump(exclude_unset=True)
+    for field, value in update_data.items():
+        setattr(agent, field, value)
+
+    agent.updated_at = datetime.utcnow()
+    db.commit()
+    db.refresh(agent)
+    return agent_to_response(agent)
+
+
+@router.delete("/{agent_id}", status_code=status.HTTP_204_NO_CONTENT)
+@instrument_endpoint("agents.delete_agent")
+def delete_agent(agent_id: str, db: Session = Depends(get_db)):
+    """Delete an agent"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    db.delete(agent)
+    db.commit()
+    return None
+
+
+@router.post("/{agent_id}/start", response_model=AgentResponse)
+@instrument_endpoint("agents.start_agent")
+def start_agent(agent_id: str, db: Session = Depends(get_db)):
+    """Start an agent"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    if agent.status == AgentStatus.RUNNING:
+        raise HTTPException(status_code=400, detail="Agent is already running")
+
+    agent.status = AgentStatus.RUNNING
+    agent.last_active_at = datetime.utcnow()
+    agent.error_message = None
+    db.commit()
+    db.refresh(agent)
+    return agent_to_response(agent)
+
+
+@router.post("/{agent_id}/stop", response_model=AgentResponse)
+@instrument_endpoint("agents.stop_agent")
+def stop_agent(agent_id: str, db: Session = Depends(get_db)):
+    """Stop an agent"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    agent.status = AgentStatus.STOPPED
+    agent.last_active_at = datetime.utcnow()
+    db.commit()
+    db.refresh(agent)
+    return agent_to_response(agent)
+
+
+@router.post("/{agent_id}/pause", response_model=AgentResponse)
+@instrument_endpoint("agents.pause_agent")
+def pause_agent(agent_id: str, db: Session = Depends(get_db)):
+    """Pause an agent"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    if agent.status != AgentStatus.RUNNING:
+        raise HTTPException(status_code=400, detail="Can only pause running agents")
+
+    agent.status = AgentStatus.PAUSED
+    agent.last_active_at = datetime.utcnow()
+    db.commit()
+    db.refresh(agent)
+    return agent_to_response(agent)
+
+
+@router.post("/{agent_id}/resume", response_model=AgentResponse)
+@instrument_endpoint("agents.resume_agent")
+def resume_agent(agent_id: str, db: Session = Depends(get_db)):
+    """Resume a paused agent"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    if agent.status != AgentStatus.PAUSED:
+        raise HTTPException(status_code=400, detail="Can only resume paused agents")
+
+    agent.status = AgentStatus.RUNNING
+    agent.last_active_at = datetime.utcnow()
+    db.commit()
+    db.refresh(agent)
+    return agent_to_response(agent)
+
+
+@router.post("/{agent_id}/execute", response_model=AgentExecutionResponse)
+@instrument_endpoint("agents.execute_agent")
+def execute_agent(
+    agent_id: str,
+    execution_request: AgentExecutionRequest,
+    db: Session = Depends(get_db),
+):
+    """Execute an agent with a prompt using LiteLLM"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    # Extract user metadata from config if provided
+    user_metadata = {}
+    if execution_request.config and "user_metadata" in execution_request.config:
+        user_metadata = execution_request.config["user_metadata"]
+
+    # Create execution record
+    execution = Execution(
+        execution_type=ExecutionType.AGENT,
+        entity_id=agent.id,
+        entity_name=agent.name,
+        prompt=execution_request.prompt,
+        system_prompt=execution_request.system_prompt,
+        status=ExecutionStatus.RUNNING,
+        execution_metadata={
+            "user_id": user_metadata.get("user_id"),
+            "user_name": user_metadata.get("user_name"),
+            "user_email": user_metadata.get("user_email"),
+            "user_avatar": user_metadata.get("user_avatar"),
+        },
+    )
+    db.add(execution)
+    db.flush()  # Get execution ID without committing
+    execution.started_at = datetime.utcnow()
+
+    # Extract model configuration
+    model = agent.model_id
+    model_config = agent.model_config or {}
+
+    # Get system prompt from configuration if not provided
+    system_prompt = execution_request.system_prompt
+    if not system_prompt and "system_prompt" in agent.configuration:
+        system_prompt = agent.configuration["system_prompt"]
+
+    try:
+        # Execute using LiteLLM service
+        result = litellm_service.execute_agent(
+            prompt=execution_request.prompt,
+            model=model,
+            system_prompt=system_prompt,
+            **model_config,
+        )
+
+        # Update execution record with results
+        execution.response = result.get("response")
+        execution.usage = result.get("usage", {})
+        execution.execution_metadata = {
+            "model": result.get("model"),
+            "finish_reason": result.get("finish_reason"),
+        }
+
+        if result.get("success"):
+            execution.status = ExecutionStatus.COMPLETED
+            agent.status = AgentStatus.COMPLETED
+        else:
+            execution.status = ExecutionStatus.FAILED
+            execution.error_message = result.get("error")
+            agent.status = AgentStatus.FAILED
+            agent.error_message = result.get("error")
+
+        execution.completed_at = datetime.utcnow()
+
+    except Exception as e:
+        # Handle execution errors
+        execution.status = ExecutionStatus.FAILED
+        execution.error_message = str(e)
+        execution.completed_at = datetime.utcnow()
+        agent.status = AgentStatus.FAILED
+        agent.error_message = str(e)
+        result = {
+            "success": False,
+            "error": str(e),
+            "model": model,
+            "usage": None,
+            "finish_reason": None,
+        }
+
+    # Update agent state
+    agent.last_active_at = datetime.utcnow()
+    db.commit()
+
+    return AgentExecutionResponse(
+        execution_id=execution.id,
+        **result
+    )
+
+
+@router.post("/{agent_id}/execute/stream")
+@instrument_endpoint("agents.execute_agent_stream")
+def execute_agent_stream(
+    agent_id: str,
+    execution_request: AgentExecutionRequest,
+    db: Session = Depends(get_db),
+):
+    """Execute an agent with a prompt using LiteLLM (streaming response)"""
+    agent = db.query(Agent).filter(Agent.id == agent_id).first()
+    if not agent:
+        raise HTTPException(status_code=404, detail="Agent not found")
+
+    # Extract model configuration
+    model = agent.model_id
+    model_config = agent.model_config or {}
+
+    # Get system prompt from configuration if not provided
+    system_prompt = execution_request.system_prompt
+    if not system_prompt and "system_prompt" in agent.configuration:
+        system_prompt = agent.configuration["system_prompt"]
+
+    # Update agent state
+    agent.last_active_at = datetime.utcnow()
+    agent.status = AgentStatus.RUNNING
+    db.commit()
+
+    # Execute using LiteLLM service (streaming)
+    return StreamingResponse(
+        litellm_service.execute_agent_stream(
+            prompt=execution_request.prompt,
+            model=model,
+            system_prompt=system_prompt,
+            **model_config,
+        ),
+        media_type="text/event-stream",
+    )