Feature/chroma memory service #4197
Conversation
Summary of ChangesHello @ayman3000, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a robust Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
adk-bot
added
the
services
- Add ChromaMemoryService for semantic search over session memories - Add BaseEmbeddingProvider abstract base class - Add OllamaEmbeddingProvider using Ollama's /api/embed endpoint - Add chromadb as optional dependency - Add comprehensive unit tests
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a ChromaMemoryService to provide semantic memory search capabilities using ChromaDB, which is a great alternative to InMemoryMemoryService for local development and self-hosted deployments requiring persistence and semantic search. The implementation is well-structured, with a pluggable BaseEmbeddingProvider and an initial OllamaEmbeddingProvider. The addition of unit tests and a sample application is also excellent. My review includes a few suggestions to improve maintainability and consistency.
ayman3000
force-pushed
the
feature/chroma-memory-service
branch
from
90c145b to
c352247
Compare
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
ryanaiagent
added
the
request clarification
|
Hi @ayman3000 , Thank you for your contribution! We appreciate you taking the time to submit this pull request. Can you fix the formatting errors and address my comment i left before we can proceed with the review. |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a ChromaMemoryService for semantic memory search using ChromaDB, along with a pluggable embedding provider architecture demonstrated with an OllamaEmbeddingProvider. The changes are well-structured, including the core service, an embedding provider base class, an Ollama implementation, unit tests, and a sample application. My review focuses on ensuring correctness, maintainability, and robustness. I've found a critical issue in the Ollama provider's network request implementation that needs to be addressed, along with a couple of medium-severity suggestions to improve code quality and user experience.
| data = json.dumps(payload).encode("utf-8") | ||
| request = urllib.request.Request( | ||
| url, | ||
| data=data, | ||
| headers={"Content-Type": "application/json"}, | ||
| method="POST", | ||
| ) | ||
|
|
||
| try: | ||
| with urllib.request.urlopen( | ||
| request, timeout=self._request_timeout | ||
| ) as response: | ||
| response_body = response.read().decode("utf-8") | ||
| except urllib.error.URLError as exc: | ||
| raise RuntimeError(f"Failed to connect to Ollama: {exc.reason}") from exc | ||
| except urllib.error.HTTPError as exc: | ||
| message = exc.read().decode("utf-8", errors="ignore") | ||
| raise RuntimeError(f"Ollama API error {exc.code}: {message}") from exc | ||
|
|
||
| return json.loads(response_body) |
There was a problem hiding this comment.
This block has a critical issue: it attempts to use urllib.request and urllib.error without importing the urllib module, which will cause a NameError at runtime.
Additionally, the requests library is imported at the top of the file but is not used.
I suggest replacing this entire block with an implementation that uses the requests library. This will fix the bug, make use of the existing import, and result in more concise and readable code.
try:
response = requests.post(
url,
json=payload,
timeout=self._request_timeout,
)
response.raise_for_status() # Raises an HTTPError for bad responses
return response.json()
except requests.exceptions.RequestException as exc:
raise RuntimeError(f"Failed to connect to Ollama: {exc}") from exc| Returns: | ||
| A SearchMemoryResponse containing the matching memories. | ||
| """ | ||
| from google.genai import types |
There was a problem hiding this comment.
For consistency and to adhere to standard Python style (PEP 8), imports should be at the top of the file. Please move this import to the top-level imports section of the file.
src/google/adk/memory/__init__.py
Outdated
| 'chromadb is not installed. If you want to use the ChromaMemoryService' | ||
| ' please install it with: pip install chromadb. If not, you can ignore' | ||
| ' this warning.' |
There was a problem hiding this comment.
The error message suggests pip install chromadb, which is helpful. However, to ensure users install a version compatible with this library as defined in pyproject.toml, it's better to guide them to install the optional dependency group. This prevents potential version conflicts.
| 'chromadb is not installed. If you want to use the ChromaMemoryService' | |
| ' please install it with: pip install chromadb. If not, you can ignore' | |
| ' this warning.' | |
| 'chromadb is not installed. If you want to use the ChromaMemoryService' | |
| ' please install it with: pip install \'google-adk[chroma]\'. If not, you can' | |
| ' ignore this warning.' |
3 participants
Problem:
The current
InMemoryMemoryServiceuses simple keyword matching to search agent memories. This means:The
VertexAiRagMemoryServiceprovides semantic search, but requires Google Cloud Platform infrastructure, which isn't suitable for local development or self-hosted deployments.Solution:
Add a new
ChromaMemoryServicethat provides semantic search capabilities using ChromaDB with pluggable embedding providers. This gives developers:Testing Plan
Unit Tests:
Manual End-to-End (E2E) Tests:
To manually test:
Start Ollama server:
Run the example:
cd contributing/samples/memory_chroma python main.pyExpected behavior:
Checklist
Additional context
Files Added:
src/google/adk/memory/chroma_memory_service.py- Main memory servicesrc/google/adk/memory/embeddings/base_embedding_provider.py- Abstract base classsrc/google/adk/memory/embeddings/ollama_embedding_provider.py- Ollama integrationtests/unittests/memory/test_chroma_memory_service.py- Unit testscontributing/samples/memory_chroma/- Example projectFiles Modified:
pyproject.toml- Added optionalchromadependencysrc/google/adk/memory/__init__.py- Added exportsUsage Example: