Add vertex search tool with integration tests (#73)

* Add vertex search tool with integration tests

* Clear notebook cells

* Address Marcelo review comments

Author: amrit110

.env.example

@@ -24,5 +24,11 @@ REPORT_GENERATION_DB__DRIVER="sqlite"
 REPORT_GENERATION_DB__DATABASE="implementations/report_generation/data/OnlineRetail.db"
 REPORT_GENERATION_DB__QUERY__MODE="ro"
 
+# Vertex AI Search (custom knowledge base) - no API key needed, uses ADC
+# On Coder/GCE workspaces the attached service account handles auth automatically.
+# Required IAM roles on the service account: roles/discoveryengine.viewer, roles/aiplatform.user
+GOOGLE_CLOUD_LOCATION="us-central1"
+VERTEX_AI_DATASTORE_ID="projects/{project}/locations/global/collections/default_collection/dataStores/{datastore-id}"
+
 # Report Generation (all optional, defaults are in implementations/report_generation/env_vars.py)
 REPORT_GENERATION_OUTPUT_PATH="..."

aieng-eval-agents/aieng/agent_evals/configs.py

@@ -160,6 +160,21 @@ class Configs(BaseSettings):
     web_search_base_url: str | None = Field(default=None, description="Base URL for web search service.")
     web_search_api_key: SecretStr | None = Field(default=None, description="API key for web search service.")
 
+    # === Vertex AI Search (custom knowledge base) ===
+    google_cloud_location: str = Field(
+        default="us-central1",
+        description="GCP region for Vertex AI model calls. Must match a region that supports Gemini.",
+    )
+    vertex_datastore_id: str | None = Field(
+        default=None,
+        validation_alias="VERTEX_AI_DATASTORE_ID",
+        description=(
+            "Full Vertex AI Search data store resource name. "
+            "Format: projects/{project}/locations/global/collections/default_collection/dataStores/{id}. "
+            "Authentication uses Application Default Credentials (ADC) — no API key required."
+        ),
+    )
+
     # === Report Generation ===
     # Defaults are set in the implementations/report_generation/env_vars.py file
     report_generation_output_path: str | None = Field(

aieng-eval-agents/aieng/agent_evals/knowledge_qa/__init__.py

@@ -21,14 +21,13 @@
     format_response_with_citations,
 )
 
-from .agent import KnowledgeAgentManager, KnowledgeGroundedAgent
+from .agent import KnowledgeGroundedAgent
 from .data import DeepSearchQADataset, DSQAExample
 
 
 __all__ = [
     # Agent
     "KnowledgeGroundedAgent",
-    "KnowledgeAgentManager",
     # Grounding tool
     "create_google_search_tool",
     "format_response_with_citations",

aieng-eval-agents/aieng/agent_evals/knowledge_qa/agent.py

@@ -747,80 +747,3 @@ def answer(
         """
         logger.info(f"Answering question (sync): {question[:100]}...")
         return asyncio.run(self.answer_async(question, session_id))
-
-
-class KnowledgeAgentManager:
-    """Manages KnowledgeGroundedAgent lifecycle with lazy initialization.
-
-    This class provides convenient lifecycle management for the knowledge agent,
-    with lazy initialization and state tracking.
-
-    Parameters
-    ----------
-    config : Configs, optional
-        Configuration object for client setup. If not provided, creates default.
-
-    Examples
-    --------
-    >>> manager = KnowledgeAgentManager()
-    >>> agent = manager.agent
-    >>> response = await agent.answer_async("What is quantum computing?")
-    >>> print(response.text)
-    >>> manager.close()
-    """
-
-    def __init__(
-        self,
-        config: Configs | None = None,
-        enable_caching: bool = True,
-        enable_planning: bool = True,
-        enable_compaction: bool = True,
-    ) -> None:
-        """Initialize the client manager.
-
-        Parameters
-        ----------
-        config : Configs, optional
-            Configuration object. If not provided, creates default config.
-        enable_caching : bool, default True
-            Whether to enable context caching.
-        enable_planning : bool, default True
-            Whether to enable built-in planning (Gemini thinking mode).
-        enable_compaction : bool, default True
-            Whether to enable context compaction.
-        """
-        self._config = config
-        self._enable_caching = enable_caching
-        self._enable_planning = enable_planning
-        self._enable_compaction = enable_compaction
-        self._agent: KnowledgeGroundedAgent | None = None
-        self._initialized = False
-
-    @property
-    def config(self) -> Configs:
-        """Get or create the config instance."""
-        if self._config is None:
-            self._config = Configs()  # type: ignore[call-arg]
-        return self._config
-
-    @property
-    def agent(self) -> KnowledgeGroundedAgent:
-        """Get or create the knowledge-grounded agent."""
-        if self._agent is None:
-            self._agent = KnowledgeGroundedAgent(
-                config=self.config,
-                enable_caching=self._enable_caching,
-                enable_planning=self._enable_planning,
-                enable_compaction=self._enable_compaction,
-            )
-            self._initialized = True
-        return self._agent
-
-    def close(self) -> None:
-        """Close all initialized clients and reset state."""
-        self._agent = None
-        self._initialized = False
-
-    def is_initialized(self) -> bool:
-        """Check if any clients have been initialized."""
-        return self._initialized

aieng-eval-agents/aieng/agent_evals/tools/__init__.py

@@ -2,6 +2,7 @@
 
 This package provides modular tools for:
 - Google Search (search.py)
+- Vertex AI Search / custom knowledge base (vertex_search.py)
 - Web content fetching - HTML and PDF (web.py)
 - File downloading and searching - CSV, XLSX, text (file.py)
 - SQL Database access (sql_database.py)
@@ -24,20 +25,23 @@
     google_search,
 )
 from .sql_database import ReadOnlySqlDatabase, ReadOnlySqlPolicy
+from .vertex_search import create_vertex_search_tool, vertex_search
 from .web import (
     create_web_fetch_tool,
     web_fetch,
 )
 
 
 __all__ = [
-    # Search tools
+    # Google Search tools
     "create_google_search_tool",
     "google_search",
     "format_response_with_citations",
-    "google_search",
     "GroundedResponse",
     "GroundingChunk",
+    # Vertex AI Search tools (custom knowledge base)
+    "create_vertex_search_tool",
+    "vertex_search",
     # Web tools (HTML pages and PDFs)
     "web_fetch",
     "create_web_fetch_tool",