100 lines
2.6 KiB
Go
100 lines
2.6 KiB
Go
package ai
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
)
|
|
|
|
type OpenAIResponse struct {
|
|
ID string `json:"id"`
|
|
Object string `json:"object"`
|
|
Created int64 `json:"created"`
|
|
Model string `json:"model"`
|
|
Choices []Choice `json:"choices"`
|
|
Usage Usage `json:"usage"`
|
|
ServiceTier string `json:"service_tier"`
|
|
}
|
|
|
|
type Choice struct {
|
|
Index int64 `json:"index"`
|
|
Message Message `json:"message"`
|
|
Logprobs interface{} `json:"logprobs"`
|
|
FinishReason string `json:"finish_reason"`
|
|
}
|
|
|
|
type Message struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
Refusal interface{} `json:"refusal"`
|
|
Annotations []interface{} `json:"annotations"`
|
|
}
|
|
|
|
type Usage struct {
|
|
PromptTokens int64 `json:"prompt_tokens"`
|
|
CompletionTokens int64 `json:"completion_tokens"`
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
PromptTokensDetails PromptTokensDetails `json:"prompt_tokens_details"`
|
|
CompletionTokensDetails CompletionTokensDetails `json:"completion_tokens_details"`
|
|
}
|
|
|
|
type CompletionTokensDetails struct {
|
|
ReasoningTokens int64 `json:"reasoning_tokens"`
|
|
AudioTokens int64 `json:"audio_tokens"`
|
|
AcceptedPredictionTokens int64 `json:"accepted_prediction_tokens"`
|
|
RejectedPredictionTokens int64 `json:"rejected_prediction_tokens"`
|
|
}
|
|
|
|
type PromptTokensDetails struct {
|
|
CachedTokens int64 `json:"cached_tokens"`
|
|
AudioTokens int64 `json:"audio_tokens"`
|
|
}
|
|
|
|
type OpenAIAPI struct {
|
|
Token string
|
|
Model string
|
|
BaseURL string
|
|
}
|
|
|
|
func NewOpenAIAPI(baseURL, token, model string) *OpenAIAPI {
|
|
return &OpenAIAPI{
|
|
Token: token,
|
|
Model: model,
|
|
BaseURL: baseURL,
|
|
}
|
|
}
|
|
|
|
type CreateCompletionReq struct {
|
|
Model string `json:"model"`
|
|
Messages []Message `json:"messages"`
|
|
}
|
|
|
|
func (o *OpenAIAPI) CreateCompletion(request CreateCompletionReq) (*OpenAIResponse, error) {
|
|
url := fmt.Sprintf("%s/v1/chat/completions", o.BaseURL)
|
|
data, err := json.Marshal(request)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
buf := bytes.NewBuffer(data)
|
|
req, err := http.NewRequest("POST", url, buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", o.Token))
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
response := new(OpenAIResponse)
|
|
err = json.Unmarshal(body, response)
|
|
return response, err
|
|
}
|