Skip to content

API Reference

This page provides the auto-generated API reference for the MiniLLM library, created directly from the source code's docstrings.

Core Components

safeagent.config

Simple configuration loader with environment variable defaults.

safeagent.governance

DataGovernanceError

Bases: Exception

Exception raised when governance policies are violated.

Source code in src/safeagent/governance.py 
30
31
32
class DataGovernanceError(Exception):
    """Exception raised when governance policies are violated."""
    pass

GovernanceManager

Manages data governance policies, including encryption, auditing, retention policies, and run ID management.

Source code in src/safeagent/governance.py 
 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
class GovernanceManager:
    """
    Manages data governance policies, including encryption, auditing,
    retention policies, and run ID management.
    """

    def __init__(self, audit_log_path: str = "audit", retention_days: int = 30, audit_log_extension: str = "json"):
        self.audit_log_path = f"{audit_log_path}.{audit_log_extension}"
        self.retention_days = retention_days
        log_dir = os.path.dirname(self.audit_log_path)
        if log_dir: 
            os.makedirs(log_dir, exist_ok=True)
        open(self.audit_log_path, "a").close() 
        self.current_run_id = None

    def start_new_run(self) -> str:
        """Generates a new unique ID for a single, complete run of an orchestrator."""
        self.current_run_id = str(uuid.uuid4())
        return self.current_run_id

    def get_current_run_id(self) -> str:
        """Returns the ID for the current run, creating one if it doesn't exist."""
        if not self.current_run_id:
            return self.start_new_run()
        return self.current_run_id

    def encrypt(self, plaintext: str) -> str:
        """Encrypt sensitive data before storage."""
        return fernet.encrypt(plaintext.encode()).decode()

    def decrypt(self, token: str) -> str:
        """Decrypt sensitive data when needed."""
        return fernet.decrypt(token.encode()).decode()

    def audit(self, user_id: str, action: str, resource: str, metadata: Dict[str, Any] = None) -> None:
        """Write an audit log entry for data actions, including the current run_id."""
        entry = {
            "timestamp": time.time(),
            "run_id": self.get_current_run_id(), 
            "user_id": user_id,
            "action": action,
            "resource": resource,
            "metadata": metadata or {}
        }
        with open(self.audit_log_path, "a") as f:
            f.write(json.dumps(entry) + "\n")

    def tag_lineage(self, record: Dict[str, Any], source: str) -> Dict[str, Any]:
        """Attach lineage metadata to a record."""
        if "_lineage" not in record:
            record["_lineage"] = []
        record["_lineage"].append({
            "timestamp": time.time(),
            "source": source
        })
        return record

    def purge_old_logs(self) -> None:
        """Purge audit log entries older than retention period."""
        cutoff = time.time() - self.retention_days * 86400
        retained = []
        try:
            with open(self.audit_log_path, "r") as f:
                for line in f:
                    try:
                        entry = json.loads(line)
                        if entry.get("timestamp", 0) >= cutoff:
                            retained.append(line)
                    except json.JSONDecodeError:
                        logging.warning(f"Skipping malformed line in audit log: {line.strip()}")
                        continue 
        except FileNotFoundError:
            logging.info(f"Audit log file not found at {self.audit_log_path} during purge. No purging needed.")
            return

        with open(self.audit_log_path, "w") as f:
            f.writelines(retained)

audit(user_id, action, resource, metadata=None)

Write an audit log entry for data actions, including the current run_id.

Source code in src/safeagent/governance.py 
68
69
70
71
72
73
74
75
76
77
78
79
def audit(self, user_id: str, action: str, resource: str, metadata: Dict[str, Any] = None) -> None:
    """Write an audit log entry for data actions, including the current run_id."""
    entry = {
        "timestamp": time.time(),
        "run_id": self.get_current_run_id(), 
        "user_id": user_id,
        "action": action,
        "resource": resource,
        "metadata": metadata or {}
    }
    with open(self.audit_log_path, "a") as f:
        f.write(json.dumps(entry) + "\n")

decrypt(token)

Decrypt sensitive data when needed.

Source code in src/safeagent/governance.py 
64
65
66
def decrypt(self, token: str) -> str:
    """Decrypt sensitive data when needed."""
    return fernet.decrypt(token.encode()).decode()

encrypt(plaintext)

Encrypt sensitive data before storage.

Source code in src/safeagent/governance.py 
60
61
62
def encrypt(self, plaintext: str) -> str:
    """Encrypt sensitive data before storage."""
    return fernet.encrypt(plaintext.encode()).decode()

get_current_run_id()

Returns the ID for the current run, creating one if it doesn't exist.

Source code in src/safeagent/governance.py 
54
55
56
57
58
def get_current_run_id(self) -> str:
    """Returns the ID for the current run, creating one if it doesn't exist."""
    if not self.current_run_id:
        return self.start_new_run()
    return self.current_run_id

purge_old_logs()

Purge audit log entries older than retention period.

Source code in src/safeagent/governance.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def purge_old_logs(self) -> None:
    """Purge audit log entries older than retention period."""
    cutoff = time.time() - self.retention_days * 86400
    retained = []
    try:
        with open(self.audit_log_path, "r") as f:
            for line in f:
                try:
                    entry = json.loads(line)
                    if entry.get("timestamp", 0) >= cutoff:
                        retained.append(line)
                except json.JSONDecodeError:
                    logging.warning(f"Skipping malformed line in audit log: {line.strip()}")
                    continue 
    except FileNotFoundError:
        logging.info(f"Audit log file not found at {self.audit_log_path} during purge. No purging needed.")
        return

    with open(self.audit_log_path, "w") as f:
        f.writelines(retained)

start_new_run()

Generates a new unique ID for a single, complete run of an orchestrator.

Source code in src/safeagent/governance.py 
49
50
51
52
def start_new_run(self) -> str:
    """Generates a new unique ID for a single, complete run of an orchestrator."""
    self.current_run_id = str(uuid.uuid4())
    return self.current_run_id

tag_lineage(record, source)

Attach lineage metadata to a record.

Source code in src/safeagent/governance.py 
81
82
83
84
85
86
87
88
89
def tag_lineage(self, record: Dict[str, Any], source: str) -> Dict[str, Any]:
    """Attach lineage metadata to a record."""
    if "_lineage" not in record:
        record["_lineage"] = []
    record["_lineage"].append({
        "timestamp": time.time(),
        "source": source
    })
    return record

safeagent.llm_client

FrameworkError

Bases: Exception

Custom exception for framework-related errors.

Source code in src/safeagent/llm_client.py 
13
14
15
class FrameworkError(Exception):
    """Custom exception for framework-related errors."""
    pass

LLMClient

Thin wrapper around any LLM provider with retries, error handling, and structured JSON logging.

Source code in src/safeagent/llm_client.py 
 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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
class LLMClient:
    """Thin wrapper around any LLM provider with retries, error handling, and structured JSON logging."""

    def __init__(self, provider: str, api_key: str, model: str, base_url: str = None):
        """
        Initialize the LLM client.

        Args:
            provider (str): Name of the provider (e.g., 'openai', 'anthropic').
            api_key (str): API key or token for authentication.
            model (str): Model identifier (e.g., 'gpt-4', 'claude-3-opus').
            base_url (str, optional): Custom endpoint URL; defaults to provider-specific default.
        """
        self.provider = provider
        self.api_key = api_key
        self.model = model
        self.base_url = base_url or self._default_url()
        if requests is not None:
            self.session = requests.Session()
        else:
            class _DummySession:
                def __init__(self):
                    self.headers = {}

                def post(self, *_, **__):
                    raise FrameworkError("requests package is required for HTTP calls")

            self.session = _DummySession()
        self.session.headers.update({
            "Content-Type": "application/json"
        })
        if self.provider != "gemini":
            self.session.headers["Authorization"] = f"Bearer {self.api_key}"
        self.gov = GovernanceManager()

    def _default_url(self) -> str:
        """Return default endpoint URL based on provider."""
        if self.provider == "openai":
            return "https://api.openai.com/v1/chat/completions"
        if self.provider == "anthropic":
            return "https://api.anthropic.com/v1/complete"
        if self.provider == "gemini":
            return f"https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent?key={self.api_key}"
        raise FrameworkError(f"No default URL configured for provider '{self.provider}'")

    def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> Dict:
        """
        Call the underlying LLM API, with up to 3 retries.

        Args:
            prompt (str): The textual prompt to send to the model.
            max_tokens (int): Maximum number of tokens in the response.
            temperature (float): Sampling temperature.

        Returns:
            Dict: A dictionary containing keys 'text', 'usage', and 'metadata'.

        Raises:
            FrameworkError: If the API fails after retries.
        """
        # Encrypt the prompt before logging
        encrypted_prompt = self.gov.encrypt(prompt)
        self.gov.audit(user_id="system", action="encrypt_prompt", resource="llm_client", metadata={"prompt_enc": encrypted_prompt[:50]})
        payload = self._build_payload(prompt, max_tokens, temperature)

        # Log start of LLM call and audit
        req_id = get_request_id()
        log_entry_start = {
            "event": "llm_call_start",
            "provider": self.provider,
            "model": self.model,
            "prompt_snippet": prompt[:100],
            "request_id": req_id,
            "timestamp": time.time(),
        }
        logging.info(json.dumps(log_entry_start))
        self.gov.audit(
            user_id="system",
            action="llm_call_start",
            resource=self.provider,
            metadata={"model": self.model, "request_id": req_id},
        )

        # Attempt with exponential backoff
        for attempt in range(3):
            try:
                resp = self.session.post(self.base_url, json=payload, timeout=30)
                if resp.status_code != 200:
                    raise FrameworkError(f"LLM returned status {resp.status_code}: {resp.text}")
                data = resp.json()
                text, usage = self._parse_response(data)

                # Log end of LLM call and audit
                log_entry_end = {
                    "event": "llm_call_end",
                    "provider": self.provider,
                    "model": self.model,
                    "usage": usage,
                    "request_id": req_id,
                    "timestamp": time.time(),
                }
                logging.info(json.dumps(log_entry_end))
                self.gov.audit(
                    user_id="system",
                    action="llm_call_end",
                    resource=self.provider,
                    metadata={"model": self.model, "usage": usage, "request_id": req_id},
                )

                return {"text": text, "usage": usage, "metadata": {"provider": self.provider, "model": self.model}}

            except Exception as e:
                wait = 2 ** attempt
                logging.warning(f"LLM call failed (attempt {attempt + 1}): {e}. Retrying in {wait}s")
                time.sleep(wait)

        raise FrameworkError("LLM generate() failed after 3 attempts")

    def _build_payload(self, prompt: str, max_tokens: int, temperature: float) -> Dict:
        """Construct provider-specific payload for the API call."""
        if self.provider == "openai":
            return {
                "model": self.model,
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": max_tokens,
                "temperature": temperature
            }
        if self.provider == "anthropic":
            return {
                "model": self.model,
                "prompt": prompt,
                "max_tokens_to_sample": max_tokens,
                "temperature": temperature
            }
        if self.provider == "gemini":
            return {
                "contents": [{"parts": [{"text": prompt}]}],
                "generationConfig": {"maxOutputTokens": max_tokens, "temperature": temperature}
            }
        raise FrameworkError(f"Payload builder not implemented for '{self.provider}'")

    def _parse_response(self, data: Dict) -> (str, Dict):
        """Extract generated text and usage info from API response."""
        if self.provider == "openai":
            choice = data.get("choices", [])[0]
            return choice.get("message", {}).get("content", ""), data.get("usage", {})
        if self.provider == "anthropic":
            return data.get("completion", ""), {
                "prompt_tokens": data.get("prompt_tokens"),
                "completion_tokens": data.get("completion_tokens")
            }
        if self.provider == "gemini":
            text = (
                data.get("candidates", [{}])[0]
                .get("content", {})
                .get("parts", [{}])[0]
                .get("text", "")
            )
            usage = data.get("usageMetadata", {})
            return text, {
                "prompt_tokens": usage.get("promptTokenCount"),
                "completion_tokens": usage.get("candidatesTokenCount"),
            }
        raise FrameworkError(f"Response parser not implemented for '{self.provider}'")

