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.
| Returns: | ||
| A SearchMemoryResponse containing the matching memories. | ||
| """ | ||
| from google.genai import types |
There was a problem hiding this comment.
For consistency and better code organization, module-level imports should be placed at the top of the file. Please move this import statement to the top with other imports.
| import urllib.error | ||
| import urllib.request |
There was a problem hiding this comment.
For making HTTP requests, it's better to use the requests library, which is already a project dependency. It provides a simpler and more modern API compared to urllib and will make the code more consistent with other parts of the codebase that might use it.
| import urllib.error | |
| import urllib.request | |
| import requests |
| def _post_embed(self, texts: list[str]) -> dict: | ||
| """Perform a blocking POST /api/embed call to Ollama. | ||
|
|
||
| Args: | ||
| texts: A list of strings to embed. | ||
|
|
||
| Returns: | ||
| The JSON response from Ollama. | ||
|
|
||
| Raises: | ||
| RuntimeError: If the request fails. | ||
| """ | ||
| url = self._host.rstrip("/") + _EMBED_ENDPOINT | ||
| payload = { | ||
| "model": self._model, | ||
| "input": texts, | ||
| } | ||
| 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.
Using the requests library simplifies this method significantly. It handles JSON serialization, setting the Content-Type header, and provides a more straightforward way to handle HTTP errors via raise_for_status().
def _post_embed(self, texts: list[str]) -> dict:
"""Perform a blocking POST /api/embed call to Ollama.
Args:
texts: A list of strings to embed.
Returns:
The JSON response from Ollama.
Raises:
RuntimeError: If the request fails.
"""
url = self._host.rstrip("/") + _EMBED_ENDPOINT
payload = {
"model": self._model,
"input": texts,
}
try:
response = requests.post(
url,
json=payload,
timeout=self._request_timeout,
)
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as exc:
raise RuntimeError(f"Failed to connect to Ollama: {exc}") from exc
ayman3000
force-pushed
the
feature/chroma-memory-service
branch
from
90c145b to
c352247
Compare
2 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: