Skip to content

Commit

Permalink
ollama: add tool testing
Browse files Browse the repository at this point in the history
  • Loading branch information
treywelsh committed Sep 13, 2024
1 parent cff3953 commit ead7477
Showing 1 changed file with 75 additions and 0 deletions.
75 changes: 75 additions & 0 deletions llms/ollama/ollama_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tmc/langchaingo/llms"
"github.com/tmc/langchaingo/llms/ollama/internal/ollamaclient"
)

func newTestClient(t *testing.T, opts ...Option) *LLM {
Expand Down Expand Up @@ -129,3 +130,77 @@ func TestWithKeepAlive(t *testing.T) {
require.NoError(t, err)
assert.NotEmpty(t, vector)
}

func TestWithTools(t *testing.T) {
t.Parallel()
llm := newTestClient(t, WithTools([]ollamaclient.Tool{
{
Type: "function",
Function: ollamaclient.ToolFunction{
Name: "getCurrentWeather",
Description: "Get the current weather in a given location",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"location": map[string]any{
"type": "string",
"description": "The city, e.g. San Francisco",
},
},
"required": []string{"location"},
},
},
},
}))

parts := []llms.ContentPart{
llms.TextContent{Text: "What is the weather like in London ?"},
}
msgs := []llms.MessageContent{
{
Role: llms.ChatMessageTypeHuman,
Parts: parts,
},
}

resp, err := llm.GenerateContent(context.Background(), msgs)
require.NoError(t, err)

assert.NotEmpty(t, resp.Choices)
c1 := resp.Choices[0]

assert.NotEmpty(t, c1.ToolCalls)

tool := c1.ToolCalls[0]

assert.Equal(t, "getCurrentWeather", tool.FunctionCall.Name)
assert.Equal(t, "{\"location\":\"london\"}", strings.ToLower(tool.FunctionCall.Arguments))

toolCallingResponse := llms.TextParts(llms.ChatMessageTypeAI, c1.Content)
for _, tc := range c1.ToolCalls {
toolCallingResponse.Parts = append(toolCallingResponse.Parts, tc)
}
msgs = append(msgs, toolCallingResponse)

toolExecResult := "{\"location\":\"london\", \"forecast\":\"63°F and cloudy\"}"

weatherCallResponse := llms.MessageContent{
Role: llms.ChatMessageTypeTool,
Parts: []llms.ContentPart{
llms.ToolCallResponse{
Name: tool.FunctionCall.Name,
Content: toolExecResult,
},
},
}
msgs = append(msgs, weatherCallResponse)

resp, err = llm.GenerateContent(context.Background(), msgs)
require.NoError(t, err)
assert.NotEmpty(t, resp.Choices)
c1Lower := strings.ToLower(resp.Choices[0].Content)

assert.Equal(t, "63", c1Lower)
assert.Equal(t, "cloudy", c1Lower)
assert.Equal(t, "london", c1Lower)
}

0 comments on commit ead7477

Please sign in to comment.