Commit Graph

45 Commits

Author SHA1 Message Date
Yuan
a0ac5334d1 chore: return error for untested fields in tools.yaml (#239)
This only checks within `SourceConfig`, `ToolConfig`, and
`AuthSourceConfig`.

Error when an unknown field is provided:
`2025-01-27T22:43:46.988401-08:00 ERROR "unable to parse tool file at
\"tools.yaml\": unable to parse as \"cloud-sql-postgres\": [2:1] unknown
field \"extra\"\n 1 | database: test_database\n> 2 | extra: here\n ^\n 3
| instance: toolbox-cloudsql\n 4 | kind: cloud-sql-postgres\n 5 |
password: postgres\n 6 | "`

Error when a required field is not provided:
`2025-01-27T17:49:47.584846-08:00 ERROR "unable to parse tool file at
\"tools.yaml\": validation failed: Key: 'Config.Region' Error:Field
validation for 'Region' failed on the 'required' tag"`

---------

Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
2025-02-03 15:30:27 -08:00
Shivaji Kharse
617cc872d1 feat: add dgraph tool and source (#233)
* add dgraph tool and source
2025-02-01 15:32:06 -07:00
Yuan
8fca0a95ee feat(sqlserver): add mssql source (#255)
Add `mssql` source for non-cloud sql server.
2025-01-31 19:13:54 +00:00
Yuan
378692ab50 feat(mysql): add source (#250)
Add `mysql` source.
2025-01-31 10:48:05 -08:00
Yuan
8452f8eb44 feat: add user agent to cloud databases (#244)
Add user agent to cloud databases that provides us anonymized data
request count, number of users, number of projects, and other
environment settings.

User agent is using the format: `genai-toolbox/$version+metadata`
2025-01-29 17:19:52 -08:00
Yuan
1357be2569 fix(mssql): fix mssql tool kind to mssql-sql (#249)
Fixing `mssql` tool to `mssql-sql`.
2025-01-29 23:46:01 +00:00
Yuan
669d6b7239 fix(mysql): fix mysql tool kind to mysql-sql (#248)
Update `mysql` tool to `mysql-sql`.
2025-01-29 23:09:13 +00:00
Wenxin Du
4639cc6560 fix: Auth token verification failure should not throw error immedicately (#234)
Currently, we are throwing 401 error immediately after auth token
verification failure. This is not expected in the following situations:
1. Non-auth tool invocation with auth token that is invalid.
2. Auth tool invocation with all the required auth token, but the header
contains extra non-required token that is invalid
These requests should pass the authorization check but fail under the
current implementation.

Change made in this PR:
1. Do not throw error immediately after auth token verification failure.
Instead only log it and continue to the next header iteration.
2. In the parseParams() method, if an auth parameter is missing, we
should error with the message telling the user that either the auth
header is missing or is invalid.
2025-01-24 08:49:51 -07:00
Wenxin Du
9bad952060 feat: Add Cloud SQL for SQL Server Source and Tool (#223)
1. `sql/database` provides a `Scan()`interface to scan query results
into typed variables. Therefore we have to create a slice of typed
variables (types retrieved from rows.ColumnTypes()) to pass them into
`Scan()`. Using []byte works but makes the printing result different
from other tools (e.g [1] instead of %!s(int32=1)]
2. MS SQL supports both named (e.g @name) and positional args (e.g @p2),
so we have to check if the name is contained in the original statement
before passing them into `db.Query()` as either named arg or as values.
2025-01-23 21:21:12 +08:00
Yuan
f1f61d7087 feat: add cloudsql mysql source and tool (#221)
Add CloudSQL for MySQL source and tool.

CloudSQLMySQL source is initialize with the following config:
```
sources:
    my-cloudsqlmysql-source:
        kind: cloud-sql-mysql
        project: my-project-name
        region: my-region
        instance: my-instance-name
        user: my_user
        password: my_pass
        database: my_db
        # ipType: public # The default dialect is public.
```

MySQL tool is initialize with the following config.
```
tools:
    test_tool:
        kind: mysql
        source: my-cloudsqlmysql-source
        description: >
            Testing tool.
        statement: "SELECT 1;"
```
2025-01-21 13:53:37 +00:00
Yuan
45fc046437 chore: migrate yaml library to goccy/go-yaml (#211)
Migrate yaml library from
[gopkg.in/yaml.v3](https://pkg.go.dev/gopkg.in/yaml.v3) to
[goccy/go-yaml](https://github.com/goccy/go-yaml).

We are trying to add validation feature to the unmarshaling of yaml.
However, `yaml.v3` currently does not support that and the library had
not been actively maintained for awhile. This migration will allow us to
use validation processing such as
[go-playground/validator](https://github.com/go-playground/validator)

---------

Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
2025-01-16 09:15:18 -08:00
Yuan
346c57da23 fix: improve return error for parameters (#206) 2025-01-14 16:35:33 +00:00
Michael Hunger
8a1224b9e0 feat: Added Neo4j Source and Tool (#189)
- configure neo4j source with url, username, password, database
- configure neo4j tools with cypher statement and paramters
- tests based on the postgres tests
- neo4j.yaml for integration tests
---------

Co-authored-by: duwenxin <duwenxin@google.com>
2025-01-14 09:17:18 -07:00
Yuan
1fcc20a846 feat: add support for OpenTelemetry (#205)
Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
2025-01-13 15:38:38 -08:00
Yuan
d65cefc600 chore: Add debug logs (#195)
Add debug logs to Toolbox. 

For example when a http fail, it will just show Error at the http level,
but not log with actual error message. err message are returned to the
api as following `{"status":"Internal Server Error","error":"error while
invoking tool: unable to execute client: spanner: code =
\"InvalidArgument\", desc = \"invalid session pool\""}`.

After adding this, if user/dev run toolbox with `--log-level=debug`, it
will output the following (debug log in addition to the error for http
request):

```
2025-01-08T14:16:25.040824-08:00 DEBUG "error while invoking tool: unable to execute client: spanner: code = \"InvalidArgument\", desc = \"invalid session pool\""

2025-01-08T14:16:25.040968-08:00 ERROR Response: 500 Server Error service: "httplog" httpRequest: {url: "http://127.0.0.1:5000/api/tool/test_tool_two/invoke" method: "POST" path: "/api/tool/test_tool_two/invoke" remoteIP: "127.0.0.1:51708" proto: "HTTP/1.1" requestID: "yuanteoh-macbookpro.roam.internal/N7LNMcLIUH-000001" scheme: "http" header: {user-agent: "curl/8.7.1" accept: "*/*" content-type: "application/json" content-length: "2"}} httpResponse: {status: 500 bytes: 167 elapsed: 0.301917}
```
2025-01-13 05:47:46 +00:00
Kurtis Van Gent
387a5b56b5 fix: correct parsing of floats/ints from json (#180)
Corrects an issue caused by Go defaulting to parsing JSON Numbers as
float64s. This caused some numbers to be incorrectly parsed as floats
when they were integers. This defaults to parsing using json.Number,
which allows us to parse between Int/Float more accurately.
2025-01-03 10:09:45 -07:00
Kurtis Van Gent
66ab70f702 fix: handle shutdown gracefully (#178)
Adds logic to make the server shutdown gracefully, including better
respecting cancelled contexts and providing up to 10 seconds to finish
current connections.
2024-12-31 09:53:53 -07:00
Yuan
809e547a48 feat: update log with given context (#147)
Update logging with the given context.
2024-12-23 19:55:25 -08:00
Twisha Bansal
938f7b41d7 chore: fix comments (#163) 2024-12-23 09:28:14 -07:00
Wenxin Du
3118104ae1 feat: add auth_required to tools (#123)
Check if a tool invocation request contains required auth token.
2024-12-16 22:41:13 -05:00
Wenxin Du
380a6fbbd5 feat: Add Tool authenticated parameters (#80)
1. Add []ParamAuthSource to every Parameter type implementation to
support authenticated configs. Create new constructors for types with
auth.
2. Tool invocation API changes to parse auth header and authentecated
parameters.
3. Add authSources to Tool manifest.
2024-12-16 17:46:06 -05:00
Wenxin Du
77b0d43175 feat: Add Auth Source configuration (#71)
1. Add authentication source configs to Toolbox server.
2. Add Google auth implementation.
2024-12-13 12:22:42 -05:00
Yuan
7e8d751d3f ci(cloudsql-pg): add end to end integration test (#113)
End to end integration test for cloudsql postgres.

Include checks for one tool's get (manifest) and post (invoke) endpoint.

Integration tests are excluded from regular unit tests.
2024-12-12 13:56:04 -08:00
Yuan
cb3357f305 fix!(spanner): use googlesql for dialect name (#129)
Google updated the official name from `google_standard_sql` to
`googlesql`.
2024-12-11 14:11:34 -08:00
Yuan
45a38ca9f0 chore: separate listen and server into two functions (#121)
Separate listen and server into two distinct functions to prevent race
condition during testing.
2024-12-09 22:57:53 +00:00
Yuan
890914aae0 feat: add Spanner source and tool (#90)
Add Spanner source and tool.

Spanner source is initialize with the following config:
```
sources:
    my-spanner-source:
        kind: spanner
        project: my-project-name
        instance: my-instance-name
        database: my_db
        # dialect: postgresql # The default dialect is google_standard_sql.
```

Spanner tool (with gsql dialect) is initialize with the following
config.
```
tools:
    get_flight_by_id:
        kind: spanner
        source: my-cloud-sql-source
        description: >
            Use this tool to list all airports matching search criteria. Takes 
            at least one of country, city, name, or all and returns all matching
            airports. The agent can decide to return the results directly to 
            the user.
        statement: "SELECT * FROM flights WHERE id = @id"
        parameters:
        - name: id
          type: int
          description: 'id' represents the unique ID for each flight. 
```

Spanner tool (with postgresql dialect) is initialize with the following
config.
```
tools:
    get_flight_by_id:
        kind: spanner
        source: my-cloud-sql-source
        description: >
            Use this tool to list all airports matching search criteria. Takes 
            at least one of country, city, name, or all and returns all matching
            airports. The agent can decide to return the results directly to 
            the user.
        statement: "SELECT * FROM flights WHERE id = $1"
        parameters:
        - name: id
          type: int
          description: 'id' represents the unique ID for each flight. 
```

Note: the only difference in config for both dialects is the sql
statement.

---------

Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
2024-12-06 16:38:03 -08:00
Yuan
6479c1dbe2 feat(source/cloudsql-pg): add configuration for public and private IP (#114)
Allow user to set if their database uses private or public ip. The
reason we add this is because the dialer require different
initialization with private and public ip.

By default, toolbox will use public ip.
2024-12-06 00:08:15 +00:00
Yuan
e815dc49f4 chore: update params type (#98)
Different databases require different types for `Params` field when
adding parameters to their statement. e.g. alloydb, cloudsql, and
postgres uses `pgxpool` to query and build sql statement, whereas
spanner uses `Spanner` library.

Added a new `ParamValue` struct. `ParseParams` helper function parses
arbitraryJSON object into `[]ParamValue`, and the tool's invoke will
convert `[]ParamValue` into it's required type.

---------

Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
2024-12-05 16:05:09 -08:00
Yuan
e88ec409d1 feat(source/alloydb-pg): add configuration for public and private IP (#103)
Allow user to set if their database uses private or public ip. The
reason we add this is because the dialer require different
initialization with private and public ip.

By default, toolbox will use public ip.

---------

Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
2024-12-05 23:05:53 +00:00
Yuan
9a0f618efc feat: add --log-level and --logging-format flags (#97)
Logging support 4 different types of logging (debug, info, warn, error).
The default logging level is Info.

User will be able to set flag for log level (allowed values: "debug",
"info", "warn", "error"), example:
`go run . --log-level debug`

User will be able to set flag for logging format (allowed values:
"standard", "JSON"), example:
`go run . --logging-format json`

**sample http request log - std:**
server
```
2024-11-12T15:08:11.451377-08:00 INFO "Initalized 0 sources.\n"
```
httplog
```
2024-11-26T15:15:53.947287-08:00 INFO Response: 200 OK service: "httplog" httpRequest: {url: "http://127.0.0.1:5000/" method: "GET" path: "/" remoteIP: "127.0.0.1:64216" proto: "HTTP/1.1" requestID: "macbookpro.roam.interna/..." scheme: "http" header: {user-agent: "curl/8.7.1" accept: "*/*"}} httpResponse: {status: 200 bytes: 22 elapsed: 0.012417}
```

**sample http request log - structured:**
server
```
{
  "timestamp":"2024-11-04T16:45:11.987299-08:00",
  "severity":"ERROR",
  "logging.googleapis.com/sourceLocation":{
    "function":"github.com/googleapis/genai-toolbox/internal/log.(*StructuredLogger).Errorf",
    "file":"/Users/yuanteoh/github/genai-toolbox/internal/log/log.go","line":157
  },
  "message":"unable to parse tool file at \"tools.yaml\": \"cloud-sql-postgres1\" is not a valid kind of data source"
}
```
httplog
```
{
  "timestamp":"2024-11-26T15:12:49.290974-08:00",
  "severity":"INFO",
  "logging.googleapis.com/sourceLocation":{
      "function":"github.com/go-chi/httplog/v2.(*RequestLoggerEntry).Write",
      "file":"/Users/yuanteoh/go/pkg/mod/github.com/go-chi/httplog/v2@v2.1.1/httplog.go","line":173
  },
  "message":"Response: 200 OK",
  "service":"httplog",
  "httpRequest":{
      "url":"http://127.0.0.1:5000/",
      "method":"GET",
      "path":"/",
      "remoteIP":"127.0.0.1:64140",
      "proto":"HTTP/1.1",
      "requestID":"yuanteoh-macbookpro.roam.internal/NBrtYBu3q9-000001",
      "scheme":"http",
      "header":{"user-agent":"curl/8.7.1","accept":"*/*"}
  },
  "httpResponse":{"status":200,"bytes":22,"elapsed":0.0115}
}
```
2024-12-03 16:52:19 +00:00
Wenxin Du
751228a774 ci: add e2e testing with cloud build (#27) 2024-11-13 17:59:41 -05:00
Kurtis Van Gent
efafba9033 feat: support requesting a single tool (#56)
Adds support for getting a ToolsManifest with a single tool when a GET
`/tools/$toolname` request is sent.
2024-11-11 19:35:57 +00:00
Kurtis Van Gent
f630965937 feat!: consolidate "x-postgres-generic" tools to "postgres-sql" tool (#43)
This PR introduces the following breaking change: The
`alloydb-pg-generic`, `cloud-sql-pg-generic`, and
`postgres-generic-tool` have been replaced by the `postgres-sql` tool,
which works with all 3 Postgres sources.

If you were using of the the previous tools, you will need to update it
as follows:
```diff
example_tool:
-    kind: cloud-sql-pg-generic
+    kind: postgres-sql
     source: my-cloud-sql-pg-instance
     description: some description
        statement: |
            SELECT * FROM SQL_STATEMENT;
        parameters:
        - name: country
          type: string
          description: some description
```

I'm proposing this change for the following reasons:
1. It provides greater flexibility between postgres-compatible sources
-- you can change between "postgres" and "alloydb-postgres" without
issue
2. The name "postgres-sql" is more clear that "postgres-generic" -- it
indicates it's a tool that runs SQL on the source
3. It's easier for us to maintain feature compatibility across a single
"postgres-sql" tool
2024-11-01 16:17:18 -07:00
Kurtis Van Gent
efe9e8bc97 chore: refactor sources/tools each into their own package (#42)
Moves all of the "source" and "tool" implementations into their own
packages. This layout makes it a bit more clear where the
implementations are, and seems likely to scale more cleanly as more
sources and tools are added.
2024-11-01 16:44:13 -06:00
Kurtis Van Gent
cf8112f856 fix: handle content-type correctly (#33)
Fixes a bug where the content-type wasn't being set correctly in the
response type.
2024-10-29 20:00:22 +00:00
Kurtis Van Gent
3903e860bc feat: add support for array type parameters (#26)
Adds support for "array" type parameters. Uses a subet of JSONSchema for
specification, in that arrays can be specified in the following way:

```yaml
parameters:
    name: "my_array"
    type: "array"
    description: "some description"
    items:
       type: "integer"
```
2024-10-25 21:54:14 +00:00
Kurtis Van Gent
de14c6f669 chore: refactor parameters logic to allow for more complex functionality (#24)
Refactors parameter logic to allow for more complex functionality (e.g.
arrays) in the future.

---------

Co-authored-by: Wenxin Du <117315983+duwenxin99@users.noreply.github.com>
2024-10-23 13:04:10 -06:00
Wenxin Du
61e7b78ad8 feat: Add Toolset manifest endpoint (#11)
1. Calculate tool manifests when server starts.
2. Add toolset manifest endpoints.

---------

Co-authored-by: Kurtis Van Gent <31518063+kurtisvg@users.noreply.github.com>
2024-10-17 18:43:58 -04:00
Kurtis Van Gent
3703176fce feat: add initial cloud sql pg invocation (#14)
Makes the following changes:
- Fills in previously stubbed "Invoke()" function for the Cloud SQL
Generic Tool
- Updates API to /{tool_name}/invoke to for invocation of said tool 
- Updates response to use JSON 
- Correctly returns error messages for invalid http codes
2024-10-15 09:53:01 -06:00
Wenxin Du
59b4bc07f4 feat: add toolset configuration (#12)
Add `Toolset` implementation to the `tools` package:
- struct and configs.
- Custom `UnmarshalYAML` function.
- Initialization function that validates if tools specified for the
toolset exist.
2024-10-11 14:11:05 -04:00
Kurtis Van Gent
27edd3b5f6 feat: add preliminary parsing of parameters (#13)
This PR adds preliminary parsing of parameters. Currently it only
supports 4 types: string, int, float32, and bool. Almost certainly we
will need to introduce more complicated parsing configuration (to handle
objects and arrays), but my initial attempts got quickly complicated, so
I simplified in the short term.

This also makes 2 breaking changes to config.yaml:
- changes "parameters" to be a list over object -- this is because
parameter ordering is important, and needs to be preserved
- removed the "required" field from parameter objects -- we need to
determine how to handle optional parameters in SQL queries
2024-10-08 15:18:49 -06:00
Kurtis Van Gent
336bdc4d56 feat: stub basic control plane functionality (#9)
Stub's out some basic control plane functionality. This also required
setting up some Source and Tools initialization.
2024-09-20 14:12:16 -06:00
Kurtis Van Gent
b9ba364fb6 feat: add basic parsing from tools file (#8) 2024-08-05 09:33:07 -06:00
Kurtis Van Gent
df9ad9e33f feat: add address and port flags (#7)
Add flags for setting the address and port of the server.
2024-07-30 10:55:13 -05:00
Kurtis Van Gent
e09ae30a90 feat: add basic http server (#6) 2024-07-26 16:59:01 -05:00