diff --git a/src/va_agent/agent.py b/src/va_agent/agent.py index 3ffcc13..cbccff9 100644 --- a/src/va_agent/agent.py +++ b/src/va_agent/agent.py @@ -10,17 +10,20 @@ from google.cloud.firestore_v1.async_client import AsyncClient from va_agent.auth import auth_headers_provider from va_agent.config import settings from va_agent.session import FirestoreSessionService +from va_agent.governance import GovernancePlugin toolset = McpToolset( connection_params=StreamableHTTPConnectionParams(url=settings.mcp_remote_url), header_provider=auth_headers_provider, ) +governance = GovernancePlugin() agent = Agent( model=settings.agent_model, name=settings.agent_name, instruction=settings.agent_instructions, tools=[toolset], + after_model_callback=governance.after_model_callback, ) session_service = FirestoreSessionService(