Skip to content

Commit ca04fcc

Browse files
committed
feat(gemini): Add generic tools parameter for all Gemini built-in tools
- Add tools.py with wrapper classes for all 6 Gemini tools: - FileSearch: RAG over documents - GoogleSearch: Ground responses with web data - CodeExecution: Run Python code - URLContext: Read specific web pages - GoogleMaps: Location-aware queries - ComputerUse: Browser automation - Replace hardcoded file_search_store param with generic tools list - Update _build_config() to handle multiple tools - Update phone example to use new tools API - Add unit tests for all tool wrappers Usage: llm = gemini.LLM(tools=[ gemini.tools.FileSearch(store), gemini.tools.GoogleSearch(), gemini.tools.CodeExecution(), ])
1 parent 2778018 commit ca04fcc

File tree

6 files changed

+317
-16
lines changed

6 files changed

+317
-16
lines changed

examples/03_phone_and_rag_example/inbound_phone_and_rag_example.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,10 @@ async def search_knowledge(query: str) -> str:
180180
return await rag.search(query, top_k=3)
181181

182182
else:
183-
llm = gemini.LLM("gemini-2.5-flash-lite", file_search_store=file_search_store)
183+
llm = gemini.LLM(
184+
"gemini-2.5-flash-lite",
185+
tools=[gemini.tools.FileSearch(file_search_store)],
186+
)
184187

185188
return Agent(
186189
edge=getstream.Edge(),

plugins/gemini/tests/test_gemini_file_search.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
async def rag():
2323
"""Create a RAG instance for testing, clean up after."""
2424
# Use unique name to avoid conflicts
25-
rag = GeminiFilesearchRAG(name=f"test-rag-123")
25+
rag = GeminiFilesearchRAG(name="test-rag-123")
2626
await rag.create()
2727
yield rag
2828
await rag.clear()
Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
"""Tests for Gemini built-in tools."""
2+
3+
from unittest.mock import MagicMock
4+
5+
from google.genai import types
6+
7+
from vision_agents.plugins.gemini import tools
8+
9+
10+
class TestGoogleSearch:
11+
"""Tests for GoogleSearch tool."""
12+
13+
def test_to_tool(self):
14+
"""Test GoogleSearch converts to Tool correctly."""
15+
tool = tools.GoogleSearch()
16+
result = tool.to_tool()
17+
18+
assert isinstance(result, types.Tool)
19+
assert result.google_search is not None
20+
21+
22+
class TestCodeExecution:
23+
"""Tests for CodeExecution tool."""
24+
25+
def test_to_tool(self):
26+
"""Test CodeExecution converts to Tool correctly."""
27+
tool = tools.CodeExecution()
28+
result = tool.to_tool()
29+
30+
assert isinstance(result, types.Tool)
31+
assert result.code_execution is not None
32+
33+
34+
class TestURLContext:
35+
"""Tests for URLContext tool."""
36+
37+
def test_to_tool(self):
38+
"""Test URLContext converts to Tool correctly."""
39+
tool = tools.URLContext()
40+
result = tool.to_tool()
41+
42+
assert isinstance(result, types.Tool)
43+
assert result.url_context is not None
44+
45+
46+
class TestGoogleMaps:
47+
"""Tests for GoogleMaps tool."""
48+
49+
def test_to_tool(self):
50+
"""Test GoogleMaps converts to Tool correctly."""
51+
tool = tools.GoogleMaps()
52+
result = tool.to_tool()
53+
54+
assert isinstance(result, types.Tool)
55+
assert result.google_maps is not None
56+
57+
58+
class TestComputerUse:
59+
"""Tests for ComputerUse tool."""
60+
61+
def test_to_tool(self):
62+
"""Test ComputerUse converts to Tool correctly."""
63+
tool = tools.ComputerUse()
64+
result = tool.to_tool()
65+
66+
assert isinstance(result, types.Tool)
67+
assert result.computer_use is not None
68+
69+
def test_custom_environment(self):
70+
"""Test ComputerUse with custom environment."""
71+
tool = tools.ComputerUse(environment=types.Environment.ENVIRONMENT_BROWSER)
72+
result = tool.to_tool()
73+
74+
assert result.computer_use.environment == types.Environment.ENVIRONMENT_BROWSER
75+
76+
77+
class TestFileSearch:
78+
"""Tests for FileSearch tool."""
79+
80+
def test_to_tool_with_created_store(self):
81+
"""Test FileSearch converts to Tool when store is created."""
82+
# Mock the store
83+
mock_store = MagicMock()
84+
mock_store.is_created = True
85+
mock_store.get_tool.return_value = types.Tool(
86+
file_search=types.FileSearch(file_search_store_names=["test-store"])
87+
)
88+
89+
tool = tools.FileSearch(mock_store)
90+
result = tool.to_tool()
91+
92+
assert isinstance(result, types.Tool)
93+
assert result.file_search is not None
94+
mock_store.get_tool.assert_called_once()
95+
96+
def test_to_tool_raises_when_store_not_created(self):
97+
"""Test FileSearch raises error when store is not created."""
98+
mock_store = MagicMock()
99+
mock_store.is_created = False
100+
101+
tool = tools.FileSearch(mock_store)
102+
103+
try:
104+
tool.to_tool()
105+
assert False, "Should have raised ValueError"
106+
except ValueError as e:
107+
assert "not created" in str(e)
Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,18 @@
11
from .gemini_llm import GeminiLLM as LLM
22
from .gemini_realtime import GeminiRealtime as Realtime
33
from .file_search import GeminiFilesearchRAG, FileSearchStore, create_file_search_store
4+
from . import tools
45
from google.genai.types import ThinkingLevel, MediaResolution
56

67
__all__ = [
78
"Realtime",
89
"LLM",
910
"ThinkingLevel",
1011
"MediaResolution",
12+
# Tools
13+
"tools",
14+
# File Search (convenience exports)
1115
"GeminiFilesearchRAG",
12-
"FileSearchStore", # Backwards compatibility alias
16+
"FileSearchStore",
1317
"create_file_search_store",
1418
]

plugins/gemini/vision_agents/plugins/gemini/gemini_llm.py

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
)
2020

