Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@ classifiers = [
"Programming Language :: Python :: 3.14",
"Typing :: Typed",
]
dependencies = []
dependencies = [
"pydantic==2.12.5",
]

[build-system]
requires = ["hatchling"]
Expand Down Expand Up @@ -40,7 +42,6 @@ dev = [
"mypy==1.19.1",
"pre-commit==4.5.1",
"pre-commit-uv==4.2.1",
"pydantic==2.12.5",
"ruff==0.15.4",
{include-group = "test"},
]
Expand Down
75 changes: 75 additions & 0 deletions src/ai_company/core/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
"""Core domain models for the AI company framework."""

from ai_company.core.agent import (
AgentIdentity,
MemoryConfig,
ModelConfig,
PersonalityConfig,
SkillSet,
ToolPermissions,
)
from ai_company.core.company import (
Company,
CompanyConfig,
Department,
HRRegistry,
Team,
)
from ai_company.core.enums import (
AgentStatus,
CompanyType,
CostTier,
CreativityLevel,
DepartmentName,
MemoryType,
ProficiencyLevel,
RiskTolerance,
SeniorityLevel,
SkillCategory,
)
from ai_company.core.role import (
Authority,
CustomRole,
Role,
SeniorityInfo,
Skill,
)
from ai_company.core.role_catalog import (
BUILTIN_ROLES,
SENIORITY_INFO,
get_builtin_role,
get_seniority_info,
)

__all__ = [
"BUILTIN_ROLES",
"SENIORITY_INFO",
"AgentIdentity",
"AgentStatus",
"Authority",
"Company",
"CompanyConfig",
"CompanyType",
"CostTier",
"CreativityLevel",
"CustomRole",
"Department",
"DepartmentName",
"HRRegistry",
"MemoryConfig",
"MemoryType",
"ModelConfig",
"PersonalityConfig",
"ProficiencyLevel",
"RiskTolerance",
"Role",
"SeniorityInfo",
"SeniorityLevel",
"Skill",
"SkillCategory",
"SkillSet",
"Team",
"ToolPermissions",
"get_builtin_role",
"get_seniority_info",
]
Comment on lines +44 to +75
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For better readability and maintainability, it's a good practice to sort the __all__ list. A common convention is to group symbols by type (e.g., constants, classes, functions) and then sort alphabetically within each group. This makes it easier to navigate as the number of exported symbols grows.

__all__ = [
    # Constants
    "BUILTIN_ROLES",
    "SENIORITY_INFO",
    # Enums & Models
    "AgentIdentity",
    "AgentStatus",
    "Authority",
    "Company",
    "CompanyConfig",
    "CompanyType",
    "CostTier",
    "CreativityLevel",
    "CustomRole",
    "Department",
    "DepartmentName",
    "HRRegistry",
    "MemoryConfig",
    "MemoryType",
    "ModelConfig",
    "PersonalityConfig",
    "ProficiencyLevel",
    "RiskTolerance",
    "Role",
    "SeniorityInfo",
    "SeniorityLevel",
    "Skill",
    "SkillCategory",
    "SkillSet",
    "Team",
    "ToolPermissions",
    # Functions
    "get_builtin_role",
    "get_seniority_info",
]

280 changes: 280 additions & 0 deletions src/ai_company/core/agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,280 @@
"""Agent identity and configuration models."""

from datetime import date # noqa: TC003 — required at runtime by Pydantic
from typing import Self
from uuid import UUID, uuid4

from pydantic import BaseModel, ConfigDict, Field, model_validator
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To use the Self type hint for model validators as per Pydantic v2 best practices, you'll need to import it. This improves maintainability by not hardcoding class names in return types.

Suggested change
from pydantic import BaseModel, ConfigDict, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, model_validator, Self


from ai_company.core.enums import (
AgentStatus,
CreativityLevel,
MemoryType,
RiskTolerance,
SeniorityLevel,
)
from ai_company.core.role import Authority


class PersonalityConfig(BaseModel):
"""Personality traits and communication style for an agent.

Attributes:
traits: Personality trait keywords.
communication_style: Free-text style description.
risk_tolerance: Risk tolerance level.
creativity: Creativity level.
description: Extended personality description.
"""