__init__(provider, api_key, model, base_url=None)

Initialize the LLM client.

Parameters:

Name Type Description Default
provider str

Name of the provider (e.g., 'openai', 'anthropic').

 required
api_key str

API key or token for authentication.

 required
model str

Model identifier (e.g., 'gpt-4', 'claude-3-opus').

 required
base_url str

Custom endpoint URL; defaults to provider-specific default.

 None
Source code in src/safeagent/llm_client.py 
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
def __init__(self, provider: str, api_key: str, model: str, base_url: str = None):
    """
    Initialize the LLM client.

    Args:
        provider (str): Name of the provider (e.g., 'openai', 'anthropic').
        api_key (str): API key or token for authentication.
        model (str): Model identifier (e.g., 'gpt-4', 'claude-3-opus').
        base_url (str, optional): Custom endpoint URL; defaults to provider-specific default.
    """
    self.provider = provider
    self.api_key = api_key
    self.model = model
    self.base_url = base_url or self._default_url()
    if requests is not None:
        self.session = requests.Session()
    else:
        class _DummySession:
            def __init__(self):
                self.headers = {}

            def post(self, *_, **__):
                raise FrameworkError("requests package is required for HTTP calls")

        self.session = _DummySession()
    self.session.headers.update({
        "Content-Type": "application/json"
    })
    if self.provider != "gemini":
        self.session.headers["Authorization"] = f"Bearer {self.api_key}"
    self.gov = GovernanceManager()

generate(prompt, max_tokens=512, temperature=0.7)

Call the underlying LLM API, with up to 3 retries.

Parameters:

Name Type Description Default
prompt str

The textual prompt to send to the model.

 required
max_tokens int

Maximum number of tokens in the response.

 512
temperature float

Sampling temperature.

 0.7

Returns:

Name Type Description
Dict Dict

A dictionary containing keys 'text', 'usage', and 'metadata'.

Raises:

Type Description
FrameworkError

If the API fails after retries.
Source code in src/safeagent/llm_client.py 
 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
def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> Dict:
    """
    Call the underlying LLM API, with up to 3 retries.

    Args:
        prompt (str): The textual prompt to send to the model.
        max_tokens (int): Maximum number of tokens in the response.
        temperature (float): Sampling temperature.

    Returns:
        Dict: A dictionary containing keys 'text', 'usage', and 'metadata'.

    Raises:
        FrameworkError: If the API fails after retries.
    """
    # Encrypt the prompt before logging
    encrypted_prompt = self.gov.encrypt(prompt)
    self.gov.audit(user_id="system", action="encrypt_prompt", resource="llm_client", metadata={"prompt_enc": encrypted_prompt[:50]})
    payload = self._build_payload(prompt, max_tokens, temperature)

    # Log start of LLM call and audit
    req_id = get_request_id()
    log_entry_start = {
        "event": "llm_call_start",
        "provider": self.provider,
        "model": self.model,
        "prompt_snippet": prompt[:100],
        "request_id": req_id,
        "timestamp": time.time(),
    }
    logging.info(json.dumps(log_entry_start))
    self.gov.audit(
        user_id="system",
        action="llm_call_start",
        resource=self.provider,
        metadata={"model": self.model, "request_id": req_id},
    )

    # Attempt with exponential backoff
    for attempt in range(3):
        try:
            resp = self.session.post(self.base_url, json=payload, timeout=30)
            if resp.status_code != 200:
                raise FrameworkError(f"LLM returned status {resp.status_code}: {resp.text}")
            data = resp.json()
            text, usage = self._parse_response(data)

            # Log end of LLM call and audit
            log_entry_end = {
                "event": "llm_call_end",
                "provider": self.provider,
                "model": self.model,
                "usage": usage,
                "request_id": req_id,
                "timestamp": time.time(),
            }
            logging.info(json.dumps(log_entry_end))
            self.gov.audit(
                user_id="system",
                action="llm_call_end",
                resource=self.provider,
                metadata={"model": self.model, "usage": usage, "request_id": req_id},
            )

            return {"text": text, "usage": usage, "metadata": {"provider": self.provider, "model": self.model}}

        except Exception as e:
            wait = 2 ** attempt
            logging.warning(f"LLM call failed (attempt {attempt + 1}): {e}. Retrying in {wait}s")
            time.sleep(wait)

    raise FrameworkError("LLM generate() failed after 3 attempts")

safeagent.memory_manager

MemoryManager

Minimal key-value memory store. Supports 'inmemory' or 'redis' backends and logs each read/write. Optionally, can summarize entire memory via an LLM.

Source code in src/safeagent/memory_manager.py 
  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
class MemoryManager:
    """
    Minimal key-value memory store.
    Supports 'inmemory' or 'redis' backends and logs each read/write.
    Optionally, can summarize entire memory via an LLM.
    """

    def __init__(self, backend: str = "inmemory", redis_url: str = None):
        """
        backend: "inmemory" (default) or "redis".
        redis_url: e.g., "redis://localhost:6379" if backend="redis".
        """
        global _redis
        self.backend = backend

        if self.backend == "redis":
            if _redis is None:
                try:
                    import redis
                    _redis = redis
                except ModuleNotFoundError:
                    logging.error("Redis backend selected, but 'redis' package not found. Falling back to in-memory.")
                    self.backend = "inmemory" 
                    self.store = {}
                    return

            if _redis: 
                self.client = _redis.from_url(redis_url)
                try:
                    self.client.ping()
                    logging.info("Successfully connected to Redis.")
                except Exception as e:
                    logging.error(f"Failed to connect to Redis at {redis_url}: {e}. Falling back to in-memory.")
                    self.backend = "inmemory"
                    self.store = {}
            else:
                logging.error("Redis package not available. Falling back to in-memory.")
                self.backend = "inmemory"
                self.store = {}

        if self.backend == "inmemory":
            self.store = {} 

    def save(self, user_id: str, key: str, value: str) -> None:
        """Saves value under (user_id, key)."""
        if self.backend == "redis":
            self.client.hset(user_id, key, value)
        else:
            self.store.setdefault(user_id, {})[key] = value

        logging.info(json.dumps({
            "event": "memory_save",
            "user_id": user_id,
            "key": key,
            "request_id": get_request_id(),
            "timestamp": time.time(),
        }))

    def load(self, user_id: str, key: str) -> str:
        """Loads value for (user_id, key). Returns empty string if missing."""
        if self.backend == "redis":
            raw = self.client.hget(user_id, key)
            if isinstance(raw, bytes):
                value = raw.decode("utf-8")
            elif raw is None:
                value = ""
            else:
                value = str(raw)
        else:
            value = self.store.get(user_id, {}).get(key, "")

        logging.info(json.dumps({
            "event": "memory_load",
            "user_id": user_id,
            "key": key,
            "request_id": get_request_id(),
            "timestamp": time.time(),
        }))
        return value

    def summarize(self, user_id: str, embed_fn, llm_client, max_tokens: int = 256) -> str:
        """
        Reads all entries for user_id, concatenates them, and calls LLM to generate a summary.
        Stores the summary under key="summary" and returns it.
        """
        if self.backend == "redis":
            # Ensure proper handling if client failed to initialize or connection dropped
            try:
                all_vals = [v.decode("utf-8") for v in self.client.hvals(user_id)]
            except Exception as e:
                logging.warning(f"Could not retrieve from Redis during summarize: {e}. Using empty history.")
                all_vals = []
        else:
            all_vals = list(self.store.get(user_id, {}).values())

        full_text = "\n".join(all_vals)
        if not full_text:
            return ""

        summary_prompt = f"Summarize the following conversation history:\n\n{full_text}"
        resp = llm_client.generate(summary_prompt, max_tokens=max_tokens)
        summary = resp["text"]

        # Save summary back to memory
        self.save(user_id, "summary", summary)
        return summary

__init__(backend='inmemory', redis_url=None)

backend: "inmemory" (default) or "redis". redis_url: e.g., "redis://localhost:6379" if backend="redis".

Source code in src/safeagent/memory_manager.py 
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
def __init__(self, backend: str = "inmemory", redis_url: str = None):
    """
    backend: "inmemory" (default) or "redis".
    redis_url: e.g., "redis://localhost:6379" if backend="redis".
    """
    global _redis
    self.backend = backend

    if self.backend == "redis":
        if _redis is None:
            try:
                import redis
                _redis = redis
            except ModuleNotFoundError:
                logging.error("Redis backend selected, but 'redis' package not found. Falling back to in-memory.")
                self.backend = "inmemory" 
                self.store = {}
                return

        if _redis: 
            self.client = _redis.from_url(redis_url)
            try:
                self.client.ping()
                logging.info("Successfully connected to Redis.")
            except Exception as e:
                logging.error(f"Failed to connect to Redis at {redis_url}: {e}. Falling back to in-memory.")
                self.backend = "inmemory"
                self.store = {}
        else:
            logging.error("Redis package not available. Falling back to in-memory.")
            self.backend = "inmemory"
            self.store = {}

    if self.backend == "inmemory":
        self.store = {}

load(user_id, key)

Loads value for (user_id, key). Returns empty string if missing.

Source code in src/safeagent/memory_manager.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def load(self, user_id: str, key: str) -> str:
    """Loads value for (user_id, key). Returns empty string if missing."""
    if self.backend == "redis":
        raw = self.client.hget(user_id, key)
        if isinstance(raw, bytes):
            value = raw.decode("utf-8")
        elif raw is None:
            value = ""
        else:
            value = str(raw)
    else:
        value = self.store.get(user_id, {}).get(key, "")

    logging.info(json.dumps({
        "event": "memory_load",
        "user_id": user_id,
        "key": key,
        "request_id": get_request_id(),
        "timestamp": time.time(),
    }))
    return value

save(user_id, key, value)

Saves value under (user_id, key).

Source code in src/safeagent/memory_manager.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def save(self, user_id: str, key: str, value: str) -> None:
    """Saves value under (user_id, key)."""
    if self.backend == "redis":
        self.client.hset(user_id, key, value)
    else:
        self.store.setdefault(user_id, {})[key] = value

    logging.info(json.dumps({
        "event": "memory_save",
        "user_id": user_id,
        "key": key,
        "request_id": get_request_id(),
        "timestamp": time.time(),
    }))

summarize(user_id, embed_fn, llm_client, max_tokens=256)

Reads all entries for user_id, concatenates them, and calls LLM to generate a summary. Stores the summary under key="summary" and returns it.

Source code in src/safeagent/memory_manager.py 
 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
