Skip to content

Commit 0333d1f

Browse files
jhrozekjhrozek
committed
Add an example with a tool talking to Trusty
Adds an example of tool function for trusty and an example program that calls the model with tools.
1 parent 8405307 commit 0333d1f

File tree

2 files changed

+227
-0
lines changed

2 files changed

+227
-0
lines changed

‎examples/tools/main.go

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
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+
"github.com/stackloklabs/gollm/examples/tools/trusty"
20+
"log"
21+
"os"
22+
"strings"
23+
"time"
24+
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 = `You are a helpful AI assistant that provides recommendations to users about software packages.
36+
Your job is to provide a recommendation based on the user's prompt.
37+
You might be provided a JSON summary along with the user's prompt. Do not summarize the JSON back to the user.
38+
Focus on whether the package is malicious or deprecated based on the provided tool input you get as JSON.
39+
Focus less on the number of stars or forks.
40+
If the user does not specify the ecosystem (one of npm, pypi, crates, maven, go), ask the user, NEVER assume the ecosystem.
41+
If the package is malicious or deprecated, recommend a safer alternative.
42+
If the package is safe, recommend the package.
43+
`
44+
summarizeMessage = `Summarize the tool response for me in plain speech. If the package is either malicious, deprecated
45+
or no longer maintained, recommend a bulleted list of two-three safer alternative packages that do the same. If the package is safe, recommend the package.`
46+
)
47+
48+
func main() {
49+
var generationBackend backend.Backend
50+
51+
beSelection := os.Getenv("BACKEND")
52+
if beSelection == "" {
53+
log.Println("No backend selected with the BACKEND env variable. Defaulting to Ollama.")
54+
beSelection = "ollama"
55+
}
56+
modelSelection := os.Getenv("MODEL")
57+
if modelSelection == "" {
58+
switch beSelection {
59+
case "ollama":
60+
modelSelection = ollamaGenModel
61+
case "openai":
62+
modelSelection = openaiModel
63+
}
64+
log.Println("No model selected with the MODEL env variable. Defaulting to ", modelSelection)
65+
}
66+
67+
switch beSelection {
68+
case "ollama":
69+
generationBackend = backend.NewOllamaBackend(ollamaHost, ollamaGenModel)
70+
log.Println("Using Ollama backend: ", ollamaGenModel)
71+
case "openai":
72+
openaiKey := os.Getenv("OPENAI_API_KEY")
73+
if openaiKey == "" {
74+
log.Fatalf("OPENAI_API_KEY is required for OpenAI backend")
75+
}
76+
generationBackend = backend.NewOpenAIBackend(openaiKey, openaiModel)
77+
log.Println("Using OpenAI backend: ", openaiModel)
78+
default:
79+
log.Fatalf("Unknown backend: %s", beSelection)
80+
}
81+
82+
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
83+
defer cancel()
84+
85+
userPrompt := os.Args[1:]
86+
if len(userPrompt) == 0 {
87+
log.Fatalf("Please provide a prompt")
88+
}
89+
90+
convo := backend.NewConversation()
91+
convo.Tools.RegisterTool(trusty.Tool())
92+
// start the conversation. We add a system message to tune the output
93+
// and add the trusty tool to the conversation so that the model knows to call it.
94+
convo.AddSystemMessage(systemMessage, nil)
95+
convo.AddUserMessage(strings.Join(userPrompt, " "), nil)
96+
97+
// generate the response
98+
resp, err := generationBackend.Converse(ctx, convo)
99+
if err != nil {
100+
log.Fatalf("Error generating response: %v", err)
101+
}
102+
103+
if len(resp.ToolCalls) == 0 {
104+
log.Println("No tool calls in response.")
105+
log.Println("Response:", convo.Messages[len(convo.Messages)-1]["content"])
106+
return
107+
}
108+
109+
log.Println("Tool called")
110+
111+
// if there was a tool response, first just feed it back to the model so it makes sense of it
112+
_, err = generationBackend.Converse(ctx, convo)
113+
if err != nil {
114+
log.Fatalf("Error generating response: %v", err)
115+
}
116+
117+
log.Println("Summarizing tool response")
118+
119+
// summarize the tool response
120+
convo.AddSystemMessage(summarizeMessage, nil)
121+
_, err = generationBackend.Converse(ctx, convo)
122+
if err != nil {
123+
log.Fatalf("Error generating response: %v", err)
124+
}
125+
126+
log.Println("Response:")
127+
log.Println(convo.Messages[len(convo.Messages)-1]["content"])
128+
}

‎examples/tools/trusty/trusty_tool.go

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
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 trusty
16+
17+
import (
18+
"encoding/json"
19+
"fmt"
20+
"github.com/stackloklabs/gollm/pkg/backend"
21+
"io"
22+
"net/http"
23+
)
24+
25+
// Tool returns a backend.Tool object that can be used to interact with the trusty tool.
26+
func Tool() backend.Tool {
27+
return backend.Tool{
28+
Type: "function",
29+
Function: backend.ToolFunction{
30+
Name: "trusty",
31+
Description: "Evaluate the trustworthiness of a package",
32+
Parameters: map[string]any{
33+
"type": "object",
34+
"properties": map[string]any{
35+
"package_name": map[string]any{
36+
"type": "string",
37+
"description": "The name of the package",
38+
},
39+
"ecosystem": map[string]any{
40+
"type": "string",
41+
"description": "The ecosystem of the package",
42+
"enum": []string{"npm", "pypi", "crates", "maven", "go"},
43+
"default": "pypi",
44+
},
45+
},
46+
"required": []string{"package_name", "ecosystem"},
47+
},
48+
Wrapper: trustyReportWrapper,
49+
},
50+
}
51+
}
52+
53+
func trustyReportWrapper(params map[string]any) (string, error) {
54+
packageName, ok := params["package_name"].(string)
55+
if !ok {
56+
return "", fmt.Errorf("package_name must be a string")
57+
}
58+
ecosystem, ok := params["ecosystem"].(string)
59+
if !ok {
60+
ecosystem = "PyPi"
61+
}
62+
return trustyReport(packageName, ecosystem)
63+
}
64+
65+
func trustyReport(packageName string, ecosystem string) (string, error) {
66+
url := fmt.Sprintf("https://api.trustypkg.dev/v1/report?package_name=%s&package_type=%s", packageName, ecosystem)
67+
68+
req, err := http.NewRequest("GET", url, nil)
69+
if err != nil {
70+
return "", err
71+
}
72+
req.Header.Set("accept", "application/json")
73+
74+
// Perform the request
75+
client := &http.Client{}
76+
resp, err := client.Do(req)
77+
if err != nil {
78+
return "", err
79+
}
80+
defer resp.Body.Close()
81+
82+
body, err := io.ReadAll(resp.Body)
83+
if err != nil {
84+
return "", err
85+
}
86+
87+
var prettyJSON map[string]interface{}
88+
err = json.Unmarshal(body, &prettyJSON)
89+
if err != nil {
90+
return "", err
91+
}
92+
93+
// Convert the JSON back to string
94+
jsonString, err := json.MarshalIndent(prettyJSON, "", " ")
95+
if err != nil {
96+
return "", err
97+
}
98+
return string(jsonString), nil
99+
}

0 commit comments

Comments
 (0)