model_config = ConfigDict(frozen=True)

traits: tuple[str, ...] = Field(
default=(),
description="Personality traits",
)
communication_style: str = Field(
default="neutral",
min_length=1,
description="Communication style description",
)
risk_tolerance: RiskTolerance = Field(
default=RiskTolerance.MEDIUM,
description="Risk tolerance level",
)
creativity: CreativityLevel = Field(
default=CreativityLevel.MEDIUM,
description="Creativity level",
)
description: str = Field(
default="",
description="Extended personality description",
)

@model_validator(mode="after")
def _validate_no_empty_traits(self) -> Self:
"""Ensure no empty or whitespace-only traits or communication_style."""
if not self.communication_style.strip():
msg = "communication_style must not be whitespace-only"
raise ValueError(msg)
for trait in self.traits:
if not trait.strip():
msg = "Empty or whitespace-only entry in traits"
raise ValueError(msg)
return self
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Comment on lines +32 to +64
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

There's a recurring pattern of using model_validator across several models in this file (PersonalityConfig, SkillSet, ModelConfig, etc.) to prevent fields from being empty or containing only whitespace. This leads to boilerplate code.

You can simplify this significantly by using Pydantic's Annotated types with StringConstraints. This approach is more declarative and reduces code duplication.

First, define a reusable NonEmptyStr type at the top of the file (after imports):

from typing import Annotated
from pydantic import StringConstraints

NonEmptyStr = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]

Then, you can apply this type to the model fields and remove the now-redundant _validate_no_empty_traits validator. This pattern can be applied to the other models in this file as well.

    traits: tuple[NonEmptyStr, ...] = Field(
        default=(),
        description="Personality traits",
    )
    communication_style: NonEmptyStr = Field(
        default="neutral",
        description="Communication style description",
    )
    risk_tolerance: RiskTolerance = Field(
        default=RiskTolerance.MEDIUM,
        description="Risk tolerance level",
    )
    creativity: CreativityLevel = Field(
        default=CreativityLevel.MEDIUM,
        description="Creativity level",
    )
    description: str = Field(
        default="",
        description="Extended personality description",
    )



class SkillSet(BaseModel):
"""Primary and secondary skills for an agent.

Attributes:
primary: Core competency skill names.
secondary: Supporting skill names.
"""

model_config = ConfigDict(frozen=True)

primary: tuple[str, ...] = Field(
default=(),
description="Primary skills",
)
secondary: tuple[str, ...] = Field(
default=(),
description="Secondary skills",
)

@model_validator(mode="after")
def _validate_no_empty_skills(self) -> Self:
"""Ensure no empty or whitespace-only skill names."""
for field_name in ("primary", "secondary"):
for skill in getattr(self, field_name):
if not skill.strip():
msg = f"Empty or whitespace-only skill name in {field_name}"
raise ValueError(msg)
return self


class ModelConfig(BaseModel):
"""LLM model configuration for an agent.

Attributes:
provider: LLM provider name (e.g. ``"anthropic"``).
model_id: Model identifier (e.g. ``"claude-sonnet-4-6"``).
temperature: Sampling temperature (0.0 to 2.0).
max_tokens: Maximum output tokens.
fallback_model: Optional fallback model identifier.
"""

model_config = ConfigDict(frozen=True)

