iris-security-vertexai 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.4
2
+ Name: iris-security-vertexai
3
+ Version: 0.1.0
4
+ Summary: IRIS governance for Vertex AI via google-cloud-aiplatform
5
+ Author-email: IRIS Platform <sdk@iris.ai>
6
+ License: Apache-2.0
7
+ Project-URL: Homepage, https://github.com/gimartinb/iris-sdk
8
+ Project-URL: Repository, https://github.com/gimartinb/iris-sdk
9
+ Requires-Python: >=3.10
10
+ Requires-Dist: iris-security-core>=0.1.0
11
+ Requires-Dist: iris-security-sdk>=0.1.0
12
+ Provides-Extra: vertexai
13
+ Requires-Dist: google-cloud-aiplatform>=1.60; extra == "vertexai"
14
+ Provides-Extra: dev
15
+ Requires-Dist: pytest>=8.0; extra == "dev"
16
+ Requires-Dist: ruff>=0.4; extra == "dev"
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.4
2
+ Name: iris-security-vertexai
3
+ Version: 0.1.0
4
+ Summary: IRIS governance for Vertex AI via google-cloud-aiplatform
5
+ Author-email: IRIS Platform <sdk@iris.ai>
6
+ License: Apache-2.0
7
+ Project-URL: Homepage, https://github.com/gimartinb/iris-sdk
8
+ Project-URL: Repository, https://github.com/gimartinb/iris-sdk
9
+ Requires-Python: >=3.10
10
+ Requires-Dist: iris-security-core>=0.1.0
11
+ Requires-Dist: iris-security-sdk>=0.1.0
12
+ Provides-Extra: vertexai
13
+ Requires-Dist: google-cloud-aiplatform>=1.60; extra == "vertexai"
14
+ Provides-Extra: dev
15
+ Requires-Dist: pytest>=8.0; extra == "dev"
16
+ Requires-Dist: ruff>=0.4; extra == "dev"
@@ -0,0 +1,11 @@
1
+ pyproject.toml
2
+ iris_security_vertexai.egg-info/PKG-INFO
3
+ iris_security_vertexai.egg-info/SOURCES.txt
4
+ iris_security_vertexai.egg-info/dependency_links.txt
5
+ iris_security_vertexai.egg-info/requires.txt
6
+ iris_security_vertexai.egg-info/top_level.txt
7
+ iris_vertexai/__init__.py
8
+ iris_vertexai/client.py
9
+ iris_vertexai/fedramp.py
10
+ iris_vertexai/init.py
11
+ tests/test_vertexai_integration.py
@@ -0,0 +1,9 @@
1
+ iris-security-core>=0.1.0
2
+ iris-security-sdk>=0.1.0
3
+
4
+ [dev]
5
+ pytest>=8.0
6
+ ruff>=0.4
7
+
8
+ [vertexai]
9
+ google-cloud-aiplatform>=1.60
@@ -0,0 +1,29 @@
1
+ """
2
+ IRIS Vertex AI integration - governed wrapper for Vertex AI GenerativeModel.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ from iris import IrisViolationError
8
+ from iris_core.models.passport import (
9
+ AgentPassport,
10
+ ComplianceTag,
11
+ DataClassification,
12
+ Environment,
13
+ )
14
+ from iris_core.models.policy import Violation
15
+
16
+ from iris_vertexai.client import IrisGenerativeModel, IrisVertexAI
17
+
18
+ __version__ = "0.1.0"
19
+
20
+ __all__ = [
21
+ "IrisVertexAI",
22
+ "IrisGenerativeModel",
23
+ "IrisViolationError",
24
+ "AgentPassport",
25
+ "ComplianceTag",
26
+ "DataClassification",
27
+ "Environment",
28
+ "Violation",
29
+ ]
@@ -0,0 +1,314 @@
1
+ """Governed Vertex AI wrapper with IRIS policy enforcement."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import importlib
6
+ import json
7
+ import os
8
+ import uuid
9
+ from datetime import datetime
10
+ from typing import Any, Optional
11
+
12
+ from iris import IrisViolationError
13
+ from iris_core.dlp import DLPScanner
14
+ from iris_core.dlp.enforcement import (
15
+ enforce_prompt_dlp,
16
+ extract_gemini_response_text,
17
+ handle_response_dlp,
18
+ )
19
+ from iris_core.engine.cedar import CedarEngine, EvaluationContext
20
+ from iris_core.rbac.context import UserContext
21
+ from iris_core.evidence.vault import EvidenceVault
22
+ from iris_core.models.passport import AgentPassport
23
+ from iris_core.models.policy import PolicyResult, Severity, Violation
24
+
25
+ from iris_gemini.guardrails import _extract_text
26
+
27
+ from iris_vertexai.fedramp import check_fedramp_location, check_fedramp_model
28
+
29
+
30
+ def _lazy_vertexai():
31
+ try:
32
+ return importlib.import_module("vertexai")
33
+ except ModuleNotFoundError as exc:
34
+ raise ImportError(
35
+ "google-cloud-aiplatform is required for IrisVertexAI. "
36
+ "Install with: pip install google-cloud-aiplatform"
37
+ ) from exc
38
+
39
+
40
+ def _lazy_vertexai_generative_models():
41
+ try:
42
+ return importlib.import_module("vertexai.generative_models")
43
+ except ModuleNotFoundError as exc:
44
+ raise ImportError(
45
+ "vertexai.generative_models is unavailable. "
46
+ "Install or upgrade: pip install google-cloud-aiplatform>=1.60"
47
+ ) from exc
48
+
49
+
50
+ def _is_fedramp_enabled(passport: AgentPassport) -> bool:
51
+ for tag in passport.compliance_tags:
52
+ value = getattr(tag, "value", str(tag)).lower()
53
+ if value == "fedramp":
54
+ return True
55
+ return False
56
+
57
+
58
+ def _load_passport_policy(engine: CedarEngine, passport: AgentPassport) -> None:
59
+ if not passport.policy_ref:
60
+ return
61
+ from pathlib import Path
62
+
63
+ policy_path = Path(passport.policy_ref)
64
+ if not policy_path.is_absolute():
65
+ policy_path = Path.cwd() / policy_path
66
+ if policy_path.exists():
67
+ engine.load_policy_file(passport.agent_id, policy_path)
68
+
69
+
70
+ def _current_environment():
71
+ from iris_core.models.passport import Environment
72
+
73
+ return Environment(os.environ.get("IRIS_ENV", "dev"))
74
+
75
+
76
+ def _enforce_result(result: PolicyResult) -> None:
77
+ if result.decision == "DENY":
78
+ raise IrisViolationError(result)
79
+
80
+
81
+ class IrisVertexAI:
82
+ """Top-level Vertex AI client wrapper with governance-aware model creation."""
83
+
84
+ def __init__(
85
+ self,
86
+ passport: AgentPassport,
87
+ project: Optional[str] = None,
88
+ location: Optional[str] = None,
89
+ user_email: Optional[str] = None,
90
+ user_role: Optional[str] = None,
91
+ **kwargs: Any,
92
+ ):
93
+ from iris_core.dev_trust import print_dev_trust_message
94
+
95
+ print_dev_trust_message()
96
+ evidence_vault_dir = kwargs.pop("evidence_vault_dir", None)
97
+ vertexai = _lazy_vertexai()
98
+ vertexai.init(project=project, location=location, **kwargs)
99
+ self._passport = passport
100
+ self._project = project
101
+ self._location = location
102
+ self._user_email = user_email
103
+ self._user_role = user_role
104
+ self._engine = CedarEngine()
105
+ self._vault = EvidenceVault(agent_id=passport.agent_id, vault_dir=evidence_vault_dir)
106
+ self._dlp = DLPScanner(passport)
107
+ _load_passport_policy(self._engine, passport)
108
+
109
+ def get_model(self, model_name: str) -> "IrisGenerativeModel":
110
+ return IrisGenerativeModel(
111
+ model_name=model_name,
112
+ passport=self._passport,
113
+ location=self._location,
114
+ project=self._project,
115
+ engine=self._engine,
116
+ vault=self._vault,
117
+ dlp=self._dlp,
118
+ user_email=self._user_email,
119
+ user_role=self._user_role,
120
+ )
121
+
122
+
123
+ class IrisGenerativeModel:
124
+ """Governed wrapper around vertexai.generative_models.GenerativeModel."""
125
+
126
+ def __init__(
127
+ self,
128
+ model_name: str,
129
+ passport: AgentPassport,
130
+ location: Optional[str] = None,
131
+ project: Optional[str] = None,
132
+ engine: Optional[CedarEngine] = None,
133
+ vault: Optional[EvidenceVault] = None,
134
+ dlp: Optional[DLPScanner] = None,
135
+ user_email: Optional[str] = None,
136
+ user_role: Optional[str] = None,
137
+ ):
138
+ generative_models = _lazy_vertexai_generative_models()
139
+ self._model_name = model_name
140
+ self._passport = passport
141
+ self._location = location
142
+ self._project = project
143
+ self._engine = engine or CedarEngine()
144
+ self._vault = vault or EvidenceVault(agent_id=passport.agent_id)
145
+ self._dlp = dlp or DLPScanner(passport)
146
+ self._user_email = user_email
147
+ self._user_role = user_role
148
+ self._model = generative_models.GenerativeModel(model_name)
149
+
150
+ def _record_with_metadata(
151
+ self, ctx: EvaluationContext, result: PolicyResult, location_violations: list[Violation]
152
+ ) -> None:
153
+ entry = {
154
+ "event_id": str(uuid.uuid4()),
155
+ "timestamp": datetime.utcnow().isoformat(),
156
+ "agent_id": self._passport.agent_id,
157
+ "action": ctx.action,
158
+ "resource": ctx.resource,
159
+ "environment": ctx.environment.value,
160
+ "decision": result.decision,
161
+ "gcp_project": self._project,
162
+ "gcp_location": self._location,
163
+ "additional": {
164
+ **ctx.additional,
165
+ "gcp_project": self._project,
166
+ "gcp_location": self._location,
167
+ },
168
+ "violations": [
169
+ {
170
+ "rule_id": v.rule_id,
171
+ "severity": v.severity.value,
172
+ "message": v.message,
173
+ "compliance_refs": v.compliance_refs,
174
+ }
175
+ for v in (list(result.violations) + list(location_violations))
176
+ ],
177
+ }
178
+ with open(self._vault._log_file, "a") as handle:
179
+ handle.write(json.dumps(entry) + "\n")
180
+
181
+ def _govern(self, action: str, contents: Any = None) -> None:
182
+ env = _current_environment()
183
+ prompt_text = "\n".join(_extract_text(contents))
184
+ dlp_result = enforce_prompt_dlp(
185
+ self._dlp,
186
+ self._vault,
187
+ self._passport,
188
+ env,
189
+ prompt_text,
190
+ resource=f"vertexai/{self._model_name}",
191
+ )
192
+ violations: list[Violation] = []
193
+
194
+ if self._location:
195
+ if self._passport.allowed_regions and self._location not in self._passport.allowed_regions:
196
+ violations.append(
197
+ Violation(
198
+ rule_id="IRIS-REGION-001",
199
+ severity=Severity.CRITICAL,
200
+ message=(
201
+ f"Location '{self._location}' is not in passport.allowed_regions."
202
+ ),
203
+ compliance_refs=["iris:region-policy"],
204
+ remediation=(
205
+ "Use an allowed location or update passport.allowed_regions "
206
+ "with security approval."
207
+ ),
208
+ )
209
+ )
210
+ if self._location.startswith("us-gov-"):
211
+ violations.append(
212
+ Violation(
213
+ rule_id="IRIS-REGION-002",
214
+ severity=Severity.HIGH,
215
+ message=(
216
+ f"Location '{self._location}' is a restricted government region."
217
+ ),
218
+ compliance_refs=["fedramp:high:region-review"],
219
+ remediation=(
220
+ "Ensure FedRAMP High authorization is documented for this location."
221
+ ),
222
+ )
223
+ )
224
+ else:
225
+ violations.append(
226
+ Violation(
227
+ rule_id="IRIS-REGION-003",
228
+ severity=Severity.CRITICAL,
229
+ message="Vertex AI location is required but missing.",
230
+ compliance_refs=["iris:region-policy"],
231
+ remediation="Provide a valid Vertex AI location before calling the model.",
232
+ )
233
+ )
234
+
235
+ if _is_fedramp_enabled(self._passport):
236
+ fedramp_location_violation = check_fedramp_location(self._location or "")
237
+ if fedramp_location_violation:
238
+ violations.append(fedramp_location_violation)
239
+ fedramp_model_violation = check_fedramp_model(self._model_name)
240
+ if fedramp_model_violation:
241
+ violations.append(fedramp_model_violation)
242
+
243
+ user_ctx = UserContext.from_params(self._user_email, self._user_role)
244
+ ctx = EvaluationContext(
245
+ agent_id=self._passport.agent_id,
246
+ action=action,
247
+ resource=f"vertexai/{self._model_name}",
248
+ resource_type="api",
249
+ environment=env,
250
+ data_region=self._location,
251
+ data_classification=self._passport.data_classification.value,
252
+ dlp_prompt_findings=dlp_result.findings,
253
+ additional={
254
+ "model": self._model_name,
255
+ "gcp_project": self._project,
256
+ "gcp_location": self._location,
257
+ },
258
+ **user_ctx.evaluation_fields(),
259
+ )
260
+ result = self._engine.evaluate(self._passport, ctx)
261
+ if violations:
262
+ result = PolicyResult(
263
+ decision="DENY",
264
+ violations=list(result.violations) + violations,
265
+ agent_id=result.agent_id,
266
+ action=result.action,
267
+ resource=result.resource,
268
+ environment=result.environment,
269
+ )
270
+
271
+ self._record_with_metadata(ctx, result, [])
272
+ _enforce_result(result)
273
+
274
+ def _scan_response(self, response: Any) -> Any:
275
+ env = _current_environment()
276
+ response_text = extract_gemini_response_text(response)
277
+ blocked, _ = handle_response_dlp(
278
+ self._dlp,
279
+ self._vault,
280
+ self._passport,
281
+ env,
282
+ response_text,
283
+ response,
284
+ resource=f"vertexai/{self._model_name}",
285
+ )
286
+ return blocked
287
+
288
+ def generate_content(self, contents: Any, **kwargs: Any) -> Any:
289
+ self._govern(action="generate_content", contents=contents)
290
+ response = self._model.generate_content(contents, **kwargs)
291
+ return self._scan_response(response)
292
+
293
+ def generate_content_stream(self, **kwargs: Any) -> Any:
294
+ self._govern(action="generate_content_stream", contents=kwargs.get("contents"))
295
+ return self._model.generate_content_stream(**kwargs)
296
+
297
+ def start_chat(self) -> "IrisChatSession":
298
+ return IrisChatSession(
299
+ chat_session=self._model.start_chat(),
300
+ model=self,
301
+ )
302
+
303
+
304
+ class IrisChatSession:
305
+ """Governed wrapper around Vertex AI chat sessions."""
306
+
307
+ def __init__(self, chat_session: Any, model: IrisGenerativeModel):
308
+ self._chat_session = chat_session
309
+ self._model = model
310
+
311
+ def send_message(self, content: Any, **kwargs: Any) -> Any:
312
+ self._model._govern(action="chat.send_message", contents=content)
313
+ response = self._chat_session.send_message(content, **kwargs)
314
+ return self._model._scan_response(response)
@@ -0,0 +1,89 @@
1
+ """FedRAMP-specific governance checks for Vertex AI integrations."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Optional
6
+
7
+ from iris_core.models.policy import Severity, Violation
8
+
9
+ FEDRAMP_AUTHORIZED_VERTEX_REGIONS = {
10
+ "us-central1",
11
+ "us-east1",
12
+ "us-east4",
13
+ "us-west1",
14
+ "us-west2",
15
+ "us-west3",
16
+ "us-west4",
17
+ "us-south1",
18
+ }
19
+
20
+ FEDRAMP_AUTHORIZED_VERTEX_MODELS = {
21
+ "gemini-1.5-pro",
22
+ "gemini-1.5-flash",
23
+ }
24
+
25
+
26
+ def check_fedramp_location(location: str) -> Optional[Violation]:
27
+ """Return violation when Vertex AI location is not FedRAMP-authorized."""
28
+ normalized = (location or "").strip().lower()
29
+ if not normalized:
30
+ return Violation(
31
+ rule_id="FEDRAMP-001",
32
+ severity=Severity.CRITICAL,
33
+ message="FedRAMP location is missing for this Vertex AI request.",
34
+ compliance_refs=["fedramp:moderate:residency"],
35
+ remediation=(
36
+ "Set a FedRAMP-authorized Vertex AI region and retry the request."
37
+ ),
38
+ )
39
+
40
+ if normalized in FEDRAMP_AUTHORIZED_VERTEX_REGIONS:
41
+ return None
42
+
43
+ if normalized.startswith("us-gov-"):
44
+ return Violation(
45
+ rule_id="FEDRAMP-001",
46
+ severity=Severity.CRITICAL,
47
+ message=(
48
+ f"Location '{location}' requires additional FedRAMP High authorization."
49
+ ),
50
+ compliance_refs=["fedramp:high:authorization-boundary"],
51
+ remediation=(
52
+ "Obtain documented FedRAMP High authorization for this region "
53
+ "or use an authorized FedRAMP Moderate region."
54
+ ),
55
+ )
56
+
57
+ return Violation(
58
+ rule_id="FEDRAMP-001",
59
+ severity=Severity.CRITICAL,
60
+ message=(
61
+ f"Location '{location}' is not in the FedRAMP-authorized Vertex AI region list."
62
+ ),
63
+ compliance_refs=["fedramp:moderate:residency"],
64
+ remediation=(
65
+ "Use one of the authorized regions: "
66
+ + ", ".join(sorted(FEDRAMP_AUTHORIZED_VERTEX_REGIONS))
67
+ + "."
68
+ ),
69
+ )
70
+
71
+
72
+ def check_fedramp_model(model_name: str) -> Optional[Violation]:
73
+ """Return violation when Vertex AI model authorization cannot be verified."""
74
+ normalized = (model_name or "").strip().lower()
75
+ if normalized in FEDRAMP_AUTHORIZED_VERTEX_MODELS:
76
+ return None
77
+
78
+ return Violation(
79
+ rule_id="FEDRAMP-002",
80
+ severity=Severity.HIGH,
81
+ message=(
82
+ f"Model '{model_name}' is not on the current FedRAMP-authorized Vertex AI list."
83
+ ),
84
+ compliance_refs=["fedramp:ai-workload:authorization"],
85
+ remediation=(
86
+ "Use a known authorized model (gemini-1.5-pro or gemini-1.5-flash) "
87
+ "or complete an authorization review for this model."
88
+ ),
89
+ )
@@ -0,0 +1,5 @@
1
+ """Compatibility exports for requested init.py module."""
2
+
3
+ from iris_vertexai.client import IrisGenerativeModel, IrisVertexAI
4
+
5
+ __all__ = ["IrisVertexAI", "IrisGenerativeModel"]
@@ -0,0 +1,36 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "iris-security-vertexai"
7
+ version = "0.1.0"
8
+ description = "IRIS governance for Vertex AI via google-cloud-aiplatform"
9
+ license = { text = "Apache-2.0" }
10
+ requires-python = ">=3.10"
11
+ authors = [{ name = "IRIS Platform", email = "sdk@iris.ai" }]
12
+ dependencies = [
13
+ "iris-security-core>=0.1.0",
14
+ "iris-security-sdk>=0.1.0",
15
+ ]
16
+
17
+ [project.optional-dependencies]
18
+ vertexai = ["google-cloud-aiplatform>=1.60"]
19
+ dev = [
20
+ "pytest>=8.0",
21
+ "ruff>=0.4",
22
+ ]
23
+
24
+ [project.urls]
25
+ Homepage = "https://github.com/gimartinb/iris-sdk"
26
+ Repository = "https://github.com/gimartinb/iris-sdk"
27
+
28
+ [tool.setuptools]
29
+ packages = ["iris_vertexai"]
30
+
31
+ [tool.ruff]
32
+ line-length = 100
33
+ target-version = "py310"
34
+
35
+ [tool.pytest.ini_options]
36
+ testpaths = ["tests"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,167 @@
1
+ """Integration tests for iris-vertexai with mocked Vertex AI SDK."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import types
7
+ from pathlib import Path
8
+ from unittest.mock import MagicMock, patch
9
+
10
+ import pytest
11
+
12
+ from iris import IrisViolationError
13
+ from iris_core.engine.cedar import CedarEngine
14
+ from iris_core.models.passport import AgentPassport, ComplianceTag, DataClassification, Environment
15
+ from iris_vertexai import IrisVertexAI
16
+ from iris_vertexai.fedramp import check_fedramp_location, check_fedramp_model
17
+
18
+
19
+ def _mock_vertex_modules():
20
+ vertexai_module = types.ModuleType("vertexai")
21
+ vertexai_module.init = MagicMock()
22
+
23
+ gm_module = types.ModuleType("vertexai.generative_models")
24
+ response = MagicMock()
25
+ response.text = "Vertex response"
26
+ chat_response = MagicMock()
27
+ chat_response.text = "chat response"
28
+ chat_session = MagicMock()
29
+ chat_session.send_message.return_value = chat_response
30
+
31
+ model_instance = MagicMock()
32
+ model_instance.generate_content.return_value = response
33
+ model_instance.generate_content_stream.return_value = iter([response])
34
+ model_instance.start_chat.return_value = chat_session
35
+
36
+ gm_module.GenerativeModel = MagicMock(return_value=model_instance)
37
+ return vertexai_module, gm_module, model_instance, chat_session
38
+
39
+
40
+ class _FedrampTag:
41
+ value = "fedramp"
42
+
43
+
44
+ @pytest.fixture
45
+ def passport():
46
+ return AgentPassport(
47
+ name="vertex-agent",
48
+ owner="team@gov-agency.gov",
49
+ data_classification=DataClassification.INTERNAL,
50
+ compliance_tags=[ComplianceTag.COLORADO_AI_ACT],
51
+ environments=[Environment.DEV, Environment.PRODUCTION],
52
+ allowed_regions=["us-central1", "us-east1"],
53
+ )
54
+
55
+
56
+ def _permit_engine(passport: AgentPassport) -> CedarEngine:
57
+ engine = CedarEngine()
58
+ engine.load_policy(passport.agent_id, "permit(principal, action, resource);")
59
+ return engine
60
+
61
+
62
+ def test_authorized_location_permitted(passport, tmp_path, monkeypatch):
63
+ monkeypatch.setenv("IRIS_ENV", "dev")
64
+ vertexai_module, gm_module, model_instance, _ = _mock_vertex_modules()
65
+
66
+ with patch.dict(
67
+ "sys.modules",
68
+ {"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
69
+ ):
70
+ client = IrisVertexAI(
71
+ passport=passport,
72
+ project="proj-1",
73
+ location="us-central1",
74
+ evidence_vault_dir=tmp_path,
75
+ )
76
+ client._engine = _permit_engine(passport)
77
+ model = client.get_model("gemini-1.5-pro")
78
+ result = model.generate_content("hello")
79
+
80
+ assert result.text == "Vertex response"
81
+ model_instance.generate_content.assert_called_once()
82
+
83
+
84
+ def test_unauthorized_location_blocked(passport, tmp_path, monkeypatch):
85
+ monkeypatch.setenv("IRIS_ENV", "production")
86
+ vertexai_module, gm_module, model_instance, _ = _mock_vertex_modules()
87
+
88
+ with patch.dict(
89
+ "sys.modules",
90
+ {"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
91
+ ):
92
+ client = IrisVertexAI(
93
+ passport=passport,
94
+ project="proj-1",
95
+ location="europe-west1",
96
+ evidence_vault_dir=tmp_path,
97
+ )
98
+ client._engine = _permit_engine(passport)
99
+ model = client.get_model("gemini-1.5-pro")
100
+ with pytest.raises(IrisViolationError):
101
+ model.generate_content("hello")
102
+
103
+ model_instance.generate_content.assert_not_called()
104
+
105
+
106
+ def test_fedramp_location_check():
107
+ assert check_fedramp_location("us-central1") is None
108
+ violation = check_fedramp_location("europe-west1")
109
+ assert violation is not None
110
+ assert violation.rule_id == "FEDRAMP-001"
111
+
112
+
113
+ def test_fedramp_model_check():
114
+ assert check_fedramp_model("gemini-1.5-pro") is None
115
+ violation = check_fedramp_model("gemini-2.0-flash")
116
+ assert violation is not None
117
+ assert violation.rule_id == "FEDRAMP-002"
118
+
119
+
120
+ def test_chat_session_intercept(passport, tmp_path, monkeypatch):
121
+ monkeypatch.setenv("IRIS_ENV", "dev")
122
+ vertexai_module, gm_module, _, chat_session = _mock_vertex_modules()
123
+
124
+ with patch.dict(
125
+ "sys.modules",
126
+ {"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
127
+ ):
128
+ client = IrisVertexAI(
129
+ passport=passport,
130
+ project="proj-1",
131
+ location="us-central1",
132
+ evidence_vault_dir=tmp_path,
133
+ )
134
+ client._engine = _permit_engine(passport)
135
+ model = client.get_model("gemini-1.5-pro")
136
+ chat = model.start_chat()
137
+ chat.send_message("hello")
138
+
139
+ chat_session.send_message.assert_called_once()
140
+
141
+
142
+ def test_evidence_vault_includes_gcp_metadata(passport, tmp_path, monkeypatch):
143
+ monkeypatch.setenv("IRIS_ENV", "dev")
144
+ vertexai_module, gm_module, _, _ = _mock_vertex_modules()
145
+ passport.compliance_tags = [ComplianceTag.COLORADO_AI_ACT, _FedrampTag()]
146
+
147
+ with patch.dict(
148
+ "sys.modules",
149
+ {"vertexai": vertexai_module, "vertexai.generative_models": gm_module},
150
+ ):
151
+ client = IrisVertexAI(
152
+ passport=passport,
153
+ project="proj-1",
154
+ location="us-central1",
155
+ evidence_vault_dir=tmp_path,
156
+ )
157
+ client._engine = _permit_engine(passport)
158
+ model = client.get_model("gemini-1.5-pro")
159
+ model.generate_content("hello")
160
+
161
+ events_file = Path(tmp_path) / passport.agent_id / "events.jsonl"
162
+ lines = events_file.read_text().strip().splitlines()
163
+ event = json.loads(lines[-1])
164
+ assert event["gcp_project"] == "proj-1"
165
+ assert event["gcp_location"] == "us-central1"
166
+ assert event["additional"]["gcp_project"] == "proj-1"
167
+ assert event["additional"]["gcp_location"] == "us-central1"