Skip to content

Comments

add adk and vertex ai integrations#35

Merged
johnymontana merged 31 commits intomainfrom
adk
Feb 11, 2026
Merged

add adk and vertex ai integrations#35
johnymontana merged 31 commits intomainfrom
adk

Conversation

@johnymontana
Copy link
Collaborator

@johnymontana johnymontana commented Jan 27, 2026

This PR implements integrations with Google Cloud's AI ecosystem for neo4j-agent-memory. It adds native support for Google's Agent Development Kit (ADK), Vertex AI embeddings, and exposes memory capabilities via MCP server for Cloud API Registry integration.

Features

1. Vertex AI Embeddings (VTX-001)

  • New VertexAIEmbedder class supporting:
    • text-embedding-004 (768 dimensions, recommended)
    • textembedding-gecko@003 (768 dimensions)
    • textembedding-gecko-multilingual@001 (768 dimensions)
  • Batch embedding with 250 item limit (Vertex AI constraint)
  • Multiple authentication methods (ADC, service account, explicit credentials)
  • Task type support for retrieval and semantic similarity

2. Google ADK Integration (ADK-001)

  • Neo4jMemoryService implementing ADK's MemoryService interface
  • Session memory persistence with entity extraction
  • Semantic memory search across all memory types
  • Type conversion utilities between neo4j-agent-memory and ADK types

3. MCP Server (MCP-001)

  • Neo4jMemoryMCPServer with 5 tools for Cloud API Registry:
    • memory_search - Hybrid vector + graph search
    • memory_store - Store messages, facts, preferences
    • entity_lookup - Get entity with graph neighbors
    • conversation_history - Retrieve session messages
    • graph_query - Execute read-only Cypher queries
  • Support for stdio and SSE transports
  • CLI command: neo4j-memory mcp serve

4. Cloud Run Template (RUN-001)

  • Production-ready Dockerfile
  • Cloud Build configuration
  • Terraform templates for infrastructure

5. Documentation & Diagrams

  • Comprehensive AsciiDoc documentation
  • 5 Excalidraw architecture diagrams
  • Updated README with Google Cloud section

@vercel
Copy link

vercel bot commented Jan 27, 2026

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
agent-memory Ready Ready Preview, Comment Feb 11, 2026 6:04am

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR implements comprehensive Google Cloud integration for neo4j-agent-memory, adding native support for Vertex AI embeddings, Google's Agent Development Kit (ADK), and a Model Context Protocol (MCP) server for Cloud API Registry integration.

Changes:

  • Vertex AI Embeddings provider supporting text-embedding-004 and gecko models with 768-dimensional embeddings
  • Neo4jMemoryService implementing ADK's MemoryService interface for session persistence and semantic search
  • MCP server with 5 tools (memory_search, memory_store, entity_lookup, conversation_history, graph_query) supporting stdio and SSE transports
  • Cloud Run deployment templates with Dockerfile, Cloud Build configuration, and infrastructure setup
  • Comprehensive documentation and examples demonstrating all Google Cloud features

Reviewed changes

Copilot reviewed 41 out of 42 changed files in this pull request and generated 20 comments.

Show a summary per file
File Description
src/neo4j_agent_memory/embeddings/vertex_ai.py Vertex AI embedder implementation with batch processing
src/neo4j_agent_memory/integrations/google_adk/ ADK MemoryService and type conversions
src/neo4j_agent_memory/mcp/ MCP server, handlers, and tool definitions
tests/unit/ Comprehensive unit tests for all new components
tests/integration/ Integration tests with real Neo4j and Vertex AI
examples/google_cloud_integration/ Full-featured examples and demos
deploy/cloudrun/ Production deployment templates
docs/ Updated documentation with Google Cloud integration guide

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