provider: str = Field(min_length=1, description="LLM provider name")
model_id: str = Field(min_length=1, description="Model identifier")
temperature: float = Field(
default=0.7,
ge=0.0,
le=2.0,
description="Sampling temperature",
)
max_tokens: int = Field(
default=4096,
gt=0,
description="Maximum output tokens",
)
fallback_model: str | None = Field(
default=None,
min_length=1,
description="Fallback model identifier",
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.

@model_validator(mode="after")
def _validate_non_blank_identifiers(self) -> Self:
"""Ensure identifier fields are not whitespace-only."""
for field_name in ("provider", "model_id", "fallback_model"):
value = getattr(self, field_name)
if value is not None and not value.strip():
msg = f"{field_name} must not be whitespace-only"
raise ValueError(msg)
return self


class MemoryConfig(BaseModel):
"""Memory configuration for an agent.

Attributes:
type: Memory persistence type.
retention_days: Days to retain memories (``None`` means forever).
"""

model_config = ConfigDict(frozen=True)

type: MemoryType = Field(
default=MemoryType.SESSION,
description="Memory persistence type",
)
retention_days: int | None = Field(
default=None,
ge=1,
description="Days to retain memories (None = forever)",
)

@model_validator(mode="after")
def _validate_retention_consistency(self) -> Self:
"""Ensure retention_days is None when memory type is MemoryType.NONE."""
if self.type is MemoryType.NONE and self.retention_days is not None:
msg = "retention_days must be None when memory type is 'none'"
raise ValueError(msg)
return self


class ToolPermissions(BaseModel):
"""Tool access permissions for an agent.

Attributes:
allowed: Explicitly allowed tool names.
denied: Explicitly denied tool names.
"""

model_config = ConfigDict(frozen=True)

allowed: tuple[str, ...] = Field(
default=(),
description="Explicitly allowed tools",
)
denied: tuple[str, ...] = Field(
default=(),
description="Explicitly denied tools",
)

@model_validator(mode="after")
def _validate_no_empty_tools(self) -> Self:
"""Ensure no empty or whitespace-only tool names."""
for field_name in ("allowed", "denied"):
for tool in getattr(self, field_name):
if not tool.strip():
msg = f"Empty or whitespace-only tool name in {field_name}"
raise ValueError(msg)
return self

@model_validator(mode="after")
def _validate_no_overlap(self) -> Self:
"""Ensure no tool appears in both allowed and denied lists.

Comparison is case-insensitive.
"""
allowed_normalized = {t.strip().casefold() for t in self.allowed}
denied_normalized = {t.strip().casefold() for t in self.denied}
overlap = allowed_normalized & denied_normalized
if overlap:
msg = f"Tools appear in both allowed and denied lists: {sorted(overlap)}"
raise ValueError(msg)
return self


class AgentIdentity(BaseModel):
"""Complete agent identity card.

Every agent in the company is represented by an ``AgentIdentity``
containing its role, personality, model backend, memory settings,
tool permissions, and authority configuration.

Attributes:
id: Unique agent identifier.
name: Agent display name.
role: Role name (string reference to :class:`~ai_company.core.role.Role`).
department: Department name (string reference).
level: Seniority level.
personality: Personality configuration.
skills: Primary and secondary skill set.
model: LLM model configuration.
memory: Memory configuration.
tools: Tool permissions.
authority: Authority configuration for this agent.
hiring_date: Date the agent was hired.
status: Current lifecycle status.
"""

model_config = ConfigDict(frozen=True)

id: UUID = Field(default_factory=uuid4, description="Unique agent identifier")
name: str = Field(min_length=1, description="Agent display name")
role: str = Field(min_length=1, description="Role name")
department: str = Field(min_length=1, description="Department name")
level: SeniorityLevel = Field(
default=SeniorityLevel.MID,
description="Seniority level",
)
personality: PersonalityConfig = Field(
default_factory=PersonalityConfig,
description="Personality configuration",
)
skills: SkillSet = Field(
default_factory=SkillSet,
description="Skill set",
)
model: ModelConfig = Field(description="LLM model configuration")
memory: MemoryConfig = Field(
default_factory=MemoryConfig,
description="Memory configuration",
)
tools: ToolPermissions = Field(
default_factory=ToolPermissions,
description="Tool permissions",
)
authority: Authority = Field(
default_factory=Authority,
description="Authority scope",
)
hiring_date: date = Field(description="Date the agent was hired")
status: AgentStatus = Field(
default=AgentStatus.ACTIVE,
description="Current lifecycle status",
)

@model_validator(mode="after")
def _validate_non_blank_identifiers(self) -> Self:
"""Ensure name, role, and department are not whitespace-only."""
for field_name in ("name", "role", "department"):
if not getattr(self, field_name).strip():
msg = f"{field_name} must not be whitespace-only"
raise ValueError(msg)
return self
Loading