push of current changes

anthropic.go

@@ -0,0 +1,133 @@
+package go_llm
+
+import (
+	"context"
+	"fmt"
+	anth "github.com/liushuangls/go-anthropic/v2"
+	"log"
+)
+
+type anthropic struct {
+	key   string
+	model string
+}
+
+var _ LLM = anthropic{}
+
+func (a anthropic) ModelVersion(modelVersion string) (ChatCompletion, error) {
+	a.model = modelVersion
+
+	// TODO: model verification?
+	return a, nil
+}
+
+func (a anthropic) requestToAnthropicRequest(req Request) anth.MessagesRequest {
+	res := anth.MessagesRequest{
+		Model:     anth.Model(a.model),
+		MaxTokens: 1000,
+	}
+
+	msgs := []anth.Message{}
+
+	// we gotta convert messages into anthropic messages, however
+	// anthropic does not have a "system" message type, so we need to
+	// append it to the res.System field instead
+
+	for _, msg := range req.Messages {
+		if msg.Role == RoleSystem {
+			if len(res.System) > 0 {
+				res.System += "\n"
+			}
+			res.System += msg.Text
+		} else {
+			role := anth.RoleUser
+
+			if msg.Role == RoleAssistant {
+				role = anth.RoleAssistant
+			}
+
+			m := anth.Message{
+				Role:    role,
+				Content: []anth.MessageContent{},
+			}
+
+			if msg.Text != "" {
+				m.Content = append(m.Content, anth.MessageContent{
+					Type: anth.MessagesContentTypeText,
+					Text: &msg.Text,
+				})
+			}
+
+			if msg.ImageBase64 != "" {
+				m.Content = append(m.Content, anth.NewImageMessageContent(anth.MessageContentImageSource{
+					Type:      "base64",
+					MediaType: "image/png",
+					Data:      msg.ImageBase64,
+				}))
+			}
+			msgs = append(msgs, m)
+		}
+	}
+
+	for _, tool := range req.Toolbox {
+		res.Tools = append(res.Tools, anth.ToolDefinition{
+			Name:        tool.Name,
+			Description: tool.Description,
+			InputSchema: tool.Parameters,
+		})
+	}
+
+	res.Messages = msgs
+
+	if req.Temperature != nil {
+		res.Temperature = req.Temperature
+	}
+
+	log.Println("llm request to anthropic request", res)
+
+	return res
+}
+
+func (a anthropic) responseToLLMResponse(in anth.MessagesResponse) Response {
+	res := Response{}
+
+	for _, msg := range in.Content {
+		choice := ResponseChoice{}
+
+		switch msg.Type {
+		case anth.MessagesContentTypeText:
+			if msg.Text != nil {
+				choice.Content = *msg.Text
+			}
+
+		case anth.MessagesContentTypeToolUse:
+			if msg.MessageContentToolUse != nil {
+				choice.Calls = append(choice.Calls, ToolCall{
+					ID: msg.MessageContentToolUse.ID,
+					FunctionCall: FunctionCall{
+						Name:      msg.MessageContentToolUse.Name,
+						Arguments: msg.MessageContentToolUse.Input,
+					},
+				})
+			}
+		}
+
+		res.Choices = append(res.Choices, choice)
+	}
+
+	log.Println("anthropic response to llm response", res)
+
+	return res
+}
+
+func (a anthropic) ChatComplete(ctx context.Context, req Request) (Response, error) {
+	cl := anth.NewClient(a.key)
+
+	res, err := cl.CreateMessages(ctx, a.requestToAnthropicRequest(req))
+
+	if err != nil {
+		return Response{}, fmt.Errorf("failed to chat complete: %w", err)
+	}
+
+	return a.responseToLLMResponse(res), nil
+}