"""

__all__ = [
"Neo4jMemoryMCPServer",
Copy link

Copilot AI Jan 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The name 'Neo4jMemoryMCPServer' is exported by all but is not defined.

Copilot uses AI. Check for mistakes.
"""Unit tests for MCP handlers."""

import json
from unittest.mock import AsyncMock, MagicMock, patch
Copy link

Copilot AI Jan 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Import of 'patch' is not used.

Copilot uses AI. Check for mistakes.
@@ -0,0 +1,164 @@
"""Unit tests for MCP tool definitions."""

import pytest
Copy link

Copilot AI Jan 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Import of 'pytest' is not used.

Copilot uses AI. Check for mistakes.

# Check if Vertex AI is available
try:
import vertexai # noqa: F401
Copy link

Copilot AI Jan 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Import of 'vertexai' is not used.

Copilot uses AI. Check for mistakes.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Copy link

Copilot AI commented Feb 5, 2026

@johnymontana I've opened a new pull request, #52, to work on those changes. Once the pull request is ready, I'll request review from you.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Copy link

Copilot AI commented Feb 5, 2026

@johnymontana I've opened a new pull request, #53, to work on those changes. Once the pull request is ready, I'll request review from you.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 98 out of 99 changed files in this pull request and generated 17 comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 117 out of 124 changed files in this pull request and generated 18 comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

self._project_id = project_id
self._location = location
self._credentials = credentials
self._batch_size = min(batch_size, DEFAULT_BATCH_SIZE)
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

batch_size can be set to 0 or a negative number, which will make range(..., step=self._batch_size) raise at runtime. Consider validating batch_size >= 1 (raise a ValueError or clamp to 1) before using it.

Suggested change
self._batch_size = min(batch_size, DEFAULT_BATCH_SIZE)
self._batch_size = min(max(batch_size, 1), DEFAULT_BATCH_SIZE)

Copilot uses AI. Check for mistakes.
Comment on lines +195 to +197
# Process in batches (Vertex AI limit is 250 per request)
for i in range(0, len(texts), self._batch_size):
batch = texts[i : i + self._batch_size]
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

batch_size can be set to 0 or a negative number, which will make range(..., step=self._batch_size) raise at runtime. Consider validating batch_size >= 1 (raise a ValueError or clamp to 1) before using it.

Suggested change
# Process in batches (Vertex AI limit is 250 per request)
for i in range(0, len(texts), self._batch_size):
batch = texts[i : i + self._batch_size]
# Determine and validate batch size
batch_size = self._batch_size
if batch_size is None or batch_size < 1:
raise EmbeddingError(f"Invalid batch_size {batch_size!r}; batch_size must be an integer >= 1.")
# Process in batches (Vertex AI limit is 250 per request)
for i in range(0, len(texts), batch_size):
batch = texts[i : i + batch_size]

Copilot uses AI. Check for mistakes.
anthropic = ["anthropic>=0.20.0"]
sentence-transformers = ["sentence-transformers>=2.2.0"]
vertex-ai = ["google-cloud-aiplatform>=1.38.0"]
google = ["google-cloud-aiplatform>=1.38.0"] # Alias for vertex-ai
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docs/examples install instructions use neo4j-agent-memory[google,mcp] to imply Vertex AI + ADK + MCP, but the google extra currently only installs google-cloud-aiplatform and does not include google-adk. Either (a) include google-adk (and any other required Google packages) in the google extra, or (b) update doc strings/examples to instruct installing [vertex-ai,google-adk,mcp] explicitly to avoid missing dependency surprises.

Suggested change
google = ["google-cloud-aiplatform>=1.38.0"] # Alias for vertex-ai
google = ["google-cloud-aiplatform>=1.38.0", "google-adk"] # Vertex AI + Google ADK

Copilot uses AI. Check for mistakes.
Comment on lines +87 to +90
mcp = ["mcp>=1.0.0"]

# Google ADK
google-adk = ["google-adk>=0.1.0"]
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docs/examples install instructions use neo4j-agent-memory[google,mcp] to imply Vertex AI + ADK + MCP, but the google extra currently only installs google-cloud-aiplatform and does not include google-adk. Either (a) include google-adk (and any other required Google packages) in the google extra, or (b) update doc strings/examples to instruct installing [vertex-ai,google-adk,mcp] explicitly to avoid missing dependency surprises.

Copilot uses AI. Check for mistakes.
Comment on lines +61 to +62
# Google Cloud (Vertex AI + ADK + MCP)
pip install neo4j-agent-memory[google,mcp]
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This install command currently doesn’t align with the package extras as defined in pyproject.toml (the google extra doesn’t include google-adk). Update the docs to match the actual extras, or adjust the extras so this command installs all components promised here.

Copilot uses AI. Check for mistakes.
Comment on lines +215 to +217
for tool in MEMORY_TOOLS:
if tool["name"] == name:
return tool
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MEMORY_TOOLS.copy() is a shallow list copy, so callers can still mutate the inner dicts and accidentally affect global tool definitions; get_tool_by_name also returns the original dict. Consider returning deep copies (or immutable structures) to prevent surprising cross-call mutation.

Copilot uses AI. Check for mistakes.
Comment on lines 42 to 47
async def get_initialized_memory_service() -> FinancialMemoryService:
"""Get the initialized memory service."""
service = get_memory_service()
if not service._initialized:
await service.initialize()
return service
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This route relies on private attributes (service._initialized, client._driver, client._database). Since the library now exposes MemoryClient.graph for Neo4j access, prefer using that public API (e.g., memory_service.client.graph.execute_read(...)) and/or add a public readiness method on FinancialMemoryService to avoid tight coupling to internals that may change.

Copilot uses AI. Check for mistakes.
Comment on lines +83 to +85
async with client._driver.session(database=client._database) as session:
result = await session.run(request.query, request.parameters)
records = await result.data()
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This route relies on private attributes (service._initialized, client._driver, client._database). Since the library now exposes MemoryClient.graph for Neo4j access, prefer using that public API (e.g., memory_service.client.graph.execute_read(...)) and/or add a public readiness method on FinancialMemoryService to avoid tight coupling to internals that may change.

Copilot uses AI. Check for mistakes.
)

client = MemoryClient(settings)
await client.initialize()
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MemoryClient usage elsewhere in this PR uses connect() / close() (and async with) rather than initialize(). If initialize() isn’t part of the public API, this script will fail at runtime. Consider switching to await client.connect() (or using async with MemoryClient(settings) as client:) to match the actual client lifecycle API.

Suggested change
await client.initialize()
await client.connect()

Copilot uses AI. Check for mistakes.
Comment on lines +47 to +48
# Mount Google Cloud credentials for local development
- ${GOOGLE_APPLICATION_CREDENTIALS:-~/.config/gcloud/application_default_credentials.json}:/app/credentials.json:ro
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docker Compose does not reliably expand ~ in bind-mount paths, so the default value here can produce an invalid mount and break local startup. Prefer requiring an absolute path (or document/export it) and/or switch to ${HOME}/.config/... to make the default portable.

Suggested change
# Mount Google Cloud credentials for local development
- ${GOOGLE_APPLICATION_CREDENTIALS:-~/.config/gcloud/application_default_credentials.json}:/app/credentials.json:ro
# Mount Google Cloud credentials for local development.
# GOOGLE_APPLICATION_CREDENTIALS must be set to an absolute path on the host.
- ${GOOGLE_APPLICATION_CREDENTIALS:?GOOGLE_APPLICATION_CREDENTIALS must be set to an absolute path on the host}:/app/credentials.json:ro

Copilot uses AI. Check for mistakes.
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 110 out of 130 changed files in this pull request and generated 11 comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +3 to +4
from neo4j_agent_memory.embeddings.base import BaseEmbedder, Embedder
from neo4j_agent_memory.embeddings.openai import OpenAIEmbedder
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Importing OpenAIEmbedder unconditionally can break base installs where the OpenAI dependency is optional (ImportError when importing neo4j_agent_memory.embeddings). Consider making OpenAIEmbedder a lazy import via __getattr__ (like VertexAIEmbedder), or ensure OpenAI is a required dependency of the base package.

Copilot uses AI. Check for mistakes.
Comment on lines +17 to +20
"description": (
"Search across all memory types using hybrid vector + graph search. "
"Finds relevant messages, entities, preferences, and reasoning traces."
),
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The description says the search is across all memory types and includes reasoning traces, but the default memory_types omits traces. Either update the description to match the default behavior, or include traces in the default list.

Copilot uses AI. Check for mistakes.
Comment on lines +33 to +41
"memory_types": {
"type": "array",
"items": {
"type": "string",
"enum": ["messages", "entities", "preferences", "traces"],
},
"description": "Types of memory to search (defaults to all)",
"default": ["messages", "entities", "preferences"],
},
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The description says the search is across all memory types and includes reasoning traces, but the default memory_types omits traces. Either update the description to match the default behavior, or include traces in the default list.

Copilot uses AI. Check for mistakes.
Returns:
List of tool definitions in MCP format.
"""
return MEMORY_TOOLS.copy()
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MEMORY_TOOLS.copy() is only a shallow copy of the list; callers can still mutate the underlying tool dicts (and impact future requests). Return a deep copy (e.g., via copy.deepcopy) or construct new dicts to keep tool definitions immutable for callers.

