{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "Header_01" }, "source": [ "# 🛡️ Agentic GraphRAG: Cybersecurity Threat Intelligence\n", "\n", "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Google
Open in Colab\n", " \n", " 		\n", " \n", " \"Google
Open in Colab Enterprise\n", " \n", " 		\n", " \n", " \"Vertex
Open in Vertex AI Workbench\n", " \n", " 		\n", " \n", " \"GitHub
View on GitHub\n", " \n", "
\n", "\n", "
" ] }, { "cell_type": "markdown", "source": [ "| Author | Architecture |\n", "| --- | --- |\n", "| [Aniket Agrawal](https://github.com/aniketagrawal2012) | **Vertex AI Agent Engine + Google ADK + Neo4j** |" ], "metadata": { "id": "OFzR-URUlNf0" } }, { "cell_type": "markdown", "metadata": { "id": "Share_Buttons" }, "source": [ "Share to:\n", "\n", "\n", " \"LinkedIn\n", "\n", "\n", "\n", " \"Bluesky\n", "\n", "\n", "\n", " \"X\n", "\n", "\n", "\n", " \"Reddit\n", "\n", "\n", "\n", " \"Facebook\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "Overview_02" }, "source": [ "## Overview\n", "\n", "This notebook builds a **Cybersecurity Threat Analysis Agent**. Unlike standard RAG which retrieves documents, this agent uses **GraphRAG** to understand the complex relationships between Threat Actors, Malware, and Vulnerabilities (CVEs).\n", "\n", "### The \"Neo4j Flavour\" + ADK Upgrade\n", "We are evolving the architecture from Dialogflow CX to the **Google Agent Development Kit (ADK)** and **Vertex AI Reasoning Engine**.\n", "\n", "**System Components:**\n", "1. **Graph Database (Neo4j):** Stores the threat landscape (e.g., `(APT28)-[:USES]->(Zebrocy)`).\n", "2. **Reasoning Engine (Vertex AI):** A managed service to deploy the agent logic.\n", "3. **Framework (Google ADK):** Defines the agent's persona, tools, and routing logic programmatically.\n", "\n", "\n", "### Objectives\n", "In this tutorial, you will:\n", "* Install the necessary libraries (Google ADK, LangChain, Neo4j drivers).\n", "* Set up your Google Cloud environment and authenticate.\n", "* Connect to a Neo4j database and seed it with cybersecurity threat data.\n", "* Define a GraphRAG tool using LangChain to query the knowledge graph.\n", "* Build an agent using the Google ADK that utilizes the GraphRAG tool.\n", "* Test the agent locally to verify its reasoning capabilities.\n", "* Deploy the agent to the Vertex AI Agent Engine as a managed service.\n", "* Visualize the threat graph interactively within the notebook.\n", "* Clean up resources to avoid incurring costs.\n", "\n", "## Before you begin\n", "\n", "1. In the Google Cloud console, on the project selector page, select or [create a Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects).\n", "2. [Make sure that billing is enabled for your Google Cloud project](https://cloud.google.com/billing/docs/how-to/verify-billing-enabled#console).\n", "3. [Make sure the Vertex AI API is enabled](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n", "\n", "### Required roles\n", "To get the permissions that you need to complete the tutorial, ask your administrator to grant you the [Vertex AI User](https://cloud.google.com/iam/docs/understanding-roles#aiplatform.user) (`roles/aiplatform.user`) IAM role on your project. For more information about granting roles, see [Manage access](https://cloud.google.com/iam/docs/granting-changing-revoking-access)." ] }, { "cell_type": "markdown", "metadata": { "id": "Install_03" }, "source": [ "### 1. Installation and Prerequisites\n", "We install the **Google ADK**, **LangChain**, and **Neo4j** drivers." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "PipInstall_04" }, "outputs": [], "source": [ "%pip install --quiet google-adk>=1.0.0\n", "%pip install --quiet google-cloud-aiplatform>=1.97.0\n", "%pip install --quiet langchain-google-vertexai\n", "%pip install --quiet langchain-community neo4j" ] }, { "cell_type": "markdown", "metadata": { "id": "DoK1gW9dsRxB" }, "source": [ "#### Authenticating your notebook environment\n", "* If you are using **Colab** to run this notebook, uncomment the cell below and continue.\n", "* If you are using **Vertex AI Workbench**, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)." ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "id": "os3H39sGXugN", "outputId": "6c8c5162-4c4b-44d4-ac18-7907efa27051", "colab": { "base_uri": "https://localhost:8080/" } }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "3.12.12 (main, Oct 10 2025, 08:52:57) [GCC 11.4.0]\n", "✅ Authenticated\n" ] } ], "source": [ "import sys\n", "print(sys.version)\n", "\n", "if \"google.colab\" in sys.modules:\n", " from google.colab import auth\n", " auth.authenticate_user()\n", " print(\"✅ Authenticated\")" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "id": "Config_05" }, "outputs": [], "source": [ "import os\n", "\n", "# --- CONFIGURATION ---\n", "PROJECT_ID = \"your-project-id\" # @param {type:\"string\"}\n", "REGION = \"us-central1\" # @param {type:\"string\"}\n", "NEO4J_URI = \"neo4j+s://your-instance.databases.neo4j.io\" # @param {type:\"string\"}\n", "NEO4J_USER = \"neo4j\" # @param {type:\"string\"}\n", "NEO4J_PASSWORD = \"your-password\" # @param {type:\"string\"}\n", "\n", "# Set Environment\n", "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = PROJECT_ID\n", "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = REGION\n", "os.environ[\"NEO4J_URI\"] = NEO4J_URI\n", "os.environ[\"NEO4J_USER\"] = NEO4J_USER\n", "os.environ[\"NEO4J_PASSWORD\"] = NEO4J_PASSWORD\n", "\n", "import vertexai\n", "vertexai.init(project=PROJECT_ID, location=REGION)" ] }, { "cell_type": "markdown", "metadata": { "id": "Seeding_06" }, "source": [ "### 2. Hydrate the Knowledge Graph\n", "We will seed the database with a sample **Cybersecurity Schema**. This creates nodes for `ThreatActor`, `Malware`, and `Vulnerability`." ] }, { "cell_type": "code", "execution_count": 31, "metadata": { "id": "SeedCode_07", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "33c8a9c9-0406-41ee-f755-026b69dfb010" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "🌱 Seeding Cyber Threat Data...\n", "✅ Database populated with Threat Intel graph.\n" ] } ], "source": [ "from langchain_community.graphs import Neo4jGraph\n", "\n", "def seed_database():\n", " print(\"🌱 Seeding Cyber Threat Data...\")\n", " graph = Neo4jGraph(url=NEO4J_URI, username=NEO4J_USER, password=NEO4J_PASSWORD)\n", "\n", " cypher = \"\"\"\n", " MERGE (a:ThreatActor {name: 'APT29', alias: 'Cozy Bear'})\n", " MERGE (m:Malware {name: 'WellMess'})\n", " MERGE (v:Vulnerability {cve: 'CVE-2023-1234', severity: 'High'})\n", " MERGE (t:Target {sector: 'Pharmaceuticals'})\n", "\n", " MERGE (a)-[:USES]->(m)\n", " MERGE (m)-[:EXPLOITS]->(v)\n", " MERGE (a)-[:TARGETS]->(t)\n", " \"\"\"\n", " graph.query(cypher)\n", " print(\"✅ Database populated with Threat Intel graph.\")\n", "\n", "if \"your-password\" not in NEO4J_PASSWORD:\n", " seed_database()" ] }, { "cell_type": "markdown", "metadata": { "id": "ToolDef_08" }, "source": [ "### 3. Define the GraphRAG Tool\n", "We define a Python function `query_threat_graph`. The **Vertex AI Agent** will call this function when it needs to answer questions about security threats. It uses **LangChain** to translate the question into Cypher." ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "id": "ToolCode_09" }, "outputs": [], "source": [ "from langchain_google_vertexai import VertexAI\n", "from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain\n", "\n", "def query_threat_graph(question: str) -> str:\n", " \"\"\"Queries the cybersecurity knowledge graph to answer questions about actors, malware, and CVEs.\"\"\"\n", " try:\n", " # Re-initialize inside function for serialization contexts\n", " import os\n", " from langchain_community.graphs import Neo4jGraph\n", " from langchain_google_vertexai import VertexAI\n", " from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain\n", "\n", " # Credentials must be explicit for remote execution\n", " neo4j_uri = \"neo4j+s://7d50da77.databases.neo4j.io\"\n", " neo4j_user = \"neo4j\"\n", " neo4j_password = \"nZasSfGzac_mgTApAqrZd_Yqty9I4HOXKyj8qeNKdYg\"\n", "\n", " graph = Neo4jGraph(\n", " url=neo4j_uri,\n", " username=neo4j_user,\n", " password=neo4j_password\n", " )\n", "\n", " llm = VertexAI(model_name=\"gemini-2.0-flash\", temperature=0)\n", "\n", " chain = GraphCypherQAChain.from_llm(\n", " llm=llm,\n", " graph=graph,\n", " verbose=True,\n", " allow_dangerous_requests=True\n", " )\n", "\n", " result = chain.invoke(question)\n", " return result['result']\n", "\n", " except Exception as e:\n", " # Catch ALL errors (imports, init, connection) to prevent Engine crash\n", " return f\"DEBUG ERROR in tool: {str(e)}\"" ] }, { "cell_type": "markdown", "metadata": { "id": "AgentDef_10" }, "source": [ "### 4. Create the ADK Agent\n", "Using the **Agent Development Kit**, we define the `CyberSecurityAgent`. We give it instructions to always check the graph first before assuming answers." ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "id": "AgentCode_11", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "ad350b95-d685-433e-e92d-83ec842b9089" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "✅ Agent Definition Created\n" ] } ], "source": [ "# Install missing dependency required by google-adk\n", "%pip install --quiet deprecated\n", "\n", "from google.adk.agents import Agent\n", "# from google.adk.tools import google_search <-- Removed to avoid API conflict\n", "\n", "cyber_agent = Agent(\n", " name=\"CyberThreatIntel\",\n", " model=\"gemini-2.0-flash\",\n", " description=\"An expert in cybersecurity threat intelligence and graph analysis.\",\n", " instruction=\"\"\"\n", " You are a Cybersecurity Analyst.\n", " 1. If the user asks about Threats, Actors, or CVEs, ALWAYS use the 'query_threat_graph' tool first.\n", " 2. Be concise and actionable in your reporting.\n", " \"\"\",\n", " tools=[query_threat_graph] # Removed google_search\n", ")\n", "\n", "print(\"✅ Agent Definition Created\")" ] }, { "cell_type": "markdown", "metadata": { "id": "LocalTest_12" }, "source": [ "### 5. Local Testing\n", "We instantiate the **ADK App** locally to test the reasoning loop." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "TestCode_13" }, "outputs": [], "source": [ "from vertexai.preview import reasoning_engines\n", "\n", "# Initialize Local App\n", "app = reasoning_engines.AdkApp(\n", " agent=cyber_agent,\n", " enable_tracing=True\n", ")\n", "\n", "print(\"💬 Querying Agent: 'Which malware does APT29 use?'\")\n", "print(\"-\" * 40)\n", "\n", "if \"your-password\" not in NEO4J_PASSWORD:\n", " # Streaming response\n", " for event in app.stream_query(user_id=\"analyst_01\", message=\"Which malware does APT29 use?\"):\n", " if 'content' in event and 'parts' in event['content']:\n", " part = event['content']['parts'][0]\n", " if 'text' in part:\n", " print(part['text'], end=\"\")\n", " elif 'function_call' in part:\n", " print(f\"\\n[🛠️ Tool Call: {part['function_call']['name']}]\", end=\"\\n\")\n", "else:\n", " print(\"⚠️ Neo4j credentials missing. Skipping test.\")" ] }, { "cell_type": "markdown", "metadata": { "id": "Deploy_14" }, "source": [ "### 6. Deploy to Vertex AI Agent Engine\n", "Finally, we package and deploy the agent as a scalable, managed service on Google Cloud." ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "id": "DeployCode_15", "colab": { "base_uri": "https://localhost:8080/" }, "outputId": "cc4a242c-7b87-4084-81dc-6ec5cb106afc" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "🚀 Deploying to Vertex AI Agent Engine... (Bucket: gs://aniket-personal-vertex-staging)\n" ] }, { "output_type": "stream", "name": "stderr", "text": [ "INFO:vertexai.reasoning_engines._reasoning_engines:Using bucket aniket-personal-vertex-staging\n", "INFO:vertexai.reasoning_engines._reasoning_engines:Writing to gs://aniket-personal-vertex-staging/reasoning_engine/reasoning_engine.pkl\n", "INFO:vertexai.reasoning_engines._reasoning_engines:Writing to gs://aniket-personal-vertex-staging/reasoning_engine/requirements.txt\n", "INFO:vertexai.reasoning_engines._reasoning_engines:Creating in-memory tarfile of extra_packages\n", "INFO:vertexai.reasoning_engines._reasoning_engines:Writing to gs://aniket-personal-vertex-staging/reasoning_engine/dependencies.tar.gz\n", "INFO:vertexai.reasoning_engines._reasoning_engines:Creating ReasoningEngine\n", "INFO:vertexai.reasoning_engines._reasoning_engines:Create ReasoningEngine backing LRO: projects/551887116707/locations/us-central1/reasoningEngines/8343273451959615488/operations/5496245304816566272\n", "INFO:vertexai.reasoning_engines._reasoning_engines:ReasoningEngine created. Resource name: projects/551887116707/locations/us-central1/reasoningEngines/8343273451959615488\n", "INFO:vertexai.reasoning_engines._reasoning_engines:To use this ReasoningEngine in another session:\n", "INFO:vertexai.reasoning_engines._reasoning_engines:reasoning_engine = vertexai.preview.reasoning_engines.ReasoningEngine('projects/551887116707/locations/us-central1/reasoningEngines/8343273451959615488')\n" ] }, { "output_type": "stream", "name": "stdout", "text": [ "✅ Deployed! Resource Name: projects/551887116707/locations/us-central1/reasoningEngines/8343273451959615488\n" ] } ], "source": [ "from vertexai.preview import reasoning_engines\n", "import vertexai\n", "\n", "STAGING_BUCKET = f\"gs://{PROJECT_ID}-vertex-staging\"\n", "\n", "# Configure the staging bucket globally\n", "vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n", "\n", "print(f\"🚀 Deploying to Vertex AI Agent Engine... (Bucket: {STAGING_BUCKET})\")\n", "\n", "# --- Wrapper to fix API registration and Context issues ---\n", "class AgentWrapper:\n", " \"\"\"Wrapper to expose stream_query and pass through all arguments.\"\"\"\n", " def __init__(self, app, project_id: str, location: str):\n", " self.app = app\n", " self.project_id = project_id\n", " self.location = location\n", "\n", " def stream_query(self, **kwargs):\n", " # 1. Re-initialize Vertex AI context inside the remote execution\n", " import vertexai\n", " vertexai.init(project=self.project_id, location=self.location)\n", "\n", " # 2. Delegate to the app with all arguments (e.g. message, user_id)\n", " return self.app.stream_query(**kwargs)\n", "# ----------------------------------------------\n", "\n", "# Deployment\n", "# Note: This requires the staging bucket to exist.\n", "try:\n", " # 1. Create a FRESH app instance\n", " clean_app = reasoning_engines.AdkApp(agent=cyber_agent, enable_tracing=False)\n", "\n", " # 2. Wrap the clean app with explicit project/location\n", " wrapped_app = AgentWrapper(clean_app, project_id=PROJECT_ID, location=REGION)\n", "\n", " remote_app = reasoning_engines.ReasoningEngine.create(\n", " wrapped_app,\n", " requirements=[\n", " \"google-adk>=1.0.0\",\n", " \"langchain-community\",\n", " \"langchain-google-vertexai\",\n", " \"langchain\",\n", " \"neo4j\",\n", " \"google-cloud-aiplatform>=1.97.0\"\n", " ],\n", " )\n", " print(f\"✅ Deployed! Resource Name: {remote_app.resource_name}\")\n", "except Exception as e:\n", " print(f\"ℹ️ Deployment skipped (Ensure GCS bucket exists): {e}\")" ] }, { "cell_type": "code", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "678d36ef", "outputId": "79de1914-05f5-488c-86ff-9d4403dcdd48" }, "source": [ "import inspect\n", "\n", "# Inspect the 'app' object to see its methods and if they are async\n", "print(\"App Type:\", type(app))\n", "print(\"Methods in app:\", dir(app))\n", "\n", "if hasattr(app, 'query'):\n", " print(\"Is 'query' async?\", inspect.iscoroutinefunction(app.query))\n", "\n", "if hasattr(app, 'stream_query'):\n", " print(\"Is 'stream_query' async?\", inspect.iscoroutinefunction(app.stream_query))" ], "execution_count": 16, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "App Type: \n", "Methods in app: ['__class__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__slotnames__', '__str__', '__subclasshook__', '__weakref__', '_convert_response_events', '_init_session', '_telemetry_enabled', '_tmpl_attrs', '_tracing_enabled', '_warn_if_telemetry_api_disabled', 'agent_framework', 'async_add_session_to_memory', 'async_create_session', 'async_delete_session', 'async_get_session', 'async_list_sessions', 'async_search_memory', 'async_stream_query', 'bidi_stream_query', 'clone', 'create_session', 'delete_session', 'get_session', 'list_sessions', 'project_id', 'register_operations', 'set_up', 'stream_query', 'streaming_agent_run_with_events']\n", "Is 'stream_query' async? False\n" ] } ] }, { "cell_type": "markdown", "metadata": { "id": "4aeb6e9d" }, "source": [ "### 7. Query the Deployed Agent\n", "Now that the agent is running on Vertex AI, we can send it queries.\n", "\n", "**Note:** If you restart your session, you can reconnect to this agent using:\n", "```python\n", "remote_app = reasoning_engines.ReasoningEngine('YOUR_RESOURCE_NAME')\n", "```" ] }, { "cell_type": "code", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "e9cf1704", "outputId": "bbb6b5be-ebe8-4ae1-ed0e-0b83f5f26b72" }, "source": [ "print(\"💬 Querying Remote Agent: 'What CVEs are associated with WellMess?'\")\n", "print(\"-\" * 40)\n", "\n", "# Query the deployed agent\n", "# We provide user_id as AdkApp often requires session context\n", "response = remote_app.stream_query(\n", " message=\"What CVEs are associated with WellMess?\",\n", " user_id=\"remote_analyst_01\"\n", ")\n", "\n", "for chunk in response:\n", " print(chunk)" ], "execution_count": 17, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "💬 Querying Remote Agent: 'What CVEs are associated with WellMess?'\n", "----------------------------------------\n", "{'model_version': 'gemini-2.0-flash', 'content': {'parts': [{'function_call': {'id': 'adk-9b45e639-f6cb-473c-a12b-4ebc87dc4d2a', 'args': {'question': 'What CVEs are associated with WellMess?'}, 'name': 'query_threat_graph'}}], 'role': 'model'}, 'finish_reason': 'STOP', 'usage_metadata': {'candidates_token_count': 15, 'candidates_tokens_details': [{'modality': 'TEXT', 'token_count': 15}], 'prompt_token_count': 119, 'prompt_tokens_details': [{'modality': 'TEXT', 'token_count': 119}], 'total_token_count': 134, 'traffic_type': 'ON_DEMAND'}, 'avg_logprobs': -4.4566400659581026e-05, 'invocation_id': 'e-e66d3e09-f7ef-488e-9a39-21fbc5133012', 'author': 'CyberThreatIntel', 'actions': {'state_delta': {}, 'artifact_delta': {}, 'requested_auth_configs': {}, 'requested_tool_confirmations': {}}, 'long_running_tool_ids': [], 'id': '067a91ca-e1da-4c6b-a9c3-c66b7c5821b3', 'timestamp': 1770299736.027741}\n", "{'content': {'parts': [{'function_response': {'id': 'adk-9b45e639-f6cb-473c-a12b-4ebc87dc4d2a', 'name': 'query_threat_graph', 'response': {'result': 'CVE-2023-1234 is associated with WellMess.\\n'}}}], 'role': 'user'}, 'invocation_id': 'e-e66d3e09-f7ef-488e-9a39-21fbc5133012', 'author': 'CyberThreatIntel', 'actions': {'state_delta': {}, 'artifact_delta': {}, 'requested_auth_configs': {}, 'requested_tool_confirmations': {}}, 'id': '6bbc55c1-5bc8-42bb-95bd-98f8ad77122b', 'timestamp': 1770299742.142327}\n", "{'model_version': 'gemini-2.0-flash', 'content': {'parts': [{'text': 'CVE-2023-1234 is associated with WellMess.\\n'}], 'role': 'model'}, 'finish_reason': 'STOP', 'usage_metadata': {'candidates_token_count': 18, 'candidates_tokens_details': [{'modality': 'TEXT', 'token_count': 18}], 'prompt_token_count': 158, 'prompt_tokens_details': [{'modality': 'TEXT', 'token_count': 158}], 'total_token_count': 176, 'traffic_type': 'ON_DEMAND'}, 'avg_logprobs': -0.0005405148387783103, 'invocation_id': 'e-e66d3e09-f7ef-488e-9a39-21fbc5133012', 'author': 'CyberThreatIntel', 'actions': {'state_delta': {}, 'artifact_delta': {}, 'requested_auth_configs': {}, 'requested_tool_confirmations': {}}, 'id': '5252ebca-08fd-440c-bfd5-124d7ff93e1a', 'timestamp': 1770299742.269534}\n" ] } ] }, { "cell_type": "markdown", "metadata": { "id": "2dd883a0" }, "source": [ "### 8. Next Steps & Cleanup\n", "\n", "**1. Try a Complex Query**\n", "Let's see if the agent can traverse multiple hops in the graph (Actor -> Target)." ] }, { "cell_type": "code", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "b9d05b0a", "outputId": "570f4f9e-d987-4379-b596-078112770895" }, "source": [ "# Ask about targets\n", "response = remote_app.stream_query(\n", " message=\"Which threat actors target the Pharmaceuticals sector?\",\n", " user_id=\"analyst_02\"\n", ")\n", "\n", "for chunk in response:\n", " print(chunk)" ], "execution_count": 18, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "{'model_version': 'gemini-2.0-flash', 'content': {'parts': [{'function_call': {'id': 'adk-3c052b1b-5333-4aa9-9a9a-02463d302cec', 'args': {'question': 'Which threat actors target the Pharmaceuticals sector?'}, 'name': 'query_threat_graph'}}], 'role': 'model'}, 'finish_reason': 'STOP', 'usage_metadata': {'candidates_token_count': 14, 'candidates_tokens_details': [{'modality': 'TEXT', 'token_count': 14}], 'prompt_token_count': 118, 'prompt_tokens_details': [{'modality': 'TEXT', 'token_count': 118}], 'total_token_count': 132, 'traffic_type': 'ON_DEMAND'}, 'avg_logprobs': -9.134594750191485e-05, 'invocation_id': 'e-aa0661ca-89d6-4386-8013-ce2324f5e69a', 'author': 'CyberThreatIntel', 'actions': {'state_delta': {}, 'artifact_delta': {}, 'requested_auth_configs': {}, 'requested_tool_confirmations': {}}, 'long_running_tool_ids': [], 'id': '12305720-9f31-4f22-b299-b0212c7091db', 'timestamp': 1770299883.442706}\n", "{'content': {'parts': [{'function_response': {'id': 'adk-3c052b1b-5333-4aa9-9a9a-02463d302cec', 'name': 'query_threat_graph', 'response': {'result': 'APT29 targets the Pharmaceuticals sector.\\n'}}}], 'role': 'user'}, 'invocation_id': 'e-aa0661ca-89d6-4386-8013-ce2324f5e69a', 'author': 'CyberThreatIntel', 'actions': {'state_delta': {}, 'artifact_delta': {}, 'requested_auth_configs': {}, 'requested_tool_confirmations': {}}, 'id': '85c3667b-1736-43a4-8e5e-31b643c8df17', 'timestamp': 1770299891.266381}\n", "{'model_version': 'gemini-2.0-flash', 'content': {'parts': [{'text': 'APT29 targets the Pharmaceuticals sector.\\n'}], 'role': 'model'}, 'finish_reason': 'STOP', 'usage_metadata': {'candidates_token_count': 9, 'candidates_tokens_details': [{'modality': 'TEXT', 'token_count': 9}], 'prompt_token_count': 147, 'prompt_tokens_details': [{'modality': 'TEXT', 'token_count': 147}], 'total_token_count': 156, 'traffic_type': 'ON_DEMAND'}, 'avg_logprobs': -0.0003546736358354489, 'invocation_id': 'e-aa0661ca-89d6-4386-8013-ce2324f5e69a', 'author': 'CyberThreatIntel', 'actions': {'state_delta': {}, 'artifact_delta': {}, 'requested_auth_configs': {}, 'requested_tool_confirmations': {}}, 'id': 'f247f585-53d4-47df-9419-502c8d750aa0', 'timestamp': 1770299891.383679}\n" ] } ] }, { "cell_type": "markdown", "metadata": { "id": "af07d9a5" }, "source": [ "**2. How to Connect from Another App**\n", "Use the `resource_name` printed in the deployment step to connect to this agent from any Python script." ] }, { "cell_type": "code", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "6fe565a1", "outputId": "85b5694f-b2b4-4a19-a572-cd66ae79b2f8" }, "source": [ "# Example of connecting to an existing agent\n", "# Replace RESOURCE_NAME with the output from DeployCode_15 (starts with projects/...)\n", "agent = reasoning_engines.ReasoningEngine(remote_app.resource_name)\n", "\n", "# We must use stream_query and PROVIDE A USER_ID for session context\n", "response = agent.stream_query(\n", " message=\"Hello\",\n", " user_id=\"external_user_01\" # Required by AdkApp\n", ")\n", "\n", "for chunk in response:\n", " print(chunk)" ], "execution_count": 23, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "{'model_version': 'gemini-2.0-flash', 'content': {'parts': [{'text': \"Hello, I'm ready to assist you with cybersecurity threat intelligence. What threats, actors, or CVEs are you interested in today?\\n\"}], 'role': 'model'}, 'finish_reason': 'STOP', 'usage_metadata': {'candidates_token_count': 29, 'candidates_tokens_details': [{'modality': 'TEXT', 'token_count': 29}], 'prompt_token_count': 111, 'prompt_tokens_details': [{'modality': 'TEXT', 'token_count': 111}], 'total_token_count': 140, 'traffic_type': 'ON_DEMAND'}, 'avg_logprobs': -0.10933220797571643, 'invocation_id': 'e-e5e6ee44-e3ec-43bc-9325-0e8a1028d372', 'author': 'CyberThreatIntel', 'actions': {'state_delta': {}, 'artifact_delta': {}, 'requested_auth_configs': {}, 'requested_tool_confirmations': {}}, 'id': 'e174f6b3-e314-47f2-8a0d-c4baecec0ceb', 'timestamp': 1770300079.152798}\n" ] } ] }, { "cell_type": "markdown", "source": [ "**3. Interactive visualization:**\n", "Following code block does the following:\n", "\n", "* Installs the `pyvis` visualization library.\n", "* Creates a function to query the raw graph data.\n", "* Renders an interactive network diagram of APT29 and its connections right here in the notebook.\n", "\n", "Graph Visualization Active! 🕸️\n", "\n", "You can now see the interactive network of threats directly in your notebook.\n", "\n", "* **Red Nodes:** The Threat Actor (e.g., APT29).\n", "* **Blue Nodes:** The connected entities (Malware, CVEs, Targets).\n", "* **Edges:** The relationships (e.g., USES, EXPLOITS, TARGETS)." ], "metadata": { "id": "OIk5tHj_IhgQ" } }, { "cell_type": "code", "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "4a2119f4", "outputId": "ae1089e8-7403-4283-99c4-e9a6c671f9a2" }, "source": [ "%pip install --quiet pyvis\n", "print(\"✅ PyVis Installed\")" ], "execution_count": 24, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "\u001b[?25l \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/756.0 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K \u001b[91m━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m399.4/756.0 kB\u001b[0m \u001b[31m11.7 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m756.0/756.0 kB\u001b[0m \u001b[31m12.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25h\u001b[?25l \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/1.6 MB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m48.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25h✅ PyVis Installed\n" ] } ] }, { "cell_type": "code", "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 575 }, "id": "3dc67aea", "outputId": "c82b16cd-4220-4d73-83c8-bfa5924129f6" }, "source": [ "from pyvis.network import Network\n", "from IPython.display import HTML, display\n", "from neo4j import GraphDatabase\n", "\n", "def visualize_attack_graph(root_node_name):\n", " \"\"\"Queries Neo4j and creates an interactive graph visualization around a node.\"\"\"\n", " # 1. Direct Connection (using the same credentials as before)\n", " uri = \"neo4j+s://7d50da77.databases.neo4j.io\"\n", " user = \"neo4j\"\n", " pwd = \"nZasSfGzac_mgTApAqrZd_Yqty9I4HOXKyj8qeNKdYg\"\n", "\n", " driver = GraphDatabase.driver(uri, auth=(user, pwd))\n", "\n", " # 2. Setup PyVis Network\n", " net = Network(notebook=True, cdn_resources='in_line', height=\"500px\", width=\"100%\", bgcolor=\"#222222\", font_color=\"white\")\n", "\n", " cypher = f\"\"\"\n", " MATCH (n)-[r]-(m)\n", " WHERE n.name = '{root_node_name}' OR n.alias = '{root_node_name}'\n", " RETURN n, r, m\n", " LIMIT 50\n", " \"\"\"\n", "\n", " with driver.session() as session:\n", " result = session.run(cypher)\n", " for record in result:\n", " src = record['n']\n", " dst = record['m']\n", " rel = record['r']\n", "\n", " # Helper to get a label for the node (handling different schemas)\n", " def get_label(node):\n", " return node.get('name') or node.get('cve') or node.get('sector') or 'Unknown'\n", "\n", " # Add Nodes\n", " # We use element_id to uniquely identify nodes in the visualizer\n", " net.add_node(src.element_id, label=get_label(src), title=str(src.labels), color='#ff4b4b') # Red for Source\n", " net.add_node(dst.element_id, label=get_label(dst), title=str(dst.labels), color='#4b94ff') # Blue for Target\n", "\n", " # Add Edge\n", " net.add_edge(src.element_id, dst.element_id, title=rel.type, label=rel.type, color='white')\n", "\n", " driver.close()\n", "\n", " # 3. Render\n", " net.show('threat_graph.html')\n", " return HTML('threat_graph.html')\n", "\n", "print(\"✨ Visualization Tool Created. Rendering graph for 'APT29'...\")\n", "visualize_attack_graph(\"APT29\")" ], "execution_count": 25, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "✨ Visualization Tool Created. Rendering graph for 'APT29'...\n", "threat_graph.html\n" ] }, { "output_type": "execute_result", "data": { "text/plain": [ "" ], "text/html": [ "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
\n", "

\n", "
\n", "\n", "\n", " \n", " \n", "\n", "\n", "
\n", "

\n", "
\n", " \n", " \n", "\n", "\n", " \n", "
\n", " \n", " \n", "
\n", "
\n", "\n", " \n", " \n", "\n", " \n", " \n", "" ] }, "metadata": {}, "execution_count": 25 } ] }, { "cell_type": "markdown", "source": [ "This feature turns the \"black box\" of the graph database into a clear, visual map for analysts.\n", "\n", "Mission Accomplished! We have:\n", "\n", "* Built a GraphRAG Agent.\n", "* Deployed it to Vertex AI.\n", "* Visualized the intelligence data.\n", "\n", "You are ready to hunt threats! 🛡️" ], "metadata": { "id": "7I20siP5J1jS" } }, { "cell_type": "markdown", "metadata": { "id": "70309767" }, "source": [ "**4. Cleanup**\n", "Delete the agent to avoid incurring costs." ] }, { "cell_type": "code", "metadata": { "id": "3be25141" }, "source": [ "# Uncomment to delete the agent when finished\n", "remote_app.delete()\n", "print(\"🗑️ Agent Deleted\")" ], "execution_count": 20, "outputs": [] } ], "metadata": { "colab": { "toc_visible": true, "provenance": [], "collapsed_sections": [ "Install_03", "Seeding_06", "ToolDef_08", "AgentDef_10", "LocalTest_12", "Deploy_14", "4aeb6e9d" ] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }