Skip to content

Add GeminiAgent for Google Gemini model support#451

Open
gavksingh wants to merge 2 commits into
2FastLabs:mainfrom
gavksingh:feat/gemini-agent
Open

Add GeminiAgent for Google Gemini model support#451
gavksingh wants to merge 2 commits into
2FastLabs:mainfrom
gavksingh:feat/gemini-agent

Conversation

@gavksingh

@gavksingh gavksingh commented Mar 17, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Adds GeminiAgent and GeminiAgentOptions for integrating Google Gemini models via the google-genai SDK
  • Follows the same pattern as OpenAIAgent — supports streaming/non-streaming, system prompts, retriever integration, and pre-configured client injection
  • Adds gemini optional dependency group (pip install agent-squad[gemini])
  • Gracefully degrades when google-genai is not installed (conditional import)
  • Added comprehensive unit tests and a Vertex AI integration test

Files Changed

  • python/src/agent_squad/agents/gemini_agent.pynew agent implementation
  • python/src/agent_squad/agents/__init__.py — conditional import for GeminiAgent
  • python/setup.cfggemini extras group with google-genai>=1.0.0
  • python/src/agent_squad/types/types.pyGEMINI_MODEL_ID_GEMINI_2_0_FLASH constant
  • python/src/agent_squad/types/__init__.py — export the new constant
  • python/src/tests/agents/test_gemini_agent.pynew test file
  • python/src/tests/pytest.ini — registered integration marker

Tests added

New file: python/src/tests/agents/test_gemini_agent.py (8 unit tests + 1 integration test)

Unit tests:

  • test_gemini_agent_requires_api_key_or_client — ValueError when neither provided
  • test_gemini_agent_client_injection — pre-configured client bypasses api_key requirement
  • test_custom_system_prompt_with_variable — template {{variable}} resolved correctly
  • test_process_request_success — non-streaming returns ConversationMessage
  • test_process_request_streaming — streaming yields AgentStreamResponse chunks + final_message
  • test_process_request_with_retriever — retriever called, context appended
  • test_process_request_api_error — exception propagation
  • test_is_streaming_enabled — False default, True when set

Integration test (requires credentials):

  • test_gemini_vertex_ai_integration — end-to-end via Vertex AI with GOOGLE_APPLICATION_CREDENTIALS

Run integration test with:

GOOGLE_APPLICATION_CREDENTIALS=/path/to/credentials.json pytest -m integration python/src/tests/agents/test_gemini_agent.py -v

Test plan

  • from agent_squad.agents import GeminiAgent works with google-genai installed
  • Import doesn't break when google-genai is NOT installed
  • All 8 unit tests pass (pytest -v -m "not integration")
  • Integration test with valid Vertex AI project

Fixes #376

Implement GeminiAgent and GeminiAgentOptions following the same pattern
as OpenAIAgent, using the google-genai SDK. Supports both streaming and
non-streaming responses, system prompts, retriever integration, and
pre-configured client injection.

- New file: agents/gemini_agent.py
- Conditional import in agents/__init__.py (graceful when google-genai not installed)
- New 'gemini' optional dependency group in setup.cfg
- GEMINI_MODEL_ID_GEMINI_2_0_FLASH constant in types

Fixes 2FastLabs#376
Unit tests cover initialization, client injection, custom prompts,
streaming/non-streaming requests, retriever integration, and error
handling. Integration test validates Vertex AI connectivity (requires
GOOGLE_APPLICATION_CREDENTIALS). Registers 'integration' pytest marker.
@cornelcroi

Copy link
Copy Markdown
Collaborator

Hey @gavksingh, a few things need to be fixed before this can merge.

Three are required: (1) there is no docs/src/content/docs/agents/built-in/gemini-agent.mdx page — every other built-in agent has one, and project rules require a usage example plus a description of each GeminiAgentOptions field; (2) the integration test has a hardcoded GCP project ID (project="project-cf964f7d-d79b-4b69-81c") that looks like a real private ID — replace it with something like os.environ.get("GOOGLE_CLOUD_PROJECT"); (3) CI has not run on this branch, it must pass before merge. There is also a logic bug in test_gemini_agent_requires_api_key_or_client — the first GeminiAgentOptions(...) call before the pytest.raises block is dead code and the test passes by accident, not because it actually validates the error path.

Two smaller things that are not blocking: GeminiAgentOptions uses model_id while OpenAIAgentOptions uses model — worth aligning for consistency. Also, handle_streaming_response iterates a sync iterator inside an async generator which blocks the event loop, same pattern as OpenAIAgent so it may be acceptable, but good to be aware of.

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.

Feature request: Implementing gemini agent out of the box

2 participants