def summarize(self, user_id: str, embed_fn, llm_client, max_tokens: int = 256) -> str:
    """
    Reads all entries for user_id, concatenates them, and calls LLM to generate a summary.
    Stores the summary under key="summary" and returns it.
    """
    if self.backend == "redis":
        # Ensure proper handling if client failed to initialize or connection dropped
        try:
            all_vals = [v.decode("utf-8") for v in self.client.hvals(user_id)]
        except Exception as e:
            logging.warning(f"Could not retrieve from Redis during summarize: {e}. Using empty history.")
            all_vals = []
    else:
        all_vals = list(self.store.get(user_id, {}).values())

    full_text = "\n".join(all_vals)
    if not full_text:
        return ""

    summary_prompt = f"Summarize the following conversation history:\n\n{full_text}"
    resp = llm_client.generate(summary_prompt, max_tokens=max_tokens)
    summary = resp["text"]

    # Save summary back to memory
    self.save(user_id, "summary", summary)
    return summary

safeagent.prompt_renderer

PromptRenderer

Jinja2-based templating engine with structured logging and lineage tagging.

Source code in src/safeagent/prompt_renderer.py 
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
class PromptRenderer:
    """Jinja2-based templating engine with structured logging and lineage tagging."""

    def __init__(self, template_dir: Path):
        """
        Args:
            template_dir (Path): Path to the directory containing Jinja2 templates.
        """
        self.env = jinja2.Environment(
            loader=jinja2.FileSystemLoader(str(template_dir)),
            autoescape=False
        )
        self.gov = GovernanceManager()

    def render(self, template_name: str, **context) -> str:
        """
        Render a Jinja2 template with provided context, logging the event and tagging lineage.

        Args:
            template_name (str): Filename of the template (e.g., 'qa_prompt.j2').
            **context: Key-value pairs to pass into the template rendering.

        Returns:
            str: The rendered template as a string.
        """
        # Audit prompt render
        lineage_metadata = {"template": template_name, "context_keys": list(context.keys())}
        self.gov.audit(user_id="system", action="prompt_render", resource=template_name, metadata=lineage_metadata)

        template = self.env.get_template(template_name)
        rendered = template.render(**context)
        log_entry = {
            "event": "prompt_render",
            "template": template_name,
            "context_keys": list(context.keys()),
            "output_length": len(rendered),
            "timestamp": time.time()
        }
        logging.info(json.dumps(log_entry))
        return rendered

__init__(template_dir)

Parameters:

Name Type Description Default
template_dir Path

Path to the directory containing Jinja2 templates.

 required
Source code in src/safeagent/prompt_renderer.py 
45
46
47
48
49
50
51
52
53
54
def __init__(self, template_dir: Path):
    """
    Args:
        template_dir (Path): Path to the directory containing Jinja2 templates.
    """
    self.env = jinja2.Environment(
        loader=jinja2.FileSystemLoader(str(template_dir)),
        autoescape=False
    )
    self.gov = GovernanceManager()

render(template_name, **context)

Render a Jinja2 template with provided context, logging the event and tagging lineage.

Parameters:

Name Type Description Default
template_name str

Filename of the template (e.g., 'qa_prompt.j2').

 required
**context

Key-value pairs to pass into the template rendering.

 {}

Returns:

Name Type Description
str str

The rendered template as a string.
Source code in src/safeagent/prompt_renderer.py 
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
def render(self, template_name: str, **context) -> str:
    """
    Render a Jinja2 template with provided context, logging the event and tagging lineage.

    Args:
        template_name (str): Filename of the template (e.g., 'qa_prompt.j2').
        **context: Key-value pairs to pass into the template rendering.

    Returns:
        str: The rendered template as a string.
    """
    # Audit prompt render
    lineage_metadata = {"template": template_name, "context_keys": list(context.keys())}
    self.gov.audit(user_id="system", action="prompt_render", resource=template_name, metadata=lineage_metadata)

    template = self.env.get_template(template_name)
    rendered = template.render(**context)
    log_entry = {
        "event": "prompt_render",
        "template": template_name,
        "context_keys": list(context.keys()),
        "output_length": len(rendered),
        "timestamp": time.time()
    }
    logging.info(json.dumps(log_entry))
    return rendered

safeagent.embeddings

EmbeddingError

Bases: Exception

Custom exception for embedding-related failures.

Source code in src/safeagent/embeddings.py 
20
21
22
class EmbeddingError(Exception):
    """Custom exception for embedding-related failures."""
    pass

gemini_embed(text, api_key, model='embedding-001')

Generates embeddings using the Google Gemini API.

This function now correctly formats the request for the embedding model, passing the API key as a URL parameter and avoiding conflicting headers.

Parameters:

Name Type Description Default
text str

The text to embed.

 required
api_key str

The Google API key.

 required
model str

The embedding model to use.

 'embedding-001'

Returns:

Type Description
Optional[List[float]]

A list of floats representing the embedding, or None on failure.

Raises:

Type Description
EmbeddingError

If the API call fails after retries.
Source code in src/safeagent/embeddings.py 
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
def gemini_embed(text: str, api_key: str, model: str = "embedding-001") -> Optional[List[float]]:
    """
    Generates embeddings using the Google Gemini API.

    This function now correctly formats the request for the embedding model,
    passing the API key as a URL parameter and avoiding conflicting headers.

    Args:
        text (str): The text to embed.
        api_key (str): The Google API key.
        model (str): The embedding model to use.

    Returns:
        A list of floats representing the embedding, or None on failure.

    Raises:
        EmbeddingError: If the API call fails after retries.
    """
    if not api_key:
        raise EmbeddingError("Gemini API key is required for embeddings.")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:embedContent?key={api_key}"

    payload = {"model": f"models/{model}", "content": {"parts": [{"text": text}]}}

    headers = {"Content-Type": "application/json"}

    try:
        resp = _session.post(url, json=payload, headers=headers, timeout=30)

        if resp.status_code != 200:
            logging.error(f"Gemini embed API request failed with status {resp.status_code}: {resp.text}")
            raise EmbeddingError(f"Gemini embed failed: {resp.text}")

        data = resp.json()
        embedding = data.get("embedding", {}).get("values")

        if not embedding:
            raise EmbeddingError("Embedding not found in Gemini API response.")

        return embedding

    except requests.exceptions.RequestException as e:
        logging.error(f"A network error occurred while calling Gemini embed API: {e}")
        raise EmbeddingError(f"Network error during embedding: {e}") from e

Orchestrators

safeagent.orchestrator

SimpleOrchestrator

Minimal DAG runner: each node is a function, edges define dependencies, with audit and lineage tagging.

Source code in src/safeagent/orchestrator.py 
 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
