134 lines
2.8 KiB
Go
134 lines
2.8 KiB
Go
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
|
|
}
|