backstory/src/utils/agents/__init__.py

51 lines
1.6 KiB
Python

from __future__ import annotations
import importlib
import pathlib
import inspect
from typing import TypeAlias, Dict, Tuple
from . types import registry
from . base import Agent
from .. setup_logging import setup_logging
from .. import defines
logger = setup_logging(defines.logging_level)
# Type alias for Agent or any subclass
AnyAgent: TypeAlias = Agent # BaseModel covers Agent and subclasses
package_dir = pathlib.Path(__file__).parent
package_name = __name__
__all__ = [ "AnyAgent", "registry"]
class_registry: Dict[str, Tuple[str, str]] = {} # Maps class_name to (module_name, class_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} from {full_module_name}")
__all__.append(name)
except ImportError as e:
logger.error(f"Error importing {full_module_name}: {e}")
continue
except Exception as e:
logger.error(f"Error processing {full_module_name}: {e}")
raise e