312 lines
7.9 KiB
Go
312 lines
7.9 KiB
Go
package openai
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/openai/openai-go"
|
|
"github.com/openai/openai-go/azure"
|
|
"github.com/openai/openai-go/option"
|
|
|
|
"github.com/ajac-zero/latticelm/internal/api"
|
|
"github.com/ajac-zero/latticelm/internal/config"
|
|
)
|
|
|
|
const Name = "openai"
|
|
|
|
// Provider implements the OpenAI SDK integration.
|
|
// It supports both direct OpenAI API and Azure-hosted endpoints.
|
|
type Provider struct {
|
|
cfg config.ProviderConfig
|
|
client *openai.Client
|
|
azure bool
|
|
}
|
|
|
|
// New constructs a Provider for the direct OpenAI API.
|
|
func New(cfg config.ProviderConfig) *Provider {
|
|
var client *openai.Client
|
|
if cfg.APIKey != "" {
|
|
c := openai.NewClient(option.WithAPIKey(cfg.APIKey))
|
|
client = &c
|
|
}
|
|
return &Provider{
|
|
cfg: cfg,
|
|
client: client,
|
|
}
|
|
}
|
|
|
|
// NewAzure constructs a Provider targeting Azure OpenAI.
|
|
// Azure OpenAI uses the OpenAI SDK with the azure subpackage for proper
|
|
// endpoint routing, api-version query parameter, and API key header.
|
|
func NewAzure(azureCfg config.AzureOpenAIConfig) *Provider {
|
|
var client *openai.Client
|
|
if azureCfg.APIKey != "" && azureCfg.Endpoint != "" {
|
|
apiVersion := azureCfg.APIVersion
|
|
if apiVersion == "" {
|
|
apiVersion = "2024-12-01-preview"
|
|
}
|
|
c := openai.NewClient(
|
|
azure.WithEndpoint(azureCfg.Endpoint, apiVersion),
|
|
azure.WithAPIKey(azureCfg.APIKey),
|
|
)
|
|
client = &c
|
|
}
|
|
return &Provider{
|
|
cfg: config.ProviderConfig{
|
|
APIKey: azureCfg.APIKey,
|
|
},
|
|
client: client,
|
|
azure: true,
|
|
}
|
|
}
|
|
|
|
// Name returns the provider identifier.
|
|
func (p *Provider) Name() string { return Name }
|
|
|
|
// Generate routes the request to OpenAI and returns a ProviderResult.
|
|
func (p *Provider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) {
|
|
if p.cfg.APIKey == "" {
|
|
return nil, fmt.Errorf("openai api key missing")
|
|
}
|
|
if p.client == nil {
|
|
return nil, fmt.Errorf("openai client not initialized")
|
|
}
|
|
|
|
// Convert messages to OpenAI format
|
|
oaiMessages := make([]openai.ChatCompletionMessageParamUnion, 0, len(messages))
|
|
for _, msg := range messages {
|
|
var content string
|
|
for _, block := range msg.Content {
|
|
if block.Type == "input_text" || block.Type == "output_text" {
|
|
content += block.Text
|
|
}
|
|
}
|
|
|
|
switch msg.Role {
|
|
case "user":
|
|
oaiMessages = append(oaiMessages, openai.UserMessage(content))
|
|
case "assistant":
|
|
oaiMessages = append(oaiMessages, openai.AssistantMessage(content))
|
|
case "system":
|
|
oaiMessages = append(oaiMessages, openai.SystemMessage(content))
|
|
case "developer":
|
|
oaiMessages = append(oaiMessages, openai.SystemMessage(content))
|
|
case "tool":
|
|
oaiMessages = append(oaiMessages, openai.ToolMessage(content, msg.CallID))
|
|
}
|
|
}
|
|
|
|
params := openai.ChatCompletionNewParams{
|
|
Model: openai.ChatModel(req.Model),
|
|
Messages: oaiMessages,
|
|
}
|
|
if req.MaxOutputTokens != nil {
|
|
params.MaxTokens = openai.Int(int64(*req.MaxOutputTokens))
|
|
}
|
|
if req.Temperature != nil {
|
|
params.Temperature = openai.Float(*req.Temperature)
|
|
}
|
|
if req.TopP != nil {
|
|
params.TopP = openai.Float(*req.TopP)
|
|
}
|
|
|
|
// Add tools if present
|
|
if req.Tools != nil && len(req.Tools) > 0 {
|
|
tools, err := parseTools(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("parse tools: %w", err)
|
|
}
|
|
params.Tools = tools
|
|
}
|
|
|
|
// Add tool_choice if present
|
|
if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
|
|
toolChoice, err := parseToolChoice(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("parse tool_choice: %w", err)
|
|
}
|
|
params.ToolChoice = toolChoice
|
|
}
|
|
|
|
// Add parallel_tool_calls if specified
|
|
if req.ParallelToolCalls != nil {
|
|
params.ParallelToolCalls = openai.Bool(*req.ParallelToolCalls)
|
|
}
|
|
|
|
// Call OpenAI API
|
|
resp, err := p.client.Chat.Completions.New(ctx, params)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("openai api error: %w", err)
|
|
}
|
|
|
|
var combinedText string
|
|
var toolCalls []api.ToolCall
|
|
|
|
for _, choice := range resp.Choices {
|
|
combinedText += choice.Message.Content
|
|
if len(choice.Message.ToolCalls) > 0 {
|
|
toolCalls = append(toolCalls, extractToolCalls(choice.Message)...)
|
|
}
|
|
}
|
|
|
|
return &api.ProviderResult{
|
|
ID: resp.ID,
|
|
Model: resp.Model,
|
|
Text: combinedText,
|
|
ToolCalls: toolCalls,
|
|
Usage: api.Usage{
|
|
InputTokens: int(resp.Usage.PromptTokens),
|
|
OutputTokens: int(resp.Usage.CompletionTokens),
|
|
TotalTokens: int(resp.Usage.TotalTokens),
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// GenerateStream handles streaming requests to OpenAI.
|
|
func (p *Provider) GenerateStream(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (<-chan *api.ProviderStreamDelta, <-chan error) {
|
|
deltaChan := make(chan *api.ProviderStreamDelta)
|
|
errChan := make(chan error, 1)
|
|
|
|
go func() {
|
|
defer close(deltaChan)
|
|
defer close(errChan)
|
|
|
|
if p.cfg.APIKey == "" {
|
|
errChan <- fmt.Errorf("openai api key missing")
|
|
return
|
|
}
|
|
if p.client == nil {
|
|
errChan <- fmt.Errorf("openai client not initialized")
|
|
return
|
|
}
|
|
|
|
// Convert messages to OpenAI format
|
|
oaiMessages := make([]openai.ChatCompletionMessageParamUnion, 0, len(messages))
|
|
for _, msg := range messages {
|
|
var content string
|
|
for _, block := range msg.Content {
|
|
if block.Type == "input_text" || block.Type == "output_text" {
|
|
content += block.Text
|
|
}
|
|
}
|
|
|
|
switch msg.Role {
|
|
case "user":
|
|
oaiMessages = append(oaiMessages, openai.UserMessage(content))
|
|
case "assistant":
|
|
oaiMessages = append(oaiMessages, openai.AssistantMessage(content))
|
|
case "system":
|
|
oaiMessages = append(oaiMessages, openai.SystemMessage(content))
|
|
case "developer":
|
|
oaiMessages = append(oaiMessages, openai.SystemMessage(content))
|
|
case "tool":
|
|
oaiMessages = append(oaiMessages, openai.ToolMessage(content, msg.CallID))
|
|
}
|
|
}
|
|
|
|
params := openai.ChatCompletionNewParams{
|
|
Model: openai.ChatModel(req.Model),
|
|
Messages: oaiMessages,
|
|
}
|
|
if req.MaxOutputTokens != nil {
|
|
params.MaxTokens = openai.Int(int64(*req.MaxOutputTokens))
|
|
}
|
|
if req.Temperature != nil {
|
|
params.Temperature = openai.Float(*req.Temperature)
|
|
}
|
|
if req.TopP != nil {
|
|
params.TopP = openai.Float(*req.TopP)
|
|
}
|
|
|
|
// Add tools if present
|
|
if req.Tools != nil && len(req.Tools) > 0 {
|
|
tools, err := parseTools(req)
|
|
if err != nil {
|
|
errChan <- fmt.Errorf("parse tools: %w", err)
|
|
return
|
|
}
|
|
params.Tools = tools
|
|
}
|
|
|
|
// Add tool_choice if present
|
|
if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
|
|
toolChoice, err := parseToolChoice(req)
|
|
if err != nil {
|
|
errChan <- fmt.Errorf("parse tool_choice: %w", err)
|
|
return
|
|
}
|
|
params.ToolChoice = toolChoice
|
|
}
|
|
|
|
// Add parallel_tool_calls if specified
|
|
if req.ParallelToolCalls != nil {
|
|
params.ParallelToolCalls = openai.Bool(*req.ParallelToolCalls)
|
|
}
|
|
|
|
// Create streaming request
|
|
stream := p.client.Chat.Completions.NewStreaming(ctx, params)
|
|
|
|
// Process stream
|
|
for stream.Next() {
|
|
chunk := stream.Current()
|
|
|
|
for _, choice := range chunk.Choices {
|
|
// Handle text content
|
|
if choice.Delta.Content != "" {
|
|
select {
|
|
case deltaChan <- &api.ProviderStreamDelta{
|
|
ID: chunk.ID,
|
|
Model: chunk.Model,
|
|
Text: choice.Delta.Content,
|
|
}:
|
|
case <-ctx.Done():
|
|
errChan <- ctx.Err()
|
|
return
|
|
}
|
|
}
|
|
|
|
// Handle tool call deltas
|
|
if len(choice.Delta.ToolCalls) > 0 {
|
|
delta := extractToolCallDelta(choice)
|
|
if delta != nil {
|
|
select {
|
|
case deltaChan <- &api.ProviderStreamDelta{
|
|
ID: chunk.ID,
|
|
Model: chunk.Model,
|
|
ToolCallDelta: delta,
|
|
}:
|
|
case <-ctx.Done():
|
|
errChan <- ctx.Err()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if err := stream.Err(); err != nil {
|
|
errChan <- fmt.Errorf("openai stream error: %w", err)
|
|
return
|
|
}
|
|
|
|
// Send final delta
|
|
select {
|
|
case deltaChan <- &api.ProviderStreamDelta{Done: true}:
|
|
case <-ctx.Done():
|
|
errChan <- ctx.Err()
|
|
}
|
|
}()
|
|
|
|
return deltaChan, errChan
|
|
}
|
|
|
|
func chooseModel(requested, defaultModel string) string {
|
|
if requested != "" {
|
|
return requested
|
|
}
|
|
if defaultModel != "" {
|
|
return defaultModel
|
|
}
|
|
return "gpt-4o-mini"
|
|
}
|