diff --git a/DESIGN_SPEC.md b/DESIGN_SPEC.md index 65b388c3b7..4e6f860143 100644 --- a/DESIGN_SPEC.md +++ b/DESIGN_SPEC.md @@ -2149,45 +2149,81 @@ The human can interact as: Templates are YAML/JSON files defining a complete company setup: ```yaml -# templates/startup.yaml +# templates/startup.yaml (simplified — real templates also declare +# variables, departments, min_agents/max_agents, and tags) template: name: "Tech Startup" description: "Small team for building MVPs and prototypes" version: "1.0" company: - name: "{{ company_name }}" type: "startup" budget_monthly: "{{ budget | default(50.00) }}" - autonomy: "semi" + autonomy: 0.5 agents: - - role: "ceo" + - role: "CEO" name: "{{ ceo_name | auto }}" model: "large" personality_preset: "visionary_leader" - - role: "full_stack_developer" + - role: "Full-Stack Developer" + merge_id: "fullstack-senior" name: "{{ dev1_name | auto }}" level: "senior" model: "medium" personality_preset: "pragmatic_builder" - - role: "full_stack_developer" + - role: "Full-Stack Developer" + merge_id: "fullstack-mid" name: "{{ dev2_name | auto }}" level: "mid" model: "small" personality_preset: "eager_learner" - - role: "product_manager" + - role: "Product Manager" name: "{{ pm_name | auto }}" model: "medium" personality_preset: "strategic_planner" workflow: "agile_kanban" communication: "hybrid" + + workflow_handoffs: + - from_department: "engineering" + to_department: "qa" + trigger: "pr_ready" + + escalation_paths: + - from_department: "engineering" + to_department: "security" + condition: "vulnerability_found" ``` +**Template Inheritance** — Templates can extend other templates using `extends`: + +```yaml +template: + name: "Extended Startup" + extends: "startup" # inherits all agents, departments, config + agents: + - role: "QA Engineer" # appended to parent agents + level: "mid" + - role: "Full-Stack Developer" + merge_id: "fullstack-mid" + department: "engineering" + _remove: true # removes matching parent agent by key +``` + +Inheritance resolves parent→child chains up to 10 levels deep. Merge semantics: +- **Scalars** (`company_name`, `company_type`): child wins if present. +- **`config`** dict: deep-merged (child keys override parent). +- **`agents`** list: merged by `(role, department, merge_id)` key. When `merge_id` is omitted, it defaults to an empty string, making the key `(role, department, "")`. Child can override, append, or remove (`_remove: true`) parent agents. +- **`departments`** list: merged by name (case-insensitive). Child dept replaces parent entirely. +- **`workflow_handoffs`**, **`escalation_paths`**: child replaces entirely if present. + +Circular inheritance is detected via chain tracking and raises `TemplateInheritanceError`. + ### 14.2 Company Builder Interactive CLI/web wizard for creating custom companies: @@ -2334,6 +2370,7 @@ ai-company/ │ │ ├── metrics.py # TaskCompletionMetrics proxy overhead model │ │ ├── react_loop.py # ReAct loop implementation │ │ ├── plan_models.py # Plan step, plan, and plan-execute config models +│ │ ├── plan_parsing.py # Plan response parsing utilities │ │ ├── plan_execute_loop.py # Plan-and-Execute loop implementation │ │ ├── loop_helpers.py # Shared stateless helpers for all loop implementations │ │ ├── recovery.py # Crash recovery strategies (RecoveryStrategy protocol) @@ -2371,6 +2408,7 @@ ai-company/ │ │ ├── events/ # Per-domain event constants │ │ │ ├── __init__.py # Package marker with usage docs; no re-exports │ │ │ ├── budget.py # BUDGET_* constants +│ │ │ ├── company.py # COMPANY_* constants │ │ │ ├── communication.py # COMM_* constants │ │ │ ├── config.py # CONFIG_* constants │ │ │ ├── correlation.py # CORRELATION_* constants @@ -2473,6 +2511,7 @@ ai-company/ │ ├── schema.py # Template schema models │ ├── loader.py # Template loader │ ├── renderer.py # Template renderer +│ ├── merge.py # Template config merging for inheritance │ ├── presets.py # Personality presets + auto-name generation │ ├── errors.py # Template errors │ └── builtins/ # Pre-built company templates @@ -2531,6 +2570,8 @@ These conventions were established during the M0–M2+ review cycle. **Adopted** | **State coordination** | Planned (M4) | Centralized single-writer: `TaskEngine` owns all task/project mutations via `asyncio.Queue`. Agents submit requests, engine applies `model_copy(update=...)` sequentially and publishes snapshots. `version: int` field on state models for future optimistic concurrency if multi-process scaling is needed. | Prevents lost updates by design. Trivial in single-threaded asyncio (no locks). Perfect audit trail. Industry consensus: MetaGPT, CrewAI, AutoGen all use prevention-by-design, not conflict resolution. See §6.8 State Coordination table. | | **Workspace isolation** | Planned (M4) | Pluggable `WorkspaceIsolationStrategy` protocol. Default: planner + git worktrees. Each agent works in an isolated worktree; sequential merge on completion. Textual conflicts detected by git; semantic conflicts reviewed by agent or human. | Industry standard (Codex, Cursor, Claude Code, VS Code). Maximum parallelism. Leverages mature git infrastructure. See §6.8. | | **Graceful shutdown** | Adopted (M3) | Pluggable `ShutdownStrategy` protocol. Default: cooperative with 30s timeout. Agents check shutdown event at turn boundaries. Force-cancel after timeout. `INTERRUPTED` status for force-cancelled tasks. M4/M5: upgrade to checkpoint-and-stop. | Cross-platform (Windows `signal.signal()` fallback). Bounded shutdown time. Mirrors cooperative shutdown in §6.7. | +| **Template inheritance** | Adopted (M2.5) | `extends` field on `CompanyTemplate` triggers parent resolution at render time. `merge.py` merges configs by field type: scalars (child wins), config dicts (deep merge), agents (by `(role, department)` key with `_remove` support), departments (by name). `_ParentEntry` dataclass tracks merge state. `DEFAULT_MERGE_DEPARTMENT = "engineering"` shared between merge and renderer. Circular chains detected via `frozenset` tracking; max depth = 10. | Enables template composition without copy-paste. Merge-by-key preserves parent order. `_remove` directive enables clean agent removal without workarounds. | +| **Pydantic alias for YAML directives** | Adopted (M2.5) | `Field(alias="_remove")` in `TemplateAgentConfig` — YAML uses `_remove: true`, Python accesses `agent.remove`. Keeps the YAML-facing name (underscore prefix signals internal directive) separate from the Python attribute name. | Underscore-prefixed YAML keys signal merge directives vs regular fields. Pydantic alias bridges the naming convention gap cleanly. | | **Communication foundation** | Adopted (M4) | `MessageBus` protocol with `InMemoryMessageBus` backend (asyncio queues, pull-model `receive()` with shutdown signaling via `asyncio.Event`). `MessageDispatcher` routes to concurrent handlers via `asyncio.TaskGroup` with pre-allocated error collection. `AgentMessenger` per-agent facade auto-fills sender/timestamp/ID; deterministic direct-channel naming `@{sorted_a}:{sorted_b}`. `DeliveryEnvelope` for delivery tracking. `NotBlankStr` validation on all protocol boundary identifiers. | Pull-model avoids callback complexity and enables agents to consume at their own pace. Protocol + backend split enables future persistent/distributed bus implementations. Deterministic DM channel names prevent duplicates. See §5. | --- diff --git a/src/ai_company/observability/events/template.py b/src/ai_company/observability/events/template.py index e0bd50bc16..3d240edfcf 100644 --- a/src/ai_company/observability/events/template.py +++ b/src/ai_company/observability/events/template.py @@ -15,3 +15,15 @@ TEMPLATE_RENDER_VALIDATION_ERROR: Final[str] = "template.render.validation_error" TEMPLATE_PERSONALITY_PRESET_UNKNOWN: Final[str] = "template.personality_preset.unknown" TEMPLATE_PASS1_FLOAT_FALLBACK: Final[str] = "template.pass1.float_fallback" +TEMPLATE_INHERIT_RESOLVE_START: Final[str] = "template.inherit.resolve_start" +TEMPLATE_INHERIT_RESOLVE_SUCCESS: Final[str] = "template.inherit.resolve_success" +TEMPLATE_INHERIT_CIRCULAR: Final[str] = "template.inherit.circular" +TEMPLATE_INHERIT_DEPTH_EXCEEDED: Final[str] = "template.inherit.depth_exceeded" +TEMPLATE_INHERIT_MERGE: Final[str] = "template.inherit.merge" +TEMPLATE_INHERIT_MERGE_ERROR: Final[str] = "template.inherit.merge_error" +TEMPLATE_RENDER_TYPE_ERROR: Final[str] = "template.render.type_error" +TEMPLATE_LOAD_NOT_FOUND: Final[str] = "template.load.not_found" +TEMPLATE_LOAD_READ_ERROR: Final[str] = "template.load.read_error" +TEMPLATE_LOAD_PARSE_ERROR: Final[str] = "template.load.parse_error" +TEMPLATE_LOAD_STRUCTURE_ERROR: Final[str] = "template.load.structure_error" +TEMPLATE_LOAD_INVALID_NAME: Final[str] = "template.load.invalid_name" diff --git a/src/ai_company/templates/__init__.py b/src/ai_company/templates/__init__.py index f47b6e9ab6..f2564ef4e1 100644 --- a/src/ai_company/templates/__init__.py +++ b/src/ai_company/templates/__init__.py @@ -16,6 +16,7 @@ TemplateAgentConfig TemplateDepartmentConfig TemplateError + TemplateInheritanceError TemplateNotFoundError TemplateRenderError TemplateValidationError @@ -23,6 +24,7 @@ from ai_company.templates.errors import ( TemplateError, + TemplateInheritanceError, TemplateNotFoundError, TemplateRenderError, TemplateValidationError, @@ -51,6 +53,7 @@ "TemplateDepartmentConfig", "TemplateError", "TemplateInfo", + "TemplateInheritanceError", "TemplateMetadata", "TemplateNotFoundError", "TemplateRenderError", diff --git a/src/ai_company/templates/builtins/agency.yaml b/src/ai_company/templates/builtins/agency.yaml index 2c2ef240e7..d33fa91006 100644 --- a/src/ai_company/templates/builtins/agency.yaml +++ b/src/ai_company/templates/builtins/agency.yaml @@ -2,6 +2,8 @@ template: name: "Agency" description: "Client-focused agency with project management and creative roles" version: "1.0.0" + min_agents: 10 + max_agents: 15 tags: - "agency" - "client-work" @@ -23,30 +25,84 @@ template: departments: - name: "operations" - budget_percent: 30 + budget_percent: 20 head_role: "Project Manager" - name: "engineering" - budget_percent: 40 + budget_percent: 35 head_role: "Full-Stack Developer" - name: "design" - budget_percent: 30 - head_role: "UI Designer" + budget_percent: 20 + head_role: "UX Designer" + - name: "quality_assurance" + budget_percent: 10 + head_role: "QA Engineer" + - name: "creative_marketing" + budget_percent: 15 + head_role: "Brand Strategist" agents: - role: "Project Manager" level: "senior" model: "medium" - personality_preset: "visionary_leader" + personality_preset: "process_optimizer" department: "operations" + - role: "Product Manager" + level: "mid" + model: "medium" + personality_preset: "strategic_planner" + department: "operations" + - role: "Scrum Master" + level: "senior" + model: "medium" + personality_preset: "process_optimizer" + department: "operations" + - role: "UX Designer" + level: "mid" + model: "medium" + personality_preset: "user_advocate" + department: "design" - role: "UI Designer" level: "mid" model: "medium" + personality_preset: "creative_innovator" department: "design" + - role: "Frontend Developer" + level: "mid" + model: "small" + personality_preset: "creative_innovator" + department: "engineering" + - role: "Backend Developer" + level: "senior" + model: "medium" + personality_preset: "pragmatic_builder" + department: "engineering" - role: "Full-Stack Developer" + merge_id: "fullstack-senior" level: "senior" model: "medium" personality_preset: "pragmatic_builder" department: "engineering" + - role: "QA Engineer" + level: "mid" + model: "small" + personality_preset: "quality_guardian" + department: "quality_assurance" + - role: "Content Writer" + level: "mid" + model: "small" + personality_preset: "technical_communicator" + department: "creative_marketing" + - role: "Brand Strategist" + level: "senior" + model: "medium" + personality_preset: "growth_hacker" + department: "creative_marketing" + - role: "Full-Stack Developer" + merge_id: "fullstack-mid" + level: "mid" + model: "small" + personality_preset: "eager_learner" + department: "engineering" workflow: "kanban" communication: "hybrid" diff --git a/src/ai_company/templates/builtins/dev_shop.yaml b/src/ai_company/templates/builtins/dev_shop.yaml index ff90b7b366..0b90087e75 100644 --- a/src/ai_company/templates/builtins/dev_shop.yaml +++ b/src/ai_company/templates/builtins/dev_shop.yaml @@ -2,6 +2,8 @@ template: name: "Dev Shop" description: "Software development focused team with QA and DevOps" version: "1.0.0" + min_agents: 6 + max_agents: 10 tags: - "development" - "engineering" @@ -36,28 +38,46 @@ template: - role: "Software Architect" level: "principal" model: "large" - personality_preset: "methodical_analyst" + personality_preset: "systems_thinker" department: "engineering" - role: "Backend Developer" + merge_id: "backend-senior-1" level: "senior" model: "medium" personality_preset: "pragmatic_builder" department: "engineering" - role: "Backend Developer" + merge_id: "backend-mid" level: "mid" model: "small" personality_preset: "eager_learner" department: "engineering" + - role: "Frontend Developer" + level: "mid" + model: "small" + personality_preset: "creative_innovator" + department: "engineering" - role: "QA Lead" level: "lead" model: "medium" + personality_preset: "quality_guardian" + department: "quality_assurance" + - role: "QA Engineer" + level: "mid" + model: "small" personality_preset: "methodical_analyst" department: "quality_assurance" - role: "DevOps/SRE Engineer" level: "mid" model: "medium" - personality_preset: "pragmatic_builder" + personality_preset: "process_optimizer" department: "operations" + - role: "Backend Developer" + merge_id: "backend-senior-2" + level: "senior" + model: "medium" + personality_preset: "disciplined_executor" + department: "engineering" workflow: "agile_kanban" communication: "hybrid" diff --git a/src/ai_company/templates/builtins/full_company.yaml b/src/ai_company/templates/builtins/full_company.yaml index bff871ead0..8cfcc234df 100644 --- a/src/ai_company/templates/builtins/full_company.yaml +++ b/src/ai_company/templates/builtins/full_company.yaml @@ -2,6 +2,8 @@ template: name: "Full Company" description: "Enterprise simulation with all departments and full hierarchy" version: "1.0.0" + min_agents: 1 + max_agents: 50 tags: - "enterprise" - "full-hierarchy" @@ -15,6 +17,14 @@ template: description: "Monthly budget in USD" var_type: "float" default: 200.0 + - name: "num_backend_devs" + description: "Number of backend developers" + var_type: "int" + default: 3 + - name: "num_frontend_devs" + description: "Number of frontend developers" + var_type: "int" + default: 2 company: type: "full_company" @@ -23,22 +33,35 @@ template: departments: - name: "executive" - budget_percent: 15 + budget_percent: 10 head_role: "CEO" - name: "engineering" - budget_percent: 50 + budget_percent: 30 head_role: "CTO" - name: "product" - budget_percent: 15 + budget_percent: 10 head_role: "Product Manager" + - name: "design" + budget_percent: 8 + head_role: "UX Designer" - name: "quality_assurance" budget_percent: 10 - head_role: "QA Engineer" + head_role: "QA Lead" + - name: "data_analytics" + budget_percent: 8 + head_role: "Data Analyst" - name: "operations" - budget_percent: 10 - head_role: "CFO" + budget_percent: 8 + head_role: "COO" + - name: "security" + budget_percent: 8 + head_role: "Security Engineer" + - name: "creative_marketing" + budget_percent: 8 + head_role: "Brand Strategist" agents: + # C-Suite (5) - role: "CEO" level: "c_suite" model: "large" @@ -47,22 +70,151 @@ template: - role: "CTO" level: "c_suite" model: "large" - personality_preset: "methodical_analyst" + personality_preset: "systems_thinker" department: "executive" - role: "CFO" level: "c_suite" model: "large" + personality_preset: "data_driven_optimizer" + department: "executive" + - role: "COO" + level: "c_suite" + model: "large" + personality_preset: "process_optimizer" department: "operations" + - role: "CPO" + level: "c_suite" + model: "large" + personality_preset: "user_advocate" + department: "product" + # Product (2) + - role: "Product Manager" + level: "senior" + model: "medium" + personality_preset: "strategic_planner" + department: "product" + - role: "Technical Writer" + level: "mid" + model: "small" + personality_preset: "technical_communicator" + department: "product" + # Design (2) + - role: "UX Designer" + level: "mid" + model: "medium" + personality_preset: "user_advocate" + department: "design" + - role: "UX Researcher" + level: "mid" + model: "medium" + personality_preset: "user_advocate" + department: "design" + # Engineering — static (3) + - role: "Software Architect" + level: "principal" + model: "large" + personality_preset: "systems_thinker" + department: "engineering" + - role: "DevOps/SRE Engineer" + level: "mid" + model: "medium" + personality_preset: "process_optimizer" + department: "engineering" + - role: "Database Engineer" + level: "mid" + model: "medium" + personality_preset: "systems_thinker" + department: "engineering" + # Engineering — dynamic backend devs +{% for i in range(num_backend_devs | default(3) | int) %} - role: "Backend Developer" + merge_id: "backend-{{ i }}" +{% if i == 0 %} level: "senior" model: "medium" personality_preset: "pragmatic_builder" +{% else %} + level: "mid" + model: "small" + personality_preset: "eager_learner" +{% endif %} + department: "engineering" +{% endfor %} + # Engineering — dynamic frontend devs +{% for i in range(num_frontend_devs | default(2) | int) %} + - role: "Frontend Developer" + merge_id: "frontend-{{ i }}" +{% if i == 0 %} + level: "senior" + model: "medium" + personality_preset: "creative_innovator" +{% else %} + level: "mid" + model: "small" + personality_preset: "eager_learner" +{% endif %} department: "engineering" +{% endfor %} + # Quality Assurance (3) + - role: "QA Lead" + level: "lead" + model: "medium" + personality_preset: "quality_guardian" + department: "quality_assurance" - role: "QA Engineer" level: "mid" model: "small" personality_preset: "methodical_analyst" department: "quality_assurance" + - role: "Automation Engineer" + level: "mid" + model: "small" + personality_preset: "disciplined_executor" + department: "quality_assurance" + # Data & Analytics (2) + - role: "Data Analyst" + level: "mid" + model: "medium" + personality_preset: "data_driven_optimizer" + department: "data_analytics" + - role: "Data Engineer" + level: "mid" + model: "medium" + personality_preset: "pragmatic_builder" + department: "data_analytics" + # Operations (2) + - role: "Project Manager" + level: "senior" + model: "medium" + personality_preset: "process_optimizer" + department: "operations" + - role: "Scrum Master" + level: "senior" + model: "medium" + personality_preset: "process_optimizer" + department: "operations" + # Security (2) + - role: "Security Engineer" + level: "senior" + model: "medium" + personality_preset: "security_sentinel" + department: "security" + - role: "Security Operations" + level: "senior" + model: "medium" + personality_preset: "security_sentinel" + department: "security" + # Creative & Marketing (2) + - role: "Content Writer" + level: "mid" + model: "small" + personality_preset: "technical_communicator" + department: "creative_marketing" + - role: "Brand Strategist" + level: "senior" + model: "medium" + personality_preset: "growth_hacker" + department: "creative_marketing" workflow: "agile_kanban" communication: "hybrid" diff --git a/src/ai_company/templates/builtins/product_team.yaml b/src/ai_company/templates/builtins/product_team.yaml index fc87c197c8..9dddc8899a 100644 --- a/src/ai_company/templates/builtins/product_team.yaml +++ b/src/ai_company/templates/builtins/product_team.yaml @@ -2,6 +2,8 @@ template: name: "Product Team" description: "Product-focused development with design and QA" version: "1.0.0" + min_agents: 8 + max_agents: 12 tags: - "product" - "design" @@ -23,25 +25,54 @@ template: departments: - name: "product" - budget_percent: 30 + budget_percent: 25 head_role: "Product Manager" - name: "engineering" - budget_percent: 50 + budget_percent: 35 head_role: "Backend Developer" + - name: "design" + budget_percent: 15 + head_role: "UX Designer" - name: "quality_assurance" - budget_percent: 20 + budget_percent: 15 head_role: "QA Engineer" + - name: "data_analytics" + budget_percent: 10 + head_role: "Data Analyst" agents: - role: "Product Manager" level: "senior" model: "medium" - personality_preset: "visionary_leader" + personality_preset: "strategic_planner" department: "product" - role: "UX Designer" level: "mid" model: "medium" - department: "product" + personality_preset: "user_advocate" + department: "design" + - role: "UX Researcher" + level: "mid" + model: "medium" + personality_preset: "user_advocate" + department: "design" + - role: "Backend Developer" + merge_id: "backend-senior" + level: "senior" + model: "medium" + personality_preset: "pragmatic_builder" + department: "engineering" + - role: "Backend Developer" + merge_id: "backend-mid" + level: "mid" + model: "small" + personality_preset: "eager_learner" + department: "engineering" + - role: "Frontend Developer" + level: "mid" + model: "small" + personality_preset: "creative_innovator" + department: "engineering" - role: "Full-Stack Developer" level: "senior" model: "medium" @@ -50,8 +81,18 @@ template: - role: "QA Engineer" level: "mid" model: "small" - personality_preset: "methodical_analyst" + personality_preset: "quality_guardian" department: "quality_assurance" + - role: "Automation Engineer" + level: "mid" + model: "small" + personality_preset: "disciplined_executor" + department: "quality_assurance" + - role: "Data Analyst" + level: "mid" + model: "medium" + personality_preset: "data_driven_optimizer" + department: "data_analytics" workflow: "agile_kanban" communication: "hybrid" diff --git a/src/ai_company/templates/builtins/research_lab.yaml b/src/ai_company/templates/builtins/research_lab.yaml index 6a0bb47b3f..c93ba72229 100644 --- a/src/ai_company/templates/builtins/research_lab.yaml +++ b/src/ai_company/templates/builtins/research_lab.yaml @@ -2,6 +2,8 @@ template: name: "Research Lab" description: "Research and analysis focused team" version: "1.0.0" + min_agents: 5 + max_agents: 10 tags: - "research" - "analysis" @@ -23,17 +25,25 @@ template: departments: - name: "engineering" - budget_percent: 40 + budget_percent: 30 head_role: "Software Architect" - name: "data_analytics" - budget_percent: 60 + budget_percent: 50 head_role: "Data Analyst" + - name: "product" + budget_percent: 20 + head_role: "Technical Writer" agents: - role: "Software Architect" level: "principal" model: "large" - personality_preset: "methodical_analyst" + personality_preset: "systems_thinker" + department: "engineering" + - role: "Backend Developer" + level: "mid" + model: "medium" + personality_preset: "pragmatic_builder" department: "engineering" - role: "Data Engineer" level: "senior" @@ -41,10 +51,27 @@ template: personality_preset: "pragmatic_builder" department: "data_analytics" - role: "Data Analyst" + merge_id: "analyst-primary" level: "mid" model: "medium" + personality_preset: "data_driven_optimizer" + department: "data_analytics" + - role: "ML Engineer" + level: "senior" + model: "medium" + personality_preset: "independent_researcher" + department: "data_analytics" + - role: "Data Analyst" + merge_id: "analyst-secondary" + level: "mid" + model: "small" personality_preset: "methodical_analyst" department: "data_analytics" + - role: "Technical Writer" + level: "mid" + model: "small" + personality_preset: "technical_communicator" + department: "product" workflow: "kanban" communication: "hybrid" diff --git a/src/ai_company/templates/builtins/solo_founder.yaml b/src/ai_company/templates/builtins/solo_founder.yaml index d87c8f040e..5ce403188b 100644 --- a/src/ai_company/templates/builtins/solo_founder.yaml +++ b/src/ai_company/templates/builtins/solo_founder.yaml @@ -2,6 +2,8 @@ template: name: "Solo Founder" description: "Minimal setup for quick prototypes and solo projects" version: "1.0.0" + min_agents: 2 + max_agents: 3 tags: - "minimal" - "solo" diff --git a/src/ai_company/templates/builtins/startup.yaml b/src/ai_company/templates/builtins/startup.yaml index c60869906c..f8acb77ba8 100644 --- a/src/ai_company/templates/builtins/startup.yaml +++ b/src/ai_company/templates/builtins/startup.yaml @@ -2,6 +2,8 @@ template: name: "Tech Startup" description: "Small team for building MVPs and prototypes" version: "1.0.0" + min_agents: 3 + max_agents: 5 tags: - "startup" - "mvp" @@ -45,11 +47,13 @@ template: personality_preset: "methodical_analyst" department: "executive" - role: "Full-Stack Developer" + merge_id: "fullstack-senior" level: "senior" model: "medium" personality_preset: "pragmatic_builder" department: "engineering" - role: "Full-Stack Developer" + merge_id: "fullstack-mid" level: "mid" model: "small" personality_preset: "eager_learner" @@ -57,6 +61,7 @@ template: - role: "Product Manager" level: "senior" model: "medium" + personality_preset: "strategic_planner" department: "product" workflow: "agile_kanban" diff --git a/src/ai_company/templates/errors.py b/src/ai_company/templates/errors.py index 42e1873e66..9175eb4999 100644 --- a/src/ai_company/templates/errors.py +++ b/src/ai_company/templates/errors.py @@ -20,6 +20,14 @@ class TemplateRenderError(TemplateError): """ +class TemplateInheritanceError(TemplateRenderError): + """Raised when template inheritance fails. + + Covers circular inheritance chains, excessive depth, + and merge conflicts. + """ + + class TemplateValidationError(TemplateError): """Raised when a rendered template fails validation. diff --git a/src/ai_company/templates/loader.py b/src/ai_company/templates/loader.py index 2ee1baf5fa..efa458eed7 100644 --- a/src/ai_company/templates/loader.py +++ b/src/ai_company/templates/loader.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from importlib import resources from pathlib import Path +from types import MappingProxyType from typing import Any, Literal import yaml @@ -25,7 +26,12 @@ TEMPLATE_BUILTIN_DEFECT, TEMPLATE_LIST_SKIP_INVALID, TEMPLATE_LOAD_ERROR, + TEMPLATE_LOAD_INVALID_NAME, + TEMPLATE_LOAD_NOT_FOUND, + TEMPLATE_LOAD_PARSE_ERROR, + TEMPLATE_LOAD_READ_ERROR, TEMPLATE_LOAD_START, + TEMPLATE_LOAD_STRUCTURE_ERROR, TEMPLATE_LOAD_SUCCESS, TEMPLATE_PASS1_FLOAT_FALLBACK, ) @@ -40,16 +46,17 @@ _USER_TEMPLATES_DIR = Path.home() / ".ai-company" / "templates" -# Registry of built-in template names -> resource filenames. -BUILTIN_TEMPLATES: dict[str, str] = { - "solo_founder": "solo_founder.yaml", - "startup": "startup.yaml", - "dev_shop": "dev_shop.yaml", - "product_team": "product_team.yaml", - "agency": "agency.yaml", - "full_company": "full_company.yaml", - "research_lab": "research_lab.yaml", -} +BUILTIN_TEMPLATES: MappingProxyType[str, str] = MappingProxyType( + { + "solo_founder": "solo_founder.yaml", + "startup": "startup.yaml", + "dev_shop": "dev_shop.yaml", + "product_team": "product_team.yaml", + "agency": "agency.yaml", + "full_company": "full_company.yaml", + "research_lab": "research_lab.yaml", + } +) @dataclass(frozen=True) @@ -169,6 +176,7 @@ def load_template(name: str) -> LoadedTemplate: # Sanitize to prevent path traversal (OS-independent). if "/" in name_clean or "\\" in name_clean or ".." in name_clean: msg = f"Invalid template name {name!r}: must not contain path separators" + logger.warning(TEMPLATE_LOAD_INVALID_NAME, template_name=name) raise TemplateNotFoundError( msg, locations=(ConfigLocation(file_path=f""),), @@ -225,6 +233,7 @@ def load_template_file(path: Path | str) -> LoadedTemplate: path = Path(path) if not path.is_file(): msg = f"Template file not found: {path}" + logger.warning(TEMPLATE_LOAD_NOT_FOUND, path=str(path)) raise TemplateNotFoundError( msg, locations=(ConfigLocation(file_path=str(path)),), @@ -242,13 +251,22 @@ def _load_builtin(name: str) -> LoadedTemplate: filename = BUILTIN_TEMPLATES.get(name) if filename is None: msg = f"Unknown built-in template: {name!r}" + logger.warning(TEMPLATE_LOAD_NOT_FOUND, template_name=name) raise TemplateNotFoundError( msg, locations=(ConfigLocation(file_path=f""),), ) - ref = resources.files("ai_company.templates.builtins") / filename - yaml_text = ref.read_text(encoding="utf-8") source_name = f"" + try: + ref = resources.files("ai_company.templates.builtins") / filename + yaml_text = ref.read_text(encoding="utf-8") + except (OSError, ImportError, TypeError) as exc: + msg = f"Failed to read built-in template resource {filename!r}: {exc}" + logger.exception(TEMPLATE_LOAD_READ_ERROR, source=source_name, error=str(exc)) + raise TemplateRenderError( + msg, + locations=(ConfigLocation(file_path=source_name),), + ) from exc template = _parse_template_yaml(yaml_text, source_name=source_name) return LoadedTemplate( template=template, @@ -261,7 +279,8 @@ def _load_from_file(path: Path) -> LoadedTemplate: """Load a template from a file path. Raises: - TemplateRenderError: If the file cannot be read. + TemplateRenderError: If the file cannot be read or YAML + parsing fails. TemplateValidationError: If validation fails. """ source_name = str(path) @@ -269,12 +288,14 @@ def _load_from_file(path: Path) -> LoadedTemplate: yaml_text = path.read_text(encoding="utf-8") except OSError as exc: msg = f"Unable to read template file: {path}" + logger.warning(TEMPLATE_LOAD_READ_ERROR, path=str(path), error=str(exc)) raise TemplateRenderError( msg, locations=(ConfigLocation(file_path=source_name),), ) from exc except UnicodeDecodeError as exc: msg = f"Template file is not valid UTF-8: {path}" + logger.warning(TEMPLATE_LOAD_READ_ERROR, path=str(path), error=str(exc)) raise TemplateRenderError( msg, locations=(ConfigLocation(file_path=source_name),), @@ -334,6 +355,7 @@ def _parse_template_yaml( data = yaml.safe_load(safe_text) except yaml.YAMLError as exc: msg = f"Template YAML syntax error in {source_name}: {exc}" + logger.warning(TEMPLATE_LOAD_PARSE_ERROR, source=source_name, error=str(exc)) raise TemplateRenderError( msg, locations=(ConfigLocation(file_path=source_name),), @@ -343,14 +365,9 @@ def _parse_template_yaml( try: normalized = _normalize_template_data(template_data) return CompanyTemplate(**normalized) - except ValidationError as exc: - msg = f"Template validation failed for {source_name}: {exc}" - raise TemplateValidationError( - msg, - locations=(ConfigLocation(file_path=source_name),), - ) from exc - except (ValueError, TypeError) as exc: + except (ValidationError, ValueError, TypeError) as exc: msg = f"Template validation failed for {source_name}: {exc}" + logger.warning(TEMPLATE_LOAD_PARSE_ERROR, source=source_name, error=str(exc)) raise TemplateValidationError( msg, locations=(ConfigLocation(file_path=source_name),), @@ -368,6 +385,7 @@ def _validate_template_structure( """ if not isinstance(data, dict) or "template" not in data: msg = f"Template YAML must have a top-level 'template' key in {source_name}" + logger.warning(TEMPLATE_LOAD_STRUCTURE_ERROR, source=source_name, error=msg) raise TemplateValidationError( msg, locations=(ConfigLocation(file_path=source_name),), @@ -375,6 +393,7 @@ def _validate_template_structure( template_data = data["template"] if not isinstance(template_data, dict): msg = f"Template 'template' key must map to an object in {source_name}" + logger.warning(TEMPLATE_LOAD_STRUCTURE_ERROR, source=source_name, error=msg) raise TemplateValidationError( msg, locations=(ConfigLocation(file_path=source_name),), @@ -394,13 +413,17 @@ def _normalize_template_data(data: dict[str, Any]) -> dict[str, Any]: Returns: Dict suitable for ``CompanyTemplate(**result)``. """ - company_raw = data.get("company", {}) - if company_raw is None: - company_raw = {} - if not isinstance(company_raw, dict): + company = data.get("company") + if company is None: + company = {} + elif not isinstance(company, dict): msg = "Template field 'template.company' must be a mapping" + logger.warning( + TEMPLATE_LOAD_STRUCTURE_ERROR, + source="template.company", + error=msg, + ) raise TypeError(msg) - company: dict[str, Any] = company_raw metadata: dict[str, Any] = { "description": data.get("description", ""), @@ -410,8 +433,12 @@ def _normalize_template_data(data: dict[str, Any]) -> dict[str, Any]: } if "name" in data: metadata["name"] = data["name"] + if "min_agents" in data: + metadata["min_agents"] = data["min_agents"] + if "max_agents" in data: + metadata["max_agents"] = data["max_agents"] - return { + result: dict[str, Any] = { "metadata": metadata, "variables": data.get("variables", ()), "agents": data.get("agents", ()), @@ -420,7 +447,12 @@ def _normalize_template_data(data: dict[str, Any]) -> dict[str, Any]: "communication": data.get("communication", "hybrid"), "budget_monthly": _to_float(company.get("budget_monthly", 50.0)), "autonomy": _to_float(company.get("autonomy", 0.5)), + "workflow_handoffs": data.get("workflow_handoffs", ()), + "escalation_paths": data.get("escalation_paths", ()), } + if "extends" in data: + result["extends"] = data["extends"] + return result def _to_float(value: Any) -> float: diff --git a/src/ai_company/templates/merge.py b/src/ai_company/templates/merge.py new file mode 100644 index 0000000000..dfd994b637 --- /dev/null +++ b/src/ai_company/templates/merge.py @@ -0,0 +1,278 @@ +"""Template config merging for inheritance. + +Provides ``merge_template_configs`` which combines a parent config dict +with a child config dict, implementing the merge semantics described in +the template inheritance design. +""" + +import copy +from dataclasses import dataclass, field +from typing import Any + +from ai_company.config.utils import deep_merge +from ai_company.observability import get_logger +from ai_company.observability.events.template import ( + TEMPLATE_INHERIT_MERGE, + TEMPLATE_INHERIT_MERGE_ERROR, +) +from ai_company.templates.errors import TemplateInheritanceError + +logger = get_logger(__name__) + +# Single source of truth for the default department. +# renderer.py re-imports this value for its own use. +DEFAULT_MERGE_DEPARTMENT = "engineering" + + +@dataclass +class _ParentEntry: + """Tracking record for a parent agent during merge.""" + + index: int + agent: dict[str, Any] | None + matched: bool = field(default=False) + + +def merge_template_configs( + parent: dict[str, Any], + child: dict[str, Any], +) -> dict[str, Any]: + """Merge a parent config dict with a child config dict. + + Merge strategies by field: + + - ``company_name``, ``company_type``: child wins if present. + - ``config`` (dict): deep-merged; child keys override parent. + - ``agents`` (list): merged by ``(role, department, merge_id)`` key. + - ``departments`` (list): merged by ``name`` (case-insensitive). + - ``workflow_handoffs``, ``escalation_paths``: child replaces + entirely if present. + + Args: + parent: Rendered parent config dict (post-Jinja2, pre-defaults). + child: Rendered child config dict (post-Jinja2, pre-defaults). + + Returns: + New merged config dict. + """ + logger.debug(TEMPLATE_INHERIT_MERGE, action="start") + + result: dict[str, Any] = {} + + # Scalars: child wins if present. + for key in ("company_name", "company_type"): + if key in child and child[key] is not None: + result[key] = child[key] + elif key in parent: + result[key] = parent[key] + + # Config dict: deep merge. + parent_config = parent.get("config", {}) + child_config = child.get("config", {}) + if parent_config or child_config: + result["config"] = deep_merge( + parent_config if isinstance(parent_config, dict) else {}, + child_config if isinstance(child_config, dict) else {}, + ) + + # Agents: merge by (role, department) key. + parent_agents = parent.get("agents", []) + child_agents = child.get("agents", []) + if parent_agents or child_agents: + result["agents"] = _merge_agents(parent_agents, child_agents) + + # Departments: merge by name. + parent_depts = parent.get("departments", []) + child_depts = child.get("departments", []) + if parent_depts or child_depts: + result["departments"] = _merge_departments(parent_depts, child_depts) + + # Replace-if-present fields (deep-copied to prevent reference sharing). + for key in ("workflow_handoffs", "escalation_paths"): + if key in child and child[key] is not None: + result[key] = copy.deepcopy(child[key]) + elif key in parent: + result[key] = copy.deepcopy(parent[key]) + + logger.debug(TEMPLATE_INHERIT_MERGE, action="done") + return result + + +def _merge_agents( + parent_agents: list[dict[str, Any]], + child_agents: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Merge agent lists by ``(role, department, merge_id)`` key. + + Algorithm: + 1. Index parent agents by ``(role.lower(), department.lower(), + merge_id.lower())``. Duplicate keys maintain an ordered list + per key. + 2. Walk child agents: + - ``_remove: true``: find first unmatched parent with same key, + remove it. Child entry is discarded. + - Otherwise: match against first unmatched parent with same key, + replace. No match -> append. + 3. Discard ``_remove`` entries; strip ``_remove`` key from + replacement/appended dicts. + 4. Result: parent agents (with replacements/removals) + appended. + + Args: + parent_agents: Parent agent dicts, each expected to have at + least ``role`` and optionally ``department`` keys. + child_agents: Child agent dicts; may include ``_remove: True`` + to remove a matching parent agent. + + Returns: + Merged agent list. + + Raises: + TemplateInheritanceError: If ``_remove`` has no matching parent. + """ + parent_entries: dict[tuple[str, str, str], list[_ParentEntry]] = {} + for idx, agent in enumerate(parent_agents): + key = _agent_key(agent) + parent_entries.setdefault(key, []).append( + _ParentEntry(index=idx, agent=copy.deepcopy(agent)), + ) + + appended: list[dict[str, Any]] = [] + for child_agent in child_agents: + _apply_child_agent(child_agent, parent_entries, appended) + + return _collect_merged_agents(parent_entries, appended) + + +def _apply_child_agent( + child_agent: dict[str, Any], + parent_entries: dict[tuple[str, str, str], list[_ParentEntry]], + appended: list[dict[str, Any]], +) -> None: + """Apply a single child agent against parent entries. + + Updates *parent_entries* and *appended* as a local mutation + scoped to the enclosing ``_merge_agents`` call. + """ + key = _agent_key(child_agent) + is_remove = child_agent.get("_remove", False) + entries = parent_entries.get(key, []) + + matched_entry = _find_unmatched(entries) + clean = copy.deepcopy( + {k: v for k, v in child_agent.items() if k not in ("_remove", "merge_id")} + ) + + if is_remove: + if matched_entry is None: + msg = f"Cannot remove agent with key {key}: no matching parent agent found" + logger.error( + TEMPLATE_INHERIT_MERGE_ERROR, + action="remove_failed", + key=key, + ) + raise TemplateInheritanceError(msg) + matched_entry.matched = True + matched_entry.agent = None # mark for removal + elif matched_entry is not None: + matched_entry.matched = True + matched_entry.agent = clean + else: + appended.append(clean) + + +def _find_unmatched( + entries: list[_ParentEntry], +) -> _ParentEntry | None: + """Find first unmatched entry in a parent entries list.""" + return next((e for e in entries if not e.matched), None) + + +def _collect_merged_agents( + parent_entries: dict[tuple[str, str, str], list[_ParentEntry]], + appended: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Collect surviving parent agents (in order) + appended.""" + all_entries = sorted( + (entry for entries in parent_entries.values() for entry in entries), + key=lambda e: e.index, + ) + _strip_keys = {"merge_id", "_remove"} + result: list[dict[str, Any]] = [ + {k: v for k, v in entry.agent.items() if k not in _strip_keys} + for entry in all_entries + if entry.agent is not None + ] + result.extend(appended) + return result + + +def _merge_departments( + parent_depts: list[dict[str, Any]], + child_depts: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Merge department lists by name (case-insensitive). + + Child dept with matching name replaces parent entirely. + Unmatched child depts are appended. + + Args: + parent_depts: Parent department dicts. + child_depts: Child department dicts. + + Returns: + Merged department list. + """ + # Build child overrides index (skip nameless departments). + child_by_name: dict[str, dict[str, Any]] = {} + nameless_child: list[dict[str, Any]] = [] + for child_dept in child_depts: + name = str(child_dept.get("name", "")).lower() + if not name: + logger.warning( + TEMPLATE_INHERIT_MERGE, + action="department_no_name", + ) + nameless_child.append(child_dept) + continue + child_by_name[name] = copy.deepcopy(child_dept) + + # Walk parent depts: apply child override if it exists. + result: list[dict[str, Any]] = [] + seen_names: set[str] = set() + for dept in parent_depts: + name = str(dept.get("name", "")).lower() + if not name: + logger.warning( + TEMPLATE_INHERIT_MERGE, + action="department_no_name", + ) + result.append(copy.deepcopy(dept)) + continue + if name in child_by_name: + result.append(child_by_name[name]) + seen_names.add(name) + else: + result.append(copy.deepcopy(dept)) + seen_names.add(name) + + # Append unmatched child depts + nameless children. + for name, child_dept in child_by_name.items(): + if name not in seen_names: + result.append(child_dept) + result.extend(copy.deepcopy(nameless_child)) + + return result + + +def _agent_key(agent: dict[str, Any]) -> tuple[str, str, str]: + """Compute the merge key for an agent dict. + + Uses ``(role, department, merge_id)`` when ``merge_id`` is present, + otherwise ``(role, department, "")`` for backwards compatibility. + """ + role = str(agent.get("role", "")).lower() + dept = agent.get("department") + if not dept: + dept = DEFAULT_MERGE_DEPARTMENT + merge_id = str(agent.get("merge_id", "")).lower() + return (role, str(dept).lower(), merge_id) diff --git a/src/ai_company/templates/presets.py b/src/ai_company/templates/presets.py index cb1918a000..c2ee5e7c87 100644 --- a/src/ai_company/templates/presets.py +++ b/src/ai_company/templates/presets.py @@ -8,6 +8,8 @@ from types import MappingProxyType from typing import Any +from pydantic import ValidationError + from ai_company.core.agent import PersonalityConfig from ai_company.observability import get_logger from ai_company.observability.events.template import ( @@ -16,8 +18,7 @@ logger = get_logger(__name__) -# Preset name -> frozen dict compatible with PersonalityConfig constructor. -# Both the outer mapping and each inner mapping are read-only. +# Mutable construction helper; frozen into PERSONALITY_PRESETS below. _RAW_PRESETS: dict[str, dict[str, Any]] = { "visionary_leader": { "traits": ("strategic", "decisive", "inspiring"), @@ -259,7 +260,88 @@ "verbosity": "balanced", "conflict_approach": "compromise", }, + "user_advocate": { + "traits": ("empathetic", "user-focused", "observant"), + "communication_style": "warm", + "risk_tolerance": "medium", + "creativity": "medium", + "description": "A user-focused advocate who champions end-user needs.", + "openness": 0.7, + "conscientiousness": 0.65, + "extraversion": 0.6, + "agreeableness": 0.85, + "stress_response": 0.6, + "decision_making": "consultative", + "collaboration": "team", + "verbosity": "balanced", + "conflict_approach": "collaborate", + }, + "process_optimizer": { + "traits": ("systematic", "efficiency-driven", "organized"), + "communication_style": "structured", + "risk_tolerance": "low", + "creativity": "medium", + "description": "A systematic optimizer who streamlines processes.", + "openness": 0.45, + "conscientiousness": 0.9, + "extraversion": 0.5, + "agreeableness": 0.55, + "stress_response": 0.75, + "decision_making": "directive", + "collaboration": "team", + "verbosity": "balanced", + "conflict_approach": "compromise", + }, + "growth_hacker": { + "traits": ("experimental", "data-informed", "ambitious"), + "communication_style": "enthusiastic", + "risk_tolerance": "high", + "creativity": "high", + "description": "An experimental growth hacker who drives rapid expansion.", + "openness": 0.85, + "conscientiousness": 0.5, + "extraversion": 0.75, + "agreeableness": 0.45, + "stress_response": 0.5, + "decision_making": "intuitive", + "collaboration": "pair", + "verbosity": "terse", + "conflict_approach": "compete", + }, + "technical_communicator": { + "traits": ("clear", "structured", "precise"), + "communication_style": "formal", + "risk_tolerance": "low", + "creativity": "medium", + "description": "A clear communicator who makes complex topics accessible.", + "openness": 0.55, + "conscientiousness": 0.85, + "extraversion": 0.4, + "agreeableness": 0.6, + "stress_response": 0.7, + "decision_making": "analytical", + "collaboration": "independent", + "verbosity": "verbose", + "conflict_approach": "avoid", + }, + "systems_thinker": { + "traits": ("holistic", "principled", "consensus-oriented"), + "communication_style": "structured", + "risk_tolerance": "medium", + "creativity": "high", + "description": "A holistic thinker who sees the big picture in systems.", + "openness": 0.8, + "conscientiousness": 0.75, + "extraversion": 0.45, + "agreeableness": 0.65, + "stress_response": 0.7, + "decision_making": "consultative", + "collaboration": "team", + "verbosity": "balanced", + "conflict_approach": "collaborate", + }, } +# Both the outer mapping and each inner mapping are read-only. PERSONALITY_PRESETS: MappingProxyType[str, MappingProxyType[str, Any]] = ( MappingProxyType({k: MappingProxyType(v) for k, v in _RAW_PRESETS.items()}) ) @@ -287,6 +369,34 @@ "data engineer": ("Reese Gallagher", "Jordan Holt", "Taylor Crane"), "security engineer": ("Quinn Steele", "Morgan Wolfe", "Avery Knox"), "content writer": ("Harper Ellis", "Kendall Frost", "Sage Monroe"), + "scrum master": ("Rowan Calloway", "Emery Dalton", "Finley Whitmore"), + "hr manager": ("Casey Pemberton", "Drew Langford", "Morgan Ashworth"), + "ml engineer": ("Quinn Fairchild", "Sage Navarro", "Avery Thornton"), + "performance engineer": ( + "Jordan Blackwell", + "Taylor Winslow", + "Blake Prescott", + ), + "automation engineer": ( + "Riley Kendrick", + "Dakota Ellsworth", + "Skyler Hargrove", + ), + "brand strategist": ( + "Phoenix Carmichael", + "Lennox Whitfield", + "Kendall Beaumont", + ), + "growth marketer": ("Harper Kingsley", "Noel Radcliffe", "Kai Vandermeer"), + "ux researcher": ("Finley Lockwood", "Emery Ashford", "Rowan Sinclair"), + "technical writer": ("Drew Fairbanks", "Casey Ellington", "Blake Holcombe"), + "database engineer": ("Reese Northcott", "Jordan Aldridge", "Taylor Wyndham"), + "security operations": ( + "Quinn Blackwood", + "Morgan Westbrook", + "Avery Cartwright", + ), + "project manager": ("Sage Pembroke", "Harley Kensington", "Lennox Beaufort"), "_default": ("Agent Alpha", "Agent Beta", "Agent Gamma", "Agent Delta"), } ) @@ -321,10 +431,9 @@ def get_personality_preset(name: str) -> dict[str, Any]: for _preset_name, _preset_dict in PERSONALITY_PRESETS.items(): try: PersonalityConfig(**_preset_dict) - except Exception as _exc: + except (ValidationError, TypeError) as _exc: msg = f"Invalid personality preset {_preset_name!r}: {_exc}" raise ValueError(msg) from _exc -# Clean up loop variables only if the loop body executed (non-empty dict). if PERSONALITY_PRESETS: del _preset_name, _preset_dict diff --git a/src/ai_company/templates/renderer.py b/src/ai_company/templates/renderer.py index 742e29b69d..d4d5e21452 100644 --- a/src/ai_company/templates/renderer.py +++ b/src/ai_company/templates/renderer.py @@ -6,6 +6,10 @@ 2. Render the raw YAML text through a Jinja2 ``SandboxedEnvironment``. 3. YAML-parse the rendered text. 4. Build a ``RootConfig``-compatible dict and validate. + +Template inheritance (``extends``) is resolved at the renderer level: +each template's Jinja2 is rendered independently, then configs are +merged via :func:`~ai_company.templates.merge.merge_template_configs`. """ from typing import TYPE_CHECKING, Any @@ -22,17 +26,25 @@ from ai_company.core.agent import PersonalityConfig from ai_company.observability import get_logger from ai_company.observability.events.template import ( + TEMPLATE_INHERIT_CIRCULAR, + TEMPLATE_INHERIT_DEPTH_EXCEEDED, + TEMPLATE_INHERIT_RESOLVE_START, + TEMPLATE_INHERIT_RESOLVE_SUCCESS, + TEMPLATE_PERSONALITY_PRESET_UNKNOWN, TEMPLATE_RENDER_JINJA2_ERROR, TEMPLATE_RENDER_START, TEMPLATE_RENDER_SUCCESS, + TEMPLATE_RENDER_TYPE_ERROR, TEMPLATE_RENDER_VALIDATION_ERROR, TEMPLATE_RENDER_VARIABLE_ERROR, TEMPLATE_RENDER_YAML_ERROR, ) from ai_company.templates.errors import ( + TemplateInheritanceError, TemplateRenderError, TemplateValidationError, ) +from ai_company.templates.merge import DEFAULT_MERGE_DEPARTMENT, merge_template_configs from ai_company.templates.presets import ( generate_auto_name, get_personality_preset, @@ -42,7 +54,10 @@ _DEFAULT_PROVIDER = "default" # Default department when not specified in template agent config. -_DEFAULT_DEPARTMENT = "engineering" +_DEFAULT_DEPARTMENT = DEFAULT_MERGE_DEPARTMENT + +# Maximum inheritance chain depth. +_MAX_INHERITANCE_DEPTH = 10 if TYPE_CHECKING: from ai_company.templates.loader import LoadedTemplate @@ -57,6 +72,8 @@ def render_template( ) -> RootConfig: """Render a loaded template into a validated RootConfig. + Resolves template inheritance (``extends``) before validation. + Args: loaded: :class:`LoadedTemplate` from the loader. variables: User-supplied variable values (overrides defaults). @@ -67,11 +84,41 @@ def render_template( Raises: TemplateRenderError: If rendering fails. TemplateValidationError: If validation fails. + TemplateInheritanceError: If inheritance resolution fails. """ logger.info( TEMPLATE_RENDER_START, source_name=loaded.source_name, ) + config_dict = _render_to_dict(loaded, variables) + + # Merge with defaults and validate. + merged = deep_merge(default_config_dict(), config_dict) + result = _validate_as_root_config(merged, loaded.source_name) + logger.info( + TEMPLATE_RENDER_SUCCESS, + source_name=loaded.source_name, + ) + return result + + +def _render_to_dict( + loaded: LoadedTemplate, + variables: dict[str, Any] | None = None, + *, + _chain: frozenset[str] = frozenset(), +) -> dict[str, Any]: + """Render a template to a config dict, resolving inheritance. + + Args: + loaded: Loaded template. + variables: User-supplied variables. + _chain: Set of already-seen template identifiers for circular + detection (internal use). + + Returns: + Config dict suitable for merging with defaults. + """ template = loaded.template vars_dict = _collect_variables(template, variables or {}) @@ -85,16 +132,148 @@ def render_template( # Parse the rendered YAML. rendered_data = _parse_rendered_yaml(rendered_text, loaded.source_name) - # Build RootConfig dict from the rendered data. - config_dict = _build_config_dict(rendered_data, template, vars_dict) + # Build config dict from the rendered data. + child_config = _build_config_dict(rendered_data, template, vars_dict) + + # If no inheritance, return child config directly. + if template.extends is None: + return child_config + + # Resolve inheritance chain. + return _resolve_inheritance( + child_config=child_config, + loaded=loaded, + vars_dict=vars_dict, + _chain=_chain, + ) + + +def _resolve_inheritance( + *, + child_config: dict[str, Any], + loaded: LoadedTemplate, + vars_dict: dict[str, Any], + _chain: frozenset[str], +) -> dict[str, Any]: + """Resolve template inheritance for a child config. + + Loads and renders the parent, detects circular dependencies and + depth violations, then merges parent + child. + + Args: + child_config: Already-rendered child config dict. + loaded: The child's :class:`LoadedTemplate`. + vars_dict: Child's resolved variables. + _chain: Already-visited parent names for circular detection. + + Returns: + Merged config dict (parent + child). + + Raises: + TemplateInheritanceError: On circular chains or depth overflow. + """ + # Guaranteed by _render_to_dict caller. + assert loaded.template.extends is not None # noqa: S101 + parent_name: str = loaded.template.extends + child_id = loaded.source_name - # Merge with defaults and validate. - merged = deep_merge(default_config_dict(), config_dict) - result = _validate_as_root_config(merged, loaded.source_name) logger.info( - TEMPLATE_RENDER_SUCCESS, - source_name=loaded.source_name, + TEMPLATE_INHERIT_RESOLVE_START, + child=child_id, + parent=parent_name, + ) + + _validate_inheritance_chain(child_id, parent_name, _chain) + + merged = _render_and_merge_parent( + parent_name, + child_config, + vars_dict, + _chain, + ) + logger.info( + TEMPLATE_INHERIT_RESOLVE_SUCCESS, + child=child_id, + parent=parent_name, + ) + return merged + + +def _validate_inheritance_chain( + child_id: str, + parent_name: str, + _chain: frozenset[str], +) -> None: + """Check for circular inheritance and depth overflow.""" + if parent_name in _chain: + logger.error( + TEMPLATE_INHERIT_CIRCULAR, + child=child_id, + parent=parent_name, + chain=sorted(_chain), + ) + msg = ( + f"Circular template inheritance: {child_id!r} extends " + f"{parent_name!r}, which is already in the inheritance chain" + ) + raise TemplateInheritanceError(msg) + + if len(_chain) >= _MAX_INHERITANCE_DEPTH: + logger.error( + TEMPLATE_INHERIT_DEPTH_EXCEEDED, + child=child_id, + depth=len(_chain), + max_depth=_MAX_INHERITANCE_DEPTH, + ) + msg = ( + f"Template inheritance depth exceeded ({len(_chain)} >= " + f"{_MAX_INHERITANCE_DEPTH}): {child_id!r}" + ) + raise TemplateInheritanceError(msg) + + +def _render_and_merge_parent( + parent_name: str, + child_config: dict[str, Any], + vars_dict: dict[str, Any], + _chain: frozenset[str], +) -> dict[str, Any]: + """Load, render, and merge a parent template with a child config.""" + from ai_company.templates.loader import load_template # noqa: PLC0415 + + parent_loaded = load_template(parent_name) + parent_vars = _collect_parent_variables( + parent_loaded.template, + vars_dict, ) + parent_config = _render_to_dict( + parent_loaded, + parent_vars, + _chain=_chain | {parent_name}, + ) + return merge_template_configs(parent_config, child_config) + + +def _collect_parent_variables( + parent_template: CompanyTemplate, + child_vars: dict[str, Any], +) -> dict[str, Any]: + """Collect variables for a parent template. + + Child's resolved variables serve as defaults for the parent. + Parent's own defaults fill gaps. + + Args: + parent_template: The parent template. + child_vars: Child's resolved variables. + + Returns: + Variable dict for parent rendering. + """ + result: dict[str, Any] = dict(child_vars) + for var in parent_template.variables: + if var.name not in result and var.default is not None: + result[var.name] = var.default return result @@ -259,10 +438,10 @@ def _build_config_dict( Returns: Dict suitable for ``RootConfig(**deep_merge(defaults, result))``. """ - company = rendered_data.get("company", {}) + company = rendered_data.get("company") if company is None: company = {} - if not isinstance(company, dict): + elif not isinstance(company, dict): msg = "Rendered template 'company' must be a mapping" logger.error(TEMPLATE_RENDER_YAML_ERROR, error=msg) raise TemplateRenderError(msg) @@ -272,7 +451,11 @@ def _build_config_dict( template.metadata.name, ) - agents = _expand_agents(_validate_list(rendered_data, "agents")) + has_extends = template.extends is not None + agents = _expand_agents( + _validate_list(rendered_data, "agents"), + has_extends=has_extends, + ) departments = _build_departments(_validate_list(rendered_data, "departments")) autonomy, budget_monthly = _extract_numeric_config(company, template) @@ -292,12 +475,20 @@ def _build_config_dict( }, } + _attach_optional_lists(rendered_data, result) + + return result + + +def _attach_optional_lists( + rendered_data: dict[str, Any], + result: dict[str, Any], +) -> None: + """Extract optional list fields from rendered data into result.""" for key in ("workflow_handoffs", "escalation_paths"): if key in rendered_data and rendered_data[key] is not None: result[key] = _validate_list(rendered_data, key) - return result - def _validate_list( rendered_data: dict[str, Any], @@ -309,6 +500,12 @@ def _validate_list( raw = [] if not isinstance(raw, list): msg = f"Rendered template {key!r} must be a list" + logger.warning( + TEMPLATE_RENDER_TYPE_ERROR, + field=key, + expected="list", + got=type(raw).__name__, + ) raise TemplateRenderError(msg) for i, item in enumerate(raw): if not isinstance(item, dict): @@ -316,6 +513,12 @@ def _validate_list( f"Rendered template {key!r}[{i}] must be a " f"mapping, got {type(item).__name__}" ) + logger.warning( + TEMPLATE_RENDER_TYPE_ERROR, + field=f"{key}[{i}]", + expected="mapping", + got=type(item).__name__, + ) raise TemplateRenderError(msg) return raw @@ -337,27 +540,40 @@ def _extract_numeric_config( ) except ValueError as exc: msg = f"Invalid numeric value in rendered template {source_name!r}: {exc}" + logger.warning( + TEMPLATE_RENDER_TYPE_ERROR, + source=source_name, + error=str(exc), + ) raise TemplateRenderError(msg) from exc return autonomy, budget_monthly def _expand_agents( raw_agents: list[dict[str, Any]], + *, + has_extends: bool, ) -> list[dict[str, Any]]: """Expand template agent dicts into AgentConfig-compatible dicts. Args: raw_agents: List of agent dicts from rendered YAML. + has_extends: Whether the template uses inheritance. Returns: List of dicts suitable for ``AgentConfig`` construction. """ - expanded: list[dict[str, Any]] = [] used_names: set[str] = set() - + expanded: list[dict[str, Any]] = [] for idx, agent in enumerate(raw_agents): - expanded.append(_expand_single_agent(agent, idx, used_names)) - + expanded.append( + _expand_single_agent( + agent, + idx, + used_names, + has_extends=has_extends, + ), + ) return expanded @@ -365,6 +581,8 @@ def _expand_single_agent( agent: dict[str, Any], idx: int, used_names: set[str], + *, + has_extends: bool, ) -> dict[str, Any]: """Expand a single template agent dict. @@ -374,6 +592,7 @@ def _expand_single_agent( role = agent.get("role") if not role: msg = f"Agent at index {idx} is missing required 'role' field" + logger.warning(TEMPLATE_RENDER_VARIABLE_ERROR, index=idx, field="role") raise TemplateRenderError(msg) name = str(agent.get("name", "")).strip() @@ -394,6 +613,47 @@ def _expand_single_agent( "level": agent.get("level", "mid"), } + _resolve_agent_personality(agent, name, agent_dict) + + model_tier = agent.get("model", "medium") + agent_dict["model"] = {"provider": _DEFAULT_PROVIDER, "model_id": model_tier} + + # Preserve _remove merge directive for inheritance. + if agent.get("_remove"): + if not has_extends: + msg = ( + f"Agent {name!r} uses '_remove' but the template " + "has no 'extends' — directive has no effect" + ) + logger.warning( + TEMPLATE_RENDER_VARIABLE_ERROR, + agent=name, + field="_remove", + ) + raise TemplateRenderError(msg) + agent_dict["_remove"] = True + + return agent_dict + + +def _resolve_agent_personality( + agent: dict[str, Any], + name: str, + agent_dict: dict[str, Any], +) -> None: + """Resolve personality from inline config or named preset. + + Mutates *agent_dict* to add the ``personality`` key when resolved. + + Args: + agent: Raw agent dict from rendered YAML. + name: Resolved agent name for error context. + agent_dict: Partially-built agent dict (mutated in place). + + Raises: + TemplateRenderError: If personality config is invalid or preset + is unknown. + """ inline_personality = agent.get("personality") preset_name = agent.get("personality_preset") if inline_personality is not None: @@ -402,6 +662,12 @@ def _expand_single_agent( f"Personality for agent {name!r} must be a mapping, " f"got {type(inline_personality).__name__}" ) + logger.warning( + TEMPLATE_RENDER_TYPE_ERROR, + agent=name, + field="personality", + got=type(inline_personality).__name__, + ) raise TemplateRenderError(msg) _validate_inline_personality(inline_personality, name) agent_dict["personality"] = inline_personality @@ -410,12 +676,13 @@ def _expand_single_agent( agent_dict["personality"] = get_personality_preset(preset_name) except KeyError as exc: msg = f"Unknown personality preset {preset_name!r} for agent {name!r}" + logger.warning( + TEMPLATE_PERSONALITY_PRESET_UNKNOWN, + agent=name, + preset=preset_name, + ) raise TemplateRenderError(msg) from exc - model_tier = agent.get("model", "medium") - agent_dict["model"] = {"provider": _DEFAULT_PROVIDER, "model_id": model_tier} - return agent_dict - def _validate_inline_personality( personality: dict[str, Any], @@ -462,6 +729,12 @@ def _build_departments( ) except ValueError as exc: msg = f"Invalid department budget value: {exc}" + logger.warning( + TEMPLATE_RENDER_TYPE_ERROR, + department=dept.get("name", ""), + field="budget_percent", + error=str(exc), + ) raise TemplateRenderError(msg) from exc dept_name = dept.get("name", "") head_role = dept.get("head_role") @@ -482,12 +755,26 @@ def _build_departments( if reporting_lines is not None: if not isinstance(reporting_lines, list): msg = f"Department {dept_name!r} 'reporting_lines' must be a list" + logger.warning( + TEMPLATE_RENDER_TYPE_ERROR, + department=dept_name, + field="reporting_lines", + expected="list", + got=type(reporting_lines).__name__, + ) raise TemplateRenderError(msg) dept_dict["reporting_lines"] = reporting_lines policies = dept.get("policies") if policies is not None: if not isinstance(policies, dict): msg = f"Department {dept_name!r} 'policies' must be a mapping" + logger.warning( + TEMPLATE_RENDER_TYPE_ERROR, + department=dept_name, + field="policies", + expected="mapping", + got=type(policies).__name__, + ) raise TemplateRenderError(msg) dept_dict["policies"] = policies departments.append(dept_dict) diff --git a/src/ai_company/templates/schema.py b/src/ai_company/templates/schema.py index a8c5c95c72..ee47c88b4d 100644 --- a/src/ai_company/templates/schema.py +++ b/src/ai_company/templates/schema.py @@ -3,7 +3,7 @@ from collections import Counter from typing import Any, Literal, Self -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from ai_company.core.enums import CompanyType, SeniorityLevel from ai_company.core.types import NotBlankStr # noqa: TC001 @@ -34,7 +34,9 @@ class TemplateVariable(BaseModel): default="str", description="Expected value type", ) - default: Any = Field(default=None, description="Default value") + default: str | int | float | bool | None = Field( + default=None, description="Default value" + ) required: bool = Field(default=False, description="Whether required") @model_validator(mode="after") @@ -89,8 +91,14 @@ class TemplateAgentConfig(BaseModel): personality_preset: Named personality preset from the presets registry. personality: Inline personality config dict (alternative to ``personality_preset``). - department: Department override (``None`` defaults to - ``"engineering"`` during rendering). + department: Department override (``None`` uses the template + system default during rendering). + merge_id: Stable identity for inheritance merge. When a + template has multiple agents with the same ``(role, + department)`` pair, ``merge_id`` disambiguates them so + child templates can target a specific agent. + remove: Merge directive — when ``True``, removes matching + parent agent during inheritance. """ model_config = ConfigDict(frozen=True, extra="forbid") @@ -114,6 +122,15 @@ class TemplateAgentConfig(BaseModel): default=None, description="Department override", ) + merge_id: str = Field( + default="", + description="Stable identity for inheritance merge", + ) + remove: bool = Field( + default=False, + alias="_remove", + description="Merge directive: remove matching parent agent", + ) @model_validator(mode="after") def _validate_personality_mutual_exclusion(self) -> Self: @@ -221,6 +238,8 @@ class CompanyTemplate(BaseModel): 1.0 = fully autonomous). workflow_handoffs: Cross-department workflow handoff definitions. escalation_paths: Cross-department escalation path definitions. + extends: Parent template name for inheritance (``None`` for + standalone templates). """ model_config = ConfigDict(frozen=True, extra="forbid") @@ -264,10 +283,31 @@ class CompanyTemplate(BaseModel): default=(), description="Cross-department escalation paths", ) + extends: NotBlankStr | None = Field( + default=None, + description="Parent template name for inheritance", + ) + + @field_validator("extends", mode="before") + @classmethod + def _normalize_extends(cls, value: Any) -> Any: + """Normalize extends to lowercase stripped form.""" + if value is None: + return None + if not isinstance(value, str): + return value # let Pydantic's type validation reject it + return value.strip().lower() @model_validator(mode="after") def _validate_agent_count_in_range(self) -> Self: - """Agent count must be within metadata min/max.""" + """Agent count must be within metadata min/max. + + Skipped when ``extends`` is set because the child may define + zero agents (inheriting all from parent). The final merged + result is validated separately. + """ + if self.extends is not None: + return self count = len(self.agents) if count < self.metadata.min_agents: msg = ( diff --git a/tests/unit/templates/conftest.py b/tests/unit/templates/conftest.py index dffdaf5b86..d86c990164 100644 --- a/tests/unit/templates/conftest.py +++ b/tests/unit/templates/conftest.py @@ -97,6 +97,88 @@ def __call__(self, content: str, name: str = ...) -> Path: ... agents: [] """ +CHILD_EXTENDS_STARTUP_YAML = """\ +template: + name: "Child of Startup" + description: "Extends startup with extra agents" + version: "1.0.0" + min_agents: 1 + max_agents: 20 + extends: "startup" + + company: + type: "startup" + + agents: + - role: "QA Engineer" + level: "mid" + model: "small" + personality_preset: "quality_guardian" + department: "engineering" +""" + +CHILD_OVERRIDE_AGENT_YAML = """\ +template: + name: "Override Child" + description: "Overrides a parent agent" + version: "1.0.0" + min_agents: 1 + max_agents: 20 + extends: "solo_founder" + + company: + type: "solo_founder" + + agents: + - role: "Full-Stack Developer" + level: "lead" + model: "large" + personality_preset: "visionary_leader" + department: "engineering" +""" + +CHILD_REMOVE_AGENT_YAML = """\ +template: + name: "Remove Child" + description: "Removes a parent agent" + version: "1.0.0" + min_agents: 1 + max_agents: 20 + extends: "solo_founder" + + company: + type: "solo_founder" + + agents: + - role: "Full-Stack Developer" + department: "engineering" + _remove: true + - role: "Backend Developer" + level: "senior" + model: "medium" + personality_preset: "pragmatic_builder" + department: "engineering" +""" + +CIRCULAR_SELF_YAML = """\ +template: + name: "Self Loop" + description: "Extends itself" + version: "1.0.0" + min_agents: 1 + max_agents: 10 + extends: "self_loop" + + company: + type: "custom" + + agents: + - role: "Backend Developer" + level: "mid" + model: "medium" + department: "engineering" +""" + def _make_template_dict(**overrides: Any) -> dict[str, Any]: """Build a minimal valid CompanyTemplate kwargs dict with overrides.""" diff --git a/tests/unit/templates/test_inheritance.py b/tests/unit/templates/test_inheritance.py new file mode 100644 index 0000000000..079a2d3e66 --- /dev/null +++ b/tests/unit/templates/test_inheritance.py @@ -0,0 +1,512 @@ +"""Tests for template inheritance (extends) and merge logic.""" + +from typing import TYPE_CHECKING, Any +from unittest.mock import patch + +if TYPE_CHECKING: + from pathlib import Path + +import pytest + +from ai_company.config.schema import RootConfig +from ai_company.core.enums import CompanyType +from ai_company.templates.errors import TemplateInheritanceError +from ai_company.templates.loader import load_template, load_template_file +from ai_company.templates.merge import ( + _merge_agents, + _merge_departments, + merge_template_configs, +) +from ai_company.templates.renderer import _collect_parent_variables, render_template +from ai_company.templates.schema import ( + CompanyTemplate, + TemplateAgentConfig, + TemplateMetadata, + TemplateVariable, +) + +from .conftest import ( + CHILD_EXTENDS_STARTUP_YAML, + CHILD_OVERRIDE_AGENT_YAML, + CHILD_REMOVE_AGENT_YAML, + CIRCULAR_SELF_YAML, +) + +pytestmark = pytest.mark.timeout(30) + + +# ── TestMergeAgents ────────────────────────────────────────────── + + +@pytest.mark.unit +class TestMergeAgents: + def test_inherit_parent_agents(self) -> None: + """Child with no agents inherits all parent agents.""" + parent = [ + {"role": "CEO", "department": "executive"}, + {"role": "Dev", "department": "engineering"}, + ] + result = _merge_agents(parent, []) + assert len(result) == 2 + assert result[0]["role"] == "CEO" + assert result[1]["role"] == "Dev" + + def test_override_by_role_dept(self) -> None: + """Child agent replaces matching parent by (role, department).""" + parent = [{"role": "Dev", "department": "engineering", "level": "mid"}] + child = [{"role": "Dev", "department": "engineering", "level": "senior"}] + result = _merge_agents(parent, child) + assert len(result) == 1 + assert result[0]["level"] == "senior" + + def test_add_new_agent(self) -> None: + """Unmatched child agent is appended.""" + parent = [{"role": "CEO", "department": "executive"}] + child = [{"role": "QA Engineer", "department": "qa"}] + result = _merge_agents(parent, child) + assert len(result) == 2 + assert result[1]["role"] == "QA Engineer" + + def test_multiple_same_role_positional_match(self) -> None: + """Multiple parents with same key are matched positionally.""" + parent = [ + {"role": "Dev", "department": "eng", "name": "first"}, + {"role": "Dev", "department": "eng", "name": "second"}, + ] + child = [ + {"role": "Dev", "department": "eng", "name": "replaced-first"}, + ] + result = _merge_agents(parent, child) + assert len(result) == 2 + assert result[0]["name"] == "replaced-first" + assert result[1]["name"] == "second" + + def test_remove_marker(self) -> None: + """_remove: true removes matching parent agent.""" + parent = [ + {"role": "CEO", "department": "executive"}, + {"role": "Dev", "department": "engineering"}, + ] + child = [{"role": "Dev", "department": "engineering", "_remove": True}] + result = _merge_agents(parent, child) + assert len(result) == 1 + assert result[0]["role"] == "CEO" + + def test_remove_nonexistent_raises(self) -> None: + """_remove with no matching parent raises error.""" + parent = [{"role": "CEO", "department": "executive"}] + child = [{"role": "QA", "department": "qa", "_remove": True}] + with pytest.raises(TemplateInheritanceError, match="no matching parent"): + _merge_agents(parent, child) + + def test_remove_marker_stripped_from_output(self) -> None: + """_remove key is not in the output for non-remove agents.""" + parent = [{"role": "Dev", "department": "eng"}] + child = [{"role": "Dev", "department": "eng", "_remove": False, "level": "sr"}] + result = _merge_agents(parent, child) + assert "_remove" not in result[0] + + +# ── TestMergeDepartments ───────────────────────────────────────── + + +@pytest.mark.unit +class TestMergeDepartments: + def test_inherit_parent_departments(self) -> None: + """Child with no departments inherits all parent depts.""" + parent = [{"name": "engineering"}, {"name": "product"}] + result = _merge_departments(parent, []) + assert len(result) == 2 + + def test_override_by_name(self) -> None: + """Child dept with matching name replaces parent entirely.""" + parent = [{"name": "engineering", "budget_percent": 50}] + child = [{"name": "Engineering", "budget_percent": 80}] + result = _merge_departments(parent, child) + assert len(result) == 1 + assert result[0]["budget_percent"] == 80 + + def test_add_new_department(self) -> None: + """Unmatched child dept is appended.""" + parent = [{"name": "engineering"}] + child = [{"name": "marketing"}] + result = _merge_departments(parent, child) + assert len(result) == 2 + assert result[1]["name"] == "marketing" + + +# ── TestMergeTemplateConfigs ───────────────────────────────────── + + +@pytest.mark.unit +class TestMergeTemplateConfigs: + def test_scalars_child_wins(self) -> None: + """Child scalars override parent.""" + parent: dict[str, Any] = {"company_name": "Parent Co"} + child: dict[str, Any] = {"company_name": "Child Co"} + result = merge_template_configs(parent, child) + assert result["company_name"] == "Child Co" + + def test_scalars_parent_fallback(self) -> None: + """Parent scalar used when child doesn't provide it.""" + parent: dict[str, Any] = { + "company_name": "Parent Co", + "company_type": "startup", + } + child: dict[str, Any] = {} + result = merge_template_configs(parent, child) + assert result["company_name"] == "Parent Co" + assert result["company_type"] == "startup" + + def test_config_deep_merge(self) -> None: + """Config dicts are deep-merged.""" + parent: dict[str, Any] = { + "config": {"autonomy": 0.5, "budget_monthly": 100.0}, + } + child: dict[str, Any] = { + "config": {"autonomy": 0.8}, + } + result = merge_template_configs(parent, child) + assert result["config"]["autonomy"] == 0.8 + assert result["config"]["budget_monthly"] == 100.0 + + def test_full_merge_integration(self) -> None: + """Full merge with agents, departments, and config.""" + parent: dict[str, Any] = { + "company_name": "Parent", + "agents": [{"role": "CEO", "department": "exec"}], + "departments": [{"name": "exec"}], + "config": {"autonomy": 0.5}, + } + child: dict[str, Any] = { + "company_name": "Child", + "agents": [{"role": "Dev", "department": "eng"}], + "departments": [{"name": "eng"}], + "config": {"budget_monthly": 200.0}, + } + result = merge_template_configs(parent, child) + assert result["company_name"] == "Child" + assert len(result["agents"]) == 2 + assert len(result["departments"]) == 2 + assert result["config"]["autonomy"] == 0.5 + assert result["config"]["budget_monthly"] == 200.0 + + def test_workflow_handoffs_child_replaces(self) -> None: + """Child workflow_handoffs replace parent entirely.""" + parent: dict[str, Any] = { + "workflow_handoffs": [{"from": "a", "to": "b"}], + } + child: dict[str, Any] = { + "workflow_handoffs": [{"from": "x", "to": "y"}], + } + result = merge_template_configs(parent, child) + assert len(result["workflow_handoffs"]) == 1 + assert result["workflow_handoffs"][0]["from"] == "x" + + def test_escalation_paths_parent_fallback(self) -> None: + """Parent escalation_paths used when child doesn't provide them.""" + parent: dict[str, Any] = { + "escalation_paths": [{"from": "eng", "to": "security"}], + } + child: dict[str, Any] = {} + result = merge_template_configs(parent, child) + assert result["escalation_paths"] == [{"from": "eng", "to": "security"}] + + def test_none_child_scalar_uses_parent(self) -> None: + """None child scalar falls back to parent value.""" + parent: dict[str, Any] = {"company_name": "Parent Co"} + child: dict[str, Any] = {"company_name": None} + result = merge_template_configs(parent, child) + assert result["company_name"] == "Parent Co" + + +# ── TestCollectParentVariables ──────────────────────────────────── + + +@pytest.mark.unit +class TestCollectParentVariables: + def test_child_vars_override_parent_defaults(self) -> None: + """Child variables take precedence over parent defaults.""" + parent = CompanyTemplate( + metadata=TemplateMetadata(name="P", company_type=CompanyType.CUSTOM), + variables=( + TemplateVariable(name="x", default="parent_x"), + TemplateVariable(name="y", default="parent_y"), + ), + agents=(TemplateAgentConfig(role="Backend Developer"),), + ) + child_vars = {"x": "child_x", "z": "child_z"} + result = _collect_parent_variables(parent, child_vars) + assert result["x"] == "child_x" + assert result["y"] == "parent_y" + assert result["z"] == "child_z" + + def test_parent_defaults_fill_gaps(self) -> None: + """Parent defaults fill variables not in child.""" + parent = CompanyTemplate( + metadata=TemplateMetadata(name="P", company_type=CompanyType.CUSTOM), + variables=(TemplateVariable(name="a", default="default_a"),), + agents=(TemplateAgentConfig(role="Backend Developer"),), + ) + result = _collect_parent_variables(parent, {}) + assert result["a"] == "default_a" + + def test_required_parent_var_without_child_value(self) -> None: + """Required parent var with no child value or default is omitted.""" + parent = CompanyTemplate( + metadata=TemplateMetadata(name="P", company_type=CompanyType.CUSTOM), + variables=(TemplateVariable(name="req", required=True),), + agents=(TemplateAgentConfig(role="Backend Developer"),), + ) + result = _collect_parent_variables(parent, {}) + assert "req" not in result + + +# ── TestResolveInheritance ─────────────────────────────────────── + + +@pytest.mark.unit +class TestResolveInheritance: + def test_single_level_extends( + self, + tmp_path: Path, + ) -> None: + """Child extends builtin and inherits its agents.""" + child_path = tmp_path / "child.yaml" + child_path.write_text(CHILD_EXTENDS_STARTUP_YAML, encoding="utf-8") + loaded = load_template_file(child_path) + config = render_template(loaded) + assert isinstance(config, RootConfig) + # Should have startup's 5 agents + 1 new QA agent. + assert len(config.agents) == 6 + + def test_override_agent_via_extends( + self, + tmp_path: Path, + ) -> None: + """Child overrides a parent agent by (role, department).""" + child_path = tmp_path / "override.yaml" + child_path.write_text(CHILD_OVERRIDE_AGENT_YAML, encoding="utf-8") + loaded = load_template_file(child_path) + config = render_template(loaded) + assert isinstance(config, RootConfig) + # solo_founder has CEO + Full-Stack Dev. Child overrides Full-Stack Dev. + assert len(config.agents) == 2 + fs_agents = [a for a in config.agents if a.role == "Full-Stack Developer"] + assert len(fs_agents) == 1 + assert fs_agents[0].level.value == "lead" + + def test_remove_agent_via_extends( + self, + tmp_path: Path, + ) -> None: + """Child removes a parent agent and adds a new one.""" + child_path = tmp_path / "remove.yaml" + child_path.write_text(CHILD_REMOVE_AGENT_YAML, encoding="utf-8") + loaded = load_template_file(child_path) + config = render_template(loaded) + assert isinstance(config, RootConfig) + # solo_founder: CEO + FS Dev. Remove FS Dev, add Backend Dev => 2 agents. + roles = [a.role for a in config.agents] + assert "CEO" in roles + assert "Backend Developer" in roles + assert "Full-Stack Developer" not in roles + + def test_variable_flow_to_parent( + self, + tmp_path: Path, + ) -> None: + """Child variables flow to parent template.""" + child_path = tmp_path / "var_child.yaml" + child_path.write_text(CHILD_EXTENDS_STARTUP_YAML, encoding="utf-8") + loaded = load_template_file(child_path) + config = render_template( + loaded, + variables={"company_name": "Custom Name"}, + ) + # Child's company_name overrides parent's (child wins in merge). + assert config.company_name == "Custom Name" + + def test_multi_level_extends( + self, + tmp_path: Path, + ) -> None: + """A→B→C multi-level inheritance resolves correctly.""" + # B extends startup + child_b_yaml = CHILD_EXTENDS_STARTUP_YAML + child_b_path = tmp_path / "child_b.yaml" + child_b_path.write_text(child_b_yaml, encoding="utf-8") + + # A extends B (using file path won't work for name-based lookup, + # so we use user template directory patching) + child_a_yaml = """\ +template: + name: "Grandchild" + description: "Two levels deep" + version: "1.0.0" + min_agents: 1 + max_agents: 30 + extends: "child_b" + + company: + type: "startup" + + agents: + - role: "Data Analyst" + level: "mid" + model: "medium" + personality_preset: "data_driven_optimizer" + department: "engineering" +""" + child_a_path = tmp_path / "child_a.yaml" + child_a_path.write_text(child_a_yaml, encoding="utf-8") + + # Patch user templates dir so child_b is found by name. + with patch( + "ai_company.templates.loader._USER_TEMPLATES_DIR", + tmp_path, + ): + loaded = load_template_file(child_a_path) + config = render_template(loaded) + assert isinstance(config, RootConfig) + # startup(5) + child_b adds QA(1) + child_a adds Data Analyst(1) = 7 + assert len(config.agents) == 7 + + +# ── TestCircularDetection ──────────────────────────────────────── + + +@pytest.mark.unit +class TestCircularDetection: + def test_self_extends_raises( + self, + tmp_path: Path, + ) -> None: + """Template extending itself raises TemplateInheritanceError.""" + path = tmp_path / "self_loop.yaml" + path.write_text(CIRCULAR_SELF_YAML, encoding="utf-8") + + with patch( + "ai_company.templates.loader._USER_TEMPLATES_DIR", + tmp_path, + ): + loaded = load_template_file(path) + with pytest.raises( + TemplateInheritanceError, + match="Circular template inheritance", + ): + render_template(loaded) + + def test_a_b_a_cycle_raises( + self, + tmp_path: Path, + ) -> None: + """A→B→A cycle raises TemplateInheritanceError.""" + a_yaml = """\ +template: + name: "Template A" + description: "test" + version: "1.0.0" + min_agents: 1 + max_agents: 10 + extends: "template_b" + + company: + type: "custom" + + agents: + - role: "Backend Developer" + level: "mid" + model: "medium" + department: "engineering" +""" + b_yaml = """\ +template: + name: "Template B" + description: "test" + version: "1.0.0" + min_agents: 1 + max_agents: 10 + extends: "template_a" + + company: + type: "custom" + + agents: + - role: "Backend Developer" + level: "mid" + model: "medium" + department: "engineering" +""" + (tmp_path / "template_a.yaml").write_text(a_yaml, encoding="utf-8") + (tmp_path / "template_b.yaml").write_text(b_yaml, encoding="utf-8") + + with patch( + "ai_company.templates.loader._USER_TEMPLATES_DIR", + tmp_path, + ): + loaded = load_template("template_a") + with pytest.raises( + TemplateInheritanceError, + match="Circular template inheritance", + ): + render_template(loaded) + + def test_depth_limit_exceeded_raises( + self, + tmp_path: Path, + ) -> None: + """Exceeding max depth raises TemplateInheritanceError.""" + # Create a chain of 12 templates: t0 → t1 → ... → t11 + for i in range(12): + parent_ref = f"chain_{i - 1}" if i > 0 else "" + extends_line = f' extends: "{parent_ref}"' if i > 0 else "" + yaml_content = f"""\ +template: + name: "Chain {i}" + description: "test" + version: "1.0.0" + min_agents: 1 + max_agents: 100 +{extends_line} + + company: + type: "custom" + + agents: + - role: "Backend Developer" + level: "mid" + model: "medium" + department: "engineering" +""" + (tmp_path / f"chain_{i}.yaml").write_text(yaml_content, encoding="utf-8") + + with patch( + "ai_company.templates.loader._USER_TEMPLATES_DIR", + tmp_path, + ): + loaded = load_template("chain_11") + with pytest.raises( + TemplateInheritanceError, + match="depth exceeded", + ): + render_template(loaded) + + +# ── TestInheritanceIntegration ─────────────────────────────────── + + +@pytest.mark.unit +class TestInheritanceIntegration: + def test_child_extends_builtin_renders_to_valid_root_config( + self, + tmp_path: Path, + ) -> None: + """File-based child extending a builtin renders to valid RootConfig.""" + child_path = tmp_path / "custom.yaml" + child_path.write_text(CHILD_EXTENDS_STARTUP_YAML, encoding="utf-8") + loaded = load_template_file(child_path) + config = render_template(loaded) + assert isinstance(config, RootConfig) + assert len(config.agents) >= 1 + assert len(config.departments) >= 1 diff --git a/tests/unit/templates/test_loader.py b/tests/unit/templates/test_loader.py index 114c7bd26e..e4cecf41be 100644 --- a/tests/unit/templates/test_loader.py +++ b/tests/unit/templates/test_loader.py @@ -53,6 +53,54 @@ def test_count_matches_registry(self) -> None: assert len(list_builtin_templates()) == len(BUILTIN_TEMPLATES) +@pytest.mark.unit +class TestMinMaxPassthrough: + def test_min_max_agents_from_yaml( + self, + tmp_template_file: TemplateFileFactory, + ) -> None: + """min_agents/max_agents pass through from YAML to TemplateMetadata.""" + yaml_with_minmax = """\ +template: + name: "MinMax Test" + description: "test" + version: "1.0.0" + min_agents: 3 + max_agents: 10 + + company: + type: "custom" + + agents: + - role: "Backend Developer" + level: "mid" + model: "medium" + department: "engineering" + - role: "Frontend Developer" + level: "mid" + model: "medium" + department: "engineering" + - role: "QA Engineer" + level: "mid" + model: "small" + department: "engineering" +""" + path = tmp_template_file(yaml_with_minmax) + loaded = load_template_file(path) + assert loaded.template.metadata.min_agents == 3 + assert loaded.template.metadata.max_agents == 10 + + def test_defaults_when_not_specified( + self, + tmp_template_file: TemplateFileFactory, + ) -> None: + """Without min/max in YAML, defaults apply (1, 100).""" + path = tmp_template_file(MINIMAL_TEMPLATE_YAML) + loaded = load_template_file(path) + assert loaded.template.metadata.min_agents == 1 + assert loaded.template.metadata.max_agents == 100 + + # ── list_templates ─────────────────────────────────────────────── diff --git a/tests/unit/templates/test_presets.py b/tests/unit/templates/test_presets.py index dee2dc06f2..d1e0c31611 100644 --- a/tests/unit/templates/test_presets.py +++ b/tests/unit/templates/test_presets.py @@ -2,7 +2,9 @@ import pytest +from ai_company.core.agent import PersonalityConfig from ai_company.templates.presets import ( + _AUTO_NAMES, PERSONALITY_PRESETS, generate_auto_name, get_personality_preset, @@ -41,12 +43,28 @@ def test_all_presets_have_required_keys(self) -> None: preset = get_personality_preset(name) assert required_keys.issubset(preset.keys()), f"{name} missing keys" - def test_preset_count_at_least_15(self) -> None: - assert len(PERSONALITY_PRESETS) >= 15 + def test_preset_count_at_least_20(self) -> None: + assert len(PERSONALITY_PRESETS) >= 20 + + @pytest.mark.parametrize( + "preset_name", + [ + "user_advocate", + "process_optimizer", + "growth_hacker", + "technical_communicator", + "systems_thinker", + ], + ) + def test_new_presets_produce_valid_personality_config( + self, + preset_name: str, + ) -> None: + preset = get_personality_preset(preset_name) + config = PersonalityConfig(**preset) + assert isinstance(config, PersonalityConfig) def test_all_presets_produce_valid_personality_config(self) -> None: - from ai_company.core.agent import PersonalityConfig - for name in PERSONALITY_PRESETS: preset = get_personality_preset(name) config = PersonalityConfig(**preset) @@ -97,3 +115,15 @@ def test_whitespace_stripped_from_role(self) -> None: a = generate_auto_name(" CEO ", seed=0) b = generate_auto_name("CEO", seed=0) assert a == b + + +@pytest.mark.unit +class TestAutoNameCoverage: + def test_auto_names_cover_all_builtin_roles(self) -> None: + """Every role in BUILTIN_ROLES has an auto-name pool.""" + from ai_company.core.role_catalog import BUILTIN_ROLES + + pool_keys = {k for k in _AUTO_NAMES if k != "_default"} + role_keys = {r.name.lower() for r in BUILTIN_ROLES} + missing = role_keys - pool_keys + assert not missing, f"Roles missing auto-name pools: {sorted(missing)}" diff --git a/tests/unit/templates/test_renderer.py b/tests/unit/templates/test_renderer.py index 65756993bf..586fd0f9fb 100644 --- a/tests/unit/templates/test_renderer.py +++ b/tests/unit/templates/test_renderer.py @@ -12,7 +12,11 @@ TEMPLATE_RENDER_SUCCESS, ) from ai_company.templates.errors import TemplateRenderError -from ai_company.templates.loader import load_template, load_template_file +from ai_company.templates.loader import ( + BUILTIN_TEMPLATES, + load_template, + load_template_file, +) from ai_company.templates.renderer import render_template from .conftest import TEMPLATE_REQUIRED_VAR_YAML, TEMPLATE_WITH_VARIABLES_YAML @@ -42,8 +46,6 @@ def test_render_builtin_startup(self) -> None: assert len(config.agents) == 5 def test_render_all_builtins_produce_valid_root_config(self) -> None: - from ai_company.templates.loader import BUILTIN_TEMPLATES - for name in BUILTIN_TEMPLATES: loaded = load_template(name) config = render_template(loaded) @@ -310,7 +312,7 @@ def test_inline_personality_applied(self) -> None: "communication_style": "custom", }, } - result = _expand_single_agent(agent, 0, set()) + result = _expand_single_agent(agent, 0, set(), has_extends=False) assert result["personality"]["communication_style"] == "custom" assert "custom-trait" in result["personality"]["traits"] @@ -393,7 +395,7 @@ def test_invalid_inline_personality_raises_template_render_error(self) -> None: "personality": {"openness": 99.0}, } with pytest.raises(TemplateRenderError, match="Invalid inline personality"): - _expand_single_agent(agent, 0, set()) + _expand_single_agent(agent, 0, set(), has_extends=False) def test_non_dict_personality_raises_template_render_error(self) -> None: """Non-dict personality value raises TemplateRenderError.""" @@ -404,7 +406,7 @@ def test_non_dict_personality_raises_template_render_error(self) -> None: "personality": "not-a-dict", } with pytest.raises(TemplateRenderError, match="must be a mapping"): - _expand_single_agent(agent, 0, set()) + _expand_single_agent(agent, 0, set(), has_extends=False) @pytest.mark.unit @@ -414,7 +416,7 @@ def test_missing_role_raises_template_render_error(self) -> None: from ai_company.templates.renderer import _expand_single_agent with pytest.raises(TemplateRenderError, match="missing required 'role'"): - _expand_single_agent({}, 0, set()) + _expand_single_agent({}, 0, set(), has_extends=False) @pytest.mark.unit @@ -486,7 +488,7 @@ def test_unknown_preset_raises_template_render_error(self) -> None: "personality_preset": "does_not_exist", } with pytest.raises(TemplateRenderError, match="Unknown personality preset"): - _expand_single_agent(agent, 0, set()) + _expand_single_agent(agent, 0, set(), has_extends=False) @pytest.mark.unit @@ -506,3 +508,40 @@ def test_non_dict_item_raises(self) -> None: with pytest.raises(TemplateRenderError, match="must be a mapping"): _validate_list({"agents": [{"role": "Dev"}, "bad"]}, "agents") + + +# ── Roster count tests ────────────────────────────────────────── + + +@pytest.mark.unit +class TestRosterCounts: + @pytest.mark.parametrize("name", sorted(BUILTIN_TEMPLATES)) + def test_template_agent_count_in_range(self, name: str) -> None: + """Each template renders agents within its declared metadata range.""" + loaded = load_template(name) + config = render_template(loaded) + lo = loaded.template.metadata.min_agents + hi = loaded.template.metadata.max_agents + assert lo <= len(config.agents) <= hi, ( + f"{name}: expected {lo}-{hi} agents, got {len(config.agents)}" + ) + assert isinstance(config, RootConfig) + + def test_full_company_variable_override(self) -> None: + """full_company num_backend_devs override changes agent count.""" + loaded = load_template("full_company") + default_config = render_template(loaded) + override_config = render_template( + loaded, + variables={"num_backend_devs": 5}, + ) + # Default is 3 backend devs, override is 5 → +2 agents. + default_backend = sum( + 1 for a in default_config.agents if a.role == "Backend Developer" + ) + override_backend = sum( + 1 for a in override_config.agents if a.role == "Backend Developer" + ) + assert default_backend == 3 + assert override_backend == 5 + assert override_backend - default_backend == 2 diff --git a/tests/unit/templates/test_schema.py b/tests/unit/templates/test_schema.py index 78fa890f5e..f0333c2763 100644 --- a/tests/unit/templates/test_schema.py +++ b/tests/unit/templates/test_schema.py @@ -124,6 +124,14 @@ def test_both_personality_and_preset_rejected(self) -> None: personality={"openness": 0.9}, ) + def test_remove_alias(self) -> None: + a = TemplateAgentConfig(role="Dev", _remove=True) + assert a.remove is True + + def test_remove_default_false(self) -> None: + a = TemplateAgentConfig(role="Dev") + assert a.remove is False + # ── TemplateDepartmentConfig ───────────────────────────────────── @@ -406,6 +414,27 @@ def test_escalation_paths_accepted( ) assert len(t.escalation_paths) == 1 + def test_extends_field_accepted( + self, + make_template_dict: Callable[..., dict[str, Any]], + ) -> None: + t = CompanyTemplate(**make_template_dict(extends="startup", agents=())) + assert t.extends == "startup" + + def test_extends_normalizes_case( + self, + make_template_dict: Callable[..., dict[str, Any]], + ) -> None: + t = CompanyTemplate(**make_template_dict(extends=" StartUp ", agents=())) + assert t.extends == "startup" + + def test_extends_skips_agent_count_validation( + self, + make_template_dict: Callable[..., dict[str, Any]], + ) -> None: + t = CompanyTemplate(**make_template_dict(extends="startup", agents=())) + assert len(t.agents) == 0 + def test_frozen( self, make_template_dict: Callable[..., dict[str, Any]],