kubiya-control-plane-api 0.9.15__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- control_plane_api/LICENSE +676 -0
- control_plane_api/README.md +350 -0
- control_plane_api/__init__.py +4 -0
- control_plane_api/__version__.py +8 -0
- control_plane_api/alembic/README +1 -0
- control_plane_api/alembic/env.py +121 -0
- control_plane_api/alembic/script.py.mako +28 -0
- control_plane_api/alembic/versions/2613c65c3dbe_initial_database_setup.py +32 -0
- control_plane_api/alembic/versions/2df520d4927d_merge_heads.py +28 -0
- control_plane_api/alembic/versions/43abf98d6a01_add_paused_status_to_executions.py +73 -0
- control_plane_api/alembic/versions/6289854264cb_merge_multiple_heads.py +28 -0
- control_plane_api/alembic/versions/6a4d4dc3d8dc_generate_execution_transitions.py +50 -0
- control_plane_api/alembic/versions/87d11cf0a783_add_disconnected_status_to_worker_.py +44 -0
- control_plane_api/alembic/versions/add_ephemeral_queue_support.py +85 -0
- control_plane_api/alembic/versions/add_model_type_to_llm_models.py +31 -0
- control_plane_api/alembic/versions/add_plan_executions_table.py +114 -0
- control_plane_api/alembic/versions/add_trace_span_tables.py +154 -0
- control_plane_api/alembic/versions/add_user_info_to_traces.py +36 -0
- control_plane_api/alembic/versions/adjusting_foreign_keys.py +32 -0
- control_plane_api/alembic/versions/b4983d976db2_initial_tables.py +1128 -0
- control_plane_api/alembic/versions/d181a3b40e71_rename_custom_metadata_to_metadata_in_.py +50 -0
- control_plane_api/alembic/versions/df9117888e82_add_missing_columns.py +82 -0
- control_plane_api/alembic/versions/f25de6ad895a_missing_migrations.py +34 -0
- control_plane_api/alembic/versions/f71305fb69b9_fix_ephemeral_queue_deletion_foreign_key.py +54 -0
- control_plane_api/alembic/versions/mark_local_exec_queues_as_ephemeral.py +68 -0
- control_plane_api/alembic.ini +148 -0
- control_plane_api/api/index.py +12 -0
- control_plane_api/app/__init__.py +11 -0
- control_plane_api/app/activities/__init__.py +20 -0
- control_plane_api/app/activities/agent_activities.py +384 -0
- control_plane_api/app/activities/plan_generation_activities.py +499 -0
- control_plane_api/app/activities/team_activities.py +424 -0
- control_plane_api/app/activities/temporal_cloud_activities.py +588 -0
- control_plane_api/app/config/__init__.py +35 -0
- control_plane_api/app/config/api_config.py +469 -0
- control_plane_api/app/config/config_loader.py +224 -0
- control_plane_api/app/config/model_pricing.py +323 -0
- control_plane_api/app/config/storage_config.py +159 -0
- control_plane_api/app/config.py +115 -0
- control_plane_api/app/controllers/__init__.py +0 -0
- control_plane_api/app/controllers/execution_environment_controller.py +1315 -0
- control_plane_api/app/database.py +135 -0
- control_plane_api/app/exceptions.py +408 -0
- control_plane_api/app/lib/__init__.py +11 -0
- control_plane_api/app/lib/environment.py +65 -0
- control_plane_api/app/lib/event_bus/__init__.py +17 -0
- control_plane_api/app/lib/event_bus/base.py +136 -0
- control_plane_api/app/lib/event_bus/manager.py +335 -0
- control_plane_api/app/lib/event_bus/providers/__init__.py +6 -0
- control_plane_api/app/lib/event_bus/providers/http_provider.py +166 -0
- control_plane_api/app/lib/event_bus/providers/nats_provider.py +324 -0
- control_plane_api/app/lib/event_bus/providers/redis_provider.py +233 -0
- control_plane_api/app/lib/event_bus/providers/websocket_provider.py +497 -0
- control_plane_api/app/lib/job_executor.py +330 -0
- control_plane_api/app/lib/kubiya_client.py +293 -0
- control_plane_api/app/lib/litellm_pricing.py +166 -0
- control_plane_api/app/lib/mcp_validation.py +163 -0
- control_plane_api/app/lib/nats/__init__.py +13 -0
- control_plane_api/app/lib/nats/credentials_manager.py +288 -0
- control_plane_api/app/lib/nats/listener.py +374 -0
- control_plane_api/app/lib/planning_prompt_builder.py +153 -0
- control_plane_api/app/lib/planning_tools/__init__.py +41 -0
- control_plane_api/app/lib/planning_tools/agents.py +409 -0
- control_plane_api/app/lib/planning_tools/agno_toolkit.py +836 -0
- control_plane_api/app/lib/planning_tools/base.py +119 -0
- control_plane_api/app/lib/planning_tools/cognitive_memory_tools.py +403 -0
- control_plane_api/app/lib/planning_tools/context_graph_tools.py +545 -0
- control_plane_api/app/lib/planning_tools/environments.py +218 -0
- control_plane_api/app/lib/planning_tools/knowledge.py +204 -0
- control_plane_api/app/lib/planning_tools/models.py +93 -0
- control_plane_api/app/lib/planning_tools/planning_service.py +646 -0
- control_plane_api/app/lib/planning_tools/resources.py +242 -0
- control_plane_api/app/lib/planning_tools/teams.py +334 -0
- control_plane_api/app/lib/policy_enforcer_client.py +1016 -0
- control_plane_api/app/lib/redis_client.py +803 -0
- control_plane_api/app/lib/sqlalchemy_utils.py +486 -0
- control_plane_api/app/lib/state_transition_tools/__init__.py +7 -0
- control_plane_api/app/lib/state_transition_tools/execution_context.py +388 -0
- control_plane_api/app/lib/storage/__init__.py +20 -0
- control_plane_api/app/lib/storage/base_provider.py +274 -0
- control_plane_api/app/lib/storage/provider_factory.py +157 -0
- control_plane_api/app/lib/storage/vercel_blob_provider.py +468 -0
- control_plane_api/app/lib/supabase.py +71 -0
- control_plane_api/app/lib/supabase_utils.py +138 -0
- control_plane_api/app/lib/task_planning/__init__.py +138 -0
- control_plane_api/app/lib/task_planning/agent_factory.py +308 -0
- control_plane_api/app/lib/task_planning/agents.py +389 -0
- control_plane_api/app/lib/task_planning/cache.py +218 -0
- control_plane_api/app/lib/task_planning/entity_resolver.py +273 -0
- control_plane_api/app/lib/task_planning/helpers.py +293 -0
- control_plane_api/app/lib/task_planning/hooks.py +474 -0
- control_plane_api/app/lib/task_planning/models.py +503 -0
- control_plane_api/app/lib/task_planning/plan_validator.py +166 -0
- control_plane_api/app/lib/task_planning/planning_workflow.py +2911 -0
- control_plane_api/app/lib/task_planning/runner.py +656 -0
- control_plane_api/app/lib/task_planning/streaming_hook.py +213 -0
- control_plane_api/app/lib/task_planning/workflow.py +424 -0
- control_plane_api/app/lib/templating/__init__.py +88 -0
- control_plane_api/app/lib/templating/compiler.py +278 -0
- control_plane_api/app/lib/templating/engine.py +178 -0
- control_plane_api/app/lib/templating/parsers/__init__.py +29 -0
- control_plane_api/app/lib/templating/parsers/base.py +96 -0
- control_plane_api/app/lib/templating/parsers/env.py +85 -0
- control_plane_api/app/lib/templating/parsers/graph.py +112 -0
- control_plane_api/app/lib/templating/parsers/secret.py +87 -0
- control_plane_api/app/lib/templating/parsers/simple.py +81 -0
- control_plane_api/app/lib/templating/resolver.py +366 -0
- control_plane_api/app/lib/templating/types.py +214 -0
- control_plane_api/app/lib/templating/validator.py +201 -0
- control_plane_api/app/lib/temporal_client.py +232 -0
- control_plane_api/app/lib/temporal_credentials_cache.py +178 -0
- control_plane_api/app/lib/temporal_credentials_service.py +203 -0
- control_plane_api/app/lib/validation/__init__.py +24 -0
- control_plane_api/app/lib/validation/runtime_validation.py +388 -0
- control_plane_api/app/main.py +531 -0
- control_plane_api/app/middleware/__init__.py +10 -0
- control_plane_api/app/middleware/auth.py +645 -0
- control_plane_api/app/middleware/exception_handler.py +267 -0
- control_plane_api/app/middleware/prometheus_middleware.py +173 -0
- control_plane_api/app/middleware/rate_limiting.py +384 -0
- control_plane_api/app/middleware/request_id.py +202 -0
- control_plane_api/app/models/__init__.py +40 -0
- control_plane_api/app/models/agent.py +90 -0
- control_plane_api/app/models/analytics.py +206 -0
- control_plane_api/app/models/associations.py +107 -0
- control_plane_api/app/models/auth_user.py +73 -0
- control_plane_api/app/models/context.py +161 -0
- control_plane_api/app/models/custom_integration.py +99 -0
- control_plane_api/app/models/environment.py +64 -0
- control_plane_api/app/models/execution.py +125 -0
- control_plane_api/app/models/execution_transition.py +50 -0
- control_plane_api/app/models/job.py +159 -0
- control_plane_api/app/models/llm_model.py +78 -0
- control_plane_api/app/models/orchestration.py +66 -0
- control_plane_api/app/models/plan_execution.py +102 -0
- control_plane_api/app/models/presence.py +49 -0
- control_plane_api/app/models/project.py +61 -0
- control_plane_api/app/models/project_management.py +85 -0
- control_plane_api/app/models/session.py +29 -0
- control_plane_api/app/models/skill.py +155 -0
- control_plane_api/app/models/system_tables.py +43 -0
- control_plane_api/app/models/task_planning.py +372 -0
- control_plane_api/app/models/team.py +86 -0
- control_plane_api/app/models/trace.py +257 -0
- control_plane_api/app/models/user_profile.py +54 -0
- control_plane_api/app/models/worker.py +221 -0
- control_plane_api/app/models/workflow.py +161 -0
- control_plane_api/app/models/workspace.py +50 -0
- control_plane_api/app/observability/__init__.py +177 -0
- control_plane_api/app/observability/context_logging.py +475 -0
- control_plane_api/app/observability/decorators.py +337 -0
- control_plane_api/app/observability/local_span_processor.py +702 -0
- control_plane_api/app/observability/metrics.py +303 -0
- control_plane_api/app/observability/middleware.py +246 -0
- control_plane_api/app/observability/optional.py +115 -0
- control_plane_api/app/observability/tracing.py +382 -0
- control_plane_api/app/policies/README.md +149 -0
- control_plane_api/app/policies/approved_users.rego +62 -0
- control_plane_api/app/policies/business_hours.rego +51 -0
- control_plane_api/app/policies/rate_limiting.rego +100 -0
- control_plane_api/app/policies/tool_enforcement/README.md +336 -0
- control_plane_api/app/policies/tool_enforcement/bash_command_validation.rego +71 -0
- control_plane_api/app/policies/tool_enforcement/business_hours_enforcement.rego +82 -0
- control_plane_api/app/policies/tool_enforcement/mcp_tool_allowlist.rego +58 -0
- control_plane_api/app/policies/tool_enforcement/production_safeguards.rego +80 -0
- control_plane_api/app/policies/tool_enforcement/role_based_tool_access.rego +44 -0
- control_plane_api/app/policies/tool_restrictions.rego +86 -0
- control_plane_api/app/routers/__init__.py +4 -0
- control_plane_api/app/routers/agents.py +382 -0
- control_plane_api/app/routers/agents_v2.py +1598 -0
- control_plane_api/app/routers/analytics.py +1310 -0
- control_plane_api/app/routers/auth.py +59 -0
- control_plane_api/app/routers/client_config.py +57 -0
- control_plane_api/app/routers/context_graph.py +561 -0
- control_plane_api/app/routers/context_manager.py +577 -0
- control_plane_api/app/routers/custom_integrations.py +490 -0
- control_plane_api/app/routers/enforcer.py +132 -0
- control_plane_api/app/routers/environment_context.py +252 -0
- control_plane_api/app/routers/environments.py +761 -0
- control_plane_api/app/routers/execution_environment.py +847 -0
- control_plane_api/app/routers/executions/__init__.py +28 -0
- control_plane_api/app/routers/executions/router.py +286 -0
- control_plane_api/app/routers/executions/services/__init__.py +22 -0
- control_plane_api/app/routers/executions/services/demo_worker_health.py +156 -0
- control_plane_api/app/routers/executions/services/status_service.py +420 -0
- control_plane_api/app/routers/executions/services/test_worker_health.py +480 -0
- control_plane_api/app/routers/executions/services/worker_health.py +514 -0
- control_plane_api/app/routers/executions/streaming/__init__.py +22 -0
- control_plane_api/app/routers/executions/streaming/deduplication.py +352 -0
- control_plane_api/app/routers/executions/streaming/event_buffer.py +353 -0
- control_plane_api/app/routers/executions/streaming/event_formatter.py +964 -0
- control_plane_api/app/routers/executions/streaming/history_loader.py +588 -0
- control_plane_api/app/routers/executions/streaming/live_source.py +693 -0
- control_plane_api/app/routers/executions/streaming/streamer.py +849 -0
- control_plane_api/app/routers/executions.py +4888 -0
- control_plane_api/app/routers/health.py +165 -0
- control_plane_api/app/routers/health_v2.py +394 -0
- control_plane_api/app/routers/integration_templates.py +496 -0
- control_plane_api/app/routers/integrations.py +287 -0
- control_plane_api/app/routers/jobs.py +1809 -0
- control_plane_api/app/routers/metrics.py +517 -0
- control_plane_api/app/routers/models.py +82 -0
- control_plane_api/app/routers/models_v2.py +628 -0
- control_plane_api/app/routers/plan_executions.py +1481 -0
- control_plane_api/app/routers/plan_generation_async.py +304 -0
- control_plane_api/app/routers/policies.py +669 -0
- control_plane_api/app/routers/presence.py +234 -0
- control_plane_api/app/routers/projects.py +987 -0
- control_plane_api/app/routers/runners.py +379 -0
- control_plane_api/app/routers/runtimes.py +172 -0
- control_plane_api/app/routers/secrets.py +171 -0
- control_plane_api/app/routers/skills.py +1010 -0
- control_plane_api/app/routers/skills_definitions.py +140 -0
- control_plane_api/app/routers/storage.py +456 -0
- control_plane_api/app/routers/task_planning.py +611 -0
- control_plane_api/app/routers/task_queues.py +650 -0
- control_plane_api/app/routers/team_context.py +274 -0
- control_plane_api/app/routers/teams.py +1747 -0
- control_plane_api/app/routers/templates.py +248 -0
- control_plane_api/app/routers/traces.py +571 -0
- control_plane_api/app/routers/websocket_client.py +479 -0
- control_plane_api/app/routers/websocket_executions_status.py +437 -0
- control_plane_api/app/routers/websocket_gateway.py +323 -0
- control_plane_api/app/routers/websocket_traces.py +576 -0
- control_plane_api/app/routers/worker_queues.py +2555 -0
- control_plane_api/app/routers/worker_websocket.py +419 -0
- control_plane_api/app/routers/workers.py +1004 -0
- control_plane_api/app/routers/workflows.py +204 -0
- control_plane_api/app/runtimes/__init__.py +6 -0
- control_plane_api/app/runtimes/validation.py +344 -0
- control_plane_api/app/schemas/__init__.py +1 -0
- control_plane_api/app/schemas/job_schemas.py +302 -0
- control_plane_api/app/schemas/mcp_schemas.py +311 -0
- control_plane_api/app/schemas/template_schemas.py +133 -0
- control_plane_api/app/schemas/trace_schemas.py +168 -0
- control_plane_api/app/schemas/worker_queue_observability_schemas.py +165 -0
- control_plane_api/app/services/__init__.py +1 -0
- control_plane_api/app/services/agno_planning_strategy.py +233 -0
- control_plane_api/app/services/agno_service.py +838 -0
- control_plane_api/app/services/claude_code_planning_service.py +203 -0
- control_plane_api/app/services/context_graph_client.py +224 -0
- control_plane_api/app/services/custom_integration_service.py +415 -0
- control_plane_api/app/services/integration_resolution_service.py +345 -0
- control_plane_api/app/services/litellm_service.py +394 -0
- control_plane_api/app/services/plan_generator.py +79 -0
- control_plane_api/app/services/planning_strategy.py +66 -0
- control_plane_api/app/services/planning_strategy_factory.py +118 -0
- control_plane_api/app/services/policy_service.py +615 -0
- control_plane_api/app/services/state_transition_service.py +755 -0
- control_plane_api/app/services/storage_service.py +593 -0
- control_plane_api/app/services/temporal_cloud_provisioning.py +150 -0
- control_plane_api/app/services/toolsets/context_graph_skill.py +432 -0
- control_plane_api/app/services/trace_retention.py +354 -0
- control_plane_api/app/services/worker_queue_metrics_service.py +190 -0
- control_plane_api/app/services/workflow_cancellation_manager.py +135 -0
- control_plane_api/app/services/workflow_operations_service.py +611 -0
- control_plane_api/app/skills/__init__.py +100 -0
- control_plane_api/app/skills/base.py +239 -0
- control_plane_api/app/skills/builtin/__init__.py +37 -0
- control_plane_api/app/skills/builtin/agent_communication/__init__.py +8 -0
- control_plane_api/app/skills/builtin/agent_communication/skill.py +246 -0
- control_plane_api/app/skills/builtin/code_ingestion/__init__.py +4 -0
- control_plane_api/app/skills/builtin/code_ingestion/skill.py +267 -0
- control_plane_api/app/skills/builtin/cognitive_memory/__init__.py +4 -0
- control_plane_api/app/skills/builtin/cognitive_memory/skill.py +174 -0
- control_plane_api/app/skills/builtin/contextual_awareness/__init__.py +4 -0
- control_plane_api/app/skills/builtin/contextual_awareness/skill.py +387 -0
- control_plane_api/app/skills/builtin/data_visualization/__init__.py +4 -0
- control_plane_api/app/skills/builtin/data_visualization/skill.py +154 -0
- control_plane_api/app/skills/builtin/docker/__init__.py +4 -0
- control_plane_api/app/skills/builtin/docker/skill.py +104 -0
- control_plane_api/app/skills/builtin/file_generation/__init__.py +4 -0
- control_plane_api/app/skills/builtin/file_generation/skill.py +94 -0
- control_plane_api/app/skills/builtin/file_system/__init__.py +4 -0
- control_plane_api/app/skills/builtin/file_system/skill.py +110 -0
- control_plane_api/app/skills/builtin/knowledge_api/__init__.py +5 -0
- control_plane_api/app/skills/builtin/knowledge_api/skill.py +124 -0
- control_plane_api/app/skills/builtin/python/__init__.py +4 -0
- control_plane_api/app/skills/builtin/python/skill.py +92 -0
- control_plane_api/app/skills/builtin/remote_filesystem/__init__.py +5 -0
- control_plane_api/app/skills/builtin/remote_filesystem/skill.py +170 -0
- control_plane_api/app/skills/builtin/shell/__init__.py +4 -0
- control_plane_api/app/skills/builtin/shell/skill.py +161 -0
- control_plane_api/app/skills/builtin/slack/__init__.py +3 -0
- control_plane_api/app/skills/builtin/slack/skill.py +302 -0
- control_plane_api/app/skills/builtin/workflow_executor/__init__.py +4 -0
- control_plane_api/app/skills/builtin/workflow_executor/skill.py +469 -0
- control_plane_api/app/skills/business_intelligence.py +189 -0
- control_plane_api/app/skills/config.py +63 -0
- control_plane_api/app/skills/loaders/__init__.py +14 -0
- control_plane_api/app/skills/loaders/base.py +73 -0
- control_plane_api/app/skills/loaders/filesystem_loader.py +199 -0
- control_plane_api/app/skills/registry.py +125 -0
- control_plane_api/app/utils/helpers.py +12 -0
- control_plane_api/app/utils/workflow_executor.py +354 -0
- control_plane_api/app/workflows/__init__.py +11 -0
- control_plane_api/app/workflows/agent_execution.py +520 -0
- control_plane_api/app/workflows/agent_execution_with_skills.py +223 -0
- control_plane_api/app/workflows/namespace_provisioning.py +326 -0
- control_plane_api/app/workflows/plan_generation.py +254 -0
- control_plane_api/app/workflows/team_execution.py +442 -0
- control_plane_api/scripts/seed_models.py +240 -0
- control_plane_api/scripts/validate_existing_tool_names.py +492 -0
- control_plane_api/shared/__init__.py +8 -0
- control_plane_api/shared/version.py +17 -0
- control_plane_api/test_deduplication.py +274 -0
- control_plane_api/test_executor_deduplication_e2e.py +309 -0
- control_plane_api/test_job_execution_e2e.py +283 -0
- control_plane_api/test_real_integration.py +193 -0
- control_plane_api/version.py +38 -0
- control_plane_api/worker/__init__.py +0 -0
- control_plane_api/worker/activities/__init__.py +0 -0
- control_plane_api/worker/activities/agent_activities.py +1585 -0
- control_plane_api/worker/activities/approval_activities.py +234 -0
- control_plane_api/worker/activities/job_activities.py +199 -0
- control_plane_api/worker/activities/runtime_activities.py +1167 -0
- control_plane_api/worker/activities/skill_activities.py +282 -0
- control_plane_api/worker/activities/team_activities.py +479 -0
- control_plane_api/worker/agent_runtime_server.py +370 -0
- control_plane_api/worker/binary_manager.py +333 -0
- control_plane_api/worker/config/__init__.py +31 -0
- control_plane_api/worker/config/worker_config.py +273 -0
- control_plane_api/worker/control_plane_client.py +1491 -0
- control_plane_api/worker/examples/analytics_integration_example.py +362 -0
- control_plane_api/worker/health_monitor.py +159 -0
- control_plane_api/worker/metrics.py +237 -0
- control_plane_api/worker/models/__init__.py +1 -0
- control_plane_api/worker/models/error_events.py +105 -0
- control_plane_api/worker/models/inputs.py +89 -0
- control_plane_api/worker/runtimes/__init__.py +35 -0
- control_plane_api/worker/runtimes/agent_runtime/runtime.py +485 -0
- control_plane_api/worker/runtimes/agno/__init__.py +34 -0
- control_plane_api/worker/runtimes/agno/config.py +248 -0
- control_plane_api/worker/runtimes/agno/hooks.py +385 -0
- control_plane_api/worker/runtimes/agno/mcp_builder.py +195 -0
- control_plane_api/worker/runtimes/agno/runtime.py +1063 -0
- control_plane_api/worker/runtimes/agno/utils.py +163 -0
- control_plane_api/worker/runtimes/base.py +979 -0
- control_plane_api/worker/runtimes/claude_code/__init__.py +38 -0
- control_plane_api/worker/runtimes/claude_code/cleanup.py +184 -0
- control_plane_api/worker/runtimes/claude_code/client_pool.py +529 -0
- control_plane_api/worker/runtimes/claude_code/config.py +829 -0
- control_plane_api/worker/runtimes/claude_code/hooks.py +482 -0
- control_plane_api/worker/runtimes/claude_code/litellm_proxy.py +1702 -0
- control_plane_api/worker/runtimes/claude_code/mcp_builder.py +467 -0
- control_plane_api/worker/runtimes/claude_code/mcp_discovery.py +558 -0
- control_plane_api/worker/runtimes/claude_code/runtime.py +1546 -0
- control_plane_api/worker/runtimes/claude_code/tool_mapper.py +403 -0
- control_plane_api/worker/runtimes/claude_code/utils.py +149 -0
- control_plane_api/worker/runtimes/factory.py +173 -0
- control_plane_api/worker/runtimes/model_utils.py +107 -0
- control_plane_api/worker/runtimes/validation.py +93 -0
- control_plane_api/worker/services/__init__.py +1 -0
- control_plane_api/worker/services/agent_communication_tools.py +908 -0
- control_plane_api/worker/services/agent_executor.py +485 -0
- control_plane_api/worker/services/agent_executor_v2.py +793 -0
- control_plane_api/worker/services/analytics_collector.py +457 -0
- control_plane_api/worker/services/analytics_service.py +464 -0
- control_plane_api/worker/services/approval_tools.py +310 -0
- control_plane_api/worker/services/approval_tools_agno.py +207 -0
- control_plane_api/worker/services/cancellation_manager.py +177 -0
- control_plane_api/worker/services/code_ingestion_tools.py +465 -0
- control_plane_api/worker/services/contextual_awareness_tools.py +405 -0
- control_plane_api/worker/services/data_visualization.py +834 -0
- control_plane_api/worker/services/event_publisher.py +531 -0
- control_plane_api/worker/services/jira_tools.py +257 -0
- control_plane_api/worker/services/remote_filesystem_tools.py +498 -0
- control_plane_api/worker/services/runtime_analytics.py +328 -0
- control_plane_api/worker/services/session_service.py +365 -0
- control_plane_api/worker/services/skill_context_enhancement.py +181 -0
- control_plane_api/worker/services/skill_factory.py +471 -0
- control_plane_api/worker/services/system_prompt_enhancement.py +410 -0
- control_plane_api/worker/services/team_executor.py +715 -0
- control_plane_api/worker/services/team_executor_v2.py +1866 -0
- control_plane_api/worker/services/tool_enforcement.py +254 -0
- control_plane_api/worker/services/workflow_executor/__init__.py +52 -0
- control_plane_api/worker/services/workflow_executor/event_processor.py +287 -0
- control_plane_api/worker/services/workflow_executor/event_publisher.py +210 -0
- control_plane_api/worker/services/workflow_executor/executors/__init__.py +15 -0
- control_plane_api/worker/services/workflow_executor/executors/base.py +270 -0
- control_plane_api/worker/services/workflow_executor/executors/json_executor.py +50 -0
- control_plane_api/worker/services/workflow_executor/executors/python_executor.py +50 -0
- control_plane_api/worker/services/workflow_executor/models.py +142 -0
- control_plane_api/worker/services/workflow_executor_tools.py +1748 -0
- control_plane_api/worker/skills/__init__.py +12 -0
- control_plane_api/worker/skills/builtin/context_graph_search/README.md +213 -0
- control_plane_api/worker/skills/builtin/context_graph_search/__init__.py +5 -0
- control_plane_api/worker/skills/builtin/context_graph_search/agno_impl.py +808 -0
- control_plane_api/worker/skills/builtin/context_graph_search/skill.yaml +67 -0
- control_plane_api/worker/skills/builtin/contextual_awareness/__init__.py +4 -0
- control_plane_api/worker/skills/builtin/contextual_awareness/agno_impl.py +62 -0
- control_plane_api/worker/skills/builtin/data_visualization/agno_impl.py +18 -0
- control_plane_api/worker/skills/builtin/data_visualization/skill.yaml +84 -0
- control_plane_api/worker/skills/builtin/docker/agno_impl.py +65 -0
- control_plane_api/worker/skills/builtin/docker/skill.yaml +60 -0
- control_plane_api/worker/skills/builtin/file_generation/agno_impl.py +47 -0
- control_plane_api/worker/skills/builtin/file_generation/skill.yaml +64 -0
- control_plane_api/worker/skills/builtin/file_system/agno_impl.py +32 -0
- control_plane_api/worker/skills/builtin/file_system/skill.yaml +54 -0
- control_plane_api/worker/skills/builtin/knowledge_api/__init__.py +4 -0
- control_plane_api/worker/skills/builtin/knowledge_api/agno_impl.py +50 -0
- control_plane_api/worker/skills/builtin/knowledge_api/skill.yaml +66 -0
- control_plane_api/worker/skills/builtin/python/agno_impl.py +25 -0
- control_plane_api/worker/skills/builtin/python/skill.yaml +60 -0
- control_plane_api/worker/skills/builtin/schema_fix_mixin.py +260 -0
- control_plane_api/worker/skills/builtin/shell/agno_impl.py +31 -0
- control_plane_api/worker/skills/builtin/shell/skill.yaml +60 -0
- control_plane_api/worker/skills/builtin/slack/__init__.py +3 -0
- control_plane_api/worker/skills/builtin/slack/agno_impl.py +1282 -0
- control_plane_api/worker/skills/builtin/slack/skill.yaml +276 -0
- control_plane_api/worker/skills/builtin/workflow_executor/agno_impl.py +62 -0
- control_plane_api/worker/skills/builtin/workflow_executor/skill.yaml +79 -0
- control_plane_api/worker/skills/loaders/__init__.py +5 -0
- control_plane_api/worker/skills/loaders/base.py +23 -0
- control_plane_api/worker/skills/loaders/filesystem_loader.py +357 -0
- control_plane_api/worker/skills/registry.py +208 -0
- control_plane_api/worker/tests/__init__.py +1 -0
- control_plane_api/worker/tests/conftest.py +12 -0
- control_plane_api/worker/tests/e2e/__init__.py +0 -0
- control_plane_api/worker/tests/e2e/test_context_graph_real_api.py +338 -0
- control_plane_api/worker/tests/e2e/test_context_graph_templates_e2e.py +523 -0
- control_plane_api/worker/tests/e2e/test_enforcement_e2e.py +344 -0
- control_plane_api/worker/tests/e2e/test_execution_flow.py +571 -0
- control_plane_api/worker/tests/e2e/test_single_execution_mode.py +656 -0
- control_plane_api/worker/tests/integration/__init__.py +0 -0
- control_plane_api/worker/tests/integration/test_builtin_skills_fixes.py +245 -0
- control_plane_api/worker/tests/integration/test_context_graph_search_integration.py +365 -0
- control_plane_api/worker/tests/integration/test_control_plane_integration.py +308 -0
- control_plane_api/worker/tests/integration/test_hook_enforcement_integration.py +579 -0
- control_plane_api/worker/tests/integration/test_scheduled_job_workflow.py +237 -0
- control_plane_api/worker/tests/integration/test_system_prompt_enhancement_integration.py +343 -0
- control_plane_api/worker/tests/unit/__init__.py +0 -0
- control_plane_api/worker/tests/unit/test_builtin_skill_autoload.py +396 -0
- control_plane_api/worker/tests/unit/test_context_graph_search.py +450 -0
- control_plane_api/worker/tests/unit/test_context_graph_templates.py +403 -0
- control_plane_api/worker/tests/unit/test_control_plane_client.py +401 -0
- control_plane_api/worker/tests/unit/test_control_plane_client_jobs.py +345 -0
- control_plane_api/worker/tests/unit/test_job_activities.py +353 -0
- control_plane_api/worker/tests/unit/test_skill_context_enhancement.py +321 -0
- control_plane_api/worker/tests/unit/test_system_prompt_enhancement.py +415 -0
- control_plane_api/worker/tests/unit/test_tool_enforcement.py +324 -0
- control_plane_api/worker/utils/__init__.py +1 -0
- control_plane_api/worker/utils/chunk_batcher.py +330 -0
- control_plane_api/worker/utils/environment.py +65 -0
- control_plane_api/worker/utils/error_publisher.py +260 -0
- control_plane_api/worker/utils/event_batcher.py +256 -0
- control_plane_api/worker/utils/logging_config.py +335 -0
- control_plane_api/worker/utils/logging_helper.py +326 -0
- control_plane_api/worker/utils/parameter_validator.py +120 -0
- control_plane_api/worker/utils/retry_utils.py +60 -0
- control_plane_api/worker/utils/streaming_utils.py +665 -0
- control_plane_api/worker/utils/tool_validation.py +332 -0
- control_plane_api/worker/utils/workspace_manager.py +163 -0
- control_plane_api/worker/websocket_client.py +393 -0
- control_plane_api/worker/worker.py +1297 -0
- control_plane_api/worker/workflows/__init__.py +0 -0
- control_plane_api/worker/workflows/agent_execution.py +909 -0
- control_plane_api/worker/workflows/scheduled_job_wrapper.py +332 -0
- control_plane_api/worker/workflows/team_execution.py +611 -0
- kubiya_control_plane_api-0.9.15.dist-info/METADATA +354 -0
- kubiya_control_plane_api-0.9.15.dist-info/RECORD +479 -0
- kubiya_control_plane_api-0.9.15.dist-info/WHEEL +5 -0
- kubiya_control_plane_api-0.9.15.dist-info/entry_points.txt +5 -0
- kubiya_control_plane_api-0.9.15.dist-info/licenses/LICENSE +676 -0
- kubiya_control_plane_api-0.9.15.dist-info/top_level.txt +3 -0
- scripts/__init__.py +1 -0
- scripts/migrations.py +39 -0
- scripts/seed_worker_queues.py +128 -0
- scripts/setup_agent_runtime.py +142 -0
- worker_internal/__init__.py +1 -0
- worker_internal/planner/__init__.py +1 -0
- worker_internal/planner/activities.py +1499 -0
- worker_internal/planner/agent_tools.py +197 -0
- worker_internal/planner/event_models.py +148 -0
- worker_internal/planner/event_publisher.py +67 -0
- worker_internal/planner/models.py +199 -0
- worker_internal/planner/retry_logic.py +134 -0
- worker_internal/planner/worker.py +300 -0
- worker_internal/planner/workflows.py +970 -0
|
@@ -0,0 +1,1585 @@
|
|
|
1
|
+
"""Agent-related Temporal activities"""
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from typing import Optional, Any, List, Dict
|
|
5
|
+
from datetime import datetime, timezone
|
|
6
|
+
from temporalio import activity
|
|
7
|
+
from temporalio.exceptions import ApplicationError
|
|
8
|
+
import structlog
|
|
9
|
+
import os
|
|
10
|
+
import httpx
|
|
11
|
+
from pathlib import Path
|
|
12
|
+
|
|
13
|
+
from control_plane_api.worker.utils.logging_helper import execution_logger
|
|
14
|
+
|
|
15
|
+
from agno.tools.shell import ShellTools
|
|
16
|
+
from agno.tools.python import PythonTools
|
|
17
|
+
from agno.tools.file import FileTools
|
|
18
|
+
from control_plane_api.worker.control_plane_client import get_control_plane_client
|
|
19
|
+
from control_plane_api.worker.services.skill_factory import SkillFactory
|
|
20
|
+
|
|
21
|
+
logger = structlog.get_logger()
|
|
22
|
+
|
|
23
|
+
# Global registry for active Agent/Team instances to support cancellation
|
|
24
|
+
# Key: execution_id, Value: {agent: Agent, run_id: str}
|
|
25
|
+
_active_agents: Dict[str, Dict[str, Any]] = {}
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def instantiate_skill(skill_data: dict) -> Optional[Any]:
|
|
29
|
+
"""
|
|
30
|
+
Instantiate an Agno toolkit based on skill configuration from Control Plane.
|
|
31
|
+
|
|
32
|
+
Args:
|
|
33
|
+
skill_data: Skill data from Control Plane API containing:
|
|
34
|
+
- type: Skill type (file_system, shell, python, docker, etc.)
|
|
35
|
+
- name: Skill name
|
|
36
|
+
- configuration: Dict with skill-specific config
|
|
37
|
+
- enabled: Whether skill is enabled
|
|
38
|
+
|
|
39
|
+
Returns:
|
|
40
|
+
Instantiated Agno toolkit or None if type not supported/enabled
|
|
41
|
+
"""
|
|
42
|
+
if not skill_data.get("enabled", True):
|
|
43
|
+
print(f" ⊗ Skipping disabled skill: {skill_data.get('name')}")
|
|
44
|
+
return None
|
|
45
|
+
|
|
46
|
+
skill_type = skill_data.get("type", "").lower()
|
|
47
|
+
config = skill_data.get("configuration", {})
|
|
48
|
+
name = skill_data.get("name", "Unknown")
|
|
49
|
+
|
|
50
|
+
try:
|
|
51
|
+
# Map Control Plane skill types to Agno toolkit classes
|
|
52
|
+
if skill_type in ["file_system", "file", "file_generation"]:
|
|
53
|
+
# FileTools: file operations (read, write, list, search)
|
|
54
|
+
# Note: file_generation is mapped to FileTools (save_file functionality)
|
|
55
|
+
base_dir = config.get("base_dir")
|
|
56
|
+
toolkit = FileTools(
|
|
57
|
+
base_dir=Path(base_dir) if base_dir else None,
|
|
58
|
+
enable_save_file=config.get("enable_save_file", True),
|
|
59
|
+
enable_read_file=config.get("enable_read_file", True),
|
|
60
|
+
enable_list_files=config.get("enable_list_files", True),
|
|
61
|
+
enable_search_files=config.get("enable_search_files", True),
|
|
62
|
+
)
|
|
63
|
+
print(f" ✓ Instantiated FileTools: {name}")
|
|
64
|
+
if skill_type == "file_generation":
|
|
65
|
+
print(f" - Type: File Generation (using FileTools.save_file)")
|
|
66
|
+
print(f" - Base Dir: {base_dir or 'Current directory'}")
|
|
67
|
+
print(f" - Read: {config.get('enable_read_file', True)}, Write: {config.get('enable_save_file', True)}")
|
|
68
|
+
return toolkit
|
|
69
|
+
|
|
70
|
+
elif skill_type in ["shell", "bash"]:
|
|
71
|
+
# ShellTools: shell command execution
|
|
72
|
+
base_dir = config.get("base_dir")
|
|
73
|
+
toolkit = ShellTools(
|
|
74
|
+
base_dir=Path(base_dir) if base_dir else None,
|
|
75
|
+
enable_run_shell_command=config.get("enable_run_shell_command", True),
|
|
76
|
+
)
|
|
77
|
+
print(f" ✓ Instantiated ShellTools: {name}")
|
|
78
|
+
print(f" - Base Dir: {base_dir or 'Current directory'}")
|
|
79
|
+
print(f" - Run Commands: {config.get('enable_run_shell_command', True)}")
|
|
80
|
+
return toolkit
|
|
81
|
+
|
|
82
|
+
elif skill_type == "python":
|
|
83
|
+
# PythonTools: Python code execution
|
|
84
|
+
base_dir = config.get("base_dir")
|
|
85
|
+
toolkit = PythonTools(
|
|
86
|
+
base_dir=Path(base_dir) if base_dir else None,
|
|
87
|
+
safe_globals=config.get("safe_globals"),
|
|
88
|
+
safe_locals=config.get("safe_locals"),
|
|
89
|
+
)
|
|
90
|
+
print(f" ✓ Instantiated PythonTools: {name}")
|
|
91
|
+
print(f" - Base Dir: {base_dir or 'Current directory'}")
|
|
92
|
+
return toolkit
|
|
93
|
+
|
|
94
|
+
elif skill_type == "docker":
|
|
95
|
+
# DockerTools requires docker package and running Docker daemon
|
|
96
|
+
try:
|
|
97
|
+
from agno.tools.docker import DockerTools
|
|
98
|
+
import docker
|
|
99
|
+
|
|
100
|
+
# Check if Docker daemon is accessible
|
|
101
|
+
try:
|
|
102
|
+
docker_client = docker.from_env()
|
|
103
|
+
docker_client.ping()
|
|
104
|
+
|
|
105
|
+
# Docker is available, instantiate toolkit
|
|
106
|
+
toolkit = DockerTools()
|
|
107
|
+
print(f" ✓ Instantiated DockerTools: {name}")
|
|
108
|
+
print(f" - Docker daemon: Connected")
|
|
109
|
+
docker_client.close()
|
|
110
|
+
return toolkit
|
|
111
|
+
|
|
112
|
+
except Exception as docker_error:
|
|
113
|
+
print(f" ⚠ Docker daemon not available - skipping: {name}")
|
|
114
|
+
print(f" Error: {str(docker_error)}")
|
|
115
|
+
return None
|
|
116
|
+
|
|
117
|
+
except ImportError:
|
|
118
|
+
print(f" ⚠ Docker skill requires 'docker' package - skipping: {name}")
|
|
119
|
+
print(f" Install with: pip install docker")
|
|
120
|
+
return None
|
|
121
|
+
|
|
122
|
+
elif skill_type in ["data_visualization", "diagramming", "visualization"]:
|
|
123
|
+
# DataVisualizationTools: Create diagrams using Mermaid syntax
|
|
124
|
+
# This is a custom implementation that uses streaming to send diagram data
|
|
125
|
+
from services.data_visualization import DataVisualizationTools
|
|
126
|
+
|
|
127
|
+
toolkit = DataVisualizationTools(
|
|
128
|
+
max_diagram_size=config.get("max_diagram_size", 50000),
|
|
129
|
+
enable_flowchart=config.get("enable_flowchart", True),
|
|
130
|
+
enable_sequence=config.get("enable_sequence", True),
|
|
131
|
+
enable_class_diagram=config.get("enable_class_diagram", True),
|
|
132
|
+
enable_er_diagram=config.get("enable_er_diagram", True),
|
|
133
|
+
enable_gantt=config.get("enable_gantt", True),
|
|
134
|
+
enable_pie_chart=config.get("enable_pie_chart", True),
|
|
135
|
+
enable_state_diagram=config.get("enable_state_diagram", True),
|
|
136
|
+
enable_git_graph=config.get("enable_git_graph", True),
|
|
137
|
+
enable_user_journey=config.get("enable_user_journey", True),
|
|
138
|
+
enable_quadrant_chart=config.get("enable_quadrant_chart", True),
|
|
139
|
+
)
|
|
140
|
+
print(f" ✓ Instantiated DataVisualizationTools: {name}")
|
|
141
|
+
print(f" - Max diagram size: {config.get('max_diagram_size', 50000)} chars")
|
|
142
|
+
print(f" - Supported: Mermaid diagrams (flowchart, sequence, class, ER, etc.)")
|
|
143
|
+
return toolkit
|
|
144
|
+
|
|
145
|
+
else:
|
|
146
|
+
print(f" ⚠ Unsupported skill type '{skill_type}': {name}")
|
|
147
|
+
return None
|
|
148
|
+
|
|
149
|
+
except Exception as e:
|
|
150
|
+
print(f" ❌ Error instantiating skill '{name}' (type: {skill_type}): {str(e)}")
|
|
151
|
+
logger.error(
|
|
152
|
+
f"Error instantiating skill",
|
|
153
|
+
extra={
|
|
154
|
+
"skill_name": name,
|
|
155
|
+
"skill_type": skill_type,
|
|
156
|
+
"error": str(e)
|
|
157
|
+
}
|
|
158
|
+
)
|
|
159
|
+
return None
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
@dataclass
|
|
163
|
+
class ActivityExecuteAgentInput:
|
|
164
|
+
"""Input for execute_agent_llm activity"""
|
|
165
|
+
execution_id: str
|
|
166
|
+
agent_id: str
|
|
167
|
+
organization_id: str
|
|
168
|
+
prompt: str
|
|
169
|
+
system_prompt: Optional[str] = None
|
|
170
|
+
model_id: Optional[str] = None
|
|
171
|
+
model_config: dict = None
|
|
172
|
+
mcp_servers: dict = None # MCP servers configuration
|
|
173
|
+
session_id: Optional[str] = None # Session ID for Agno session management (use execution_id)
|
|
174
|
+
user_id: Optional[str] = None # User ID for multi-user support
|
|
175
|
+
user_message_id: Optional[str] = None # Message ID from workflow for deduplication
|
|
176
|
+
user_name: Optional[str] = None # User name for attribution
|
|
177
|
+
user_email: Optional[str] = None # User email for attribution
|
|
178
|
+
user_avatar: Optional[str] = None # User avatar for attribution
|
|
179
|
+
# Note: control_plane_url and api_key are read from worker environment variables (CONTROL_PLANE_URL, KUBIYA_API_KEY)
|
|
180
|
+
|
|
181
|
+
def __post_init__(self):
|
|
182
|
+
if self.model_config is None:
|
|
183
|
+
self.model_config = {}
|
|
184
|
+
if self.mcp_servers is None:
|
|
185
|
+
self.mcp_servers = {}
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
@dataclass
|
|
189
|
+
class ActivityUpdateExecutionInput:
|
|
190
|
+
"""Input for update_execution_status activity"""
|
|
191
|
+
execution_id: str
|
|
192
|
+
status: str
|
|
193
|
+
started_at: Optional[str] = None
|
|
194
|
+
completed_at: Optional[str] = None
|
|
195
|
+
response: Optional[str] = None
|
|
196
|
+
error_message: Optional[str] = None
|
|
197
|
+
usage: dict = None
|
|
198
|
+
execution_metadata: dict = None
|
|
199
|
+
|
|
200
|
+
def __post_init__(self):
|
|
201
|
+
if self.usage is None:
|
|
202
|
+
self.usage = {}
|
|
203
|
+
if self.execution_metadata is None:
|
|
204
|
+
self.execution_metadata = {}
|
|
205
|
+
|
|
206
|
+
|
|
207
|
+
@dataclass
|
|
208
|
+
class ActivityGetExecutionInput:
|
|
209
|
+
"""Input for get_execution_details activity"""
|
|
210
|
+
execution_id: str
|
|
211
|
+
|
|
212
|
+
|
|
213
|
+
@dataclass
|
|
214
|
+
class ActivityUpdateAgentInput:
|
|
215
|
+
"""Input for update_agent_status activity"""
|
|
216
|
+
agent_id: str
|
|
217
|
+
organization_id: str
|
|
218
|
+
status: str
|
|
219
|
+
last_active_at: str
|
|
220
|
+
error_message: Optional[str] = None
|
|
221
|
+
state: dict = None
|
|
222
|
+
|
|
223
|
+
def __post_init__(self):
|
|
224
|
+
if self.state is None:
|
|
225
|
+
self.state = {}
|
|
226
|
+
|
|
227
|
+
|
|
228
|
+
@activity.defn
|
|
229
|
+
async def execute_agent_llm(input: ActivityExecuteAgentInput) -> dict:
|
|
230
|
+
"""
|
|
231
|
+
Execute an agent's LLM call with Agno Teams and session management.
|
|
232
|
+
|
|
233
|
+
This activity uses Agno Teams with session support for persistent conversation history.
|
|
234
|
+
The session_id should be set to execution_id for 1:1 mapping.
|
|
235
|
+
|
|
236
|
+
Args:
|
|
237
|
+
input: Activity input with execution details
|
|
238
|
+
|
|
239
|
+
Returns:
|
|
240
|
+
Dict with response, usage, success flag, session messages, etc.
|
|
241
|
+
"""
|
|
242
|
+
print("\n" + "="*80)
|
|
243
|
+
print("🤖 AGENT EXECUTION START")
|
|
244
|
+
print("="*80)
|
|
245
|
+
print(f"Execution ID: {input.execution_id}")
|
|
246
|
+
print(f"Agent ID: {input.agent_id}")
|
|
247
|
+
print(f"Organization: {input.organization_id}")
|
|
248
|
+
print(f"Model: {input.model_id or 'default'}")
|
|
249
|
+
print(f"Session ID: {input.session_id}")
|
|
250
|
+
print(f"MCP Servers: {len(input.mcp_servers)} configured" if input.mcp_servers else "MCP Servers: None")
|
|
251
|
+
print(f"Prompt: {input.prompt[:100]}..." if len(input.prompt) > 100 else f"Prompt: {input.prompt}")
|
|
252
|
+
print("="*80 + "\n")
|
|
253
|
+
|
|
254
|
+
activity.logger.info(
|
|
255
|
+
f"Executing agent LLM call with Agno Sessions",
|
|
256
|
+
extra={
|
|
257
|
+
"execution_id": input.execution_id,
|
|
258
|
+
"agent_id": input.agent_id,
|
|
259
|
+
"organization_id": input.organization_id,
|
|
260
|
+
"model_id": input.model_id,
|
|
261
|
+
"has_mcp_servers": bool(input.mcp_servers),
|
|
262
|
+
"mcp_server_count": len(input.mcp_servers) if input.mcp_servers else 0,
|
|
263
|
+
"mcp_server_ids": list(input.mcp_servers.keys()) if input.mcp_servers else [],
|
|
264
|
+
"session_id": input.session_id,
|
|
265
|
+
}
|
|
266
|
+
)
|
|
267
|
+
|
|
268
|
+
try:
|
|
269
|
+
# Get Control Plane client for all communication with Control Plane
|
|
270
|
+
control_plane = get_control_plane_client()
|
|
271
|
+
|
|
272
|
+
# STEP 1: Load existing session history from Control Plane (if this is a continuation)
|
|
273
|
+
# This enables conversation continuity across multiple execution turns
|
|
274
|
+
# IMPORTANT: This must be non-blocking and have proper timeout/retry
|
|
275
|
+
session_history = []
|
|
276
|
+
if input.session_id:
|
|
277
|
+
print(f"\n📥 Loading session history from Control Plane...")
|
|
278
|
+
|
|
279
|
+
# Try up to 3 times with exponential backoff for transient failures
|
|
280
|
+
max_retries = 3
|
|
281
|
+
for attempt in range(max_retries):
|
|
282
|
+
try:
|
|
283
|
+
if attempt > 0:
|
|
284
|
+
print(f" 🔄 Retry attempt {attempt + 1}/{max_retries}...")
|
|
285
|
+
|
|
286
|
+
session_data = control_plane.get_session(
|
|
287
|
+
execution_id=input.execution_id,
|
|
288
|
+
session_id=input.session_id
|
|
289
|
+
)
|
|
290
|
+
if session_data and session_data.get("messages"):
|
|
291
|
+
session_history = session_data["messages"]
|
|
292
|
+
print(f" ✅ Loaded {len(session_history)} messages from previous turns")
|
|
293
|
+
|
|
294
|
+
activity.logger.info(
|
|
295
|
+
"Session history loaded from Control Plane",
|
|
296
|
+
extra={
|
|
297
|
+
"execution_id": input.execution_id,
|
|
298
|
+
"session_id": input.session_id,
|
|
299
|
+
"message_count": len(session_history),
|
|
300
|
+
"attempt": attempt + 1,
|
|
301
|
+
}
|
|
302
|
+
)
|
|
303
|
+
break # Success - exit retry loop
|
|
304
|
+
else:
|
|
305
|
+
print(f" ℹ️ No previous session found - starting new conversation")
|
|
306
|
+
break # No session exists - not an error
|
|
307
|
+
|
|
308
|
+
except httpx.TimeoutException as e:
|
|
309
|
+
print(f" ⏱️ Timeout loading session (attempt {attempt + 1}/{max_retries})")
|
|
310
|
+
activity.logger.warning(
|
|
311
|
+
"Session load timeout",
|
|
312
|
+
extra={"error": str(e), "execution_id": input.execution_id, "attempt": attempt + 1}
|
|
313
|
+
)
|
|
314
|
+
if attempt < max_retries - 1:
|
|
315
|
+
import time
|
|
316
|
+
time.sleep(2 ** attempt) # Exponential backoff: 1s, 2s, 4s
|
|
317
|
+
continue
|
|
318
|
+
else:
|
|
319
|
+
print(f" ⚠️ Session load failed after {max_retries} attempts - continuing without history")
|
|
320
|
+
|
|
321
|
+
except Exception as e:
|
|
322
|
+
error_type = type(e).__name__
|
|
323
|
+
print(f" ⚠️ Failed to load session history ({error_type}): {str(e)[:100]}")
|
|
324
|
+
activity.logger.warning(
|
|
325
|
+
"Failed to load session history from Control Plane",
|
|
326
|
+
extra={
|
|
327
|
+
"error": str(e),
|
|
328
|
+
"error_type": error_type,
|
|
329
|
+
"execution_id": input.execution_id,
|
|
330
|
+
"attempt": attempt + 1
|
|
331
|
+
}
|
|
332
|
+
)
|
|
333
|
+
# For non-timeout errors, don't retry - likely invalid session
|
|
334
|
+
break
|
|
335
|
+
|
|
336
|
+
# Always continue execution even if session loading fails
|
|
337
|
+
print(f" → Continuing with {len(session_history)} messages in context\n")
|
|
338
|
+
|
|
339
|
+
# Get LiteLLM credentials from environment (set by worker from registration)
|
|
340
|
+
litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
|
|
341
|
+
litellm_api_key = os.getenv("LITELLM_API_KEY")
|
|
342
|
+
|
|
343
|
+
if not litellm_api_key:
|
|
344
|
+
raise ValueError("LITELLM_API_KEY environment variable not set")
|
|
345
|
+
|
|
346
|
+
# Get model from input or use default
|
|
347
|
+
model = input.model_id or os.environ.get("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4")
|
|
348
|
+
|
|
349
|
+
# Fetch resolved skills from Control Plane if available
|
|
350
|
+
skills = []
|
|
351
|
+
if input.agent_id:
|
|
352
|
+
print(f"🔧 Fetching skills from Control Plane...")
|
|
353
|
+
try:
|
|
354
|
+
skills = control_plane.get_skills(input.agent_id)
|
|
355
|
+
if skills:
|
|
356
|
+
print(f"✅ Resolved {len(skills)} skills from Control Plane")
|
|
357
|
+
print(f" Skill Types: {[t.get('type') for t in skills]}")
|
|
358
|
+
print(f" Skill Sources: {[t.get('source') for t in skills]}")
|
|
359
|
+
print(f" Skill Names: {[t.get('name') for t in skills]}\n")
|
|
360
|
+
|
|
361
|
+
activity.logger.info(
|
|
362
|
+
f"Resolved skills from Control Plane",
|
|
363
|
+
extra={
|
|
364
|
+
"agent_id": input.agent_id,
|
|
365
|
+
"skill_count": len(skills),
|
|
366
|
+
"skill_types": [t.get("type") for t in skills],
|
|
367
|
+
"skill_sources": [t.get("source") for t in skills],
|
|
368
|
+
"skill_names": [t.get("name") for t in skills],
|
|
369
|
+
}
|
|
370
|
+
)
|
|
371
|
+
else:
|
|
372
|
+
print(f"⚠️ No skills found for agent\n")
|
|
373
|
+
except Exception as e:
|
|
374
|
+
print(f"❌ Error fetching skills: {str(e)}\n")
|
|
375
|
+
activity.logger.error(
|
|
376
|
+
f"Error fetching skills from Control Plane: {str(e)}",
|
|
377
|
+
extra={"error": str(e)}
|
|
378
|
+
)
|
|
379
|
+
# Continue execution without skills
|
|
380
|
+
else:
|
|
381
|
+
print(f"ℹ️ No agent_id provided - skipping skill resolution\n")
|
|
382
|
+
|
|
383
|
+
# Instantiate Agno toolkits from Control Plane skills
|
|
384
|
+
print(f"\n🔧 Instantiating Skills:")
|
|
385
|
+
agno_toolkits = []
|
|
386
|
+
if skills:
|
|
387
|
+
# Create factory instance for agno runtime
|
|
388
|
+
skill_factory = SkillFactory(runtime_type="agno")
|
|
389
|
+
skill_factory.initialize()
|
|
390
|
+
|
|
391
|
+
for skill in skills:
|
|
392
|
+
# Add execution_id to skill data for workflow streaming
|
|
393
|
+
skill['execution_id'] = input.execution_id
|
|
394
|
+
# Use SkillFactory which supports all skill types including workflow_executor
|
|
395
|
+
toolkit = skill_factory.create_skill(skill)
|
|
396
|
+
if toolkit:
|
|
397
|
+
agno_toolkits.append(toolkit)
|
|
398
|
+
|
|
399
|
+
if agno_toolkits:
|
|
400
|
+
print(f"\n✅ Successfully instantiated {len(agno_toolkits)} skill(s)")
|
|
401
|
+
else:
|
|
402
|
+
print(f"\nℹ️ No skills instantiated\n")
|
|
403
|
+
|
|
404
|
+
print(f"📦 Total Tools Available:")
|
|
405
|
+
print(f" MCP Servers: {len(input.mcp_servers)}")
|
|
406
|
+
print(f" OS-Level Skills: {len(agno_toolkits)}\n")
|
|
407
|
+
|
|
408
|
+
activity.logger.info(
|
|
409
|
+
f"Using Agno Agent with sessions and skills",
|
|
410
|
+
extra={
|
|
411
|
+
"execution_id": input.execution_id,
|
|
412
|
+
"session_id": input.session_id,
|
|
413
|
+
"has_mcp_servers": bool(input.mcp_servers),
|
|
414
|
+
"mcp_server_count": len(input.mcp_servers) if input.mcp_servers else 0,
|
|
415
|
+
"mcp_servers": list(input.mcp_servers.keys()) if input.mcp_servers else [],
|
|
416
|
+
"skill_count": len(agno_toolkits),
|
|
417
|
+
"model": model,
|
|
418
|
+
}
|
|
419
|
+
)
|
|
420
|
+
|
|
421
|
+
# Import Agno libraries
|
|
422
|
+
from agno.agent import Agent
|
|
423
|
+
from agno.models.litellm import LiteLLM
|
|
424
|
+
|
|
425
|
+
print(f"\n🤖 Creating Agno Agent:")
|
|
426
|
+
print(f" Model: {model}")
|
|
427
|
+
print(f" Skills: {len(agno_toolkits)}")
|
|
428
|
+
|
|
429
|
+
# Send heartbeat: Creating agent
|
|
430
|
+
activity.heartbeat({"status": "Creating agent with skills..."})
|
|
431
|
+
|
|
432
|
+
# Track tool executions for real-time streaming
|
|
433
|
+
tool_execution_messages = []
|
|
434
|
+
|
|
435
|
+
# Create tool hook to capture tool execution for real-time streaming
|
|
436
|
+
# Agno inspects the signature and passes matching parameters
|
|
437
|
+
def tool_hook(name: str = None, function_name: str = None, function=None, arguments: dict = None, **kwargs):
|
|
438
|
+
"""Hook to capture tool execution and add to messages for streaming
|
|
439
|
+
|
|
440
|
+
Agno passes these parameters based on our signature:
|
|
441
|
+
- name or function_name: The tool function name
|
|
442
|
+
- function: The callable being executed (this is the NEXT function in the chain)
|
|
443
|
+
- arguments: Dict of arguments passed to the tool
|
|
444
|
+
|
|
445
|
+
The hook must CALL the function and return its result.
|
|
446
|
+
"""
|
|
447
|
+
# Get tool name from Agno's parameters
|
|
448
|
+
tool_name = name or function_name or "unknown"
|
|
449
|
+
tool_args = arguments or {}
|
|
450
|
+
|
|
451
|
+
# Generate unique tool execution ID using UUID to avoid collisions
|
|
452
|
+
import uuid
|
|
453
|
+
tool_execution_id = f"{tool_name}_{uuid.uuid4().hex[:12]}"
|
|
454
|
+
|
|
455
|
+
print(f" 🔧 Tool Starting: {tool_name} (ID: {tool_execution_id})")
|
|
456
|
+
if tool_args:
|
|
457
|
+
args_preview = str(tool_args)[:200]
|
|
458
|
+
print(f" Args: {args_preview}{'...' if len(str(tool_args)) > 200 else ''}")
|
|
459
|
+
|
|
460
|
+
# Publish streaming event to Control Plane (real-time UI update)
|
|
461
|
+
control_plane.publish_event(
|
|
462
|
+
execution_id=input.execution_id,
|
|
463
|
+
event_type="tool_started",
|
|
464
|
+
data={
|
|
465
|
+
"tool_name": tool_name,
|
|
466
|
+
"tool_execution_id": tool_execution_id, # Unique ID for this execution
|
|
467
|
+
"tool_arguments": tool_args,
|
|
468
|
+
"message_id": message_id, # Link tool to assistant message turn
|
|
469
|
+
"message": f"🔧 Executing tool: {tool_name}",
|
|
470
|
+
}
|
|
471
|
+
)
|
|
472
|
+
|
|
473
|
+
tool_execution_messages.append({
|
|
474
|
+
"role": "system",
|
|
475
|
+
"content": f"🔧 Executing tool: **{tool_name}**",
|
|
476
|
+
"tool_name": tool_name,
|
|
477
|
+
"tool_event": "started",
|
|
478
|
+
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
479
|
+
})
|
|
480
|
+
|
|
481
|
+
# CRITICAL: Actually call the function and handle completion
|
|
482
|
+
result = None
|
|
483
|
+
error = None
|
|
484
|
+
try:
|
|
485
|
+
# Call the actual function (next in the hook chain)
|
|
486
|
+
if function and callable(function):
|
|
487
|
+
result = function(**tool_args) if tool_args else function()
|
|
488
|
+
else:
|
|
489
|
+
raise ValueError(f"Function not callable: {function}")
|
|
490
|
+
|
|
491
|
+
status = "success"
|
|
492
|
+
icon = "✅"
|
|
493
|
+
print(f" {icon} Tool Success: {tool_name}")
|
|
494
|
+
|
|
495
|
+
except Exception as e:
|
|
496
|
+
error = e
|
|
497
|
+
status = "failed"
|
|
498
|
+
icon = "❌"
|
|
499
|
+
print(f" {icon} Tool Failed: {tool_name} - {str(e)}")
|
|
500
|
+
|
|
501
|
+
# Publish completion event to Control Plane (real-time UI update)
|
|
502
|
+
control_plane.publish_event(
|
|
503
|
+
execution_id=input.execution_id,
|
|
504
|
+
event_type="tool_completed",
|
|
505
|
+
data={
|
|
506
|
+
"tool_name": tool_name,
|
|
507
|
+
"tool_execution_id": tool_execution_id, # Same ID to match the started event
|
|
508
|
+
"status": status,
|
|
509
|
+
"error": str(error) if error else None,
|
|
510
|
+
"tool_output": result if result is not None else None, # Include tool output for UI display
|
|
511
|
+
"message_id": message_id, # Link tool to assistant message turn
|
|
512
|
+
"message": f"{icon} Tool {status}: {tool_name}",
|
|
513
|
+
}
|
|
514
|
+
)
|
|
515
|
+
|
|
516
|
+
tool_execution_messages.append({
|
|
517
|
+
"role": "system",
|
|
518
|
+
"content": f"{icon} Tool {status}: **{tool_name}**",
|
|
519
|
+
"tool_name": tool_name,
|
|
520
|
+
"tool_event": "completed",
|
|
521
|
+
"tool_status": status,
|
|
522
|
+
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
523
|
+
})
|
|
524
|
+
|
|
525
|
+
# If there was an error, re-raise it so Agno knows the tool failed
|
|
526
|
+
if error:
|
|
527
|
+
raise error
|
|
528
|
+
|
|
529
|
+
# Return the result to continue the chain
|
|
530
|
+
return result
|
|
531
|
+
|
|
532
|
+
# Build conversation context from session history for manual session management
|
|
533
|
+
# Workers don't have database access, so we manage sessions via Control Plane API
|
|
534
|
+
conversation_context = []
|
|
535
|
+
if session_history:
|
|
536
|
+
print(f"\n📝 Building conversation context from {len(session_history)} previous messages...")
|
|
537
|
+
for msg in session_history:
|
|
538
|
+
# Convert Control Plane message format to Agno format
|
|
539
|
+
conversation_context.append({
|
|
540
|
+
"role": msg.get("role", "user"),
|
|
541
|
+
"content": msg.get("content", ""),
|
|
542
|
+
})
|
|
543
|
+
print(f" ✅ Conversation context ready\n")
|
|
544
|
+
|
|
545
|
+
# Create Agno Agent with LiteLLM configuration
|
|
546
|
+
# Note: NO database - workers use Control Plane API for session management
|
|
547
|
+
# Use openai/ prefix for custom proxy compatibility
|
|
548
|
+
agent = Agent(
|
|
549
|
+
name=f"Agent {input.agent_id}",
|
|
550
|
+
role=input.system_prompt or "You are a helpful AI assistant",
|
|
551
|
+
model=LiteLLM(
|
|
552
|
+
id=f"openai/{model}",
|
|
553
|
+
api_base=litellm_api_base,
|
|
554
|
+
api_key=litellm_api_key,
|
|
555
|
+
),
|
|
556
|
+
tools=agno_toolkits if agno_toolkits else None, # Add skills to agent
|
|
557
|
+
tool_hooks=[tool_hook], # Add hook for real-time tool updates
|
|
558
|
+
# NO db parameter - session management via Control Plane API
|
|
559
|
+
)
|
|
560
|
+
|
|
561
|
+
# Register agent for cancellation support
|
|
562
|
+
_active_agents[input.execution_id] = {
|
|
563
|
+
"agent": agent,
|
|
564
|
+
"run_id": None, # Will be set when run starts
|
|
565
|
+
"started_at": datetime.now(timezone.utc).isoformat(),
|
|
566
|
+
}
|
|
567
|
+
print(f"✅ Agent registered for cancellation support (execution_id: {input.execution_id})\n")
|
|
568
|
+
|
|
569
|
+
# Cache execution metadata in Redis for fast SSE lookups (avoid DB queries)
|
|
570
|
+
control_plane.cache_metadata(input.execution_id, "AGENT")
|
|
571
|
+
|
|
572
|
+
# Execute agent run with streaming
|
|
573
|
+
print("⚡ Executing Agent Run with Streaming...\n")
|
|
574
|
+
|
|
575
|
+
# Send heartbeat: Starting execution
|
|
576
|
+
activity.heartbeat({"status": "Agent is processing your request..."})
|
|
577
|
+
|
|
578
|
+
import asyncio
|
|
579
|
+
|
|
580
|
+
# Stream the response and collect chunks
|
|
581
|
+
response_chunks = []
|
|
582
|
+
full_response = ""
|
|
583
|
+
|
|
584
|
+
# Generate unique message ID for this turn (execution_id + timestamp)
|
|
585
|
+
import time
|
|
586
|
+
message_id = f"{input.execution_id}_{int(time.time() * 1000000)}"
|
|
587
|
+
|
|
588
|
+
def stream_agent_run():
|
|
589
|
+
"""Run agent with streaming and collect response"""
|
|
590
|
+
nonlocal full_response
|
|
591
|
+
run_id_published = False
|
|
592
|
+
|
|
593
|
+
try:
|
|
594
|
+
# Build full prompt with conversation history for context
|
|
595
|
+
# Since worker has no database, we manually prepend history
|
|
596
|
+
if conversation_context:
|
|
597
|
+
# Agno Agent supports passing messages parameter for conversation context
|
|
598
|
+
run_response = agent.run(
|
|
599
|
+
input.prompt,
|
|
600
|
+
stream=True,
|
|
601
|
+
messages=conversation_context, # Pass previous conversation
|
|
602
|
+
)
|
|
603
|
+
else:
|
|
604
|
+
# First turn - no history
|
|
605
|
+
run_response = agent.run(input.prompt, stream=True)
|
|
606
|
+
|
|
607
|
+
# Iterate over streaming chunks
|
|
608
|
+
for chunk in run_response:
|
|
609
|
+
# Capture and publish run_id from first chunk for cancellation support
|
|
610
|
+
if not run_id_published and hasattr(chunk, 'run_id') and chunk.run_id:
|
|
611
|
+
agno_run_id = chunk.run_id
|
|
612
|
+
print(f"\n🆔 Agno run_id: {agno_run_id}")
|
|
613
|
+
|
|
614
|
+
# Store run_id in registry for cancellation
|
|
615
|
+
if input.execution_id in _active_agents:
|
|
616
|
+
_active_agents[input.execution_id]["run_id"] = agno_run_id
|
|
617
|
+
|
|
618
|
+
# Publish run_id to Redis for Control Plane cancellation access
|
|
619
|
+
# This allows users to cancel via STOP button in UI
|
|
620
|
+
control_plane.publish_event(
|
|
621
|
+
execution_id=input.execution_id,
|
|
622
|
+
event_type="run_started",
|
|
623
|
+
data={
|
|
624
|
+
"run_id": agno_run_id,
|
|
625
|
+
"agent_id": input.agent_id,
|
|
626
|
+
"cancellable": True,
|
|
627
|
+
}
|
|
628
|
+
)
|
|
629
|
+
run_id_published = True
|
|
630
|
+
|
|
631
|
+
# Filter out whitespace-only chunks to prevent "(no content)" in UI
|
|
632
|
+
if hasattr(chunk, 'content') and chunk.content:
|
|
633
|
+
content = str(chunk.content)
|
|
634
|
+
# Only process and send chunks with non-whitespace content
|
|
635
|
+
if content.strip():
|
|
636
|
+
full_response += content
|
|
637
|
+
response_chunks.append(content)
|
|
638
|
+
print(content, end='', flush=True)
|
|
639
|
+
|
|
640
|
+
# Stream chunk to Control Plane for real-time UI updates
|
|
641
|
+
# Include message_id so UI knows which message these chunks belong to
|
|
642
|
+
control_plane.publish_event(
|
|
643
|
+
execution_id=input.execution_id,
|
|
644
|
+
event_type="message_chunk",
|
|
645
|
+
data={
|
|
646
|
+
"role": "assistant",
|
|
647
|
+
"content": content,
|
|
648
|
+
"is_chunk": True,
|
|
649
|
+
"message_id": message_id, # Unique ID for this turn
|
|
650
|
+
}
|
|
651
|
+
)
|
|
652
|
+
|
|
653
|
+
# Note: Cannot send heartbeat from sync context (thread pool)
|
|
654
|
+
|
|
655
|
+
print() # New line after streaming
|
|
656
|
+
|
|
657
|
+
# Return the iterator's final result
|
|
658
|
+
return run_response
|
|
659
|
+
except Exception as e:
|
|
660
|
+
print(f"\n❌ Streaming error: {str(e)}")
|
|
661
|
+
# Fall back to non-streaming
|
|
662
|
+
if conversation_context:
|
|
663
|
+
return agent.run(input.prompt, stream=False, messages=conversation_context)
|
|
664
|
+
else:
|
|
665
|
+
return agent.run(input.prompt, stream=False)
|
|
666
|
+
|
|
667
|
+
# Execute in thread pool (NO TIMEOUT - tasks can run as long as needed)
|
|
668
|
+
# Control Plane can cancel via Agno's cancel_run API if user requests it
|
|
669
|
+
result = await asyncio.to_thread(stream_agent_run)
|
|
670
|
+
|
|
671
|
+
# Send heartbeat: Completed
|
|
672
|
+
activity.heartbeat({"status": "Agent execution completed, preparing response..."})
|
|
673
|
+
|
|
674
|
+
print("✅ Agent Execution Completed!")
|
|
675
|
+
print(f" Response Length: {len(full_response)} chars\n")
|
|
676
|
+
|
|
677
|
+
activity.logger.info(
|
|
678
|
+
f"Agent LLM call completed",
|
|
679
|
+
extra={
|
|
680
|
+
"execution_id": input.execution_id,
|
|
681
|
+
"has_content": bool(full_response),
|
|
682
|
+
}
|
|
683
|
+
)
|
|
684
|
+
|
|
685
|
+
# Use the streamed response content
|
|
686
|
+
response_content = full_response if full_response else (result.content if hasattr(result, "content") else str(result))
|
|
687
|
+
|
|
688
|
+
# Extract tool call messages for UI streaming
|
|
689
|
+
tool_messages = []
|
|
690
|
+
if hasattr(result, "messages") and result.messages:
|
|
691
|
+
for msg in result.messages:
|
|
692
|
+
# Check if message has tool calls
|
|
693
|
+
if hasattr(msg, "tool_calls") and msg.tool_calls:
|
|
694
|
+
for tool_call in msg.tool_calls:
|
|
695
|
+
tool_name = getattr(tool_call, "function", {}).get("name") if hasattr(tool_call, "function") else str(tool_call)
|
|
696
|
+
tool_args = getattr(tool_call, "function", {}).get("arguments") if hasattr(tool_call, "function") else {}
|
|
697
|
+
|
|
698
|
+
print(f" 🔧 Tool Call: {tool_name}")
|
|
699
|
+
|
|
700
|
+
tool_messages.append({
|
|
701
|
+
"role": "tool",
|
|
702
|
+
"content": f"Executing {tool_name}...",
|
|
703
|
+
"tool_name": tool_name,
|
|
704
|
+
"tool_execution_id": tool_execution_id, # CRITICAL: For deduplication
|
|
705
|
+
"tool_input": tool_args,
|
|
706
|
+
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
707
|
+
})
|
|
708
|
+
|
|
709
|
+
if tool_messages:
|
|
710
|
+
print(f"\n🔧 Tool Calls Captured: {len(tool_messages)}")
|
|
711
|
+
|
|
712
|
+
# Extract usage metrics if available
|
|
713
|
+
usage = {}
|
|
714
|
+
if hasattr(result, "metrics") and result.metrics:
|
|
715
|
+
metrics = result.metrics
|
|
716
|
+
usage = {
|
|
717
|
+
"prompt_tokens": getattr(metrics, "input_tokens", 0),
|
|
718
|
+
"completion_tokens": getattr(metrics, "output_tokens", 0),
|
|
719
|
+
"total_tokens": getattr(metrics, "total_tokens", 0),
|
|
720
|
+
}
|
|
721
|
+
print(f"📊 Token Usage:")
|
|
722
|
+
print(f" Input Tokens: {usage.get('prompt_tokens', 0)}")
|
|
723
|
+
print(f" Output Tokens: {usage.get('completion_tokens', 0)}")
|
|
724
|
+
print(f" Total Tokens: {usage.get('total_tokens', 0)}\n")
|
|
725
|
+
|
|
726
|
+
print(f"📝 Response Preview:")
|
|
727
|
+
print(f" {response_content[:200]}..." if len(response_content) > 200 else f" {response_content}")
|
|
728
|
+
|
|
729
|
+
# CRITICAL: Persist COMPLETE session history to Control Plane API
|
|
730
|
+
# This includes previous history + current turn for conversation continuity
|
|
731
|
+
# IMPORTANT: Use retry logic - persistence failures shouldn't break execution
|
|
732
|
+
print("\n💾 Persisting session history to Control Plane...")
|
|
733
|
+
|
|
734
|
+
# Build complete session: previous history + current turn's messages
|
|
735
|
+
updated_session_messages = list(session_history) # Start with loaded history
|
|
736
|
+
|
|
737
|
+
# Track proper timestamps for chronological ordering
|
|
738
|
+
# User message happened at execution start, assistant response completed now
|
|
739
|
+
from datetime import timedelta
|
|
740
|
+
response_completed_at = datetime.now(timezone.utc)
|
|
741
|
+
# User message must be BEFORE assistant message (subtract execution time)
|
|
742
|
+
# Use a conservative estimate: user message was at least 1 second before response
|
|
743
|
+
execution_started_at = response_completed_at - timedelta(seconds=1)
|
|
744
|
+
|
|
745
|
+
# Add current turn messages (user prompt + assistant response)
|
|
746
|
+
# CRITICAL: User message must have EARLIER timestamp than assistant for proper ordering
|
|
747
|
+
# CRITICAL: Use message_id from workflow signal if available (for deduplication consistency)
|
|
748
|
+
user_message_id = getattr(input, "user_message_id", None) or f"{input.execution_id}_user_{int(execution_started_at.timestamp() * 1000000)}"
|
|
749
|
+
|
|
750
|
+
current_turn_messages = [
|
|
751
|
+
{
|
|
752
|
+
"role": "user",
|
|
753
|
+
"content": input.prompt,
|
|
754
|
+
"timestamp": execution_started_at.isoformat(), # When user sent message
|
|
755
|
+
"user_id": input.user_id,
|
|
756
|
+
"user_name": getattr(input, "user_name", None),
|
|
757
|
+
"user_email": getattr(input, "user_email", None),
|
|
758
|
+
"user_avatar": getattr(input, "user_avatar", None),
|
|
759
|
+
"message_id": user_message_id, # Use provided message_id or generate stable one
|
|
760
|
+
},
|
|
761
|
+
{
|
|
762
|
+
"role": "assistant",
|
|
763
|
+
"content": response_content,
|
|
764
|
+
"timestamp": response_completed_at.isoformat(), # When assistant finished responding
|
|
765
|
+
"message_id": message_id, # Use the message_id generated at start of execution
|
|
766
|
+
}
|
|
767
|
+
]
|
|
768
|
+
|
|
769
|
+
print(f" 📝 Adding {len(current_turn_messages)} messages from current turn (user + assistant)...")
|
|
770
|
+
updated_session_messages.extend(current_turn_messages)
|
|
771
|
+
|
|
772
|
+
# CRITICAL: Also add tool messages to session history for proper rendering
|
|
773
|
+
if tool_messages:
|
|
774
|
+
print(f" 🔧 Adding {len(tool_messages)} tool messages to session history...")
|
|
775
|
+
updated_session_messages.extend(tool_messages)
|
|
776
|
+
|
|
777
|
+
if updated_session_messages:
|
|
778
|
+
# Try up to 3 times to persist session history
|
|
779
|
+
max_retries = 3
|
|
780
|
+
persisted = False
|
|
781
|
+
|
|
782
|
+
for attempt in range(max_retries):
|
|
783
|
+
try:
|
|
784
|
+
if attempt > 0:
|
|
785
|
+
print(f" 🔄 Retry persistence attempt {attempt + 1}/{max_retries}...")
|
|
786
|
+
|
|
787
|
+
success = control_plane.persist_session(
|
|
788
|
+
execution_id=input.execution_id,
|
|
789
|
+
session_id=input.session_id or input.execution_id,
|
|
790
|
+
user_id=input.user_id,
|
|
791
|
+
messages=updated_session_messages, # Complete conversation history
|
|
792
|
+
metadata={
|
|
793
|
+
"agent_id": input.agent_id,
|
|
794
|
+
"organization_id": input.organization_id,
|
|
795
|
+
"turn_count": len(updated_session_messages),
|
|
796
|
+
}
|
|
797
|
+
)
|
|
798
|
+
|
|
799
|
+
if success:
|
|
800
|
+
print(f" ✅ Complete session history persisted ({len(updated_session_messages)} total messages)")
|
|
801
|
+
persisted = True
|
|
802
|
+
break
|
|
803
|
+
else:
|
|
804
|
+
print(f" ⚠️ Persistence failed (attempt {attempt + 1}/{max_retries})")
|
|
805
|
+
if attempt < max_retries - 1:
|
|
806
|
+
import time
|
|
807
|
+
time.sleep(2 ** attempt) # Exponential backoff
|
|
808
|
+
|
|
809
|
+
except Exception as session_error:
|
|
810
|
+
error_type = type(session_error).__name__
|
|
811
|
+
print(f" ⚠️ Persistence error ({error_type}, attempt {attempt + 1}/{max_retries})")
|
|
812
|
+
logger.warning(
|
|
813
|
+
"session_persistence_error",
|
|
814
|
+
extra={
|
|
815
|
+
"error": str(session_error),
|
|
816
|
+
"error_type": error_type,
|
|
817
|
+
"execution_id": input.execution_id,
|
|
818
|
+
"attempt": attempt + 1
|
|
819
|
+
}
|
|
820
|
+
)
|
|
821
|
+
if attempt < max_retries - 1:
|
|
822
|
+
import time
|
|
823
|
+
time.sleep(2 ** attempt) # Exponential backoff
|
|
824
|
+
|
|
825
|
+
if not persisted:
|
|
826
|
+
print(f" ⚠️ Session persistence failed after {max_retries} attempts")
|
|
827
|
+
logger.error(
|
|
828
|
+
"session_persistence_failed_all_retries",
|
|
829
|
+
extra={
|
|
830
|
+
"execution_id": input.execution_id,
|
|
831
|
+
"message_count": len(updated_session_messages)
|
|
832
|
+
}
|
|
833
|
+
)
|
|
834
|
+
# Don't fail execution - session loss is better than execution failure
|
|
835
|
+
else:
|
|
836
|
+
print(" ℹ️ No messages - skipping session persistence")
|
|
837
|
+
|
|
838
|
+
print("\n" + "="*80)
|
|
839
|
+
print("🏁 AGENT EXECUTION END")
|
|
840
|
+
print("="*80 + "\n")
|
|
841
|
+
|
|
842
|
+
# Cleanup: Remove agent from registry
|
|
843
|
+
if input.execution_id in _active_agents:
|
|
844
|
+
del _active_agents[input.execution_id]
|
|
845
|
+
print(f"✅ Agent unregistered (execution_id: {input.execution_id})\n")
|
|
846
|
+
|
|
847
|
+
return {
|
|
848
|
+
"success": True,
|
|
849
|
+
"response": response_content,
|
|
850
|
+
"usage": usage,
|
|
851
|
+
"model": model,
|
|
852
|
+
"finish_reason": "stop",
|
|
853
|
+
"mcp_tools_used": 0, # TODO: Track MCP tool usage
|
|
854
|
+
"tool_messages": tool_messages, # Include tool call messages for UI
|
|
855
|
+
"tool_execution_messages": tool_execution_messages, # Include real-time tool execution status
|
|
856
|
+
}
|
|
857
|
+
|
|
858
|
+
except Exception as e:
|
|
859
|
+
# Cleanup on error
|
|
860
|
+
if input.execution_id in _active_agents:
|
|
861
|
+
del _active_agents[input.execution_id]
|
|
862
|
+
|
|
863
|
+
# Ensure error message is never empty
|
|
864
|
+
error_msg = str(e) or repr(e) or f"{type(e).__name__}: No error details available"
|
|
865
|
+
|
|
866
|
+
print("\n" + "="*80)
|
|
867
|
+
print("❌ AGENT EXECUTION FAILED")
|
|
868
|
+
print("="*80)
|
|
869
|
+
print(f"Error: {error_msg}")
|
|
870
|
+
print(f"Error Type: {type(e).__name__}")
|
|
871
|
+
print("="*80 + "\n")
|
|
872
|
+
|
|
873
|
+
activity.logger.error(
|
|
874
|
+
f"Agent LLM call failed",
|
|
875
|
+
extra={
|
|
876
|
+
"execution_id": input.execution_id,
|
|
877
|
+
"error": error_msg,
|
|
878
|
+
"error_type": type(e).__name__,
|
|
879
|
+
},
|
|
880
|
+
exc_info=True,
|
|
881
|
+
)
|
|
882
|
+
|
|
883
|
+
# Raise ApplicationError so Temporal marks the workflow as FAILED
|
|
884
|
+
raise ApplicationError(
|
|
885
|
+
f"Agent execution failed: {error_msg}",
|
|
886
|
+
non_retryable=False, # Allow retries per retry policy
|
|
887
|
+
type=type(e).__name__
|
|
888
|
+
)
|
|
889
|
+
|
|
890
|
+
|
|
891
|
+
@activity.defn
|
|
892
|
+
async def update_execution_status(input: ActivityUpdateExecutionInput) -> dict:
|
|
893
|
+
"""
|
|
894
|
+
Update execution status in database via Control Plane API.
|
|
895
|
+
|
|
896
|
+
This activity calls the Control Plane API to update execution records.
|
|
897
|
+
Also records which worker processed this execution.
|
|
898
|
+
|
|
899
|
+
Args:
|
|
900
|
+
input: Activity input with update details
|
|
901
|
+
|
|
902
|
+
Returns:
|
|
903
|
+
Dict with success flag
|
|
904
|
+
"""
|
|
905
|
+
execution_logger.activity_started(
|
|
906
|
+
"Update Status",
|
|
907
|
+
input.execution_id,
|
|
908
|
+
details={"status": input.status}
|
|
909
|
+
)
|
|
910
|
+
|
|
911
|
+
try:
|
|
912
|
+
# Get Control Plane URL and Kubiya API key from environment
|
|
913
|
+
control_plane_url = os.getenv("CONTROL_PLANE_URL")
|
|
914
|
+
kubiya_api_key = os.getenv("KUBIYA_API_KEY")
|
|
915
|
+
worker_id = os.getenv("WORKER_ID", "unknown")
|
|
916
|
+
|
|
917
|
+
if not control_plane_url:
|
|
918
|
+
raise ValueError("CONTROL_PLANE_URL environment variable not set")
|
|
919
|
+
if not kubiya_api_key:
|
|
920
|
+
raise ValueError("KUBIYA_API_KEY environment variable not set")
|
|
921
|
+
|
|
922
|
+
# Collect worker system information
|
|
923
|
+
import socket
|
|
924
|
+
import platform
|
|
925
|
+
worker_info = {
|
|
926
|
+
"worker_id": worker_id,
|
|
927
|
+
"hostname": socket.gethostname(),
|
|
928
|
+
"platform": platform.platform(),
|
|
929
|
+
"python_version": platform.python_version(),
|
|
930
|
+
}
|
|
931
|
+
|
|
932
|
+
# Build update payload
|
|
933
|
+
update_payload = {}
|
|
934
|
+
|
|
935
|
+
if input.status:
|
|
936
|
+
update_payload["status"] = input.status
|
|
937
|
+
|
|
938
|
+
if input.started_at:
|
|
939
|
+
update_payload["started_at"] = input.started_at
|
|
940
|
+
|
|
941
|
+
if input.completed_at:
|
|
942
|
+
update_payload["completed_at"] = input.completed_at
|
|
943
|
+
|
|
944
|
+
if input.response is not None:
|
|
945
|
+
update_payload["response"] = input.response
|
|
946
|
+
|
|
947
|
+
if input.error_message is not None:
|
|
948
|
+
update_payload["error_message"] = input.error_message
|
|
949
|
+
|
|
950
|
+
if input.usage:
|
|
951
|
+
update_payload["usage"] = input.usage
|
|
952
|
+
|
|
953
|
+
# Merge worker info into execution_metadata
|
|
954
|
+
execution_metadata = input.execution_metadata or {}
|
|
955
|
+
if not execution_metadata.get("worker_info"):
|
|
956
|
+
execution_metadata["worker_info"] = worker_info
|
|
957
|
+
update_payload["execution_metadata"] = execution_metadata
|
|
958
|
+
|
|
959
|
+
# Call Control Plane API
|
|
960
|
+
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
961
|
+
response = await client.patch(
|
|
962
|
+
f"{control_plane_url}/api/v1/executions/{input.execution_id}",
|
|
963
|
+
json=update_payload,
|
|
964
|
+
headers={
|
|
965
|
+
"Authorization": f"Bearer {kubiya_api_key}",
|
|
966
|
+
"Content-Type": "application/json",
|
|
967
|
+
}
|
|
968
|
+
)
|
|
969
|
+
|
|
970
|
+
if response.status_code == 404:
|
|
971
|
+
# Execution not found - this is not a retryable error
|
|
972
|
+
# The execution may have been deleted, or never existed
|
|
973
|
+
execution_logger.activity_failed(
|
|
974
|
+
"Update Status",
|
|
975
|
+
input.execution_id,
|
|
976
|
+
"Execution record not found - may have been deleted or cancelled",
|
|
977
|
+
will_retry=False
|
|
978
|
+
)
|
|
979
|
+
|
|
980
|
+
# Raise ApplicationError with non_retryable=True since execution doesn't exist
|
|
981
|
+
raise ApplicationError(
|
|
982
|
+
f"Execution {input.execution_id} not found - may have been deleted",
|
|
983
|
+
non_retryable=True,
|
|
984
|
+
type="ExecutionNotFound"
|
|
985
|
+
)
|
|
986
|
+
elif response.status_code != 200:
|
|
987
|
+
raise Exception(f"Failed to update execution: {response.status_code} - {response.text}")
|
|
988
|
+
|
|
989
|
+
# IMPORTANT: Log status changes to RUNNING prominently
|
|
990
|
+
if input.status == "running":
|
|
991
|
+
logger.info(
|
|
992
|
+
"execution_status_updated_to_running",
|
|
993
|
+
execution_id=input.execution_id[:8] if input.execution_id else "unknown",
|
|
994
|
+
status=input.status,
|
|
995
|
+
worker_id=worker_id
|
|
996
|
+
)
|
|
997
|
+
|
|
998
|
+
execution_logger.activity_completed(
|
|
999
|
+
"Update Status",
|
|
1000
|
+
input.execution_id,
|
|
1001
|
+
result=f"Status: {input.status}"
|
|
1002
|
+
)
|
|
1003
|
+
|
|
1004
|
+
activity.logger.info(
|
|
1005
|
+
f"Execution status updated via API",
|
|
1006
|
+
extra={
|
|
1007
|
+
"execution_id": input.execution_id,
|
|
1008
|
+
"status": input.status,
|
|
1009
|
+
}
|
|
1010
|
+
)
|
|
1011
|
+
|
|
1012
|
+
return {"success": True, "execution_not_found": False}
|
|
1013
|
+
|
|
1014
|
+
except Exception as e:
|
|
1015
|
+
print(f"❌ Failed to update status: {str(e)}\n")
|
|
1016
|
+
|
|
1017
|
+
activity.logger.error(
|
|
1018
|
+
f"Failed to update execution status",
|
|
1019
|
+
extra={
|
|
1020
|
+
"execution_id": input.execution_id,
|
|
1021
|
+
"error": str(e),
|
|
1022
|
+
}
|
|
1023
|
+
)
|
|
1024
|
+
raise
|
|
1025
|
+
|
|
1026
|
+
|
|
1027
|
+
@activity.defn
|
|
1028
|
+
async def get_execution_details(input: ActivityGetExecutionInput) -> dict:
|
|
1029
|
+
"""
|
|
1030
|
+
Get execution details from Control Plane API.
|
|
1031
|
+
|
|
1032
|
+
This activity fetches the current execution state including status and metadata.
|
|
1033
|
+
|
|
1034
|
+
Args:
|
|
1035
|
+
input: Activity input with execution_id
|
|
1036
|
+
|
|
1037
|
+
Returns:
|
|
1038
|
+
Dict with execution details including status
|
|
1039
|
+
"""
|
|
1040
|
+
execution_logger.activity_started(
|
|
1041
|
+
"Get Execution Details",
|
|
1042
|
+
input.execution_id
|
|
1043
|
+
)
|
|
1044
|
+
|
|
1045
|
+
try:
|
|
1046
|
+
# Get Control Plane URL and Kubiya API key from environment
|
|
1047
|
+
control_plane_url = os.getenv("CONTROL_PLANE_URL")
|
|
1048
|
+
kubiya_api_key = os.getenv("KUBIYA_API_KEY")
|
|
1049
|
+
|
|
1050
|
+
if not control_plane_url:
|
|
1051
|
+
raise ValueError("CONTROL_PLANE_URL environment variable not set")
|
|
1052
|
+
if not kubiya_api_key:
|
|
1053
|
+
raise ValueError("KUBIYA_API_KEY environment variable not set")
|
|
1054
|
+
|
|
1055
|
+
# Call Control Plane API to get execution
|
|
1056
|
+
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
1057
|
+
response = await client.get(
|
|
1058
|
+
f"{control_plane_url}/api/v1/executions/{input.execution_id}",
|
|
1059
|
+
headers={
|
|
1060
|
+
"Authorization": f"Bearer {kubiya_api_key}",
|
|
1061
|
+
"Content-Type": "application/json",
|
|
1062
|
+
}
|
|
1063
|
+
)
|
|
1064
|
+
|
|
1065
|
+
if response.status_code == 404:
|
|
1066
|
+
execution_logger.activity_failed(
|
|
1067
|
+
"Get Execution Details",
|
|
1068
|
+
input.execution_id,
|
|
1069
|
+
"Execution record not found",
|
|
1070
|
+
will_retry=False
|
|
1071
|
+
)
|
|
1072
|
+
|
|
1073
|
+
raise ApplicationError(
|
|
1074
|
+
f"Execution {input.execution_id} not found",
|
|
1075
|
+
non_retryable=True,
|
|
1076
|
+
type="ExecutionNotFound"
|
|
1077
|
+
)
|
|
1078
|
+
elif response.status_code != 200:
|
|
1079
|
+
raise Exception(f"Failed to get execution: {response.status_code} - {response.text}")
|
|
1080
|
+
|
|
1081
|
+
execution_data = response.json()
|
|
1082
|
+
|
|
1083
|
+
execution_logger.activity_completed(
|
|
1084
|
+
"Get Execution Details",
|
|
1085
|
+
input.execution_id,
|
|
1086
|
+
result=f"Status: {execution_data.get('status', 'unknown')}"
|
|
1087
|
+
)
|
|
1088
|
+
|
|
1089
|
+
activity.logger.info(
|
|
1090
|
+
f"Execution details fetched via API",
|
|
1091
|
+
extra={
|
|
1092
|
+
"execution_id": input.execution_id,
|
|
1093
|
+
"status": execution_data.get("status"),
|
|
1094
|
+
}
|
|
1095
|
+
)
|
|
1096
|
+
|
|
1097
|
+
return execution_data
|
|
1098
|
+
|
|
1099
|
+
except Exception as e:
|
|
1100
|
+
activity.logger.error(
|
|
1101
|
+
f"Failed to get execution details",
|
|
1102
|
+
extra={
|
|
1103
|
+
"execution_id": input.execution_id,
|
|
1104
|
+
"error": str(e),
|
|
1105
|
+
}
|
|
1106
|
+
)
|
|
1107
|
+
raise
|
|
1108
|
+
|
|
1109
|
+
|
|
1110
|
+
@activity.defn
|
|
1111
|
+
async def update_agent_status(input: ActivityUpdateAgentInput) -> dict:
|
|
1112
|
+
"""
|
|
1113
|
+
Update agent status in database via Control Plane API.
|
|
1114
|
+
|
|
1115
|
+
This activity calls the Control Plane API to update agent records.
|
|
1116
|
+
|
|
1117
|
+
Args:
|
|
1118
|
+
input: Activity input with update details
|
|
1119
|
+
|
|
1120
|
+
Returns:
|
|
1121
|
+
Dict with success flag
|
|
1122
|
+
"""
|
|
1123
|
+
activity.logger.info(
|
|
1124
|
+
f"Updating agent status via Control Plane API",
|
|
1125
|
+
extra={
|
|
1126
|
+
"agent_id": input.agent_id,
|
|
1127
|
+
"status": input.status,
|
|
1128
|
+
}
|
|
1129
|
+
)
|
|
1130
|
+
|
|
1131
|
+
try:
|
|
1132
|
+
# Get Control Plane URL and Kubiya API key from environment
|
|
1133
|
+
control_plane_url = os.getenv("CONTROL_PLANE_URL")
|
|
1134
|
+
kubiya_api_key = os.getenv("KUBIYA_API_KEY")
|
|
1135
|
+
|
|
1136
|
+
if not control_plane_url:
|
|
1137
|
+
raise ValueError("CONTROL_PLANE_URL environment variable not set")
|
|
1138
|
+
if not kubiya_api_key:
|
|
1139
|
+
raise ValueError("KUBIYA_API_KEY environment variable not set")
|
|
1140
|
+
|
|
1141
|
+
# Build update payload
|
|
1142
|
+
update_payload = {
|
|
1143
|
+
"status": input.status,
|
|
1144
|
+
"last_active_at": input.last_active_at,
|
|
1145
|
+
}
|
|
1146
|
+
|
|
1147
|
+
if input.error_message is not None:
|
|
1148
|
+
update_payload["error_message"] = input.error_message
|
|
1149
|
+
|
|
1150
|
+
if input.state:
|
|
1151
|
+
update_payload["state"] = input.state
|
|
1152
|
+
|
|
1153
|
+
# Call Control Plane API
|
|
1154
|
+
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
1155
|
+
response = await client.patch(
|
|
1156
|
+
f"{control_plane_url}/api/v1/agents/{input.agent_id}",
|
|
1157
|
+
json=update_payload,
|
|
1158
|
+
headers={
|
|
1159
|
+
"Authorization": f"Bearer {kubiya_api_key}",
|
|
1160
|
+
"Content-Type": "application/json",
|
|
1161
|
+
}
|
|
1162
|
+
)
|
|
1163
|
+
|
|
1164
|
+
# For team executions, the "agent_id" is actually a team_id, so it won't be found in agents table
|
|
1165
|
+
# This is expected and not an error - just log and return success
|
|
1166
|
+
if response.status_code == 404:
|
|
1167
|
+
activity.logger.info(
|
|
1168
|
+
f"Agent not found (likely a team execution) - skipping agent status update",
|
|
1169
|
+
extra={
|
|
1170
|
+
"agent_id": input.agent_id,
|
|
1171
|
+
"status": input.status,
|
|
1172
|
+
}
|
|
1173
|
+
)
|
|
1174
|
+
return {"success": True, "skipped": True}
|
|
1175
|
+
elif response.status_code != 200:
|
|
1176
|
+
raise Exception(f"Failed to update agent: {response.status_code} - {response.text}")
|
|
1177
|
+
|
|
1178
|
+
activity.logger.info(
|
|
1179
|
+
f"Agent status updated via API",
|
|
1180
|
+
extra={
|
|
1181
|
+
"agent_id": input.agent_id,
|
|
1182
|
+
"status": input.status,
|
|
1183
|
+
}
|
|
1184
|
+
)
|
|
1185
|
+
|
|
1186
|
+
return {"success": True}
|
|
1187
|
+
|
|
1188
|
+
except Exception as e:
|
|
1189
|
+
activity.logger.error(
|
|
1190
|
+
f"Failed to update agent status",
|
|
1191
|
+
extra={
|
|
1192
|
+
"agent_id": input.agent_id,
|
|
1193
|
+
"error": str(e),
|
|
1194
|
+
}
|
|
1195
|
+
)
|
|
1196
|
+
raise
|
|
1197
|
+
|
|
1198
|
+
|
|
1199
|
+
@dataclass
|
|
1200
|
+
class ActivityCancelAgentInput:
|
|
1201
|
+
execution_id: str
|
|
1202
|
+
|
|
1203
|
+
|
|
1204
|
+
@activity.defn(name="cancel_agent_run")
|
|
1205
|
+
async def cancel_agent_run(input: ActivityCancelAgentInput) -> dict:
|
|
1206
|
+
"""
|
|
1207
|
+
Cancel an active agent/team run using Agno's cancel_run API.
|
|
1208
|
+
|
|
1209
|
+
This is called when a user clicks the STOP button in the UI.
|
|
1210
|
+
Uses the global registry to find the Agent instance and run_id,
|
|
1211
|
+
then calls agent.cancel_run(run_id) to stop execution.
|
|
1212
|
+
|
|
1213
|
+
Args:
|
|
1214
|
+
input: Contains execution_id to identify which run to cancel
|
|
1215
|
+
|
|
1216
|
+
Returns:
|
|
1217
|
+
Dict with success status and cancellation details
|
|
1218
|
+
"""
|
|
1219
|
+
print("\n" + "="*80)
|
|
1220
|
+
print("🛑 CANCEL AGENT RUN")
|
|
1221
|
+
print("="*80)
|
|
1222
|
+
print(f"Execution ID: {input.execution_id}\n")
|
|
1223
|
+
|
|
1224
|
+
try:
|
|
1225
|
+
# Look up agent in registry
|
|
1226
|
+
if input.execution_id not in _active_agents:
|
|
1227
|
+
print(f"⚠️ Agent not found in registry - may have already completed")
|
|
1228
|
+
activity.logger.warning(
|
|
1229
|
+
"cancel_agent_not_found",
|
|
1230
|
+
extra={"execution_id": input.execution_id}
|
|
1231
|
+
)
|
|
1232
|
+
return {
|
|
1233
|
+
"success": False,
|
|
1234
|
+
"error": "Agent not found or already completed",
|
|
1235
|
+
"execution_id": input.execution_id,
|
|
1236
|
+
}
|
|
1237
|
+
|
|
1238
|
+
agent_info = _active_agents[input.execution_id]
|
|
1239
|
+
agent = agent_info["agent"]
|
|
1240
|
+
run_id = agent_info.get("run_id")
|
|
1241
|
+
|
|
1242
|
+
if not run_id:
|
|
1243
|
+
print(f"⚠️ No run_id found - execution may not have started yet")
|
|
1244
|
+
return {
|
|
1245
|
+
"success": False,
|
|
1246
|
+
"error": "Execution not started yet",
|
|
1247
|
+
"execution_id": input.execution_id,
|
|
1248
|
+
}
|
|
1249
|
+
|
|
1250
|
+
print(f"🆔 Found run_id: {run_id}")
|
|
1251
|
+
print(f"🛑 Calling agent.cancel_run()...")
|
|
1252
|
+
|
|
1253
|
+
# Call Agno's cancel_run API
|
|
1254
|
+
success = agent.cancel_run(run_id)
|
|
1255
|
+
|
|
1256
|
+
if success:
|
|
1257
|
+
print(f"✅ Agent run cancelled successfully!\n")
|
|
1258
|
+
activity.logger.info(
|
|
1259
|
+
"agent_run_cancelled",
|
|
1260
|
+
extra={
|
|
1261
|
+
"execution_id": input.execution_id,
|
|
1262
|
+
"run_id": run_id,
|
|
1263
|
+
}
|
|
1264
|
+
)
|
|
1265
|
+
|
|
1266
|
+
# Clean up registry
|
|
1267
|
+
del _active_agents[input.execution_id]
|
|
1268
|
+
|
|
1269
|
+
return {
|
|
1270
|
+
"success": True,
|
|
1271
|
+
"execution_id": input.execution_id,
|
|
1272
|
+
"run_id": run_id,
|
|
1273
|
+
"cancelled_at": datetime.now(timezone.utc).isoformat(),
|
|
1274
|
+
}
|
|
1275
|
+
else:
|
|
1276
|
+
print(f"⚠️ Cancel failed - run may have already completed\n")
|
|
1277
|
+
return {
|
|
1278
|
+
"success": False,
|
|
1279
|
+
"error": "Cancel failed - run may be completed",
|
|
1280
|
+
"execution_id": input.execution_id,
|
|
1281
|
+
"run_id": run_id,
|
|
1282
|
+
}
|
|
1283
|
+
|
|
1284
|
+
except Exception as e:
|
|
1285
|
+
print(f"❌ Error cancelling run: {str(e)}\n")
|
|
1286
|
+
activity.logger.error(
|
|
1287
|
+
"cancel_agent_error",
|
|
1288
|
+
extra={
|
|
1289
|
+
"execution_id": input.execution_id,
|
|
1290
|
+
"error": str(e),
|
|
1291
|
+
}
|
|
1292
|
+
)
|
|
1293
|
+
return {
|
|
1294
|
+
"success": False,
|
|
1295
|
+
"error": str(e),
|
|
1296
|
+
"execution_id": input.execution_id,
|
|
1297
|
+
}
|
|
1298
|
+
|
|
1299
|
+
|
|
1300
|
+
@dataclass
|
|
1301
|
+
class ActivityPersistConversationInput:
|
|
1302
|
+
"""Input for persisting conversation history"""
|
|
1303
|
+
execution_id: str
|
|
1304
|
+
session_id: str
|
|
1305
|
+
messages: List[Dict[str, Any]]
|
|
1306
|
+
user_id: Optional[str] = None
|
|
1307
|
+
metadata: Optional[Dict[str, Any]] = None
|
|
1308
|
+
|
|
1309
|
+
|
|
1310
|
+
@activity.defn(name="persist_conversation_history")
|
|
1311
|
+
async def persist_conversation_history(input: ActivityPersistConversationInput) -> dict:
|
|
1312
|
+
"""
|
|
1313
|
+
Persist conversation history to Control Plane after each turn.
|
|
1314
|
+
|
|
1315
|
+
This ensures conversation state is saved end-to-end, making it available:
|
|
1316
|
+
- For future turns in the same conversation
|
|
1317
|
+
- For UI display and history views
|
|
1318
|
+
- For analytics and debugging
|
|
1319
|
+
- Even if the worker crashes or restarts
|
|
1320
|
+
|
|
1321
|
+
The Control Plane stores this in the database and caches it in Redis
|
|
1322
|
+
for fast retrieval on subsequent turns.
|
|
1323
|
+
|
|
1324
|
+
Args:
|
|
1325
|
+
input: Contains execution_id, session_id, messages, and optional metadata
|
|
1326
|
+
|
|
1327
|
+
Returns:
|
|
1328
|
+
Dict with success status and persistence details
|
|
1329
|
+
"""
|
|
1330
|
+
execution_id_short = input.execution_id[:8] if input.execution_id else "unknown"
|
|
1331
|
+
|
|
1332
|
+
activity.logger.info(
|
|
1333
|
+
"persisting_conversation",
|
|
1334
|
+
extra={
|
|
1335
|
+
"execution_id": execution_id_short,
|
|
1336
|
+
"session_id": input.session_id[:8] if input.session_id else "none",
|
|
1337
|
+
"message_count": len(input.messages),
|
|
1338
|
+
}
|
|
1339
|
+
)
|
|
1340
|
+
|
|
1341
|
+
try:
|
|
1342
|
+
# Get Control Plane client
|
|
1343
|
+
control_plane = get_control_plane_client()
|
|
1344
|
+
|
|
1345
|
+
# CRITICAL: Deduplicate messages by message_id before persisting
|
|
1346
|
+
# This prevents duplicate messages from being stored in the database
|
|
1347
|
+
seen_message_ids = set()
|
|
1348
|
+
deduplicated_messages = []
|
|
1349
|
+
|
|
1350
|
+
for msg in input.messages:
|
|
1351
|
+
# CRITICAL: For tool messages, use tool_execution_id as unique identifier
|
|
1352
|
+
# This prevents duplicate tool calls from being persisted
|
|
1353
|
+
if msg.get("role") == "tool" and msg.get("tool_execution_id"):
|
|
1354
|
+
msg_id = msg["tool_execution_id"]
|
|
1355
|
+
activity.logger.debug(
|
|
1356
|
+
"using_tool_execution_id_for_deduplication",
|
|
1357
|
+
extra={
|
|
1358
|
+
"execution_id": execution_id_short,
|
|
1359
|
+
"tool_execution_id": msg_id,
|
|
1360
|
+
"tool_name": msg.get("tool_name")
|
|
1361
|
+
}
|
|
1362
|
+
)
|
|
1363
|
+
else:
|
|
1364
|
+
msg_id = msg.get("message_id")
|
|
1365
|
+
|
|
1366
|
+
if not msg_id:
|
|
1367
|
+
# Generate stable message_id for messages without one
|
|
1368
|
+
# Use timestamp + role to create consistent ID
|
|
1369
|
+
timestamp_str = msg.get("timestamp", "")
|
|
1370
|
+
role = msg.get("role", "unknown")
|
|
1371
|
+
msg_id = f"{input.execution_id}_{role}_{timestamp_str}"
|
|
1372
|
+
msg["message_id"] = msg_id
|
|
1373
|
+
activity.logger.debug(
|
|
1374
|
+
"generated_missing_message_id",
|
|
1375
|
+
extra={
|
|
1376
|
+
"execution_id": execution_id_short,
|
|
1377
|
+
"role": role,
|
|
1378
|
+
"generated_id": msg_id
|
|
1379
|
+
}
|
|
1380
|
+
)
|
|
1381
|
+
|
|
1382
|
+
# Check for duplicates
|
|
1383
|
+
if msg_id in seen_message_ids:
|
|
1384
|
+
activity.logger.warning(
|
|
1385
|
+
"skipping_duplicate_message_id",
|
|
1386
|
+
extra={
|
|
1387
|
+
"execution_id": execution_id_short,
|
|
1388
|
+
"message_id": msg_id,
|
|
1389
|
+
"role": msg.get("role"),
|
|
1390
|
+
"content_preview": msg.get("content", "")[:50]
|
|
1391
|
+
}
|
|
1392
|
+
)
|
|
1393
|
+
continue
|
|
1394
|
+
|
|
1395
|
+
seen_message_ids.add(msg_id)
|
|
1396
|
+
deduplicated_messages.append(msg)
|
|
1397
|
+
|
|
1398
|
+
activity.logger.info(
|
|
1399
|
+
"deduplication_complete",
|
|
1400
|
+
extra={
|
|
1401
|
+
"execution_id": execution_id_short,
|
|
1402
|
+
"before": len(input.messages),
|
|
1403
|
+
"after": len(deduplicated_messages),
|
|
1404
|
+
"removed": len(input.messages) - len(deduplicated_messages)
|
|
1405
|
+
}
|
|
1406
|
+
)
|
|
1407
|
+
|
|
1408
|
+
# Persist conversation via Control Plane API with deduplicated messages
|
|
1409
|
+
success = control_plane.persist_session(
|
|
1410
|
+
execution_id=input.execution_id,
|
|
1411
|
+
session_id=input.session_id or input.execution_id,
|
|
1412
|
+
user_id=input.user_id,
|
|
1413
|
+
messages=deduplicated_messages, # Use deduplicated messages
|
|
1414
|
+
metadata=input.metadata or {}
|
|
1415
|
+
)
|
|
1416
|
+
|
|
1417
|
+
if success:
|
|
1418
|
+
activity.logger.info(
|
|
1419
|
+
"conversation_persisted",
|
|
1420
|
+
extra={
|
|
1421
|
+
"execution_id": execution_id_short,
|
|
1422
|
+
"message_count": len(input.messages),
|
|
1423
|
+
}
|
|
1424
|
+
)
|
|
1425
|
+
return {
|
|
1426
|
+
"success": True,
|
|
1427
|
+
"execution_id": input.execution_id,
|
|
1428
|
+
"message_count": len(input.messages),
|
|
1429
|
+
"persisted_at": datetime.now(timezone.utc).isoformat(),
|
|
1430
|
+
}
|
|
1431
|
+
else:
|
|
1432
|
+
activity.logger.warning(
|
|
1433
|
+
"conversation_persistence_failed",
|
|
1434
|
+
extra={
|
|
1435
|
+
"execution_id": execution_id_short,
|
|
1436
|
+
}
|
|
1437
|
+
)
|
|
1438
|
+
return {
|
|
1439
|
+
"success": False,
|
|
1440
|
+
"error": "Control Plane API returned failure",
|
|
1441
|
+
"execution_id": input.execution_id,
|
|
1442
|
+
}
|
|
1443
|
+
|
|
1444
|
+
except Exception as e:
|
|
1445
|
+
error_type = type(e).__name__
|
|
1446
|
+
error_msg = str(e) if str(e) else "No error message provided"
|
|
1447
|
+
|
|
1448
|
+
activity.logger.error(
|
|
1449
|
+
"conversation_persistence_error",
|
|
1450
|
+
extra={
|
|
1451
|
+
"execution_id": execution_id_short,
|
|
1452
|
+
"error_type": error_type,
|
|
1453
|
+
"error": error_msg[:500], # Truncate very long errors
|
|
1454
|
+
"message_count": len(input.messages),
|
|
1455
|
+
},
|
|
1456
|
+
exc_info=True,
|
|
1457
|
+
)
|
|
1458
|
+
return {
|
|
1459
|
+
"success": False,
|
|
1460
|
+
"error": f"{error_type}: {error_msg}",
|
|
1461
|
+
"error_type": error_type,
|
|
1462
|
+
"execution_id": input.execution_id,
|
|
1463
|
+
}
|
|
1464
|
+
|
|
1465
|
+
|
|
1466
|
+
@dataclass
|
|
1467
|
+
class AnalyticsActivityInput:
|
|
1468
|
+
"""Input for analytics submission activity"""
|
|
1469
|
+
execution_id: str
|
|
1470
|
+
turn_number: int
|
|
1471
|
+
result: Dict[str, Any] # RuntimeExecutionResult as dict
|
|
1472
|
+
turn_start_time: float
|
|
1473
|
+
|
|
1474
|
+
|
|
1475
|
+
@activity.defn
|
|
1476
|
+
async def submit_runtime_analytics_activity(input: AnalyticsActivityInput) -> dict:
|
|
1477
|
+
"""
|
|
1478
|
+
Temporal activity for submitting runtime analytics.
|
|
1479
|
+
|
|
1480
|
+
Runs independently from main execution flow, with its own timeout and retry logic.
|
|
1481
|
+
Failures are logged but do not affect execution success.
|
|
1482
|
+
"""
|
|
1483
|
+
from control_plane_api.worker.services.runtime_analytics import submit_runtime_analytics
|
|
1484
|
+
from control_plane_api.worker.services.analytics_service import AnalyticsService
|
|
1485
|
+
from control_plane_api.worker.runtimes.base import RuntimeExecutionResult
|
|
1486
|
+
import time
|
|
1487
|
+
|
|
1488
|
+
execution_id_short = input.execution_id[:8]
|
|
1489
|
+
|
|
1490
|
+
activity.logger.info(
|
|
1491
|
+
"analytics_activity_started",
|
|
1492
|
+
extra={
|
|
1493
|
+
"execution_id": execution_id_short,
|
|
1494
|
+
"turn_number": input.turn_number,
|
|
1495
|
+
"tokens": input.result.get("usage", {}).get("total_tokens", 0),
|
|
1496
|
+
}
|
|
1497
|
+
)
|
|
1498
|
+
|
|
1499
|
+
try:
|
|
1500
|
+
# Initialize analytics service with required parameters
|
|
1501
|
+
control_plane_url = os.environ.get("CONTROL_PLANE_URL", "http://localhost:8000")
|
|
1502
|
+
api_key = os.environ.get("KUBIYA_API_KEY", "")
|
|
1503
|
+
|
|
1504
|
+
if not api_key:
|
|
1505
|
+
raise ValueError("KUBIYA_API_KEY environment variable not set")
|
|
1506
|
+
|
|
1507
|
+
analytics_service = AnalyticsService(
|
|
1508
|
+
control_plane_url=control_plane_url,
|
|
1509
|
+
api_key=api_key,
|
|
1510
|
+
)
|
|
1511
|
+
|
|
1512
|
+
# Convert dict result to RuntimeExecutionResult
|
|
1513
|
+
# The workflow passes a dict, but submit_runtime_analytics expects RuntimeExecutionResult
|
|
1514
|
+
result_obj = RuntimeExecutionResult(
|
|
1515
|
+
response=input.result.get("response", ""),
|
|
1516
|
+
usage=input.result.get("usage", {}),
|
|
1517
|
+
success=input.result.get("success", True),
|
|
1518
|
+
finish_reason=input.result.get("finish_reason", "stop"),
|
|
1519
|
+
tool_execution_messages=input.result.get("tool_messages", []),
|
|
1520
|
+
tool_messages=input.result.get("tool_messages", []),
|
|
1521
|
+
model=input.result.get("model", "unknown"),
|
|
1522
|
+
metadata=input.result.get("metadata", {}),
|
|
1523
|
+
error=input.result.get("error"),
|
|
1524
|
+
)
|
|
1525
|
+
|
|
1526
|
+
# Submit analytics
|
|
1527
|
+
try:
|
|
1528
|
+
await submit_runtime_analytics(
|
|
1529
|
+
result=result_obj,
|
|
1530
|
+
execution_id=input.execution_id,
|
|
1531
|
+
turn_number=input.turn_number,
|
|
1532
|
+
turn_start_time=input.turn_start_time,
|
|
1533
|
+
analytics_service=analytics_service,
|
|
1534
|
+
turn_end_time=time.time(),
|
|
1535
|
+
)
|
|
1536
|
+
|
|
1537
|
+
activity.logger.info(
|
|
1538
|
+
"analytics_activity_completed",
|
|
1539
|
+
extra={
|
|
1540
|
+
"execution_id": execution_id_short,
|
|
1541
|
+
"turn_number": input.turn_number,
|
|
1542
|
+
}
|
|
1543
|
+
)
|
|
1544
|
+
|
|
1545
|
+
return {
|
|
1546
|
+
"success": True,
|
|
1547
|
+
"execution_id": input.execution_id,
|
|
1548
|
+
"turn_number": input.turn_number,
|
|
1549
|
+
}
|
|
1550
|
+
|
|
1551
|
+
finally:
|
|
1552
|
+
# Cleanup HTTP client
|
|
1553
|
+
try:
|
|
1554
|
+
await analytics_service.aclose()
|
|
1555
|
+
except Exception as cleanup_error:
|
|
1556
|
+
activity.logger.warning(
|
|
1557
|
+
"analytics_service_cleanup_error",
|
|
1558
|
+
extra={
|
|
1559
|
+
"execution_id": execution_id_short,
|
|
1560
|
+
"error": str(cleanup_error),
|
|
1561
|
+
}
|
|
1562
|
+
)
|
|
1563
|
+
|
|
1564
|
+
except Exception as e:
|
|
1565
|
+
# Log but return success - analytics errors are non-critical
|
|
1566
|
+
error_msg = str(e)[:500]
|
|
1567
|
+
activity.logger.warning(
|
|
1568
|
+
"analytics_activity_failed",
|
|
1569
|
+
extra={
|
|
1570
|
+
"execution_id": execution_id_short,
|
|
1571
|
+
"turn_number": input.turn_number,
|
|
1572
|
+
"error": error_msg,
|
|
1573
|
+
"error_type": type(e).__name__,
|
|
1574
|
+
"note": "Analytics failed but execution continues normally"
|
|
1575
|
+
}
|
|
1576
|
+
)
|
|
1577
|
+
|
|
1578
|
+
return {
|
|
1579
|
+
"success": False,
|
|
1580
|
+
"execution_id": input.execution_id,
|
|
1581
|
+
"turn_number": input.turn_number,
|
|
1582
|
+
"error": error_msg,
|
|
1583
|
+
"error_type": type(e).__name__,
|
|
1584
|
+
}
|
|
1585
|
+
|