Copilot uses AI. Check for mistakes.
self._batch_size = min(batch_size, DEFAULT_BATCH_SIZE)
self._task_type = task_type
self._embedding_model: TextEmbeddingModel | None = None
self._initialized = False
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

self._initialized is set but not used anywhere in this class. Either remove it to reduce state/complexity, or use it consistently (e.g., to short-circuit initialization checks or expose an initialized property).

Copilot uses AI. Check for mistakes.
Comment on lines +50 to +53
valueFrom:
secretKeyRef:
key: latest
name: neo4j-user
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This Cloud Run manifest requires a neo4j-user Secret Manager secret, but the repo’s setup instructions/scripts shown in this PR primarily create neo4j-uri and neo4j-password. Either update the setup docs/scripts to also create neo4j-user, or change the manifest to use a plain env var (common default is neo4j).

Suggested change
valueFrom:
secretKeyRef:
key: latest
name: neo4j-user
value: "neo4j"

Copilot uses AI. Check for mistakes.
Comment on lines +44 to +48
pip install neo4j-agent-memory[google,mcp]

# Or for development
cd neo4j-agent-memory
pip install -e ".[google,mcp]"
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The README advertises installing "all Google Cloud features" with [google,mcp], but ADK support is exposed as a separate extra (google-adk) in pyproject.toml. Consider updating these commands to include google-adk (e.g., [google,google-adk,mcp]) so the documented installation actually matches the described features.