class SimpleOrchestrator:
    """Minimal DAG runner: each node is a function, edges define dependencies, with audit and lineage tagging."""

    def __init__(self):
        # Map node name to function
        self.nodes: Dict[str, Callable[..., Any]] = {}
        # Map node name to list of dependent node names
        self.edges: Dict[str, List[str]] = {}
        self.gov = GovernanceManager()

    def add_node(self, name: str, func: Callable[..., Any]):
        """Register a function under the given node name."""
        self.nodes[name] = func
        self.edges.setdefault(name, [])

    def add_edge(self, src: str, dest: str):
        """Specify that 'dest' depends on 'src'."""
        if src not in self.nodes or dest not in self.nodes:
            raise ValueError(f"Either '{src}' or '{dest}' is not registered as a node.")
        self.edges[src].append(dest)

    def run(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute all nodes in topological order, audit pipeline start/end, and tag lineage on outputs.

        Args:
            inputs (Dict[str, Any]): Global inputs (e.g., 'user_input', 'user_id').

        Returns:
            Dict[str, Any]: Mapping of node name to its return value.
        """
        results: Dict[str, Any] = {}
        visited = set()

        # Audit pipeline start
        self.gov.audit(user_id=inputs.get("user_id", "system"), action="pipeline_start", resource="orchestrator")

        def execute(node: str):
            if node in visited:
                return results.get(node)
            visited.add(node)
            func = self.nodes[node]
            kwargs = {}
            import inspect
            params = inspect.signature(func).parameters
            for name in params:
                if name in results:
                    kwargs[name] = results[name]
                elif name.startswith("node_") and name[5:] in results:
                    kwargs[name] = results[name[5:]]
                elif name in inputs:
                    kwargs[name] = inputs[name]
            output = func(**kwargs)
            # Tag lineage on dict outputs
            if isinstance(output, dict):
                output = self.gov.tag_lineage(output, source=node)
            results[node] = output
            return output

        for node in self.nodes:
            execute(node)

        # Audit pipeline end
        self.gov.audit(user_id=inputs.get("user_id", "system"), action="pipeline_end", resource="orchestrator")

        return results

add_edge(src, dest)

Specify that 'dest' depends on 'src'.

Source code in src/safeagent/orchestrator.py 
19
20
21
22
23
def add_edge(self, src: str, dest: str):
    """Specify that 'dest' depends on 'src'."""
    if src not in self.nodes or dest not in self.nodes:
        raise ValueError(f"Either '{src}' or '{dest}' is not registered as a node.")
    self.edges[src].append(dest)

add_node(name, func)

Register a function under the given node name.

Source code in src/safeagent/orchestrator.py 
14
15
16
17
def add_node(self, name: str, func: Callable[..., Any]):
    """Register a function under the given node name."""
    self.nodes[name] = func
    self.edges.setdefault(name, [])

run(inputs)

Execute all nodes in topological order, audit pipeline start/end, and tag lineage on outputs.

Parameters:

Name Type Description Default
inputs Dict[str, Any]

Global inputs (e.g., 'user_input', 'user_id').

 required

Returns:

Type Description
Dict[str, Any]

Dict[str, Any]: Mapping of node name to its return value.
Source code in src/safeagent/orchestrator.py 
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
def run(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
    """
    Execute all nodes in topological order, audit pipeline start/end, and tag lineage on outputs.

    Args:
        inputs (Dict[str, Any]): Global inputs (e.g., 'user_input', 'user_id').

    Returns:
        Dict[str, Any]: Mapping of node name to its return value.
    """
    results: Dict[str, Any] = {}
    visited = set()

    # Audit pipeline start
    self.gov.audit(user_id=inputs.get("user_id", "system"), action="pipeline_start", resource="orchestrator")

    def execute(node: str):
        if node in visited:
            return results.get(node)
        visited.add(node)
        func = self.nodes[node]
        kwargs = {}
        import inspect
        params = inspect.signature(func).parameters
        for name in params:
            if name in results:
                kwargs[name] = results[name]
            elif name.startswith("node_") and name[5:] in results:
                kwargs[name] = results[name[5:]]
            elif name in inputs:
                kwargs[name] = inputs[name]
        output = func(**kwargs)
        # Tag lineage on dict outputs
        if isinstance(output, dict):
            output = self.gov.tag_lineage(output, source=node)
        results[node] = output
        return output

    for node in self.nodes:
        execute(node)

    # Audit pipeline end
    self.gov.audit(user_id=inputs.get("user_id", "system"), action="pipeline_end", resource="orchestrator")

    return results

safeagent.stateful_orchestrator

EdgeRegistrationError

Bases: OrchestratorError

Raised during an invalid attempt to register an edge.

Source code in src/safeagent/stateful_orchestrator.py 
19
20
21
22
23
class EdgeRegistrationError(OrchestratorError):
    """Raised during an invalid attempt to register an edge."""
    def __init__(self, node_name: str, message: str):
        self.node_name = node_name
        super().__init__("{}: '{}'".format(message, node_name))

NodeNotFoundError

Bases: OrchestratorError

Raised when a node name is not found in the graph.

Source code in src/safeagent/stateful_orchestrator.py 
13
14
15
16
17
class NodeNotFoundError(OrchestratorError):
    """Raised when a node name is not found in the graph."""
    def __init__(self, node_name: str):
        self.node_name = node_name
        super().__init__("Node '{}' not found in the graph.".format(node_name))

OrchestratorError

Bases: Exception

Base exception for all stateful orchestrator errors.

Source code in src/safeagent/stateful_orchestrator.py 
 9
10
11
class OrchestratorError(Exception):
    """Base exception for all stateful orchestrator errors."""
    pass

StateValidationError

Bases: OrchestratorError

Raised when the state does not conform to the defined schema.

Source code in src/safeagent/stateful_orchestrator.py 
25
26
27
28
class StateValidationError(OrchestratorError):
    """Raised when the state does not conform to the defined schema."""
    def __init__(self, message: str):
        super().__init__(message)

StatefulOrchestrator

An orchestrator that manages a central state object, allowing for complex, cyclical, and conditional workflows with integrated governance, human-in-the-loop interrupts, and optional state schema validation.

Source code in src/safeagent/stateful_orchestrator.py 
 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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
class StatefulOrchestrator:
    """
    An orchestrator that manages a central state object, allowing for complex,
    cyclical, and conditional workflows with integrated governance, human-in-the-loop
    interrupts, and optional state schema validation.
    """

    def __init__(self, entry_node: str, state_schema: Optional[Dict[str, Type]] = None):
        """
        Initializes the stateful orchestrator.

        Args:
            entry_node (str): The name of the first node to execute in the graph.
            state_schema (Optional[Dict[str, Type]]): An optional schema defining
                expected keys and their Python types in the state object.
        """
        if not isinstance(entry_node, str) or not entry_node:
            raise ValueError("entry_node must be a non-empty string.")

        self.nodes: Dict[str, Callable[[Dict], Dict]] = {}
        self.edges: Dict[str, Callable[[Dict], str]] = {}
        self.entry_node = entry_node
        self.state_schema = state_schema
        self.gov = GovernanceManager()

    def add_node(self, name: str, func: Callable[[Dict], Dict]):
        self.nodes[name] = func

    def add_edge(self, src: str, dest: str):
        if src not in self.nodes:
            raise EdgeRegistrationError(src, "Source node for edge is not registered")
        if dest not in self.nodes and dest not in ("__end__", "__interrupt__"):
             raise EdgeRegistrationError(dest, "Destination node for edge is not registered")
        self.edges[src] = lambda state: dest

    def add_conditional_edge(self, src: str, path_func: Callable[[Dict], str]):
        if src not in self.nodes:
            raise EdgeRegistrationError(src, "Source node for conditional edge is not registered")
        self.edges[src] = path_func

    def _validate_state(self, state: Dict[str, Any], keys_to_check: List[str]):
        """Validates a subset of the state against the schema if it exists."""
        if not self.state_schema:
            return

        for key in keys_to_check:
            if key not in self.state_schema:
                raise StateValidationError("Key '{}' in state is not defined in the schema.".format(key))
            if key in state and not isinstance(state[key], self.state_schema[key]):
                expected_type = self.state_schema[key].__name__
                actual_type = type(state[key]).__name__
                msg = "Type mismatch for key '{}'. Expected '{}', got '{}'.".format(key, expected_type, actual_type)
                raise StateValidationError(msg)

    def run(self, inputs: Dict[str, Any], user_id: str = "system", max_steps: int = 15) -> Tuple[str, Dict[str, Any]]:
        """
        Executes the graph starting from the entry node.

        Returns:
            A tuple containing the final status ('completed', 'paused', 'error')
            and the final state of the graph.
        """
        state = inputs.copy()
        self._validate_state(state, list(state.keys()))
        self.gov.audit(user_id, "stateful_run_start", "StatefulOrchestrator", {"initial_keys": list(state.keys())})

        return self._execute_from(self.entry_node, state, user_id, max_steps)

    def resume(self, state: Dict[str, Any], human_input: Dict[str, Any], user_id: str = "system", max_steps: int = 15) -> Tuple[str, Dict[str, Any]]:
        """
        Resumes execution of a paused graph.
        """
        if "__next_node__" not in state:
            raise OrchestratorError("Cannot resume. The provided state is not a valid paused state.")

        next_node = state.pop("__next_node__")
        state.update(human_input)

        self.gov.audit(user_id, "graph_resume", "StatefulOrchestrator", {"resuming_at_node": next_node, "human_input_keys": list(human_input.keys())})
        self._validate_state(state, list(human_input.keys()))

        return self._execute_from(next_node, state, user_id, max_steps, start_step=state.get('__step__', 0))

    def _execute_from(self, start_node: str, state: Dict[str, Any], user_id: str, max_steps: int, start_step: int = 0) -> Tuple[str, Dict[str, Any]]:
        current_node_name = start_node

        for step in range(start_step, max_steps):
            if current_node_name == "__end__":
                self.gov.audit(user_id, "graph_end_reached", "StatefulOrchestrator", {"step": step})
                return "completed", state

            if current_node_name == "__interrupt__":
                self.gov.audit(user_id, "graph_interrupt_human_input", "StatefulOrchestrator", {"step": step})
                if state['__previous_node__'] in self.edges:
                    state["__next_node__"] = self.edges[state['__previous_node__']](state)
                    state["__step__"] = step
                return "paused", state

            if current_node_name not in self.nodes:
                raise NodeNotFoundError(current_node_name)

            self.gov.audit(user_id, "node_start", current_node_name, {"step": step})
            node_func = self.nodes[current_node_name]

            try:
                updates = node_func(state)
                self._validate_state(updates, list(updates.keys()))

                for key, value in updates.items():
                    record_to_tag = value if isinstance(value, dict) else {'value': value}
                    tagged_record = self.gov.tag_lineage(record_to_tag, source=current_node_name)
                    state[key] = tagged_record.get('value', tagged_record)

                self.gov.audit(user_id, "node_end", current_node_name, {"step": step, "updated_keys": list(updates.keys())})
            except Exception as e:
                self.gov.audit(user_id, "node_error", current_node_name, {"step": step, "error": str(e)})
                raise

            if current_node_name not in self.edges:
                self.gov.audit(user_id, "graph_path_end", "StatefulOrchestrator", {"last_node": current_node_name})
                return "completed", state

            path_func = self.edges[current_node_name]
            state["__previous_node__"] = current_node_name
            next_node_name = path_func(state)

            self.gov.audit(user_id, "conditional_edge_traversed", current_node_name, {"destination": next_node_name})
            current_node_name = next_node_name
        else:
             self.gov.audit(user_id, "max_steps_reached", "StatefulOrchestrator", {"max_steps": max_steps})
             return "max_steps_reached", state

        return "completed", state

__init__(entry_node, state_schema=None)

Initializes the stateful orchestrator.

Parameters:

Name Type Description Default
entry_node str

The name of the first node to execute in the graph.

 required
state_schema Optional[Dict[str, Type]]

An optional schema defining expected keys and their Python types in the state object.

 None
Source code in src/safeagent/stateful_orchestrator.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def __init__(self, entry_node: str, state_schema: Optional[Dict[str, Type]] = None):
    """
    Initializes the stateful orchestrator.

    Args:
        entry_node (str): The name of the first node to execute in the graph.
        state_schema (Optional[Dict[str, Type]]): An optional schema defining
            expected keys and their Python types in the state object.
    """
    if not isinstance(entry_node, str) or not entry_node:
        raise ValueError("entry_node must be a non-empty string.")

    self.nodes: Dict[str, Callable[[Dict], Dict]] = {}
    self.edges: Dict[str, Callable[[Dict], str]] = {}
    self.entry_node = entry_node
    self.state_schema = state_schema
    self.gov = GovernanceManager()

resume(state, human_input, user_id='system', max_steps=15)

Resumes execution of a paused graph.

Source code in src/safeagent/stateful_orchestrator.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def resume(self, state: Dict[str, Any], human_input: Dict[str, Any], user_id: str = "system", max_steps: int = 15) -> Tuple[str, Dict[str, Any]]:
    """
    Resumes execution of a paused graph.
    """
    if "__next_node__" not in state:
        raise OrchestratorError("Cannot resume. The provided state is not a valid paused state.")

    next_node = state.pop("__next_node__")
    state.update(human_input)

    self.gov.audit(user_id, "graph_resume", "StatefulOrchestrator", {"resuming_at_node": next_node, "human_input_keys": list(human_input.keys())})
    self._validate_state(state, list(human_input.keys()))

    return self._execute_from(next_node, state, user_id, max_steps, start_step=state.get('__step__', 0))

run(inputs, user_id='system', max_steps=15)

Executes the graph starting from the entry node.

Returns:

Type Description
str

A tuple containing the final status ('completed', 'paused', 'error')
Dict[str, Any]

and the final state of the graph.
Source code in src/safeagent/stateful_orchestrator.py 
86
87
88
89
90
91
92
93
94
95
96
97
98
def run(self, inputs: Dict[str, Any], user_id: str = "system", max_steps: int = 15) -> Tuple[str, Dict[str, Any]]:
    """
    Executes the graph starting from the entry node.

    Returns:
        A tuple containing the final status ('completed', 'paused', 'error')
        and the final state of the graph.
    """
    state = inputs.copy()
    self._validate_state(state, list(state.keys()))
    self.gov.audit(user_id, "stateful_run_start", "StatefulOrchestrator", {"initial_keys": list(state.keys())})

    return self._execute_from(self.entry_node, state, user_id, max_steps)

Tooling

safeagent.tool_registry

AccessManager

A centralized class to manage Role-Based Access Control (RBAC).

This class can be initialized with a user role mapping, or it will use a default set of roles for demonstration purposes.

Source code in src/safeagent/tool_registry.py 
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
class AccessManager:
    """
    A centralized class to manage Role-Based Access Control (RBAC).

    This class can be initialized with a user role mapping, or it will
    use a default set of roles for demonstration purposes.
    """
    def __init__(self, role_config: Optional[Dict[str, List[str]]] = None):
        """
        Initializes the AccessManager.

        Args:
            role_config: A dictionary mapping user IDs to a list of their roles.
                         If None, a default demo configuration is used.
        """
        if role_config is not None:
            self._user_role_database = role_config
        else:
            self._user_role_database = {
                "billing_user_01": ["billing_agent", "support"],
                "weather_analyst_7": ["weather_forecaster"],
                "data_auditor_3": ["readonly_viewer", "guest_access"]
            }

    def check_access(self, user_id: str, required_role: str) -> bool:
        """
        Checks if a user has a required role by looking them up in the
        internal role database.
        """
        current_user_roles = self._user_role_database.get(user_id, [])
        return required_role in current_user_roles

__init__(role_config=None)

Initializes the AccessManager.

Parameters:

Name Type Description Default
role_config Optional[Dict[str, List[str]]]

A dictionary mapping user IDs to a list of their roles. If None, a default demo configuration is used.

 None
Source code in src/safeagent/tool_registry.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def __init__(self, role_config: Optional[Dict[str, List[str]]] = None):
    """
    Initializes the AccessManager.

    Args:
        role_config: A dictionary mapping user IDs to a list of their roles.
                     If None, a default demo configuration is used.
    """
    if role_config is not None:
        self._user_role_database = role_config
    else:
        self._user_role_database = {
            "billing_user_01": ["billing_agent", "support"],
            "weather_analyst_7": ["weather_forecaster"],
            "data_auditor_3": ["readonly_viewer", "guest_access"]
        }

check_access(user_id, required_role)

Checks if a user has a required role by looking them up in the internal role database.

Source code in src/safeagent/tool_registry.py 
40
41
42
43
44
45
46
def check_access(self, user_id: str, required_role: str) -> bool:
    """
    Checks if a user has a required role by looking them up in the
    internal role database.
    """
    current_user_roles = self._user_role_database.get(user_id, [])
    return required_role in current_user_roles

SimilarityMetric

Bases: Enum

Specifies the similarity metric for vector search.

Source code in src/safeagent/tool_registry.py 
76
77
78
79
80
class SimilarityMetric(Enum):
    """Specifies the similarity metric for vector search."""
    L2 = "l2"
    COSINE = "cosine"
    DOT_PRODUCT = "dot_product"

ToolExecutionError

Bases: ToolRegistryError

Raised when a tool fails to execute after all retries.

Source code in src/safeagent/tool_registry.py 
91
92
93
class ToolExecutionError(ToolRegistryError):
    """Raised when a tool fails to execute after all retries."""
    pass

ToolNotFoundError

Bases: ToolRegistryError

Raised when a tool is not found in the registry.

Source code in src/safeagent/tool_registry.py 
87
88
89
class ToolNotFoundError(ToolRegistryError):
    """Raised when a tool is not found in the registry."""
    pass

ToolRegistry

A central, governed registry for tools that includes RBAC, automatic retries, circuit breakers, cost/latency tracking, caching, async support, output sinks, and dynamic schemas.

Source code in src/safeagent/tool_registry.py 
 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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
class ToolRegistry:
    """
    A central, governed registry for tools that includes RBAC, automatic retries,
    circuit breakers, cost/latency tracking, caching, async support, output sinks,
    and dynamic schemas.
    """
    def __init__(
        self,
        governance_manager: GovernanceManager,
        access_manager: Optional[AccessManager] = None,
        embedding_config: Optional[Dict] = None,
        similarity_metric: SimilarityMetric = SimilarityMetric.L2,
        embedding_dimension: int = 768
    ):
        self._tools: Dict[str, Callable] = {}
        self._tool_metadata: Dict[str, Dict] = {}
        self.gov = governance_manager
        self.access_manager = access_manager or AccessManager()
        self.embedding_config = embedding_config or {}
        self.similarity_metric = similarity_metric
        self.embedding_dimension = embedding_dimension
        self._circuit_breaker_state: Dict[str, Dict] = {}
        self._cache: Dict[str, Dict] = {}  # In-memory cache

        self._tool_index = None
        self._index_to_tool_name: Dict[int, str] = {}
        if _EMBEDDINGS_ENABLED:
            self._initialize_faiss_index()

    def _initialize_faiss_index(self):
        """Initializes the correct FAISS index based on the chosen similarity metric."""
        if self.similarity_metric == SimilarityMetric.L2:
            self._tool_index = faiss.IndexFlatL2(self.embedding_dimension)
        elif self.similarity_metric in (SimilarityMetric.COSINE, SimilarityMetric.DOT_PRODUCT):
            self._tool_index = faiss.IndexFlatIP(self.embedding_dimension)
        else:
            raise ValueError("Unsupported similarity metric: {}".format(self.similarity_metric))

    def _index_tool(self, tool_name: str):
        """Embeds and indexes a tool's description for semantic search."""
        if not _EMBEDDINGS_ENABLED or self._tool_index is None: return
        metadata = self._tool_metadata.get(tool_name, {})
        description = "Tool: {}. Description: {}".format(tool_name, metadata.get("docstring", ""))
        api_key = self.embedding_config.get("api_key", "")
        vector = gemini_embed(text=description, api_key=api_key)
        if vector:
            vector_np = np.array([vector], dtype=np.float32)
            if self.similarity_metric == SimilarityMetric.COSINE:
                faiss.normalize_L2(vector_np)
            new_index_id = self._tool_index.ntotal
            self._tool_index.add(vector_np)
            self._index_to_tool_name[new_index_id] = tool_name

    def register(
        self,
        required_role: Optional[str] = None,
        retry_attempts: int = 0,
        retry_delay: float = 1.0,
        circuit_breaker_threshold: int = 0,
        cache_ttl_seconds: int = 0,
        cost_per_call: Optional[float] = None,
        cost_calculator: Optional[Callable[[Any], float]] = None,
        output_sinks: Optional[List[BaseOutputSink]] = None
    ) -> Callable:
        """A decorator to register a tool with advanced, governed execution policies."""
        def decorator(func: Callable) -> Callable:
            tool_name = func.__name__
            self._tools[tool_name] = func
            self._tool_metadata[tool_name] = {
                "docstring": inspect.getdoc(func),
                "signature": inspect.signature(func),
                "is_async": inspect.iscoroutinefunction(func),
                "policies": {
                    "role": required_role, "retry_attempts": retry_attempts,
                    "retry_delay": retry_delay, "circuit_breaker_threshold": circuit_breaker_threshold,
                    "cache_ttl_seconds": cache_ttl_seconds, "cost_per_call": cost_per_call,
                    "cost_calculator": cost_calculator, "output_sinks": output_sinks or []
                }
            }
            self._circuit_breaker_state[tool_name] = {'failure_count': 0, 'is_open': False, 'opened_at': 0}
            self._index_tool(tool_name)
            return func
        return decorator

    def _create_cache_key(self, tool_name: str, **kwargs) -> str:
        """Creates a stable cache key from the tool name and arguments."""
        hasher = hashlib.md5()
        encoded = json.dumps(kwargs, sort_keys=True).encode('utf-8')
        hasher.update(encoded)
        return "{}:{}".format(tool_name, hasher.hexdigest())

    def _check_pre_execution_policies(self, name: str, user_id: str, policies: Dict, **kwargs) -> Optional[Any]:
        """Handles caching, circuit breaker, and RBAC checks. Returns cached result if hit."""
        # Caching
        if policies["cache_ttl_seconds"] > 0:
            cache_key = self._create_cache_key(name, **kwargs)
            if cache_key in self._cache:
                cached_item = self._cache[cache_key]
                if time.time() - cached_item["timestamp"] < policies["cache_ttl_seconds"]:
                    self.gov.audit(user_id, "tool_cache_hit", name, {"args": kwargs})
                    return cached_item["result"]

        # Circuit Breaker
        cb_state = self._circuit_breaker_state[name]
        if cb_state['is_open']:
            if time.time() - cb_state['opened_at'] > 60:  # 1-minute cooldown
                cb_state['is_open'] = False
            else:
                msg = "Circuit breaker for tool '{}' is open.".format(name)
                self.gov.audit(user_id, "tool_circuit_breaker_open", name, {"error": msg})
                raise ToolExecutionError(msg)

        # RBAC
        if policies["role"] and not self.access_manager.check_access(user_id, policies["role"]):
            msg = "User '{}' lacks required role '{}' for tool '{}'.".format(user_id, policies["role"], name)
            self.gov.audit(user_id, "tool_access_denied", name, {"required_role": policies["role"]})
            raise RBACError(msg)

        return None

    def _handle_post_execution(self, name: str, user_id: str, policies: Dict, result: Any, latency_ms: float, **kwargs):
        """Handles auditing, cost calculation, caching, and output sinks after successful execution."""
        cost = policies["cost_per_call"]
        if policies["cost_calculator"]:
            cost = policies["cost_calculator"](result)

        audit_metadata = {"result_type": type(result).__name__, "latency_ms": round(latency_ms), "cost": cost}
        self.gov.audit(user_id, "tool_call_end", name, audit_metadata)

        if policies["cache_ttl_seconds"] > 0:
            cache_key = self._create_cache_key(name, **kwargs)
            self._cache[cache_key] = {"timestamp": time.time(), "result": result}

        run_id = self.gov.get_current_run_id()
        for sink in policies["output_sinks"]:
            try:
                sink_metadata = sink.handle(name, result, run_id, **kwargs)
                self.gov.audit(user_id, "output_sink_success", str(sink), {"tool_name": name, **sink_metadata})
            except Exception as e:
                self.gov.audit(user_id, "output_sink_failure", str(sink), {"tool_name": name, "error": str(e)})

    def _handle_execution_error(self, name: str, user_id: str, policies: Dict, e: Exception, attempt: int):
        """Handles failures, including retry logic and circuit breaker trips."""
        self.gov.audit(user_id, "tool_call_error", name, {"error": str(e), "attempt": attempt + 1})
        if attempt >= policies["retry_attempts"]:
            cb_state = self._circuit_breaker_state[name]
            cb_state['failure_count'] += 1
            if policies["circuit_breaker_threshold"] > 0 and cb_state['failure_count'] >= policies["circuit_breaker_threshold"]:
                cb_state['is_open'] = True
                cb_state['opened_at'] = time.time()
                self.gov.audit(user_id, "tool_circuit_breaker_tripped", name)
            raise ToolExecutionError("Tool '{}' failed after all retry attempts.".format(name)) from e

    def _get_governed_sync_tool(self, name: str, user_id: str, original_func: Callable, policies: Dict) -> Callable:
        """Returns the fully governed wrapper for a synchronous tool."""
        def sync_wrapper(**kwargs):
            cached_result = self._check_pre_execution_policies(name, user_id, policies, **kwargs)
            if cached_result is not None: return cached_result

            for attempt in range(policies["retry_attempts"] + 1):
                start_time = time.monotonic()
                try:
                    self.gov.audit(user_id, "tool_call_start", name, {"args": kwargs, "attempt": attempt + 1})
                    result = original_func(**kwargs)
                    latency_ms = (time.monotonic() - start_time) * 1000
                    self._handle_post_execution(name, user_id, policies, result, latency_ms, **kwargs)
                    return result
                except Exception as e:
                    self._handle_execution_error(name, user_id, policies, e, attempt)
                    time.sleep(policies["retry_delay"] * (2 ** attempt))
            # This line should be logically unreachable if retry_attempts >= 0
            raise ToolExecutionError("Tool '{}' execution logic failed unexpectedly.".format(name))
        return sync_wrapper

    def _get_governed_async_tool(self, name: str, user_id: str, original_func: Callable, policies: Dict) -> Callable:
        """Returns the fully governed wrapper for an asynchronous tool."""
        async def async_wrapper(**kwargs):
            cached_result = self._check_pre_execution_policies(name, user_id, policies, **kwargs)
            if cached_result is not None: return cached_result

            for attempt in range(policies["retry_attempts"] + 1):
                start_time = time.monotonic()
                try:
                    self.gov.audit(user_id, "tool_call_start", name, {"args": kwargs, "attempt": attempt + 1})
                    result = await original_func(**kwargs)
                    latency_ms = (time.monotonic() - start_time) * 1000
                    self._handle_post_execution(name, user_id, policies, result, latency_ms, **kwargs)
                    return result
                except Exception as e:
                    self._handle_execution_error(name, user_id, policies, e, attempt)
                    await asyncio.sleep(policies["retry_delay"] * (2 ** attempt))
            # This line should be logically unreachable if retry_attempts >= 0
            raise ToolExecutionError("Tool '{}' execution logic failed unexpectedly.".format(name))
        return async_wrapper

    def get_governed_tool(self, name: str, user_id: str) -> Callable:
        """
        Retrieves a tool by name and wraps it in all registered governance policies.
        This method correctly handles both synchronous and asynchronous tools.
        """
        if name not in self._tools:
            raise ToolNotFoundError("Tool '{}' not found in registry.".format(name))

        metadata = self._tool_metadata[name]
        original_func = self._tools[name]
        policies = metadata["policies"]

        if metadata["is_async"]:
            return self._get_governed_async_tool(name, user_id, original_func, policies)
        else:
            return self._get_governed_sync_tool(name, user_id, original_func, policies)

    def generate_tool_schema(self, tool_names: List[str]) -> List[Dict[str, Any]]:
        """Generates a JSON Schema-like description for a list of tools."""
        schema = []
        for name in tool_names:
            if name in self._tool_metadata:
                metadata = self._tool_metadata[name]
                sig = metadata["signature"]
                properties = {}
                for param in sig.parameters.values():
                    if param.name != 'self':
                        type_map = {str: 'string', int: 'number', float: 'number', bool: 'boolean'}
                        param_type = type_map.get(param.annotation, 'string')
                        properties[param.name] = {'type': param_type, 'description': ''}
                schema.append({
                    "name": name,
                    "description": metadata["docstring"],
                    "parameters": {
                        "type": "object",
                        "properties": properties,
                        "required": [p.name for p in sig.parameters.values() if p.default == inspect.Parameter.empty and p.name != 'self']
                    }
                })
        return schema

    def get_relevant_tools(self, query: str, top_k: int = 3) -> List[str]:
        """Finds the most semantically relevant tools for a given query using a vector index."""
        if not _EMBEDDINGS_ENABLED or self._tool_index is None or self._tool_index.ntotal == 0:
            return []
        api_key = self.embedding_config.get("api_key", "")
        query_vector = gemini_embed(text=query, api_key=api_key)
        if not query_vector:
            return []
        query_np = np.array([query_vector], dtype=np.float32)
        if self.similarity_metric == SimilarityMetric.COSINE:
            faiss.normalize_L2(query_np)
        distances, indices = self._tool_index.search(query_np, min(top_k, self._tool_index.ntotal))
        return [self._index_to_tool_name[i] for i in indices[0]]

generate_tool_schema(tool_names)

Generates a JSON Schema-like description for a list of tools.

Source code in src/safeagent/tool_registry.py 
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
def generate_tool_schema(self, tool_names: List[str]) -> List[Dict[str, Any]]:
    """Generates a JSON Schema-like description for a list of tools."""
    schema = []
    for name in tool_names:
        if name in self._tool_metadata:
            metadata = self._tool_metadata[name]
            sig = metadata["signature"]
            properties = {}
            for param in sig.parameters.values():
                if param.name != 'self':
                    type_map = {str: 'string', int: 'number', float: 'number', bool: 'boolean'}
                    param_type = type_map.get(param.annotation, 'string')
                    properties[param.name] = {'type': param_type, 'description': ''}
            schema.append({
                "name": name,
                "description": metadata["docstring"],
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": [p.name for p in sig.parameters.values() if p.default == inspect.Parameter.empty and p.name != 'self']
                }
            })
    return schema

