Skip to content

Commit 1427fb6

Browse files
jhrozekjhrozek
committed
Add an example with a tool pretending to do weather reports
Adds an example of tool function with pretend weather reports. Example invocations: ``` BACKEND=ollama go run ./examples/tools/main.go "What is the weather in London?" 2024/10/17 22:42:11 No model selected with the MODEL env variable. Defaulting to qwen2.5 2024/10/17 22:42:11 Using Ollama backend: qwen2.5 2024/10/17 22:42:14 Tool called 2024/10/17 22:42:14 Response: 2024/10/17 22:42:14 The current temperature in London is 15°C and the conditions are rainy. ```
1 parent c23c465 commit 1427fb6

File tree

2 files changed

+197
-0
lines changed

2 files changed

+197
-0
lines changed

‎examples/tools/main.go

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
// Copyright 2024 Stacklok, Inc
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package main
16+
17+
import (
18+
"context"
19+
"log"
20+
"os"
21+
"strings"
22+
"time"
23+
24+
"github.com/stackloklabs/gollm/examples/tools/weather"
25+
"github.com/stackloklabs/gollm/pkg/backend"
26+
)
27+
28+
var (
29+
ollamaHost = "http://localhost:11434"
30+
ollamaGenModel = "qwen2.5"
31+
openaiModel = "gpt-4o-mini"
32+
)
33+
34+
const (
35+
systemMessage = `
36+
You are an AI assistant that can retrieve weather forecasts by calling a tool
37+
that returns weather data in JSON format. You will be provided with a city
38+
name, and you will use a tool to call out to a weather forecast API that
39+
provides the weather for that city. The tool returns a JSON object with three
40+
fields: city, temperature, and conditions.
41+
`
42+
summarizeMessage = `
43+
Summarize the tool's forecast of the weather in clear, plain language for the user.
44+
`
45+
)
46+
47+
func main() {
48+
var generationBackend backend.Backend
49+
50+
beSelection := os.Getenv("BACKEND")
51+
if beSelection == "" {
52+
log.Println("No backend selected with the BACKEND env variable. Defaulting to Ollama.")
53+
beSelection = "ollama"
54+
}
55+
modelSelection := os.Getenv("MODEL")
56+
if modelSelection == "" {
57+
switch beSelection {
58+
case "ollama":
59+
modelSelection = ollamaGenModel
60+
case "openai":
61+
modelSelection = openaiModel
62+
}
63+
log.Println("No model selected with the MODEL env variable. Defaulting to ", modelSelection)
64+
}
65+
66+
switch beSelection {
67+
case "ollama":
68+
generationBackend = backend.NewOllamaBackend(ollamaHost, ollamaGenModel, 30*time.Second)
69+
log.Println("Using Ollama backend: ", ollamaGenModel)
70+
case "openai":
71+
openaiKey := os.Getenv("OPENAI_API_KEY")
72+
if openaiKey == "" {
73+
log.Fatalf("OPENAI_API_KEY is required for OpenAI backend")
74+
}
75+
generationBackend = backend.NewOpenAIBackend(openaiKey, openaiModel, 30*time.Second)
76+
log.Println("Using OpenAI backend: ", openaiModel)
77+
default:
78+
log.Fatalf("Unknown backend: %s", beSelection)
79+
}
80+
81+
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
82+
defer cancel()
83+
84+
userPrompt := os.Args[1:]
85+
if len(userPrompt) == 0 {
86+
log.Fatalf("Please provide a prompt")
87+
}
88+
89+
convo := backend.NewPrompt()
90+
convo.Tools.RegisterTool(weather.Tool())
91+
// start the conversation. We add a system message to tune the output
92+
// and add the weather tool to the conversation so that the model knows to call it.
93+
convo.AddMessage("system", systemMessage)
94+
convo.AddMessage("user", strings.Join(userPrompt, " "))
95+
96+
// generate the response
97+
resp, err := generationBackend.Converse(ctx, convo)
98+
if err != nil {
99+
log.Fatalf("Error generating response: %v", err)
100+
}
101+
102+
if len(resp.ToolCalls) == 0 {
103+
log.Println("No tool calls in response.")
104+
log.Println("Response:", convo.Messages[len(convo.Messages)-1].Content)
105+
return
106+
}
107+
108+
log.Println("Tool called")
109+
110+
// if there was a tool response, first just feed it back to the model so it makes sense of it
111+
_, err = generationBackend.Converse(ctx, convo)
112+
if err != nil {
113+
log.Fatalf("Error generating response: %v", err)
114+
}
115+
116+
log.Println("Response:")
117+
log.Println(convo.Messages[len(convo.Messages)-1].Content)
118+
}

‎examples/tools/weather/weather_tool.go

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
// Copyright 2024 Stacklok, Inc
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package weather
16+
17+
import (
18+
"encoding/json"
19+
"errors"
20+
"fmt"
21+
"github.com/stackloklabs/gollm/pkg/backend"
22+
)
23+
24+
// Tool returns a backend.Tool object that can be used to interact with the weather tool.
25+
func Tool() backend.Tool {
26+
return backend.Tool{
27+
Type: "function",
28+
Function: backend.ToolFunction{
29+
Name: "weather",
30+
Description: "Get weather report for a city",
31+
Parameters: map[string]any{
32+
"type": "object",
33+
"properties": map[string]any{
34+
"city": map[string]any{
35+
"type": "string",
36+
"description": "The city for which to get the weather report",
37+
},
38+
},
39+
"required": []string{"city"},
40+
},
41+
Wrapper: weatherReportWrapper,
42+
},
43+
}
44+
}
45+
46+
func weatherReportWrapper(params map[string]any) (string, error) {
47+
city, ok := params["city"].(string)
48+
if !ok {
49+
return "", fmt.Errorf("city must be a string")
50+
}
51+
return weatherReport(city)
52+
}
53+
54+
// WeatherReport defines the structure of the JSON response
55+
type WeatherReport struct {
56+
City string `json:"city"`
57+
Temperature string `json:"temperature"`
58+
Conditions string `json:"conditions"`
59+
}
60+
61+
// weatherReport returns a dummy weather report for the specified city in JSON format.
62+
func weatherReport(city string) (string, error) {
63+
// in a real application, this data would be fetched from an external API
64+
weatherData := map[string]WeatherReport{
65+
"London": {City: "London", Temperature: "15°C", Conditions: "Rainy"},
66+
"Stockholm": {City: "Stockholm", Temperature: "10°C", Conditions: "Sunny"},
67+
"Brno": {City: "Brno", Temperature: "18°C", Conditions: "Clear skies"},
68+
}
69+
70+
if report, ok := weatherData[city]; ok {
71+
jsonReport, err := json.Marshal(report)
72+
if err != nil {
73+
return "", err
74+
}
75+
return string(jsonReport), nil
76+
}
77+
78+
return "", errors.New("city not found")
79+
}

0 commit comments

Comments
 (0)