Suggested change
pip install neo4j-agent-memory[google,mcp]
# Or for development
cd neo4j-agent-memory
pip install -e ".[google,mcp]"
pip install neo4j-agent-memory[google,google-adk,mcp]
# Or for development
cd neo4j-agent-memory
pip install -e ".[google,google-adk,mcp]"

Copilot uses AI. Check for mistakes.
Comment on lines 60 to 78
query_upper = request.query.upper().strip()
forbidden_keywords = [
"CREATE",
"MERGE",
"DELETE",
"REMOVE",
"SET",
"DROP",
"DETACH",
"CALL",
"LOAD",
]

for keyword in forbidden_keywords:
if keyword in query_upper:
raise HTTPException(
status_code=400,
detail=f"Write operations ({keyword}) are not allowed. Use read-only queries.",
)
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The validation uses substring matching, which will incorrectly reject legitimate read-only queries containing keywords inside identifiers (e.g., n.created_at contains CREATE after .upper()). Use a token-based check or a regex with word boundaries (and ideally ignore string literals/comments) to avoid false positives.

Copilot uses AI. Check for mistakes.
Comment on lines +81 to +85
# Execute query through the memory client
client = memory_service.client
async with client._driver.session(database=client._database) as session:
result = await session.run(request.query, request.parameters)
records = await result.data()
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This route uses private MemoryClient internals (_driver, _database). Since this PR introduces MemoryClient.graph for custom queries, prefer using that public API (e.g., client.graph.execute_read(...)) to avoid tight coupling to internal attributes that can change.

Suggested change
# Execute query through the memory client
client = memory_service.client
async with client._driver.session(database=client._database) as session:
result = await session.run(request.query, request.parameters)
records = await result.data()
# Execute query through the memory client's public graph interface
client = memory_service.client
records = await client.graph.execute_read(
request.query,
request.parameters,
)

Copilot uses AI. Check for mistakes.
Comment on lines 20 to 21
if not service._initialized:
await service.initialize()
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Accessing service._initialized relies on a private attribute. Prefer a public readiness check (or always call await service.initialize() since it is already idempotent) to avoid depending on internal state.

Suggested change
if not service._initialized:
await service.initialize()
await service.initialize()

Copilot uses AI. Check for mistakes.
@johnymontana johnymontana merged commit 6847272 into main Feb 11, 2026
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants