Add cohere tools

.golangci.yml

@@ -1,17 +1,3 @@
-# Copyright 2013-2023 The Cobra Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 run:
   deadline: 5m
 

assistant/assistant.go

@@ -24,15 +24,15 @@ type observer interface {
 }
 
 const (
-	DefaultMaxIterations = 3
+	DefaultToolsMaxIterations = 1
 )
 
 type Assistant struct {
-	llm           LLM
-	rag           RAG
-	thread        *thread.Thread
-	parameters    Parameters
-	maxIterations uint
+	llm                LLM
+	rag                RAG
+	thread             *thread.Thread
+	parameters         Parameters
+	toolsMaxIterations uint
 }
 
 type LLM interface {
@@ -54,7 +54,7 @@ func New(llm LLM) *Assistant {
 			CompanyName:        defaultCompanyName,
 			CompanyDescription: defaultCompanyDescription,
 		},
-		maxIterations: DefaultMaxIterations,
+		toolsMaxIterations: DefaultToolsMaxIterations,
 	}
 
 	return assistant
@@ -94,7 +94,7 @@ func (a *Assistant) Run(ctx context.Context) error {
 		a.injectSystemMessage()
 	}
 
-	for i := 0; i < int(a.maxIterations); i++ {
+	for i := 0; i < int(a.toolsMaxIterations); i++ {
 		err = a.runIteration(ctx, i)
 		if err != nil {
 			return err
@@ -181,8 +181,8 @@ func (a *Assistant) generateRAGMessage(ctx context.Context) error {
 	return nil
 }
 
-func (a *Assistant) WithMaxIterations(maxIterations uint) *Assistant {
-	a.maxIterations = maxIterations
+func (a *Assistant) WithToolsMaxIterations(maxIterations uint) *Assistant {
+	a.toolsMaxIterations = maxIterations
 	return a
 }
 

examples/assistant/agent/main.go → examples/assistant/tools/main.go

@@ -43,10 +43,10 @@ func main() {
 	).WithThread(
 		thread.New().AddMessages(
 			thread.NewUserMessage().AddContent(
-				thread.NewTextContent("search the top 3 italian dishes and then their costs, then ask the user's budget in euros and calculate how many guests can be invited for each dish"),
+				thread.NewTextContent("search the top 3 italian meals and then their costs, then ask the user's budget in euros and calculate how many guests can be invited for each meal"),
 			),
 		),
-	).WithMaxIterations(10)
+	).WithToolsMaxIterations(10)
 
 	err = myAssistant.Run(ctx)
 	if err != nil {

go.mod

@@ -7,7 +7,7 @@ require github.com/mitchellh/mapstructure v1.5.0
 require (
 	github.com/RediSearch/redisearch-go/v2 v2.1.1
 	github.com/google/uuid v1.6.0
-	github.com/henomis/cohere-go v1.1.2
+	github.com/henomis/cohere-go v1.2.0
 	github.com/henomis/langfuse-go v0.0.3
 	github.com/henomis/milvus-go v0.0.4
 	github.com/henomis/pinecone-go/v2 v2.0.0

go.sum

@@ -7,8 +7,8 @@ github.com/gomodule/redigo v1.8.9 h1:Sl3u+2BI/kk+VEatbj0scLdrFhjPmbxOc1myhDP41ws
 github.com/gomodule/redigo v1.8.9/go.mod h1:7ArFNvsTjH8GMMzB4uy1snslv2BwmginuMs06a1uzZE=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/henomis/cohere-go v1.1.2 h1:rzEA1JRm26RnaQValoVeVQ1y1nTofuhr/z/fPyhUW/4=
-github.com/henomis/cohere-go v1.1.2/go.mod h1:z+UIgBbNCnLH5M47FYqg4h3CDWxjUv2VDfEwdTb95MU=
+github.com/henomis/cohere-go v1.2.0 h1:rHi0JDdR/LYQPtQJrP0wKD6cLg8rFdJTW7KF5KjMG0s=
+github.com/henomis/cohere-go v1.2.0/go.mod h1:z+UIgBbNCnLH5M47FYqg4h3CDWxjUv2VDfEwdTb95MU=
 github.com/henomis/langfuse-go v0.0.3 h1:Z5Mqlnj1fsok7eAT7jt+N4tegjbhY5HnZQ7wu1iVKss=
 github.com/henomis/langfuse-go v0.0.3/go.mod h1:gSRuO3nvjAvk/mgmb7b+9BcoN9s64GvXeaSN7PfVEKQ=
 github.com/henomis/milvus-go v0.0.4 h1:ArddXRJx/EGdQ75gB7TyEzD4Z4BqXzW16p8jRcjJOhs=

llm/cohere/cohere.go

@@ -2,6 +2,7 @@ package cohere
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
@@ -31,6 +32,8 @@ const (
 	ModelCommandNightly      Model = model.ModelCommandNightly
 	ModelCommandLight        Model = model.ModelCommandLight
 	ModelCommandLightNightly Model = model.ModelCommandLightNightly
+	ModelCommandR            Model = model.ModelCommandR
+	ModelCommandRPlus        Model = model.ModelCommandRPlus
 )
 
 const (
@@ -51,8 +54,7 @@ type Cohere struct {
 	cache            *cache.Cache
 	streamCallbackFn StreamCallbackFn
 	name             string
-	observer         llmobserver.LLMObserver
-	observerTraceID  string
+	functions        map[string]Function
 }
 
 func (c *Cohere) WithCache(cache *cache.Cache) *Cohere {
@@ -72,6 +74,7 @@ func New() *Cohere {
 		temperature: DefaultTemperature,
 		maxTokens:   DefaultMaxTokens,
 		name:        "cohere",
+		functions:   make(map[string]Function),
 	}
 }
 
@@ -116,13 +119,8 @@ func (c *Cohere) WithStream(callbackFn StreamCallbackFn) *Cohere {
 	return c
 }
 
-func (c *Cohere) WithObserver(observer llmobserver.LLMObserver, traceID string) *Cohere {
-	c.observer = observer
-	c.observerTraceID = traceID
-	return c
-}
-
 // Completion returns the completion for the given prompt
+// nolint:staticcheck
 func (c *Cohere) Completion(ctx context.Context, prompt string) (string, error) {
 	resp := &response.Generate{}
 	err := c.client.Generate(
@@ -219,11 +217,17 @@ func (c *Cohere) Generate(ctx context.Context, t *thread.Thread) error {
 
 	chatRequest := c.buildChatCompletionRequest(t)
 
+	if len(c.functions) > 0 {
+		chatRequest.Tools = c.buildChatCompletionRequestTools()
+	}
+
 	generation, err := c.startObserveGeneration(ctx, t)
 	if err != nil {
 		return fmt.Errorf("%w: %w", ErrCohereChat, err)
 	}
 
+	nMessageBeforeGeneration := len(t.Messages)
+
 	if c.streamCallbackFn != nil {
 		err = c.stream(ctx, t, chatRequest)
 	} else {
@@ -233,7 +237,7 @@ func (c *Cohere) Generate(ctx context.Context, t *thread.Thread) error {
 		return err
 	}
 
-	err = c.stopObserveGeneration(ctx, generation, []*thread.Message{t.LastMessage()})
+	err = c.stopObserveGeneration(ctx, generation, t.Messages[nMessageBeforeGeneration:])
 	if err != nil {
 		return fmt.Errorf("%w: %w", ErrCohereChat, err)
 	}
@@ -251,18 +255,36 @@ func (c *Cohere) Generate(ctx context.Context, t *thread.Thread) error {
 func (c *Cohere) generate(ctx context.Context, t *thread.Thread, chatRequest *request.Chat) error {
 	var response response.Chat
 
+	b, _ := json.MarshalIndent(chatRequest, "", "  ")
+	fmt.Println(string(b))
+
 	err := c.client.Chat(
 		ctx,
 		chatRequest,
 		&response,
 	)
 	if err != nil {
 		return fmt.Errorf("%w: %w", ErrCohereChat, err)
+	} else if !response.IsSuccess() {
+		if response.RawBody == nil {
+			return fmt.Errorf("%w: %d", ErrCohereChat, response.Code)
+		}
+		return fmt.Errorf("%w: %s", ErrCohereChat, *response.RawBody)
 	}
 
-	t.AddMessage(thread.NewAssistantMessage().AddContent(
-		thread.NewTextContent(response.Text),
-	))
+	var messages []*thread.Message
+	if len(response.ToolCalls) > 0 {
+		messages = append(messages, toolCallsToToolCallMessage(response.ToolCalls))
+		messages = append(messages, c.callTools(response.ToolCalls)...)
+	} else {
+		messages = []*thread.Message{
+			thread.NewAssistantMessage().AddContent(
+				thread.NewTextContent(response.Text),
+			),
+		}
+	}
+
+	t.Messages = append(t.Messages, messages...)
 
 	return nil
 }
@@ -317,3 +339,54 @@ func (c *Cohere) stopObserveGeneration(
 		messages,
 	)
 }
+
+func (c *Cohere) buildChatCompletionRequestTools() []model.Tool {
+	tools := []model.Tool{}
+
+	for _, function := range c.functions {
+		tool := buildChatCompletionRequestTool(function)
+		if tool != nil {
+			tools = append(tools, *tool)
+		}
+	}
+
+	return tools
+}
+
+func (c *Cohere) callTools(toolCalls []model.ToolCall) []*thread.Message {
+	if len(c.functions) == 0 || len(toolCalls) == 0 {
+		return nil
+	}
+
+	var messages []*thread.Message
+	for _, toolCall := range toolCalls {
+		result, err := c.callTool(toolCall)
+		if err != nil {
+			result = fmt.Sprintf("error: %s", err)
+		}
+
+		messages = append(messages, toolCallResultToThreadMessage(toolCall, result))
+	}
+
+	return messages
+}
+
+func (c *Cohere) callTool(toolCall model.ToolCall) (string, error) {
+	fn, ok := c.functions[toolCall.Name]
+	if !ok {
+		return "", fmt.Errorf("unknown function %s", toolCall.Name)
+	}
+
+	parameters, err := json.Marshal(toolCall.Parameters)
+	if err != nil {
+		return "", err
+	}
+
+	resultAsJSON, err := callFnWithArgumentAsJSON(fn.Fn, string(parameters))
+	//resultAsJSON, err := callFnWithArgumentAsJSON(fn.Fn, "{}")
+	if err != nil {
+		return "", err
+	}
+
+	return resultAsJSON, nil
+}