kubiya-control-plane-api 0.1.0__py3-none-any.whl → 0.3.4__py3-none-any.whl

control_plane_api/app/activities/team_activities.py

@@ -0,0 +1,410 @@
+"""Team-related Temporal activities"""
+
+import os
+from dataclasses import dataclass
+from typing import Optional, List
+from temporalio import activity
+import structlog
+
+from agno.agent import Agent
+from agno.team import Team
+from agno.models.litellm import LiteLLM
+
+from control_plane_api.app.lib.supabase import get_supabase
+from control_plane_api.app.activities.agent_activities import update_execution_status, ActivityUpdateExecutionInput
+
+logger = structlog.get_logger()
+
+
+@dataclass
+class ActivityGetTeamAgentsInput:
+    """Input for get_team_agents activity"""
+    team_id: str
+    organization_id: str
+
+
+@dataclass
+class ActivityExecuteTeamInput:
+    """Input for execute_team_coordination activity"""
+    execution_id: str
+    team_id: str
+    organization_id: str
+    prompt: str
+    system_prompt: Optional[str] = None
+    agents: List[dict] = None
+    team_config: dict = None
+
+    def __post_init__(self):
+        if self.agents is None:
+            self.agents = []
+        if self.team_config is None:
+            self.team_config = {}
+
+
+@activity.defn
+async def get_team_agents(input: ActivityGetTeamAgentsInput) -> dict:
+    """
+    Get all agents in a team.
+
+    This activity fetches agents belonging to the team from the database.
+
+    Args:
+        input: Activity input with team details
+
+    Returns:
+        Dict with agents list
+    """
+    print(f"\n\n=== GET_TEAM_AGENTS START ===")
+    print(f"team_id: {input.team_id} (type: {type(input.team_id).__name__})")
+    print(f"organization_id: {input.organization_id} (type: {type(input.organization_id).__name__})")
+    print(f"================================\n")
+
+    activity.logger.info(
+        f"[DEBUG] Getting team agents START",
+        extra={
+            "team_id": input.team_id,
+            "team_id_type": type(input.team_id).__name__,
+            "organization_id": input.organization_id,
+            "organization_id_type": type(input.organization_id).__name__,
+        }
+    )
+
+    try:
+        client = get_supabase()
+        print(f"Supabase client created successfully")
+
+        activity.logger.info(
+            f"[DEBUG] Supabase client created, fetching team configuration",
+            extra={
+                "team_id": input.team_id,
+                "organization_id": input.organization_id,
+            }
+        )
+
+        # First, get the team configuration to get member_ids (source of truth)
+        print(f"Fetching team configuration for team_id={input.team_id}")
+        team_result = (
+            client.table("teams")
+            .select("configuration")
+            .eq("id", input.team_id)
+            .eq("organization_id", input.organization_id)
+            .single()
+            .execute()
+        )
+
+        if not team_result.data:
+            print(f"Team not found!")
+            activity.logger.error(
+                f"[DEBUG] Team not found",
+                extra={
+                    "team_id": input.team_id,
+                    "organization_id": input.organization_id,
+                }
+            )
+            return {"agents": [], "count": 0}
+
+        team_config = team_result.data.get("configuration", {})
+        member_ids = team_config.get("member_ids", [])
+        print(f"Team configuration member_ids: {member_ids}")
+
+        activity.logger.info(
+            f"[DEBUG] Team configuration loaded",
+            extra={
+                "team_id": input.team_id,
+                "member_ids": member_ids,
+                "member_count": len(member_ids),
+            }
+        )
+
+        # Get agents by member_ids (not by FK relationship)
+        agents = []
+        if member_ids:
+            print(f"Fetching agents with IDs: {member_ids}")
+            result = (
+                client.table("agents")
+                .select("*")
+                .in_("id", member_ids)
+                .eq("organization_id", input.organization_id)
+                .execute()
+            )
+            agents = result.data or []
+        print(f"Query executed. Result data length: {len(agents)}")
+
+        activity.logger.info(
+            f"[DEBUG] Query executed, processing results",
+            extra={
+                "agents_found": len(agents),
+                "agent_ids": [a.get("id") for a in agents],
+            }
+        )
+
+        print(f"Agents found: {len(agents)}")
+        if agents:
+            for agent in agents:
+                print(f"  - {agent.get('name')} (ID: {agent.get('id')})")
+
+        activity.logger.info(
+            f"[DEBUG] Retrieved team agents",
+            extra={
+                "team_id": input.team_id,
+                "agent_count": len(agents),
+                "agent_names": [a.get("name") for a in agents],
+                "agent_ids": [a.get("id") for a in agents],
+            }
+        )
+
+        if not agents:
+            print(f"\n!!! NO AGENTS FOUND - Running verification query !!!")
+            activity.logger.warning(
+                f"[DEBUG] WARNING: No agents found for team - running verification query",
+                extra={
+                    "team_id": input.team_id,
+                    "organization_id": input.organization_id,
+                }
+            )
+
+            # Try query without org filter to debug
+            result_no_org = (
+                client.table("agents")
+                .select("id, name, team_id, organization_id")
+                .eq("team_id", input.team_id)
+                .execute()
+            )
+            print(f"Query without org filter returned {len(result_no_org.data or [])} agents")
+            if result_no_org.data:
+                for agent in result_no_org.data:
+                    print(f"  - {agent.get('name')} (team_id: {agent.get('team_id')}, org_id: {agent.get('organization_id')})")
+
+            activity.logger.warning(
+                f"[DEBUG] Query without org filter returned {len(result_no_org.data or [])} agents",
+                extra={
+                    "agents_found": result_no_org.data if result_no_org.data else [],
+                }
+            )
+
+        print(f"\n=== GET_TEAM_AGENTS END: Returning {len(agents)} agents ===\n\n")
+        return {
+            "agents": agents,
+            "count": len(agents),
+        }
+
+    except Exception as e:
+        print(f"\n!!! EXCEPTION in get_team_agents: {type(e).__name__}: {str(e)} !!!\n")
+        activity.logger.error(
+            f"[DEBUG] EXCEPTION in get_team_agents",
+            extra={
+                "team_id": input.team_id,
+                "organization_id": input.organization_id,
+                "error": str(e),
+                "error_type": type(e).__name__,
+            }
+        )
+        raise
+
+
+@activity.defn
+async def execute_team_coordination(input: ActivityExecuteTeamInput) -> dict:
+    """
+    Execute team coordination using Agno Teams.
+
+    This activity creates an Agno Team with member Agents and executes
+    the team run, allowing Agno to handle coordination.
+
+    Args:
+        input: Activity input with team execution details
+
+    Returns:
+        Dict with aggregated response, usage, success flag
+    """
+    activity.logger.info(
+        f"Executing team coordination with Agno Teams",
+        extra={
+            "execution_id": input.execution_id,
+            "team_id": input.team_id,
+            "agent_count": len(input.agents),
+        }
+    )
+
+    try:
+        # Create Agno Agent objects for each team member
+        member_agents = []
+        for agent_data in input.agents:
+            # Get model ID (default to kubiya/claude-sonnet-4 if not specified)
+            model_id = agent_data.get("model_id") or "kubiya/claude-sonnet-4"
+
+            # Get LiteLLM configuration from environment
+            litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
+            litellm_api_key = os.getenv("LITELLM_API_KEY")
+
+            # Create Agno Agent with explicit LiteLLM proxy configuration
+            # IMPORTANT: Use openai/ prefix for custom proxy compatibility (same as agno_service)
+            member_agent = Agent(
+                name=agent_data["name"],
+                role=agent_data.get("description", agent_data["name"]),
+                model=LiteLLM(
+                    id=f"openai/{model_id}",  # e.g., "openai/kubiya/claude-sonnet-4"
+                    api_base=litellm_api_base,
+                    api_key=litellm_api_key,
+                ),
+            )
+            member_agents.append(member_agent)
+
+            activity.logger.info(
+                f"Created Agno Agent",
+                extra={
+                    "agent_name": agent_data["name"],
+                    "model": model_id,
+                }
+            )
+
+        # Create Agno Team with member agents and LiteLLM model for coordination
+        litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
+        litellm_api_key = os.getenv("LITELLM_API_KEY")
+
+        # Get coordinator model from team configuration (if specified by user in UI)
+        # Falls back to default if not configured
+        team_model = (
+            input.team_config.get("llm", {}).get("model")
+            or "kubiya/claude-sonnet-4"  # Default coordinator model
+        )
+
+        # Create Team with openai/ prefix for custom proxy compatibility (same as agno_service)
+        team = Team(
+            members=member_agents,
+            name=f"Team {input.team_id}",
+            model=LiteLLM(
+                id=f"openai/{team_model}",  # e.g., "openai/kubiya/claude-sonnet-4"
+                api_base=litellm_api_base,
+                api_key=litellm_api_key,
+            ),
+        )
+
+        activity.logger.info(
+            f"Created Agno Team with {len(member_agents)} members",
+            extra={
+                "coordinator_model": team_model,
+                "member_count": len(member_agents),
+            }
+        )
+
+        # Execute team run in a thread pool since team.run() is synchronous
+        # This prevents blocking the async event loop in Temporal
+        import asyncio
+        result = await asyncio.to_thread(team.run, input.prompt)
+
+        activity.logger.info(
+            f"Agno Team execution completed",
+            extra={
+                "execution_id": input.execution_id,
+                "has_content": bool(result.content),
+            }
+        )
+
+        # Extract response content
+        response_content = result.content if hasattr(result, "content") else str(result)
+
+        # Extract usage metrics if available
+        usage = {}
+        if hasattr(result, "metrics") and result.metrics:
+            metrics = result.metrics
+            usage = {
+                "input_tokens": getattr(metrics, "input_tokens", 0),
+                "output_tokens": getattr(metrics, "output_tokens", 0),
+                "total_tokens": getattr(metrics, "total_tokens", 0),
+            }
+
+        return {
+            "success": True,
+            "response": response_content,
+            "usage": usage,
+            "coordination_type": "agno_team",
+        }
+
+    except Exception as e:
+        activity.logger.error(
+            f"Team coordination failed",
+            extra={
+                "execution_id": input.execution_id,
+                "error": str(e),
+            }
+        )
+        return {
+            "success": False,
+            "error": str(e),
+            "coordination_type": "agno_team",
+            "usage": {},
+        }
+
+
+async def _execute_sequential(input: ActivityExecuteTeamInput) -> dict:
+    """Execute agents sequentially"""
+    activity.logger.info("Executing team sequentially")
+
+    responses = []
+    total_usage = {}
+
+    for agent in input.agents:
+        try:
+            # Import here to avoid circular dependency
+            from control_plane_api.app.services.litellm_service import litellm_service
+
+            model = agent.get("model_id") or "kubiya/claude-sonnet-4"
+            model_config = agent.get("model_config", {})
+
+            # Execute agent
+            result = litellm_service.execute_agent(
+                prompt=f"{input.prompt}\n\nAgent role: {agent.get('description', agent['name'])}",
+                model=model,
+                system_prompt=input.system_prompt,
+                **model_config
+            )
+
+            if result.get("success"):
+                responses.append({
+                    "agent": agent["name"],
+                    "response": result.get("response"),
+                })
+
+            # Aggregate usage
+            usage = result.get("usage", {})
+            for key, value in usage.items():
+                total_usage[key] = total_usage.get(key, 0) + value
+
+        except Exception as e:
+            activity.logger.error(
+                f"Agent execution failed",
+                extra={"agent_id": agent["id"], "error": str(e)}
+            )
+            responses.append({
+                "agent": agent["name"],
+                "error": str(e),
+            })
+
+    # Combine responses
+    combined_response = "\n\n".join([
+        f"**{r['agent']}**: {r.get('response', r.get('error'))}"
+        for r in responses
+    ])
+
+    return {
+        "success": True,
+        "response": combined_response,
+        "usage": total_usage,
+        "coordination_type": "sequential",
+    }
+
+
+async def _execute_parallel(input: ActivityExecuteTeamInput) -> dict:
+    """Execute agents in parallel (placeholder - would use asyncio.gather in real impl)"""
+    activity.logger.info("Executing team in parallel")
+    # For now, fall back to sequential
+    # In a real implementation, this would use asyncio.gather or child workflows
+    return await _execute_sequential(input)
+
+
+async def _execute_hierarchical(input: ActivityExecuteTeamInput) -> dict:
+    """Execute agents hierarchically with a coordinator (placeholder)"""
+    activity.logger.info("Executing team hierarchically")
+    # For now, fall back to sequential
+    # In a real implementation, this would have a coordinator agent that delegates
+    return await _execute_sequential(input)