@@ -2,12 +2,15 @@ package ollama_test
22
33import (
44 "context"
5+ "encoding/json"
6+ "log"
57 "os"
68 "testing"
79
810 // Packages
911 opts "github.com/mutablelogic/go-client"
1012 ollama "github.com/mutablelogic/go-llm/pkg/ollama"
13+ "github.com/mutablelogic/go-llm/pkg/tool"
1114 assert "github.com/stretchr/testify/assert"
1215)
1316
@@ -60,13 +63,17 @@ func Test_chat_002(t *testing.T) {
6063 t .FailNow ()
6164 }
6265
66+ // Make a toolkit
67+ toolkit := tool .NewToolKit ()
68+ if err := toolkit .Register (new (weather )); err != nil {
69+ t .FailNow ()
70+ }
71+
6372 t .Run ("Tools" , func (t * testing.T ) {
6473 assert := assert .New (t )
6574 response , err := client .Chat (context .TODO (),
6675 model .UserPrompt ("what is the weather in berlin?" ),
67- ollama .WithTool (ollama .MustTool ("get_weather" , "Return weather conditions in a location" , struct {
68- Location string `help:"Location to get weather for" required:""`
69- }{})),
76+ ollama .WithToolKit (toolkit ),
7077 )
7178 if ! assert .NoError (err ) {
7279 t .FailNow ()
@@ -108,3 +115,31 @@ func Test_chat_003(t *testing.T) {
108115 t .Log (response )
109116 })
110117}
118+
119+ ////////////////////////////////////////////////////////////////////////////////
120+ // TOOLS
121+
122+ type weather struct {
123+ Location string `json:"location" name:"location" help:"The location to get the weather for" required:"true"`
124+ }
125+
126+ func (* weather ) Name () string {
127+ return "weather_in_location"
128+ }
129+
130+ func (* weather ) Description () string {
131+ return "Get the weather in a location"
132+ }
133+
134+ func (weather * weather ) String () string {
135+ data , err := json .MarshalIndent (weather , "" , " " )
136+ if err != nil {
137+ return err .Error ()
138+ }
139+ return string (data )
140+ }
141+
142+ func (weather * weather ) Run (ctx context.Context ) (any , error ) {
143+ log .Println ("weather_in_location" , "=>" , weather )
144+ return "very sunny today" , nil
145+ }
0 commit comments