agent_data_isolation/agent_isolation_plugins.py at main · weg-9000/agent_data_isolation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
from abc import ABC, abstractmethod
from typing import Dict, List, Any, Optional
from isolation_manager import UniversalIsolationManager

class AgentIsolationPlugin(ABC):
    def __init__(self, isolation_manager: UniversalIsolationManager):
        self.isolation_manager = isolation_manager
        self.plugin_name = self.__class__.__name__

    @abstractmethod
    def isolate_agent_data(self, data: Any, context: str = "") -> Any:
        pass

    @abstractmethod
    def get_supported_agent_types(self) -> List[str]:
        pass

class TextProcessingAgentPlugin(AgentIsolationPlugin):
    def isolate_agent_data(self, data: Any, context: str = "") -> Any:
        if isinstance(data, dict):
            return self._isolate_text_dict(data, context)
        elif isinstance(data, list):
            return self.isolation_manager.filter_contaminated_data(data, context)
        elif isinstance(data, str):
            return data if not self.isolation_manager.is_contaminated(data, context) else ""
        return data

    def _isolate_text_dict(self, data: dict, context: str) -> dict:
        clean_data = {}
        for key, value in data.items():
            if not self.isolation_manager.is_contaminated(value, f"{context}.{key}"):
                clean_data[key] = value

        clean_data["_isolation_metadata"] = {
            "plugin": self.plugin_name,
            "context": context,
            "isolation_applied": True
        }
        return clean_data

    def get_supported_agent_types(self) -> List[str]:
        return ["TextAgent", "ContentAgent", "WritingAgent", "AnalysisAgent"]

class ImageProcessingAgentPlugin(AgentIsolationPlugin):
    def isolate_agent_data(self, data: Any, context: str = "") -> Any:
        if isinstance(data, dict):
            return self._isolate_image_dict(data, context)
        elif isinstance(data, list):
            return self._isolate_image_list(data, context)
        return data

    def _isolate_image_dict(self, data: dict, context: str) -> dict:
        clean_data = {}

        for key, value in data.items():
            if key.lower() in ['image_url', 'image_urls', 'images']:
                clean_data[key] = self._filter_trusted_urls(value)
            elif not self.isolation_manager.is_contaminated(value, f"{context}.{key}"):
                clean_data[key] = value

        return clean_data

    def _isolate_image_list(self, data: list, context: str) -> list:
        clean_images = []
        max_items = self.isolation_manager.config_manager.config.max_items_per_section

        for i, item in enumerate(data[:max_items]):
            if not self.isolation_manager.is_contaminated(item, f"{context}[{i}]"):
                clean_images.append(item)

        return clean_images

    def _filter_trusted_urls(self, urls: Any) -> Any:
        trusted_domains = self.isolation_manager.config_manager.config.trusted_domains

        if isinstance(urls, str):
            return urls if any(domain in urls for domain in trusted_domains) else ""
        elif isinstance(urls, list):
            return [url for url in urls if any(domain in url for domain in trusted_domains)]

        return urls

    def get_supported_agent_types(self) -> List[str]:
        return ["ImageAgent", "LayoutAgent", "DesignAgent", "MediaAgent"]

class DataProcessingAgentPlugin(AgentIsolationPlugin):
    def isolate_agent_data(self, data: Any, context: str = "") -> Any:
        if isinstance(data, dict):
            return self._isolate_data_dict(data, context)
        elif isinstance(data, list):
            return self._isolate_data_list(data, context)
        return data

    def _isolate_data_dict(self, data: dict, context: str) -> dict:
        clean_data = {}

        for key, value in data.items():
            if key.startswith('_') or key.lower() in ['metadata', 'config', 'settings']:
                if not self.isolation_manager.is_contaminated(value, f"{context}.{key}"):
                    clean_data[key] = value
            else:
                clean_data[key] = self.isolate_agent_data(value, f"{context}.{key}")

        return clean_data

    def _isolate_data_list(self, data: list, context: str) -> list:
        return [
            self.isolate_agent_data(item, f"{context}[{i}]")
            for i, item in enumerate(data)
            if not self.isolation_manager.is_contaminated(item, f"{context}[{i}]")
        ]

    def get_supported_agent_types(self) -> List[str]:
        return ["DataAgent", "ProcessingAgent", "TransformAgent", "FilterAgent"]

class PluginManager:
    def __init__(self, isolation_manager: UniversalIsolationManager):
        self.isolation_manager = isolation_manager
        self.plugins = {}
        self._register_default_plugins()

    def _register_default_plugins(self):
        default_plugins = [
            TextProcessingAgentPlugin(self.isolation_manager),
            ImageProcessingAgentPlugin(self.isolation_manager),
            DataProcessingAgentPlugin(self.isolation_manager)
        ]

        for plugin in default_plugins:
            self.register_plugin(plugin)

    def register_plugin(self, plugin: AgentIsolationPlugin):
        for agent_type in plugin.get_supported_agent_types():
            self.plugins[agent_type] = plugin

    def get_plugin_for_agent(self, agent_type: str) -> Optional[AgentIsolationPlugin]:
        return self.plugins.get(agent_type)

    def isolate_agent_data(self, agent_type: str, data: Any, context: str = "") -> Any:
        plugin = self.get_plugin_for_agent(agent_type)
        if plugin:
            return plugin.isolate_agent_data(data, context)
        else:
            return self.isolation_manager.filter_contaminated_data(data, context) if isinstance(data, list) else data