mcmas.ctx

  1"""
  2mcmas.ctx.
  3"""
  4
  5import inspect
  6import typing
  7
  8import pydantic
  9
 10import mcmas
 11from mcmas import util
 12
 13LOGGER = util.get_logger(__name__)
 14
 15
 16def patch_pydantic():
 17    """
 18    Patches current and future pydantic_ai.Agent instances to
 19    have an additional `specification` attribute, which resolves
 20    to a `mcmas.ispl.Agent`.
 21    """
 22    LOGGER.warning("patching pydantic agent")
 23    pydantic_ai = get_pydantic_ai()
 24    if pydantic_ai:
 25        pydantic_ai.Agent.specification = property(
 26            lambda self: mcmas.ispl.Agent(**agent_builder[self])
 27        )
 28        oinit = pydantic_ai.Agent.__init__
 29        otool = pydantic_ai.Agent.tool
 30
 31        def agent_data(obj):
 32            return {
 33                "actions": ["none"],
 34                "metadata": {"file": inspect.getfile(obj.__class__)},
 35            }
 36
 37        def init(self, *args, **kwargs):
 38            LOGGER.warning("running patched init")
 39            oinit(self, *args, **kwargs)
 40            agent_builder.add_tag(self, **agent_data(self))
 41
 42        def tool(self, fxn, *args, **kwargs):
 43            out = otool(self, fxn, *args, **kwargs)
 44            agent_builder.add_tag(
 45                self,
 46                **{
 47                    **agent_data(self),
 48                    **{
 49                        "actions": agent_builder[self].get("actions", [])
 50                        + [f"{fxn.__name__}"],
 51                    },
 52                },
 53            )
 54            # raise Exception([args,kwargs])
 55            return out
 56
 57        LOGGER.warning(f"patching class: {pydantic_ai.Agent}")
 58        pydantic_ai.Agent.__init__ = init
 59        pydantic_ai.Agent.tool = tool
 60
 61        for obj in util.find_instances(pydantic_ai.Agent):
 62            LOGGER.warning(f"patching {obj}")
 63            obj.__init__ = init
 64            obj.tool = tool
 65            agent_builder.add_tag(obj)
 66            agent_builder.add_tag(
 67                obj,
 68                **{
 69                    **agent_data(obj),
 70                    **{
 71                        "actions": agent_builder[obj].get("actions", [])
 72                        + list(obj._function_toolset.tools.keys()),
 73                    },
 74                },
 75            )
 76
 77
 78class ModelBuilder:
 79    """
 80    
 81    """
 82
 83    def __init__(self):
 84        self.TAGS = {}
 85
 86    @property
 87    def wrapping(self):
 88        from mcmas.models import Agent
 89
 90        return Agent
 91
 92    def __call__(self, obj: typing.Any, **data) -> typing.Any:
 93        """
 94        
 95        """
 96        LOGGER.warning(f"updating {obj} with {data}")
 97        key = self.key_for(obj)
 98        orig = self.TAGS.get(key, {})
 99        orig.update(**data)
100        self.TAGS[key] = orig
101        return self.wrapping(**orig)
102
103    @pydantic.validate_call
104    def key_for(self, obj: typing.Any) -> str:
105        """
106        
107        """
108        skey = getattr(obj, "name", str(obj))
109        key = f"{obj.__class__.__module__}.{obj.__class__.__name__}.{skey}"
110        return key
111
112    @pydantic.validate_call
113    def __getitem__(self, other) -> typing.Dict:
114        return self.TAGS.get(self.key_for(other), {})
115
116    @pydantic.validate_call
117    def __setitem__(self, obj, val: typing.Dict) -> typing.Any:
118        return self.add_tag(self.key_for(obj), **val)
119
120    def add_tag(self, obj, **tags) -> typing.Any:
121        """
122        
123        """
124        LOGGER.warning(f"tagging {obj}")
125        # skey = getattr(obj, 'name', str(obj))
126        key = self.key_for(obj)
127        # f"{obj.__class__.__module__}.{obj.__class__.__name__}.{skey}"
128        _tags = self.TAGS.get(key, {})
129        _tags.update(**tags)
130        self.TAGS[key] = _tags
131        # .setdefault(key, {}).update(**tag)
132        LOGGER.warning(f"tagged {obj} with {_tags}")
133
134        return obj
135
136
137agent_builder = ModelBuilder()
138
139
140def get_pydantic_ai():
141    try:
142        import pydantic_ai
143
144        return pydantic_ai
145    except (ImportError,):
146        LOGGER.critical(
147            "`pydantic_ai` module not available.  "
148            "pip install py-mcmas[ai] or pydantic-ai"
149        )
150        return None