google-antigravity 0.1.2__py3-none-win_amd64.whl

google/antigravity/__init__.py

@@ -0,0 +1,41 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Google Antigravity SDK for building AI agents."""
+
+from google.antigravity.agent import Agent
+from google.antigravity.connections.connection import AgentConfig
+from google.antigravity.connections.local.local_connection_config import LocalAgentConfig
+from google.antigravity.tools.tool_context import ToolContext
+from google.antigravity.types import CapabilitiesConfig
+from google.antigravity.types import GeminiConfig
+from google.antigravity.types import GenerationConfig
+from google.antigravity.types import ModelConfig
+from google.antigravity.types import ModelEntry
+from google.antigravity.types import ThinkingLevel
+from google.antigravity.types import UsageMetadata
+
+__all__ = [
+    "Agent",
+    "AgentConfig",
+    "LocalAgentConfig",
+    "ToolContext",
+    "CapabilitiesConfig",
+    "GeminiConfig",
+    "GenerationConfig",
+    "ModelConfig",
+    "ModelEntry",
+    "ThinkingLevel",
+    "UsageMetadata",
+]

google/antigravity/agent.py

@@ -0,0 +1,244 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Layer 1 API for Antigravity SDK."""
+
+import contextlib
+import logging
+from typing import cast
+
+from google.antigravity import types
+from google.antigravity.connections import connection as connection_module
+from google.antigravity.conversation import conversation
+from google.antigravity.hooks import hook_runner
+from google.antigravity.hooks import hooks
+from google.antigravity.hooks import policy
+from google.antigravity.mcp import bridge
+from google.antigravity.tools import tool_context
+from google.antigravity.tools import tool_runner
+from google.antigravity.triggers import trigger_runner
+from google.antigravity.triggers import triggers as triggers_lib
+
+
+__all__ = ["Agent"]
+
+
+class Agent:
+  """High-level Agent API for simplified interaction."""
+
+  def __init__(self, config: connection_module.AgentConfig):
+    """Initializes the Agent.
+
+    Args:
+        config: Declarative agent configuration.
+    """
+    self._config = config.model_copy(deep=True)
+    if self._config.response_schema:
+      # The response_schema is validated/stringified in AgentConfig.
+      self._config.capabilities.finish_tool_schema_json = cast(
+          str, self._config.response_schema
+      )
+    self._strategy = None
+    self._conversation = None
+    self._tool_runner = None
+    self._hook_runner = None
+    self._trigger_runner = None
+    self._mcp_bridge = None
+    # Use the original config (not self._config) for hooks and triggers:
+    # model_copy(deep=True) creates new objects, breaking reference equality
+    # for user-provided hooks/triggers. The list() snapshot prevents the
+    # caller from mutating our copy, while preserving object identity.
+    self._pending_hooks = list(config.hooks)
+    self._pending_triggers = list(config.triggers)
+    self._exit_stack = contextlib.AsyncExitStack()
+
+  def register_hook(self, hook: hooks.Hook):
+    """Registers a hook by inferring its type.
+
+    Args:
+        hook: The hook to register.
+    """
+    if not self._hook_runner:
+      self._pending_hooks.append(hook)
+      return
+    self._hook_runner.register_hook(hook)
+
+  def register_trigger(self, trigger: triggers_lib.Trigger):
+    """Registers a trigger.
+
+    Cannot be called after the agent has started.
+
+    Args:
+      trigger: The trigger function to register.
+
+    Raises:
+      RuntimeError: If the agent has already started.
+    """
+    if self._trigger_runner:
+      raise RuntimeError(
+          "Cannot register triggers after the agent has started."
+      )
+    self._pending_triggers.append(trigger)
+
+  async def __aenter__(self) -> "Agent":
+    """Starts the agent session.
+
+    Returns:
+        The started Agent instance.
+    """
+    logging.info("Starting Agent session")
+    try:
+      self._hook_runner = hook_runner.HookRunner()
+
+      # Register pending hooks
+      for hook in self._pending_hooks:
+        self._hook_runner.register_hook(hook)
+      self._pending_hooks.clear()
+
+      # Apply policies
+      active_policies = list(self._config.policies)
+      cfg = self._config.capabilities
+      read_only_tools = set(types.BuiltinTools.read_only())
+      # enabled_tools and disabled_tools are mutually exclusive
+      # (enforced by CapabilitiesConfig validation).
+      if cfg.enabled_tools is not None:
+        active_tools = set(cfg.enabled_tools)
+      elif cfg.disabled_tools is not None:
+        active_tools = set(types.BuiltinTools) - set(cfg.disabled_tools)
+      else:
+        active_tools = set(types.BuiltinTools)
+      has_write_tools = bool(active_tools - read_only_tools)
+      has_mcp_servers = bool(self._config.mcp_servers)
+      has_tool_decide_hook = bool(self._hook_runner.pre_tool_call_decide_hooks)
+
+      if (
+          (has_write_tools or has_mcp_servers)
+          and not active_policies
+          and not has_tool_decide_hook
+      ):
+        raise ValueError(
+            "Write tools or MCP servers are enabled without a safety policy. "
+            "Add policies=[policy.allow_all()] to approve all tool calls, "
+            "or policies=[policy.deny_all(), policy.allow('tool_name')] "
+            "to selectively allow specific tools."
+        )
+
+      if active_policies:
+        self._hook_runner.register_hook(
+            policy.enforce(
+                active_policies, mcp_servers=self._config.mcp_servers
+            )
+        )
+
+      all_tools = list(self._config.tools)
+      # Connect MCP servers
+      if self._config.mcp_servers:
+        logging.info("Connecting to MCP servers...")
+        self._mcp_bridge = bridge.McpBridge()
+        self._exit_stack.push_async_callback(self._mcp_bridge.stop)
+
+        for server_cfg in self._config.mcp_servers:
+          await self._mcp_bridge.connect(server_cfg)
+        all_tools.extend(self._mcp_bridge.tools)
+
+      self._tool_runner = tool_runner.ToolRunner(tools=all_tools)
+
+      self._strategy = self._config.create_strategy(
+          tool_runner=self._tool_runner,
+          hook_runner=self._hook_runner,
+      )
+
+      logging.info("Starting connection and creating conversation...")
+      self._conversation = await self._exit_stack.enter_async_context(
+          conversation.Conversation.create(self._strategy)
+      )
+
+      # Start triggers via TriggerRunner.
+      if self._pending_triggers:
+        logging.info("Starting triggers...")
+        self._trigger_runner = await self._exit_stack.enter_async_context(
+            trigger_runner.TriggerRunner(
+                triggers=list(self._pending_triggers),
+                connection=self.conversation.connection,
+            )
+        )
+        self._pending_triggers.clear()
+
+      # Wire ToolContext into ToolRunner so tools can access
+      # conversation capabilities (same pattern as TriggerRunner).
+      if self._tool_runner:
+        ctx = tool_context.ToolContext(self.conversation.connection)
+        self._tool_runner.set_context(ctx)
+
+      return self
+    except Exception:
+      logging.exception("Failed to start Agent session, cleaning up...")
+      await self._exit_stack.aclose()
+      raise
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb):
+    """Stops the agent session.
+
+    Args:
+        exc_type: The exception type, if any.
+        exc_val: The exception value, if any.
+        exc_tb: The traceback, if any.
+    """
+    logging.info("Stopping Agent session")
+    return await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+
+  async def chat(self, prompt: types.Content) -> types.ChatResponse:
+    """Sends a prompt and returns the final response.
+
+    Args:
+        prompt: The user prompt or content to send.
+
+    Returns:
+        The final response from the agent.
+    """
+    return await self.conversation.chat(prompt)
+
+  @property
+  def is_started(self) -> bool:
+    """Whether the agent session has been started."""
+    return self._conversation is not None
+
+  @property
+  def conversation(self) -> conversation.Conversation:
+    """Returns the active Conversation session.
+
+    Use this for advanced session introspection: history, turn count,
+    compaction indices, usage, or direct send/receive_steps control.
+    For most use cases, prefer chat() instead.
+
+    Raises:
+      RuntimeError: If the agent session has not been started.
+    """
+    if not self._conversation:
+      raise RuntimeError(
+          "Agent session not started. Use 'async with Agent(...)'."
+      )
+    return self._conversation
+
+  @property
+  def conversation_id(self) -> str | None:
+    """Returns the conversation identifier assigned by the runtime.
+
+    Available after the session has started and at least one message has
+    been exchanged.  Pass this value back via SessionConfig.conversation_id
+    to resume from a saved session.  Returns None before the session starts.
+    """
+    if not self._conversation:
+      return None
+    return self._conversation.conversation_id or None