| 1 | from helpers import files |
| 2 | from helpers import cache |
| 3 | from helpers import yaml as yaml_helper |
| 4 | from typing import TypedDict, TYPE_CHECKING, Literal |
| 5 | from pydantic import BaseModel, model_validator |
| 6 | import json |
| 7 | import os |
| 8 | |
| 9 | GLOBAL_DIR = "." |
| 10 | USER_DIR = "usr" |
| 11 | DEFAULT_AGENTS_DIR = "agents" |
| 12 | USER_AGENTS_DIR = "usr/agents" |
| 13 | PATHS_CACHE_AREA = "subagent_paths(plugins)" |
| 14 | |
| 15 | cache.toggle_area(PATHS_CACHE_AREA, False) |
| 16 | |
| 17 | type Origin = Literal["default", "user", "project", "plugin"] |
| 18 | |
| 19 | if TYPE_CHECKING: |
| 20 | from agent import Agent |
| 21 | |
| 22 | |
| 23 | class SubAgentListItem(BaseModel): |
| 24 | name: str = "" |
| 25 | title: str = "" |
| 26 | description: str = "" |
| 27 | context: str = "" |
| 28 | path: str = "" |
| 29 | origin: list[Origin] = [] |
| 30 | enabled: bool = True |
| 31 | avatar: dict[str, str] | None = None |
| 32 | |
| 33 | @model_validator(mode="after") |
| 34 | def post_validator(self): |
| 35 | if "title" not in self.model_fields_set and self.name: |
| 36 | object.__setattr__(self, "title", self.name) |
| 37 | return self |
| 38 | |
| 39 | |
| 40 | class SubAgent(SubAgentListItem): |
| 41 | prompts: dict[str, str] = {} |
| 42 | |
| 43 | |
| 44 | def get_agents_list(project_name: str | None = None) -> list[SubAgentListItem]: |
| 45 | return list(get_agents_dict(project_name).values()) |
| 46 | |
| 47 | |
| 48 | def get_agents_dict( |
| 49 | project_name: str | None = None, |
| 50 | ) -> dict[str, SubAgentListItem]: |
| 51 | def _merge_agent_dicts( |
| 52 | base: dict[str, SubAgentListItem], |
| 53 | overrides: dict[str, SubAgentListItem], |
| 54 | ) -> dict[str, SubAgentListItem]: |
| 55 | merged: dict[str, SubAgentListItem] = dict(base) |
| 56 | for name, override in overrides.items(): |
| 57 | base_agent = merged.get(name) |
| 58 | merged[name] = ( |
| 59 | _merge_agent_list_item(base_agent, override) |
| 60 | if base_agent |
| 61 | else override |
| 62 | ) |
| 63 | return merged |
| 64 | |
| 65 | from helpers import plugins |
| 66 | |
| 67 | # load default, plugin, and custom agents and merge |
| 68 | default_agents = _get_agents_list_from_dir(DEFAULT_AGENTS_DIR, origin="default") |
| 69 | merged: dict[str, SubAgentListItem] = dict(default_agents) |
| 70 | |
| 71 | # merge with plugin agents |
| 72 | for plugin_dir in plugins.get_enabled_plugin_paths(None, "agents"): |
| 73 | plugin_agents = _get_agents_list_from_dir(plugin_dir, origin="plugin") |
| 74 | merged = _merge_agent_dicts(merged, plugin_agents) |
| 75 | |
| 76 | custom_agents = _get_agents_list_from_dir(USER_AGENTS_DIR, origin="user") |
| 77 | merged = _merge_agent_dicts(merged, custom_agents) |
| 78 | |
| 79 | # merge with project agents if possible |
| 80 | if project_name: |
| 81 | from helpers import projects |
| 82 | |
| 83 | project_agents_dir = projects.get_project_meta(project_name, "agents") |
| 84 | project_agents = _get_agents_list_from_dir(project_agents_dir, origin="project") |
| 85 | merged = _merge_agent_dicts(merged, project_agents) |
| 86 | |
| 87 | return merged |
| 88 | |
| 89 | |
| 90 | def _get_agents_list_from_dir(dir: str, origin: Origin) -> dict[str, SubAgentListItem]: |
| 91 | result: dict[str, SubAgentListItem] = {} |
| 92 | subdirs = files.get_subdirectories(dir) |
| 93 | |
| 94 | for subdir in subdirs: |
| 95 | try: |
| 96 | try: |
| 97 | raw = _read_agent_definition(dir, subdir) |
| 98 | except FileNotFoundError: |
| 99 | if origin == "default": |
| 100 | continue |
| 101 | raw = {} |
| 102 | agent_data = SubAgentListItem.model_validate(raw) |
| 103 | name = agent_data.name or subdir |
| 104 | agent_data.name = name |
| 105 | if "title" not in agent_data.model_fields_set: |
| 106 | object.__setattr__(agent_data, "title", name) |
| 107 | agent_data.path = files.get_abs_path(dir, subdir) |
| 108 | agent_data.origin = [origin] |
| 109 | result[name] = agent_data |
| 110 | except Exception: |
| 111 | continue |
| 112 | |
| 113 | return result |
| 114 | |
| 115 | |
| 116 | def load_agent_data(name: str, project_name: str | None = None) -> SubAgent: |
| 117 | from helpers import plugins |
| 118 | |
| 119 | # load default, plugin, and user agents and merge |
| 120 | default_agent = _load_agent_data_from_dir( |
| 121 | DEFAULT_AGENTS_DIR, name, origin="default" |
| 122 | ) |
| 123 | merged = default_agent |
| 124 | |
| 125 | # merge with plugin agents |
| 126 | # TODO review this |
| 127 | for plugin_dir in plugins.get_enabled_plugin_paths(None, "agents"): |
| 128 | plugin_agent = _load_agent_data_from_dir(plugin_dir, name, origin="plugin") |
| 129 | merged = _merge_agent(merged, plugin_agent) |
| 130 | |
| 131 | user_agent = _load_agent_data_from_dir(USER_AGENTS_DIR, name, origin="user") |
| 132 | merged = _merge_agent(merged, user_agent) |
| 133 | |
| 134 | # merge with project agent if possible |
| 135 | if project_name: |
| 136 | from helpers import projects |
| 137 | |
| 138 | project_agents_dir = projects.get_project_meta(project_name, "agents") |
| 139 | project_agent = _load_agent_data_from_dir( |
| 140 | project_agents_dir, name, origin="project" |
| 141 | ) |
| 142 | merged = _merge_agent(merged, project_agent) |
| 143 | |
| 144 | if merged is None: |
| 145 | raise FileNotFoundError( |
| 146 | f"Agent '{name}' not found in default, plugin, or custom directories" |
| 147 | ) |
| 148 | |
| 149 | return merged |
| 150 | |
| 151 | |
| 152 | def save_agent_data(name: str, subagent: SubAgent) -> None: |
| 153 | # write agent.json in custom directory |
| 154 | agent_dir = f"{USER_AGENTS_DIR}/{name}" |
| 155 | agent_json = { |
| 156 | "title": subagent.title, |
| 157 | "description": subagent.description, |
| 158 | "context": subagent.context, |
| 159 | "enabled": subagent.enabled, |
| 160 | } |
| 161 | files.write_file(f"{agent_dir}/agent.json", json.dumps(agent_json, indent=2)) |
| 162 | |
| 163 | # replace prompts in custom directory |
| 164 | prompts_dir = f"{agent_dir}/prompts" |
| 165 | # clear existing custom prompts directory (if any) |
| 166 | files.delete_dir(prompts_dir) |
| 167 | |
| 168 | prompts = subagent.prompts or {} |
| 169 | for name, content in prompts.items(): |
| 170 | safe_name = files.safe_file_name(name) |
| 171 | if not safe_name.endswith(".md"): |
| 172 | safe_name += ".md" |
| 173 | files.write_file(f"{prompts_dir}/{safe_name}", content) |
| 174 | |
| 175 | |
| 176 | def delete_agent_data(name: str) -> None: |
| 177 | files.delete_dir(f"{USER_AGENTS_DIR}/{name}") |
| 178 | |
| 179 | |
| 180 | def _load_agent_data_from_dir(dir: str, name: str, origin: Origin) -> SubAgent | None: |
| 181 | agent_dir = files.get_abs_path(dir, name) |
| 182 | if not os.path.isdir(agent_dir): |
| 183 | return None |
| 184 | |
| 185 | try: |
| 186 | subagent = SubAgent.model_validate(_read_agent_definition(dir, name)) |
| 187 | except Exception: |
| 188 | # backward compatibility (before agent.json existed) |
| 189 | try: |
| 190 | subagent = SubAgent( |
| 191 | context=files.read_file(files.get_abs_path(dir, name, "_context.md")) |
| 192 | ) |
| 193 | except Exception: |
| 194 | subagent = SubAgent() |
| 195 | |
| 196 | # non-stored fields |
| 197 | subagent.name = name |
| 198 | if "title" not in subagent.model_fields_set: |
| 199 | object.__setattr__(subagent, "title", name) |
| 200 | subagent.path = agent_dir |
| 201 | subagent.origin = [origin] |
| 202 | |
| 203 | prompts_dir = f"{dir}/{name}/prompts" |
| 204 | try: |
| 205 | prompts = files.read_text_files_in_dir(prompts_dir, pattern="*.md") |
| 206 | except Exception: |
| 207 | prompts = {} |
| 208 | |
| 209 | subagent.prompts = prompts or {} |
| 210 | return subagent |
| 211 | |
| 212 | |
| 213 | def _read_agent_definition(dir: str, name: str) -> dict: |
| 214 | yaml_path = files.get_abs_path(dir, name, "agent.yaml") |
| 215 | if files.exists(yaml_path): |
| 216 | return yaml_helper.loads(files.read_file(yaml_path)) or {} |
| 217 | json_path = files.get_abs_path(dir, name, "agent.json") |
| 218 | if files.exists(json_path): |
| 219 | return json.loads(files.read_file(json_path)) or {} |
| 220 | raise FileNotFoundError |
| 221 | |
| 222 | |
| 223 | def _merge_agent(base: SubAgent | None, override: SubAgent | None) -> SubAgent | None: |
| 224 | if base is None: |
| 225 | return override |
| 226 | if override is None: |
| 227 | return base |
| 228 | |
| 229 | data = _merge_agent_metadata(base, override) |
| 230 | data["prompts"] = {**(base.prompts or {}), **(override.prompts or {})} |
| 231 | return SubAgent.model_validate(data) |
| 232 | |
| 233 | |
| 234 | def _merge_agent_list_item( |
| 235 | base: SubAgentListItem, override: SubAgentListItem |
| 236 | ) -> SubAgentListItem: |
| 237 | return SubAgentListItem.model_validate(_merge_agent_metadata(base, override)) |
| 238 | |
| 239 | |
| 240 | def _merge_agent_metadata( |
| 241 | base: SubAgentListItem, override: SubAgentListItem |
| 242 | ) -> dict: |
| 243 | data = base.model_dump() |
| 244 | data.update( |
| 245 | override.model_dump( |
| 246 | exclude_unset=True, exclude={"name", "path", "origin", "prompts"} |
| 247 | ) |
| 248 | ) |
| 249 | data.update( |
| 250 | name=override.name or base.name, |
| 251 | path=override.path or base.path, |
| 252 | origin=[*base.origin, *override.origin], |
| 253 | ) |
| 254 | return data |
| 255 | |
| 256 | |
| 257 | def get_agents_roots() -> list[str]: |
| 258 | # from helpers import plugins |
| 259 | |
| 260 | plugin_agents = plugins.get_enabled_plugin_paths(None, "agents") |
| 261 | project_agents = files.find_existing_paths_by_pattern("usr/projects/*/.a0proj/agents") |
| 262 | paths = [ |
| 263 | files.get_abs_path(DEFAULT_AGENTS_DIR), |
| 264 | *plugin_agents, |
| 265 | files.get_abs_path(USER_AGENTS_DIR), |
| 266 | *project_agents, |
| 267 | ] |
| 268 | unique: list[str] = [] |
| 269 | seen = set() |
| 270 | for p in paths: |
| 271 | if not p: |
| 272 | continue |
| 273 | key = str(p) |
| 274 | if key in seen: |
| 275 | continue |
| 276 | seen.add(key) |
| 277 | if os.path.exists(p): |
| 278 | unique.append(p) |
| 279 | return unique |
| 280 | |
| 281 | |
| 282 | def get_all_agents_list() -> list[dict[str, str]]: |
| 283 | def _origin_from_root(root: str) -> Origin: |
| 284 | rel = files.deabsolute_path(root).replace("\\", "/") |
| 285 | if rel.startswith("usr/projects/"): |
| 286 | return "project" |
| 287 | if rel.startswith("usr/agents"): |
| 288 | return "user" |
| 289 | if "/plugins/" in rel or rel.startswith("plugins/"): |
| 290 | return "plugin" |
| 291 | return "default" |
| 292 | |
| 293 | merged: dict[str, SubAgentListItem] = {} |
| 294 | for root in get_agents_roots(): |
| 295 | origin = _origin_from_root(root) |
| 296 | items = _get_agents_list_from_dir(root, origin=origin) |
| 297 | for name, item in items.items(): |
| 298 | if name in merged: |
| 299 | merged[name] = _merge_agent_list_item(merged[name], item) |
| 300 | else: |
| 301 | merged[name] = item |
| 302 | |
| 303 | return [ |
| 304 | {"key": key, "label": item.title or key} |
| 305 | for key, item in sorted(merged.items()) |
| 306 | if key != "default" |
| 307 | ] |
| 308 | |
| 309 | |
| 310 | def get_default_promp_file_names() -> list[str]: |
| 311 | return files.list_files("prompts", filter="*.md") |
| 312 | |
| 313 | |
| 314 | def get_available_agents_dict( |
| 315 | project_name: str | None, |
| 316 | ) -> dict[str, SubAgentListItem]: |
| 317 | all_agents = get_agents_dict(project_name) |
| 318 | from helpers import projects |
| 319 | |
| 320 | project_settings = ( |
| 321 | projects.load_project_subagents(project_name) if project_name else {} |
| 322 | ) |
| 323 | |
| 324 | filtered_agents: dict[str, SubAgentListItem] = {} |
| 325 | for name, agent in all_agents.items(): |
| 326 | if name == "_example": |
| 327 | continue |
| 328 | if name in project_settings: |
| 329 | agent.enabled = project_settings[name]["enabled"] |
| 330 | if agent.enabled: |
| 331 | filtered_agents[name] = agent |
| 332 | return filtered_agents |
| 333 | |
| 334 | |
| 335 | def get_paths( |
| 336 | agent: "Agent|None", |
| 337 | *subpaths, |
| 338 | must_exist_completely: bool = True, |
| 339 | include_project: bool = True, |
| 340 | include_user: bool = True, |
| 341 | include_default: bool = True, |
| 342 | include_plugins: bool = True, |
| 343 | default_root: str = "", |
| 344 | ) -> list[str]: |
| 345 | """Returns list of file paths for the given agent and subpaths, searched in order of priority: |
| 346 | project/agents/, project/, usr/agents/, plugin agents/, agents/, usr/, plugins/, default.""" |
| 347 | cache_key = cache.determine_cache_key( |
| 348 | agent, |
| 349 | *subpaths, |
| 350 | must_exist_completely, |
| 351 | include_project, |
| 352 | include_user, |
| 353 | include_default, |
| 354 | include_plugins, |
| 355 | default_root, |
| 356 | ) |
| 357 | cached = cache.get(PATHS_CACHE_AREA, cache_key) |
| 358 | if cached is not None: |
| 359 | return cached |
| 360 | |
| 361 | paths: list[str] = [] |
| 362 | check_subpaths = subpaths if must_exist_completely else [] |
| 363 | profile_name = agent.config.profile if agent and agent.config.profile else "" |
| 364 | project_name = "" |
| 365 | |
| 366 | if include_project and agent: |
| 367 | from helpers import projects |
| 368 | |
| 369 | project_name = projects.get_context_project_name(agent.context) or "" |
| 370 | |
| 371 | if project_name and profile_name: |
| 372 | # project/agents/<profile>/... |
| 373 | project_agent_dir = projects.get_project_meta( |
| 374 | project_name, "agents", profile_name |
| 375 | ) |
| 376 | if files.exists(files.get_abs_path(project_agent_dir, *check_subpaths)): |
| 377 | paths.append(files.get_abs_path(project_agent_dir, *subpaths)) |
| 378 | |
| 379 | if project_name: |
| 380 | # project/.a0proj/... |
| 381 | path = projects.get_project_meta(project_name, *subpaths) |
| 382 | if (not must_exist_completely) or files.exists(path): |
| 383 | paths.append(path) |
| 384 | |
| 385 | if profile_name: |
| 386 | |
| 387 | # usr/agents/<profile>/... |
| 388 | path = files.get_abs_path(USER_AGENTS_DIR, profile_name, *subpaths) |
| 389 | if (not must_exist_completely) or files.exists(files.get_abs_path(USER_AGENTS_DIR, profile_name, *check_subpaths)): |
| 390 | paths.append(path) |
| 391 | |
| 392 | # plugin agents/<profile>/... |
| 393 | if include_plugins: |
| 394 | # from helpers import plugins |
| 395 | for plugin_dir in plugins.get_enabled_plugin_paths(agent, "agents", profile_name): |
| 396 | path = files.get_abs_path(plugin_dir, *subpaths) |
| 397 | if (not must_exist_completely) or files.exists(files.get_abs_path(plugin_dir, *check_subpaths)): |
| 398 | paths.append(path) |
| 399 | |
| 400 | # agents/<profile>/... |
| 401 | path = files.get_abs_path(DEFAULT_AGENTS_DIR, profile_name, *subpaths) |
| 402 | if (not must_exist_completely) or files.exists(files.get_abs_path(DEFAULT_AGENTS_DIR, profile_name, *check_subpaths)): |
| 403 | paths.append(path) |
| 404 | |
| 405 | if include_user: |
| 406 | # usr/... |
| 407 | path = files.get_abs_path(USER_DIR, *subpaths) |
| 408 | if (not must_exist_completely) or files.exists(path): |
| 409 | paths.append(path) |
| 410 | |
| 411 | if include_plugins: |
| 412 | # plugins/*/subpaths... |
| 413 | # from helpers import plugins |
| 414 | |
| 415 | for plugin_dir in plugins.get_enabled_plugin_paths(agent): |
| 416 | path = files.get_abs_path(plugin_dir, *subpaths) |
| 417 | if (not must_exist_completely) or files.exists(path): |
| 418 | if path not in paths: |
| 419 | paths.append(path) |
| 420 | |
| 421 | if include_default: |
| 422 | # default_root/... |
| 423 | path = files.get_abs_path(default_root, *subpaths) |
| 424 | if (not must_exist_completely) or files.exists(path): |
| 425 | paths.append(path) |
| 426 | |
| 427 | cache.add(PATHS_CACHE_AREA, cache_key, paths) |
| 428 | return paths |
| 429 | |
| 430 | |
| 431 | # end-of-file imports to prevent circular imports |
| 432 | from helpers import plugins |