2121
from . import events
22-
from .file_search import FileSearchStore
22+
from .tools import GeminiTool
2323

2424
from vision_agents.core.processors import Processor
2525

@@ -58,7 +58,7 @@ def __init__(
5858
thinking_level: Optional[ThinkingLevel] = None,
5959
media_resolution: Optional[MediaResolution] = None,
6060
config: Optional[GenerateContentConfig] = None,
61-
file_search_store: Optional[FileSearchStore] = None,
61+
tools: Optional[List[GeminiTool]] = None,
6262
**kwargs,
6363
):
6464
"""
@@ -77,17 +77,22 @@ def __init__(
7777
"low"/"medium" for general video, "high" for text-heavy video.
7878
config: Optional[GenerateContentConfig] to use as base. Any kwargs will be passed
7979
to GenerateContentConfig constructor if config is not provided.
80-
file_search_store: Optional FileSearchStore for RAG functionality. When provided,
81-
the model will use the indexed documents to answer questions.
82-
See: https://ai.google.dev/gemini-api/docs/file-search
80+
tools: Optional list of Gemini built-in tools. Available tools:
81+
- tools.FileSearch(store): RAG over your documents
82+
- tools.GoogleSearch(): Ground responses with web data
83+
- tools.CodeExecution(): Run Python code
84+
- tools.URLContext(): Read specific web pages
85+
- tools.GoogleMaps(): Location-aware queries (Preview)
86+
- tools.ComputerUse(): Browser automation (Preview)
87+
See: https://ai.google.dev/gemini-api/docs/tools
8388
**kwargs: Additional arguments passed to GenerateContentConfig constructor.
8489
"""
8590
super().__init__()
8691
self.events.register_events_from_module(events)
8792
self.model = model
8893
self.thinking_level = thinking_level
8994
self.media_resolution = media_resolution
90-
self.file_search_store = file_search_store
95+
self._builtin_tools = tools or []
9196

9297
if config is not None:
9398
self._base_config: Optional[GenerateContentConfig] = config
@@ -109,15 +114,15 @@ def _build_config(
109114
base_config: Optional[GenerateContentConfig] = None,
110115
) -> GenerateContentConfig:
111116
"""
112-
Build GenerateContentConfig with Gemini 3 features and file search.
117+
Build GenerateContentConfig with Gemini 3 features and built-in tools.
113118
114119
Args:
115120
system_instruction: Optional system instruction to include. If not provided,
116121
uses self._instructions to ensure instructions are always passed.
117122
base_config: Optional base config to extend (takes precedence over self._base_config)
118123
119124
Returns:
120-
GenerateContentConfig with thinking_level, media_resolution, and file_search if set
125+
GenerateContentConfig with thinking_level, media_resolution, and tools if set
121126
"""
122127
if base_config is not None:
123128
config = base_config
@@ -142,14 +147,14 @@ def _build_config(
142147
if self.media_resolution:
143148
config.media_resolution = self.media_resolution
144149

145-
# Add file search tool if store is configured
146-
if self.file_search_store and self.file_search_store.is_created:
147-
file_search_tool = self.file_search_store.get_tool()
150+
# Add built-in tools if configured
151+
if self._builtin_tools:
152+
builtin_tool_objects = [tool.to_tool() for tool in self._builtin_tools]
148153
if config.tools is None:
149-
config.tools = [file_search_tool]
154+
config.tools = builtin_tool_objects
150155
else:
151156
# Append to existing tools
152-
config.tools = list(config.tools) + [file_search_tool]
157+
config.tools = list(config.tools) + builtin_tool_objects
153158

154159
return config
155160

0 commit comments

Comments
 (0)