77 lines
2.1 KiB
Go
77 lines
2.1 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
|
|
openai "github.com/sashabaranov/go-openai"
|
|
|
|
coretypes "github.com/enmanuel/agents/pkg/llm"
|
|
)
|
|
|
|
// NewOpenAIComplete returns a CompleteFunc backed by the OpenAI-compatible API.
|
|
// Works with OpenAI, Ollama, vLLM, LMStudio — just change baseURL.
|
|
func NewOpenAIComplete(apiKeyEnv, baseURL string) coretypes.CompleteFunc {
|
|
return func(ctx context.Context, req coretypes.CompletionRequest) (coretypes.CompletionResponse, error) {
|
|
apiKey := os.Getenv(apiKeyEnv)
|
|
if apiKey == "" {
|
|
apiKey = "ollama" // Ollama doesn't require a real key
|
|
}
|
|
|
|
cfg := openai.DefaultConfig(apiKey)
|
|
if baseURL != "" {
|
|
cfg.BaseURL = baseURL
|
|
}
|
|
client := openai.NewClientWithConfig(cfg)
|
|
|
|
msgs := make([]openai.ChatCompletionMessage, 0, len(req.Messages)+1)
|
|
if req.SystemPrompt != "" {
|
|
msgs = append(msgs, openai.ChatCompletionMessage{
|
|
Role: openai.ChatMessageRoleSystem,
|
|
Content: req.SystemPrompt,
|
|
})
|
|
}
|
|
for _, m := range req.Messages {
|
|
role := openai.ChatMessageRoleUser
|
|
switch m.Role {
|
|
case coretypes.RoleAssistant:
|
|
role = openai.ChatMessageRoleAssistant
|
|
case coretypes.RoleSystem:
|
|
role = openai.ChatMessageRoleSystem
|
|
case coretypes.RoleTool:
|
|
role = openai.ChatMessageRoleTool
|
|
}
|
|
msgs = append(msgs, openai.ChatCompletionMessage{
|
|
Role: role,
|
|
Content: m.Content,
|
|
})
|
|
}
|
|
|
|
openReq := openai.ChatCompletionRequest{
|
|
Model: req.Model,
|
|
Messages: msgs,
|
|
MaxTokens: req.MaxTokens,
|
|
Temperature: float32(req.Temperature),
|
|
}
|
|
|
|
resp, err := client.CreateChatCompletion(ctx, openReq)
|
|
if err != nil {
|
|
return coretypes.CompletionResponse{}, fmt.Errorf("openai completion: %w", err)
|
|
}
|
|
if len(resp.Choices) == 0 {
|
|
return coretypes.CompletionResponse{}, fmt.Errorf("openai: empty choices")
|
|
}
|
|
|
|
return coretypes.CompletionResponse{
|
|
Content: resp.Choices[0].Message.Content,
|
|
FinishReason: string(resp.Choices[0].FinishReason),
|
|
Usage: coretypes.TokenUsage{
|
|
InputTokens: resp.Usage.PromptTokens,
|
|
OutputTokens: resp.Usage.CompletionTokens,
|
|
TotalTokens: resp.Usage.TotalTokens,
|
|
},
|
|
}, nil
|
|
}
|
|
}
|