52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
from __future__ import annotations
|
|
from typing import TypeAlias, Dict, Tuple, Optional
|
|
import importlib
|
|
import pathlib
|
|
import inspect
|
|
|
|
from . types import registry
|
|
from .. setup_logging import setup_logging
|
|
from .. import defines
|
|
from . base import Agent
|
|
|
|
logger = setup_logging(defines.logging_level)
|
|
|
|
__all__ = [ "AnyAgent", "Agent", "registry", "class_registry" ]
|
|
|
|
# Type alias for Agent or any subclass
|
|
AnyAgent: TypeAlias = Agent # BaseModel covers Agent and subclasses
|
|
|
|
class_registry: Dict[str, Tuple[str, str]] = {} # Maps class_name to (module_name, class_name)
|
|
|
|
package_dir = pathlib.Path(__file__).parent
|
|
package_name = __name__
|
|
|
|
for path in package_dir.glob("*.py"):
|
|
if path.name in ("__init__.py", "base.py") or path.name.startswith("_"):
|
|
continue
|
|
|
|
module_name = path.stem
|
|
full_module_name = f"{package_name}.{module_name}"
|
|
|
|
try:
|
|
module = importlib.import_module(full_module_name)
|
|
|
|
# Find all Agent subclasses in the module
|
|
for name, obj in inspect.getmembers(module, inspect.isclass):
|
|
if (
|
|
issubclass(obj, AnyAgent)
|
|
and obj is not AnyAgent
|
|
and obj is not Agent
|
|
and name not in class_registry
|
|
):
|
|
class_registry[name] = (full_module_name, name)
|
|
globals()[name] = obj
|
|
logger.info(f"Adding agent: {name}")
|
|
__all__.append(name) # type: ignore
|
|
except ImportError as e:
|
|
logger.error(f"Error importing {full_module_name}: {e}")
|
|
raise e
|
|
except Exception as e:
|
|
logger.error(f"Error processing {full_module_name}: {e}")
|
|
raise e
|