get_governed_tool(name, user_id)

Retrieves a tool by name and wraps it in all registered governance policies. This method correctly handles both synchronous and asynchronous tools.

Source code in src/safeagent/tool_registry.py 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
def get_governed_tool(self, name: str, user_id: str) -> Callable:
    """
    Retrieves a tool by name and wraps it in all registered governance policies.
    This method correctly handles both synchronous and asynchronous tools.
    """
    if name not in self._tools:
        raise ToolNotFoundError("Tool '{}' not found in registry.".format(name))

    metadata = self._tool_metadata[name]
    original_func = self._tools[name]
    policies = metadata["policies"]

    if metadata["is_async"]:
        return self._get_governed_async_tool(name, user_id, original_func, policies)
    else:
        return self._get_governed_sync_tool(name, user_id, original_func, policies)

get_relevant_tools(query, top_k=3)

Finds the most semantically relevant tools for a given query using a vector index.

Source code in src/safeagent/tool_registry.py 
332
333
334
335
336
337
338
339
340
341
342
343
344
def get_relevant_tools(self, query: str, top_k: int = 3) -> List[str]:
    """Finds the most semantically relevant tools for a given query using a vector index."""
    if not _EMBEDDINGS_ENABLED or self._tool_index is None or self._tool_index.ntotal == 0:
        return []
    api_key = self.embedding_config.get("api_key", "")
    query_vector = gemini_embed(text=query, api_key=api_key)
    if not query_vector:
        return []
    query_np = np.array([query_vector], dtype=np.float32)
    if self.similarity_metric == SimilarityMetric.COSINE:
        faiss.normalize_L2(query_np)
    distances, indices = self._tool_index.search(query_np, min(top_k, self._tool_index.ntotal))
    return [self._index_to_tool_name[i] for i in indices[0]]

