Commit Graph

18 Commits

Author SHA1 Message Date
Yuan
04cb5fbc3e feat!: deprecate authsource in favor of authservice (#297)
Rename existing `authSource` to `authService` through deprecation.
`AuthService` more clearly distinguishes it from `Sources` objects.

`authSources` will be converted into `authServices` after the
unmarshalling process. A warning log is shown if `authSources` are used
(for both within tools parameters and defining auth services):
```
2025-02-20T13:57:51.156025-08:00 WARN "`authSources` is deprecated, use `authServices` for parameters instead"
2025-02-20T13:57:51.156569-08:00 WARN "`authSources` is deprecated, use `authServices` instead"
2025-02-20T13:57:52.047584-08:00 INFO "Initialized 1 sources."
...
```

The manifest generated will continue to use `authSources` to keep
compatibility with the sdks:
```
{
"serverVersion":"0.1.0",
"tools":{
  "test_tool2":{
    "description":"Use this tool to test\n",
    "parameters":[{
      "name":"user_id",
      "type":"string",
      "description":"Auto-populated from Google login",
      "authSources":["my-google-auth"]
      }]
    }
  }
}
```



Test cases with `authSources` are kept for compatibility. Will be
removed when `authSources` are no longer supported.
2025-02-21 18:36:04 -08:00
Yuan
95d29f9e8a chore: pass logger via context through unmarshaller (#300)
Passing logger via context. With this, we will be able to add logs
during the unmarshaling process.
2025-02-21 13:02:23 -08:00
Yuan
e494d11e6e fix: validate required common fields for parameters (#298)
Check common parameters for required fields. Throw an error if name or
description is not provided.

Example of error when `description` field not provided:
```
2025-02-18T14:11:55.475101-08:00 ERROR "unable to parse tool file at \"tools.yaml\": unable to parse as \"postgres-sql\": unable to parse as \"array\": Key: 'CommonParameter.Desc' Error:Field validation for 'Desc' failed on the 'required' tag"
```
2025-02-19 10:09:57 -08:00
Yuan
541612d72d fix: add items to parameter manifest (#293)
`items` are not specified in the manifest for parameters, which means
the client doesn't parse them correctly when presented to the LLMs. This
PR adds `items` as part of manifest.

Example manifest:
```
{
    "serverVersion": "0.1.0",
    "tools": {
        "tool_name": {
            "description": "description of the tool...",
            "parameters": [{
                "name": "foos",
                "type": "array",
                "description": "A list of foo.",
                "authSources": [],
                "items": {"name": "foo", "type": "string", "description": "foobar.", "authSources": []}
            }]
        }
    }
}
```

Fixes #292 🦕
2025-02-14 15:29:16 -08:00
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
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
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
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
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
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
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
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
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