mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-01-14 01:48:29 -05:00
Compare commits
7 Commits
host-error
...
docs/toolb
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e14b5437b1 | ||
|
|
8944ffc658 | ||
|
|
eeb8f7b614 | ||
|
|
1760253462 | ||
|
|
ec8bd283ee | ||
|
|
ef0b51750e | ||
|
|
8061ef73e1 |
@@ -509,7 +509,7 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"! pip install toolbox-core --quiet\n",
|
||||
"! pip install toolbox-adk --quiet\n",
|
||||
"! pip install google-adk --quiet"
|
||||
]
|
||||
},
|
||||
@@ -525,14 +525,18 @@
|
||||
"from google.adk.runners import Runner\n",
|
||||
"from google.adk.sessions import InMemorySessionService\n",
|
||||
"from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService\n",
|
||||
"from google.adk.tools.toolbox_toolset import ToolboxToolset\n",
|
||||
"from google.genai import types\n",
|
||||
"from toolbox_core import ToolboxSyncClient\n",
|
||||
"\n",
|
||||
"import os\n",
|
||||
"# TODO(developer): replace this with your Google API key\n",
|
||||
"os.environ['GOOGLE_API_KEY'] = \"<GOOGLE_API_KEY>\"\n",
|
||||
"\n",
|
||||
"toolbox_client = ToolboxSyncClient(\"http://127.0.0.1:5000\")\n",
|
||||
"# Configure toolset\n",
|
||||
"toolset = ToolboxToolset(\n",
|
||||
" server_url=\"http://127.0.0.1:5000\",\n",
|
||||
" toolset_name=\"my-toolset\"\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"prompt = \"\"\"\n",
|
||||
" You're a helpful hotel assistant. You handle hotel searching, booking and\n",
|
||||
@@ -549,7 +553,7 @@
|
||||
" name='hotel_agent',\n",
|
||||
" description='A helpful AI assistant.',\n",
|
||||
" instruction=prompt,\n",
|
||||
" tools=toolbox_client.load_toolset(\"my-toolset\"),\n",
|
||||
" tools=[toolset],\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"session_service = InMemorySessionService()\n",
|
||||
|
||||
@@ -52,7 +52,7 @@ runtime](https://research.google.com/colaboratory/local-runtimes.html).
|
||||
{{< tabpane persist=header >}}
|
||||
{{< tab header="ADK" lang="bash" >}}
|
||||
|
||||
pip install toolbox-core
|
||||
pip install toolbox-adk
|
||||
{{< /tab >}}
|
||||
{{< tab header="Langchain" lang="bash" >}}
|
||||
|
||||
|
||||
@@ -1,15 +1,17 @@
|
||||
from google.adk import Agent
|
||||
from google.adk.apps import App
|
||||
from toolbox_core import ToolboxSyncClient
|
||||
from google.adk.tools.toolbox_toolset import ToolboxToolset
|
||||
|
||||
# TODO(developer): update the TOOLBOX_URL to your toolbox endpoint
|
||||
client = ToolboxSyncClient("http://127.0.0.1:5000")
|
||||
toolset = ToolboxToolset(
|
||||
server_url="http://127.0.0.1:5000",
|
||||
)
|
||||
|
||||
root_agent = Agent(
|
||||
name='root_agent',
|
||||
model='gemini-2.5-flash',
|
||||
instruction="You are a helpful AI assistant designed to provide accurate and useful information.",
|
||||
tools=client.load_toolset(),
|
||||
tools=[toolset],
|
||||
)
|
||||
|
||||
app = App(root_agent=root_agent, name="my_agent")
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
google-adk==1.21.0
|
||||
toolbox-core==0.5.4
|
||||
toolbox-adk>=0.1.0
|
||||
pytest==9.0.2
|
||||
@@ -49,7 +49,7 @@ with the necessary configuration for deployment to Vertex AI Agent Engine.
|
||||
4. Add `toolbox-core` as a dependency to the new project:
|
||||
|
||||
```bash
|
||||
uv add toolbox-core
|
||||
uv add toolbox-adk
|
||||
```
|
||||
|
||||
## Step 3: Configure Google Cloud Authentication
|
||||
@@ -95,22 +95,23 @@ authentication token.
|
||||
```python
|
||||
from google.adk import Agent
|
||||
from google.adk.apps import App
|
||||
from toolbox_core import ToolboxSyncClient, auth_methods
|
||||
from google.adk.tools.toolbox_toolset import ToolboxToolset
|
||||
from toolbox_adk import CredentialStrategy
|
||||
|
||||
# TODO(developer): Replace with your Toolbox Cloud Run Service URL
|
||||
TOOLBOX_URL = "https://your-toolbox-service-xyz.a.run.app"
|
||||
|
||||
# Initialize the client with the Cloud Run URL and Auth headers
|
||||
client = ToolboxSyncClient(
|
||||
TOOLBOX_URL,
|
||||
client_headers={"Authorization": auth_methods.get_google_id_token(TOOLBOX_URL)}
|
||||
# Initialize the toolset with Workload Identity (generates ID token for the URL)
|
||||
toolset = ToolboxToolset(
|
||||
server_url=TOOLBOX_URL,
|
||||
credentials=CredentialStrategy.workload_identity(target_audience=TOOLBOX_URL)
|
||||
)
|
||||
|
||||
root_agent = Agent(
|
||||
name='root_agent',
|
||||
model='gemini-2.5-flash',
|
||||
instruction="You are a helpful AI assistant designed to provide accurate and useful information.",
|
||||
tools=client.load_toolset(),
|
||||
tools=[toolset],
|
||||
)
|
||||
|
||||
app = App(root_agent=root_agent, name="my_agent")
|
||||
|
||||
@@ -365,7 +365,7 @@ pip install llama-index-llms-google-genai
|
||||
|
||||
{{< /tab >}}
|
||||
{{< tab header="ADK" lang="bash" >}}
|
||||
pip install toolbox-core
|
||||
pip install toolbox-adk
|
||||
{{< /tab >}}
|
||||
{{< /tabpane >}}
|
||||
|
||||
@@ -607,8 +607,8 @@ from google.adk.agents import Agent
|
||||
from google.adk.runners import Runner
|
||||
from google.adk.sessions import InMemorySessionService
|
||||
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
||||
from google.adk.tools.toolbox_toolset import ToolboxToolset
|
||||
from google.genai import types # For constructing message content
|
||||
from toolbox_core import ToolboxSyncClient
|
||||
|
||||
import os
|
||||
os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'True'
|
||||
@@ -623,48 +623,47 @@ os.environ['GOOGLE_CLOUD_LOCATION'] = 'us-central1'
|
||||
|
||||
# --- Load Tools from Toolbox ---
|
||||
|
||||
# TODO(developer): Ensure the Toolbox server is running at <http://127.0.0.1:5000>
|
||||
# TODO(developer): Ensure the Toolbox server is running at http://127.0.0.1:5000
|
||||
toolset = ToolboxToolset(server_url="http://127.0.0.1:5000")
|
||||
|
||||
with ToolboxSyncClient("<http://127.0.0.1:5000>") as toolbox_client:
|
||||
# TODO(developer): Replace "my-toolset" with the actual ID of your toolset as configured in your MCP Toolbox server.
|
||||
agent_toolset = toolbox_client.load_toolset("my-toolset")
|
||||
# --- Define the Agent's Prompt ---
|
||||
prompt = """
|
||||
You're a helpful hotel assistant. You handle hotel searching, booking and
|
||||
cancellations. When the user searches for a hotel, mention it's name, id,
|
||||
location and price tier. Always mention hotel ids while performing any
|
||||
searches. This is very important for any operations. For any bookings or
|
||||
cancellations, please provide the appropriate confirmation. Be sure to
|
||||
update checkin or checkout dates if mentioned by the user.
|
||||
Don't ask for confirmations from the user.
|
||||
"""
|
||||
|
||||
# --- Define the Agent's Prompt ---
|
||||
prompt = """
|
||||
You're a helpful hotel assistant. You handle hotel searching, booking and
|
||||
cancellations. When the user searches for a hotel, mention it's name, id,
|
||||
location and price tier. Always mention hotel ids while performing any
|
||||
searches. This is very important for any operations. For any bookings or
|
||||
cancellations, please provide the appropriate confirmation. Be sure to
|
||||
update checkin or checkout dates if mentioned by the user.
|
||||
Don't ask for confirmations from the user.
|
||||
"""
|
||||
# --- Configure the Agent ---
|
||||
|
||||
# --- Configure the Agent ---
|
||||
root_agent = Agent(
|
||||
model='gemini-2.0-flash-001',
|
||||
name='hotel_agent',
|
||||
description='A helpful AI assistant that can search and book hotels.',
|
||||
instruction=prompt,
|
||||
tools=[toolset], # Pass the loaded toolset
|
||||
)
|
||||
|
||||
root_agent = Agent(
|
||||
model='gemini-2.0-flash-001',
|
||||
name='hotel_agent',
|
||||
description='A helpful AI assistant that can search and book hotels.',
|
||||
instruction=prompt,
|
||||
tools=agent_toolset, # Pass the loaded toolset
|
||||
)
|
||||
# --- Initialize Services for Running the Agent ---
|
||||
session_service = InMemorySessionService()
|
||||
artifacts_service = InMemoryArtifactService()
|
||||
|
||||
# --- Initialize Services for Running the Agent ---
|
||||
session_service = InMemorySessionService()
|
||||
artifacts_service = InMemoryArtifactService()
|
||||
runner = Runner(
|
||||
app_name='hotel_agent',
|
||||
agent=root_agent,
|
||||
artifact_service=artifacts_service,
|
||||
session_service=session_service,
|
||||
)
|
||||
|
||||
async def main():
|
||||
# Create a new session for the interaction.
|
||||
session = session_service.create_session(
|
||||
session = await session_service.create_session(
|
||||
state={}, app_name='hotel_agent', user_id='123'
|
||||
)
|
||||
|
||||
runner = Runner(
|
||||
app_name='hotel_agent',
|
||||
agent=root_agent,
|
||||
artifact_service=artifacts_service,
|
||||
session_service=session_service,
|
||||
)
|
||||
|
||||
# --- Define Queries and Run the Agent ---
|
||||
queries = [
|
||||
"Find hotels in Basel with Basel in it's name.",
|
||||
@@ -687,6 +686,10 @@ with ToolboxSyncClient("<http://127.0.0.1:5000>") as toolbox_client:
|
||||
|
||||
for text in responses:
|
||||
print(text)
|
||||
|
||||
import asyncio
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
{{< /tab >}}
|
||||
{{< /tabpane >}}
|
||||
|
||||
|
||||
@@ -304,14 +304,10 @@ func hostCheck(allowedHosts map[string]struct{}) func(http.Handler) http.Handler
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, hasWildcard := allowedHosts["*"]
|
||||
hostname := r.Host
|
||||
if host, _, err := net.SplitHostPort(r.Host); err == nil {
|
||||
hostname = host
|
||||
}
|
||||
_, hostIsAllowed := allowedHosts[hostname]
|
||||
_, hostIsAllowed := allowedHosts[r.Host]
|
||||
if !hasWildcard && !hostIsAllowed {
|
||||
// Return 403 Forbidden to block the attack
|
||||
http.Error(w, "Invalid Host header", http.StatusForbidden)
|
||||
// Return 400 Bad Request or 403 Forbidden to block the attack
|
||||
http.Error(w, "Invalid Host header", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
@@ -410,11 +406,7 @@ func NewServer(ctx context.Context, cfg ServerConfig) (*Server, error) {
|
||||
}
|
||||
allowedHostsMap := make(map[string]struct{}, len(cfg.AllowedHosts))
|
||||
for _, h := range cfg.AllowedHosts {
|
||||
hostname := h
|
||||
if host, _, err := net.SplitHostPort(h); err == nil {
|
||||
hostname = host
|
||||
}
|
||||
allowedHostsMap[hostname] = struct{}{}
|
||||
allowedHostsMap[h] = struct{}{}
|
||||
}
|
||||
r.Use(hostCheck(allowedHostsMap))
|
||||
|
||||
|
||||
@@ -85,66 +85,13 @@ func initDataplexConnection(ctx context.Context) (*dataplex.CatalogClient, error
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// cleanupOldAspectTypes Deletes AspectTypes older than the specified duration.
|
||||
func cleanupOldAspectTypes(t *testing.T, ctx context.Context, client *dataplex.CatalogClient, oldThreshold time.Duration) {
|
||||
parent := fmt.Sprintf("projects/%s/locations/us", DataplexProject)
|
||||
olderThanTime := time.Now().Add(-oldThreshold)
|
||||
|
||||
listReq := &dataplexpb.ListAspectTypesRequest{
|
||||
Parent: parent,
|
||||
PageSize: 100, // Fetch up to 100 items
|
||||
OrderBy: "create_time asc", // Order by creation time
|
||||
}
|
||||
|
||||
const maxDeletes = 8 // Explicitly limit the number of deletions
|
||||
it := client.ListAspectTypes(ctx, listReq)
|
||||
var aspectTypesToDelete []string
|
||||
for len(aspectTypesToDelete) < maxDeletes {
|
||||
aspectType, err := it.Next()
|
||||
if err == iterator.Done {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
t.Logf("Warning: Failed to list aspect types during cleanup: %v", err)
|
||||
return
|
||||
}
|
||||
// Perform time-based filtering in memory
|
||||
if aspectType.CreateTime != nil {
|
||||
createTime := aspectType.CreateTime.AsTime()
|
||||
if createTime.Before(olderThanTime) {
|
||||
aspectTypesToDelete = append(aspectTypesToDelete, aspectType.GetName())
|
||||
}
|
||||
} else {
|
||||
t.Logf("Warning: AspectType %s has no CreateTime", aspectType.GetName())
|
||||
}
|
||||
}
|
||||
if len(aspectTypesToDelete) == 0 {
|
||||
t.Logf("cleanupOldAspectTypes: No aspect types found older than %s to delete.", oldThreshold.String())
|
||||
return
|
||||
}
|
||||
|
||||
for _, aspectTypeName := range aspectTypesToDelete {
|
||||
deleteReq := &dataplexpb.DeleteAspectTypeRequest{Name: aspectTypeName}
|
||||
op, err := client.DeleteAspectType(ctx, deleteReq)
|
||||
if err != nil {
|
||||
t.Logf("Warning: Failed to delete aspect type %s: %v", aspectTypeName, err)
|
||||
continue // Skip to the next item if initiation fails
|
||||
}
|
||||
|
||||
if err := op.Wait(ctx); err != nil {
|
||||
t.Logf("Warning: Failed to delete aspect type %s, operation error: %v", aspectTypeName, err)
|
||||
} else {
|
||||
t.Logf("cleanupOldAspectTypes: Successfully deleted %s", aspectTypeName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDataplexToolEndpoints(t *testing.T) {
|
||||
sourceConfig := getDataplexVars(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
var args []string
|
||||
|
||||
bigqueryClient, err := initBigQueryConnection(ctx, DataplexProject)
|
||||
if err != nil {
|
||||
t.Fatalf("unable to create Cloud SQL connection pool: %s", err)
|
||||
@@ -155,9 +102,6 @@ func TestDataplexToolEndpoints(t *testing.T) {
|
||||
t.Fatalf("unable to create Dataplex connection: %s", err)
|
||||
}
|
||||
|
||||
// Cleanup older aspecttypes
|
||||
cleanupOldAspectTypes(t, ctx, dataplexClient, 1*time.Hour)
|
||||
|
||||
// create resources with UUID
|
||||
datasetName := fmt.Sprintf("temp_toolbox_test_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
|
||||
tableName := fmt.Sprintf("param_table_%s", strings.ReplaceAll(uuid.New().String(), "-", ""))
|
||||
|
||||
Reference in New Issue
Block a user