register(required_role=None, retry_attempts=0, retry_delay=1.0, circuit_breaker_threshold=0, cache_ttl_seconds=0, cost_per_call=None, cost_calculator=None, output_sinks=None)

A decorator to register a tool with advanced, governed execution policies.

Source code in src/safeagent/tool_registry.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def register(
    self,
    required_role: Optional[str] = None,
    retry_attempts: int = 0,
    retry_delay: float = 1.0,
    circuit_breaker_threshold: int = 0,
    cache_ttl_seconds: int = 0,
    cost_per_call: Optional[float] = None,
    cost_calculator: Optional[Callable[[Any], float]] = None,
    output_sinks: Optional[List[BaseOutputSink]] = None
) -> Callable:
    """A decorator to register a tool with advanced, governed execution policies."""
    def decorator(func: Callable) -> Callable:
        tool_name = func.__name__
        self._tools[tool_name] = func
        self._tool_metadata[tool_name] = {
            "docstring": inspect.getdoc(func),
            "signature": inspect.signature(func),
            "is_async": inspect.iscoroutinefunction(func),
            "policies": {
                "role": required_role, "retry_attempts": retry_attempts,
                "retry_delay": retry_delay, "circuit_breaker_threshold": circuit_breaker_threshold,
                "cache_ttl_seconds": cache_ttl_seconds, "cost_per_call": cost_per_call,
                "cost_calculator": cost_calculator, "output_sinks": output_sinks or []
            }
        }
        self._circuit_breaker_state[tool_name] = {'failure_count': 0, 'is_open': False, 'opened_at': 0}
        self._index_tool(tool_name)
        return func
    return decorator

ToolRegistryError

Bases: Exception

Base class for tool registry exceptions.

Source code in src/safeagent/tool_registry.py 
83
84
85
class ToolRegistryError(Exception):
    """Base class for tool registry exceptions."""
    pass

Retrievers

safeagent.retriever

BaseRetriever

Base interface for retrieval. Requires implementing index and query.

Source code in src/safeagent/retriever.py 
59
60
61
62
63
64
65
class BaseRetriever:
    """Base interface for retrieval. Requires implementing index and query."""
    def index(self, embeddings: List[Any], metadata: List[Dict[str, Any]]) -> None:
        raise NotImplementedError

    def query(self, query_text: str, top_k: int = 5) -> List[Dict[str, Any]]:
        raise NotImplementedError

GraphRetriever

Bases: BaseRetriever

Neo4j-backed GraphRAG retriever using GDS k-NN, with governance integration.

Source code in src/safeagent/retriever.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
class GraphRetriever(BaseRetriever):
    """Neo4j-backed GraphRAG retriever using GDS k-NN, with governance integration."""

    def __init__(self, neo4j_uri: str, user: str, password: str, gds_graph_name: str, embed_model_fn):
        """Create the retriever. If neo4j_uri is falsy, the retriever is disabled."""
        self.driver = None
        self.gov = GovernanceManager()
        self.embed = embed_model_fn
        self.gds_graph = gds_graph_name

        if not neo4j_uri:
            logging.info("GraphRetriever is disabled because no neo4j_uri was provided.")
            return

        try:
            from neo4j import GraphDatabase, exceptions
            self.driver = GraphDatabase.driver(neo4j_uri, auth=(user, password))
            # Test the connection to fail fast
            with self.driver.session() as session:
                session.run("RETURN 1")
            logging.info("Successfully connected to Neo4j.")
        except ImportError:
            logging.warning("The 'neo4j' library is not installed. GraphRetriever will be disabled.")
            self.driver = None
        except exceptions.ServiceUnavailable:
            logging.warning(f"Could not connect to Neo4j at '{neo4j_uri}'. GraphRetriever is disabled.")
            self.driver = None
        except Exception as e:
            logging.warning(f"An unexpected error occurred while connecting to Neo4j. GraphRetriever is disabled. Error: {e}")
            self.driver = None


    def index(self, embeddings: List[List[float]], metadata: List[Dict[str, Any]]):
        """
        Ingest each document as a node with a 'vector' property and 'metadata' (with lineage tagging).
        """
        if not self.driver:
            return 

        self.gov.audit(user_id="system", action="graph_index", resource="neo4j", metadata={"count": len(embeddings)})
        with self.driver.session() as session:
            for vec, meta in zip(embeddings, metadata):
                tagged_meta = self.gov.tag_lineage(meta.copy(), source="graph_index")
                session.run(
                    "MERGE (d:Document {id: $id}) "
                    "SET d.vector = $vector, d.metadata = $meta",
                    id=meta["id"], vector=vec, meta=tagged_meta
                )
        log_entry = {
            "event": "graph_index",
            "count": len(embeddings),
            "timestamp": time.time()
        }
        logging.info(json.dumps(log_entry))

    def query(self, query_text: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        Compute embedding for query_text, run GDS K-NN, and return nearest documents (with lineage tagging).
        """
        if not self.driver:
            return []

        # Encrypt and audit query
        encrypted_query = self.gov.encrypt(query_text)
        self.gov.audit(user_id="system", action="graph_query", resource="neo4j", metadata={"query_enc": encrypted_query[:50], "top_k": top_k})

        vec = self.embed(query_text)
        cypher = f"""
            CALL gds.knn.stream(
                '{self.gds_graph}',
                {{
                    topK: $k,
                    nodeWeightProperty: 'vector',
                    queryVector: $vector
                }}
            ) YIELD nodeId, similarity
            RETURN gds.util.asNode(nodeId).id AS id, similarity
        """
        results = []
        try:
            with self.driver.session() as session:
                for record in session.run(cypher, vector=vec, k=top_k):
                    node_id = record["id"]
                    score = record["similarity"]
                    meta_record = session.run(
                        "MATCH (d:Document {id: $id}) RETURN d.metadata AS meta", id=node_id
                    ).single()
                    if meta_record:
                        meta = meta_record["meta"]
                        tagged_meta = self.gov.tag_lineage(meta.copy(), source="graph_query")
                        results.append({"id": node_id, "score": score, "metadata": tagged_meta})
        except Exception as e:
            logging.error(f"Error querying Neo4j GDS: {e}")
            return []

        log_entry = {
            "event": "graph_query",
            "top_k": top_k,
            "timestamp": time.time()
        }
        logging.info(json.dumps(log_entry))
        return results

__init__(neo4j_uri, user, password, gds_graph_name, embed_model_fn)

Create the retriever. If neo4j_uri is falsy, the retriever is disabled.

Source code in src/safeagent/retriever.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def __init__(self, neo4j_uri: str, user: str, password: str, gds_graph_name: str, embed_model_fn):
    """Create the retriever. If neo4j_uri is falsy, the retriever is disabled."""
    self.driver = None
    self.gov = GovernanceManager()
    self.embed = embed_model_fn
    self.gds_graph = gds_graph_name

    if not neo4j_uri:
        logging.info("GraphRetriever is disabled because no neo4j_uri was provided.")
        return

    try:
        from neo4j import GraphDatabase, exceptions
        self.driver = GraphDatabase.driver(neo4j_uri, auth=(user, password))
        # Test the connection to fail fast
        with self.driver.session() as session:
            session.run("RETURN 1")
        logging.info("Successfully connected to Neo4j.")
    except ImportError:
        logging.warning("The 'neo4j' library is not installed. GraphRetriever will be disabled.")
        self.driver = None
    except exceptions.ServiceUnavailable:
        logging.warning(f"Could not connect to Neo4j at '{neo4j_uri}'. GraphRetriever is disabled.")
        self.driver = None
    except Exception as e:
        logging.warning(f"An unexpected error occurred while connecting to Neo4j. GraphRetriever is disabled. Error: {e}")
        self.driver = None

index(embeddings, metadata)

Ingest each document as a node with a 'vector' property and 'metadata' (with lineage tagging).

Source code in src/safeagent/retriever.py 
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def index(self, embeddings: List[List[float]], metadata: List[Dict[str, Any]]):
    """
    Ingest each document as a node with a 'vector' property and 'metadata' (with lineage tagging).
    """
    if not self.driver:
        return 

    self.gov.audit(user_id="system", action="graph_index", resource="neo4j", metadata={"count": len(embeddings)})
    with self.driver.session() as session:
        for vec, meta in zip(embeddings, metadata):
            tagged_meta = self.gov.tag_lineage(meta.copy(), source="graph_index")
            session.run(
                "MERGE (d:Document {id: $id}) "
                "SET d.vector = $vector, d.metadata = $meta",
                id=meta["id"], vector=vec, meta=tagged_meta
            )
    log_entry = {
        "event": "graph_index",
        "count": len(embeddings),
        "timestamp": time.time()
    }
    logging.info(json.dumps(log_entry))

query(query_text, top_k=5)

Compute embedding for query_text, run GDS K-NN, and return nearest documents (with lineage tagging).

Source code in src/safeagent/retriever.py 
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def query(self, query_text: str, top_k: int = 5) -> List[Dict[str, Any]]:
    """
    Compute embedding for query_text, run GDS K-NN, and return nearest documents (with lineage tagging).
    """
    if not self.driver:
        return []

    # Encrypt and audit query
    encrypted_query = self.gov.encrypt(query_text)
    self.gov.audit(user_id="system", action="graph_query", resource="neo4j", metadata={"query_enc": encrypted_query[:50], "top_k": top_k})

    vec = self.embed(query_text)
    cypher = f"""
        CALL gds.knn.stream(
            '{self.gds_graph}',
            {{
                topK: $k,
                nodeWeightProperty: 'vector',
                queryVector: $vector
            }}
        ) YIELD nodeId, similarity
        RETURN gds.util.asNode(nodeId).id AS id, similarity
    """
    results = []
    try:
        with self.driver.session() as session:
            for record in session.run(cypher, vector=vec, k=top_k):
                node_id = record["id"]
                score = record["similarity"]
                meta_record = session.run(
                    "MATCH (d:Document {id: $id}) RETURN d.metadata AS meta", id=node_id
                ).single()
                if meta_record:
                    meta = meta_record["meta"]
                    tagged_meta = self.gov.tag_lineage(meta.copy(), source="graph_query")
                    results.append({"id": node_id, "score": score, "metadata": tagged_meta})
    except Exception as e:
        logging.error(f"Error querying Neo4j GDS: {e}")
        return []

    log_entry = {
        "event": "graph_query",
        "top_k": top_k,
        "timestamp": time.time()
    }
    logging.info(json.dumps(log_entry))
    return results

VectorRetriever

Bases: BaseRetriever

FAISS-backed vector retriever. Uses an embedding function to map text to vectors, with governance integration.

Source code in src/safeagent/retriever.py 
 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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class VectorRetriever(BaseRetriever):
    """FAISS-backed vector retriever. Uses an embedding function to map text to vectors, with governance integration."""
    def __init__(self, index_path: str, embed_model_fn):
        """
        Args:
            index_path (str): Filesystem path to store/load FAISS index.
            embed_model_fn (callable): Function that maps text (str) to a numpy ndarray vector.
        """
        self.embed = embed_model_fn
        self.gov = GovernanceManager()
        self.metadata_store: Dict[int, Dict[str, Any]] = {}
        self.next_id = 0
        self.index_path = index_path
        if _FAISS:
            if Path(index_path).exists():
                self._index = faiss.read_index(index_path)
            else:
                self._index = faiss.IndexFlatL2(768)
        else:
            self._index = []  # type: ignore

    def index(self, embeddings: List[np.ndarray], metadata: List[Dict[str, Any]]):
        """
        Add embeddings to the FAISS index and store metadata (with lineage tagging).

        Args:
            embeddings (List[np.ndarray]): List of vectors.
            metadata (List[Dict[str, Any]]): Corresponding metadata dicts (must include 'id').
        """
        if _FAISS:
            vectors = np.vstack(embeddings)
            self._index.add(vectors)
        else:
            for vec in embeddings:
                self._index.append(np.array(vec))
        for vec, meta in zip(embeddings, metadata):
            tagged_meta = self.gov.tag_lineage(meta.copy(), source="vector_index")
            self.metadata_store[self.next_id] = tagged_meta
            self.next_id += 1

        log_entry = {
            "event": "vector_index",
            "count": len(embeddings),
            "timestamp": time.time()
        }
        logging.info(json.dumps(log_entry))
        if _FAISS:
            faiss.write_index(self._index, self.index_path)

    def query(self, query_text: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        Perform KNN search on the FAISS index using the embedded query, with encryption and audit.

        Args:
            query_text (str): The query string.
            top_k (int): Number of nearest neighbors to return.

        Returns:
            List[Dict[str, Any]]: Each dict contains 'id', 'score', and 'metadata'.
        """
        # Encrypt and audit query
        encrypted_query = self.gov.encrypt(query_text)
        self.gov.audit(user_id="system", action="vector_query", resource="faiss", metadata={"query_enc": encrypted_query[:50], "top_k": top_k})

        vec = self.embed(query_text)
        if _FAISS:
            distances, indices = self._index.search(np.array([vec]), top_k)
            idx_list = indices[0]
            dist_list = distances[0]
        else:
            if not self._index:
                idx_list, dist_list = [], []
            else:
                def dist(a, b):
                    return sum((ai - bi) ** 2 for ai, bi in zip(a, b)) ** 0.5

                dists = [dist(v, vec) for v in self._index]
                sorted_idx = sorted(range(len(dists)), key=lambda i: dists[i])[:top_k]
                idx_list = sorted_idx
                dist_list = [dists[i] for i in sorted_idx]
        results = []
        for idx, dist in zip(idx_list, dist_list):
            meta = self.metadata_store.get(int(idx), {})
            results.append({"id": int(idx), "score": float(dist), "metadata": meta})

        log_entry = {
            "event": "vector_query",
            "top_k": top_k,
            "timestamp": time.time()
        }
        logging.info(json.dumps(log_entry))
        return results

__init__(index_path, embed_model_fn)

Parameters:

Name Type Description Default
index_path str

Filesystem path to store/load FAISS index.

 required
embed_model_fn callable

Function that maps text (str) to a numpy ndarray vector.

 required
Source code in src/safeagent/retriever.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def __init__(self, index_path: str, embed_model_fn):
    """
    Args:
        index_path (str): Filesystem path to store/load FAISS index.
        embed_model_fn (callable): Function that maps text (str) to a numpy ndarray vector.
    """
    self.embed = embed_model_fn
    self.gov = GovernanceManager()
    self.metadata_store: Dict[int, Dict[str, Any]] = {}
    self.next_id = 0
    self.index_path = index_path
    if _FAISS:
        if Path(index_path).exists():
            self._index = faiss.read_index(index_path)
        else:
            self._index = faiss.IndexFlatL2(768)
    else:
        self._index = []  # type: ignore

index(embeddings, metadata)

Add embeddings to the FAISS index and store metadata (with lineage tagging).

Parameters:

Name Type Description Default
embeddings List[ndarray]

List of vectors.

 required
metadata List[Dict[str, Any]]

Corresponding metadata dicts (must include 'id').

 required
Source code in src/safeagent/retriever.py 
 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
def index(self, embeddings: List[np.ndarray], metadata: List[Dict[str, Any]]):
    """
    Add embeddings to the FAISS index and store metadata (with lineage tagging).

    Args:
        embeddings (List[np.ndarray]): List of vectors.
        metadata (List[Dict[str, Any]]): Corresponding metadata dicts (must include 'id').
    """
    if _FAISS:
        vectors = np.vstack(embeddings)
        self._index.add(vectors)
    else:
        for vec in embeddings:
            self._index.append(np.array(vec))
    for vec, meta in zip(embeddings, metadata):
        tagged_meta = self.gov.tag_lineage(meta.copy(), source="vector_index")
        self.metadata_store[self.next_id] = tagged_meta
        self.next_id += 1

    log_entry = {
        "event": "vector_index",
        "count": len(embeddings),
        "timestamp": time.time()
    }
    logging.info(json.dumps(log_entry))
    if _FAISS:
        faiss.write_index(self._index, self.index_path)

query(query_text, top_k=5)

Perform KNN search on the FAISS index using the embedded query, with encryption and audit.

Parameters:

Name Type Description Default
query_text str

The query string.

 required
top_k int

Number of nearest neighbors to return.

 5

Returns:

Type Description
List[Dict[str, Any]]

List[Dict[str, Any]]: Each dict contains 'id', 'score', and 'metadata'.
Source code in src/safeagent/retriever.py 
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def query(self, query_text: str, top_k: int = 5) -> List[Dict[str, Any]]:
    """
    Perform KNN search on the FAISS index using the embedded query, with encryption and audit.

    Args:
        query_text (str): The query string.
        top_k (int): Number of nearest neighbors to return.

    Returns:
        List[Dict[str, Any]]: Each dict contains 'id', 'score', and 'metadata'.
    """
    # Encrypt and audit query
    encrypted_query = self.gov.encrypt(query_text)
    self.gov.audit(user_id="system", action="vector_query", resource="faiss", metadata={"query_enc": encrypted_query[:50], "top_k": top_k})

    vec = self.embed(query_text)
    if _FAISS:
        distances, indices = self._index.search(np.array([vec]), top_k)
        idx_list = indices[0]
        dist_list = distances[0]
    else:
        if not self._index:
            idx_list, dist_list = [], []
        else:
            def dist(a, b):
                return sum((ai - bi) ** 2 for ai, bi in zip(a, b)) ** 0.5

            dists = [dist(v, vec) for v in self._index]
            sorted_idx = sorted(range(len(dists)), key=lambda i: dists[i])[:top_k]
            idx_list = sorted_idx
            dist_list = [dists[i] for i in sorted_idx]
    results = []
    for idx, dist in zip(idx_list, dist_list):
        meta = self.metadata_store.get(int(idx), {})
        results.append({"id": int(idx), "score": float(dist), "metadata": meta})

    log_entry = {
        "event": "vector_query",
        "top_k": top_k,
        "timestamp": time.time()
    }
    logging.info(json.dumps(log_entry))
    return results

register_retriever(name, cls)

Register a retriever class for dynamic loading.

Source code in src/safeagent/retriever.py 
51
52
53
def register_retriever(name: str, cls):
    """Register a retriever class for dynamic loading."""
    RETRIEVER_REGISTRY[name] = cls

Protocols

safeagent.protocol_manager

PROTOCOLS

Bases: Enum

Defines the supported communication/execution protocols.

Source code in src/safeagent/protocol_manager.py 
23
24
25
26
class PROTOCOLS(Enum):
    """Defines the supported communication/execution protocols."""
    MCP = "mcp"  # Master/Controller/Program protocol
    AGENT2AGENT = "agent2agent"

ProtocolManager

Manages the selection and execution of different agent workflows (protocols). This class acts as the main entry point for running a complete agent system.

Source code in src/safeagent/protocol_manager.py 
 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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
class ProtocolManager:
    """
    Manages the selection and execution of different agent workflows (protocols).
    This class acts as the main entry point for running a complete agent system.
    """

    def __init__(self, protocol: str = None, access_manager: AccessManager = None):
        self.protocol = protocol or DEFAULT_PROTOCOL
        self.access_manager = access_manager or AccessManager()
        self.cfg = Config()
        if self.protocol not in (p.value for p in PROTOCOLS):
            raise ValueError(f"Unsupported protocol: {self.protocol}")
        gov.audit(
            user_id="system",
            action="protocol_selected",
            resource="ProtocolManager",
            metadata={"protocol": self.protocol},
        )

    def run(self, inputs: Dict[str, Any]) -> Any:
        """
        Executes the configured workflow based on the selected protocol.
        """
        if self.protocol == PROTOCOLS.MCP.value:
            return self._run_mcp(inputs)
        elif self.protocol == PROTOCOLS.AGENT2AGENT.value:
            return self._run_agent2agent(inputs)
        else:
            raise NotImplementedError(f"Protocol '{self.protocol}' is not implemented.")

    def _initialize_shared_resources(self):
        """Initializes all shared components needed by the protocols."""
        llm = LLMClient(
            provider=self.cfg.llm_provider,
            api_key=self.cfg.api_key,
            model=self.cfg.llm_model,
        )
        renderer = PromptRenderer(template_dir=Path(self.cfg.template_dir))
        embedding_fn = lambda text: gemini_embed(text, self.cfg.api_key)

        vector_ret = VectorRetriever(
            index_path=self.cfg.faiss_index_path, embed_model_fn=embedding_fn
        )
        graph_ret = GraphRetriever(
            neo4j_uri=self.cfg.neo4j_uri,
            user=self.cfg.neo4j_user,
            password=self.cfg.neo4j_password,
            gds_graph_name=self.cfg.gds_graph_name,
            embed_model_fn=embedding_fn,
        )
        mem_mgr = MemoryManager(
            backend=self.cfg.memory_backend, redis_url=self.cfg.redis_url
        )

        tool_registry = ToolRegistry(
            governance_manager=gov,
            access_manager=self.access_manager,
            embedding_config={"api_key": self.cfg.api_key},
            similarity_metric=SimilarityMetric(self.cfg.tool_similarity_metric),
            embedding_dimension=self.cfg.embedding_dimension,
        )
        return llm, renderer, vector_ret, graph_ret, mem_mgr, tool_registry

    def _define_tools(self, tool_registry: ToolRegistry):
        """A central place to define and register all available tools with policies."""

        @tool_registry.register(
            required_role="weather_forecaster",
            cost_per_call=0.001, 
            cache_ttl_seconds=300, 
            retry_attempts=2
        )
        def get_weather(city: str) -> str:
            """A governed tool to fetch the weather for a given city."""
            if "new york" in city.lower():
                return "It is currently 75°F and sunny in New York."
            elif "san francisco" in city.lower():
                return "It is currently 62°F and foggy in San Francisco."
            else:
                return f"Weather data for {city} is not available."

    def _build_mcp_orchestrator(self, resources: tuple) -> SimpleOrchestrator:
        """Builds the MCP orchestrator with the superior tool-use workflow."""
        llm, renderer, vector_ret, graph_ret, mem_mgr, tool_registry = resources
        self._define_tools(tool_registry)

        orch = SimpleOrchestrator()

        def retrieve_docs(user_input: str, user_id: str, **kwargs):
            if not self.access_manager.check_access(user_id, "vector_store"):
                raise RBACError(f"User {user_id} unauthorized for retrieval")
            v_docs = vector_ret.query(user_input, top_k=3)
            g_docs = graph_ret.query(user_input, top_k=3)
            combined = {d.get("id"): d for d in (v_docs + g_docs) if d.get("id")}
            return list(combined.values())

        def make_initial_prompt(
            user_input: str, retrieve_docs: List[dict], **kwargs
        ) -> str:
            relevant_tools = tool_registry.get_relevant_tools(user_input, top_k=3)
            tool_schemas = tool_registry.generate_tool_schema(relevant_tools)
            return renderer.render(
                "tool_decider_prompt.j2",
                question=user_input,
                docs=retrieve_docs,
                tools=json.dumps(tool_schemas, indent=2),
            )

        def call_llm_for_tool(make_initial_prompt: str, user_id: str, **kwargs) -> dict:
            if not self.access_manager.check_access(user_id, "llm_call"):
                raise RBACError(f"User {user_id} unauthorized for LLM calls")
            summary = mem_mgr.load(user_id, "summary") or ""
            full_prompt = f"{summary}\n\n{make_initial_prompt}"
            return llm.generate(full_prompt)

        def execute_tool(call_llm_for_tool: dict, user_id: str, **kwargs) -> dict:
            response_text = call_llm_for_tool.get("text", "")
            try:
                # Clean the response text from markdown code blocks
                cleaned_text = response_text.replace("```json", "").replace("```", "").strip()
                data = json.loads(cleaned_text)

                tool_name = data.get("tool_name")
                tool_args = data.get("arguments")

                if tool_name and isinstance(tool_args, dict):
                    governed_tool = tool_registry.get_governed_tool(tool_name, user_id)
                    result = governed_tool(**tool_args)
                    return {"status": "success", "output": result}
            except (json.JSONDecodeError, TypeError, NameError) as e:
                logging.info(f"Could not parse tool call, treating as direct answer. Error: {e}, Response: '{response_text}'")

            return {"status": "no_tool_needed", "output": response_text}

        def generate_final_answer(
            execute_tool: dict, user_input: str, **kwargs
        ) -> dict:
            if execute_tool["status"] != "success":
                return {"text": execute_tool["output"]}
            final_prompt = renderer.render(
                "synthesis_prompt.j2",
                question=user_input,
                tool_result=execute_tool["output"],
            )
            return llm.generate(final_prompt)

        # Define the graph structure
        orch.add_node("retrieve_docs", retrieve_docs)
        orch.add_node("make_initial_prompt", make_initial_prompt)
        orch.add_node("call_llm_for_tool", call_llm_for_tool)
        orch.add_node("execute_tool", execute_tool)
        orch.add_node("generate_final_answer", generate_final_answer)

        # Define the execution flow
        # orch.add_edge("user_input", "retrieve_docs")
        # orch.add_edge("user_input", "make_initial_prompt")
        orch.add_edge("retrieve_docs", "make_initial_prompt")
        orch.add_edge("make_initial_prompt", "call_llm_for_tool")
        orch.add_edge("call_llm_for_tool", "execute_tool")
        # orch.add_edge("user_id", "execute_tool")
        orch.add_edge("execute_tool", "generate_final_answer")
        # orch.add_edge("user_input", "generate_final_answer")

        return orch

    def _run_mcp(self, inputs: Dict[str, Any]) -> Any:
        """Runs the complete MCP workflow."""
        resources = self._initialize_shared_resources()
        orch = self._build_mcp_orchestrator(resources)
        gov.audit(
            user_id=inputs.get("user_id", "system"),
            action="run_mcp_start",
            resource="ProtocolManager",
        )
        results = orch.run(inputs)
        gov.audit(
            user_id=inputs.get("user_id", "system"),
            action="run_mcp_end",
            resource="ProtocolManager",
        )

        if self.cfg.log_mcp_separate:
            run_id = gov.get_current_run_id()
            filename = f"mcp_output_{run_id}.json"
            try:
                with open(filename, "w", encoding="utf-8") as f:
                    json.dump(results, f, indent=2, default=str)
                logging.info(f"MCP run output saved to {filename}")
            except Exception as e:
                logging.error(f"Failed to save MCP output to {filename}: {e}")

        return results

    def _run_agent2agent(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Runs the Agent-to-Agent simulation workflow."""
        gov.audit(
            user_id=inputs.get("user_id", "system"),
            action="run_agent2agent_start",
            resource="ProtocolManager",
        )
        llm, _, vector_ret, _, mem_mgr, _ = self._initialize_shared_resources()
        agents = {}
        agent_ids = ["analyst_agent", "manager_agent"]

        for aid in agent_ids:
            orch = SimpleOrchestrator()

            def retrieve(agent_id=aid, user_input: str = inputs["user_input"], **kwargs):
                return vector_ret.query(f"Query for {agent_id}: {user_input}", top_k=2)

            def respond(retrieve: List[dict], agent_id=aid, **kwargs) -> dict:
                doc_ids = [d.get("id", "N/A") for d in retrieve]
                prompt = (
                    f"As {agent_id}, generate a one-sentence response based on "
                    f"documents: {doc_ids}"
                )
                return llm.generate(prompt)

            orch.add_node("retrieve", retrieve)
            orch.add_node("respond", respond)
            orch.add_edge("retrieve", "respond")
            agents[aid] = orch

        outputs = {}
        for aid, orch in agents.items():
            gov.audit(
                user_id=inputs.get("user_id", "system"), action="agent_start", resource=aid
            )
            res = orch.run(inputs)
            outputs[aid] = res.get("respond", {}).get("text", "")
            mem_mgr.save(aid, "last_response", outputs[aid])
            gov.audit(
                user_id=inputs.get("user_id", "system"), action="agent_end", resource=aid
            )

        gov.audit(
            user_id=inputs.get("user_id", "system"),
            action="run_agent2agent_end",
            resource="ProtocolManager",
        )

        if self.cfg.log_a2a_separate:
            run_id = gov.get_current_run_id()
            filename = f"a2a_output_{run_id}.json"
            try:
                with open(filename, "w", encoding="utf-8") as f:
                    json.dump(outputs, f, indent=2, default=str)
                logging.info(f"Agent-to-Agent run output saved to {filename}")
            except Exception as e:
                logging.error(f"Failed to save Agent-to-Agent output to {filename}: {e}")

        return outputs

run(inputs)

Executes the configured workflow based on the selected protocol.

Source code in src/safeagent/protocol_manager.py 
54
55
56
57
58
59
60
61
62
63
def run(self, inputs: Dict[str, Any]) -> Any:
    """
    Executes the configured workflow based on the selected protocol.
    """
    if self.protocol == PROTOCOLS.MCP.value:
        return self._run_mcp(inputs)
    elif self.protocol == PROTOCOLS.AGENT2AGENT.value:
        return self._run_agent2agent(inputs)
    else:
        raise NotImplementedError(f"Protocol '{self.protocol}' is not implemented.")