127 lines
3.4 KiB
Go
127 lines
3.4 KiB
Go
package translator
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/wholetrans/d10n/config"
|
|
)
|
|
|
|
// Translator handles communication with the OpenAI API
|
|
type Translator struct {
|
|
config *config.Config
|
|
}
|
|
|
|
// Message represents a message in the OpenAI chat format
|
|
type Message struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
// ChatRequest represents a request to the OpenAI chat completion API
|
|
type ChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []Message `json:"messages"`
|
|
}
|
|
|
|
// ChatResponse represents a response from the OpenAI chat completion API
|
|
type ChatResponse struct {
|
|
Choices []struct {
|
|
Message struct {
|
|
Content string `json:"content"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
Error *struct {
|
|
Message string `json:"message"`
|
|
} `json:"error,omitempty"`
|
|
}
|
|
|
|
// NewTranslator creates a new translator
|
|
func NewTranslator(cfg *config.Config) *Translator {
|
|
return &Translator{
|
|
config: cfg,
|
|
}
|
|
}
|
|
|
|
// Translate translates content from sourceLanguage to targetLanguage
|
|
func (t *Translator) Translate(content, sourceLanguage, targetLanguage string) (string, error) {
|
|
messages := []Message{
|
|
{
|
|
Role: "system",
|
|
Content: t.getSystemPrompt(sourceLanguage, targetLanguage),
|
|
},
|
|
{
|
|
Role: "user",
|
|
Content: content,
|
|
},
|
|
}
|
|
|
|
// Create the API request
|
|
requestBody := ChatRequest{
|
|
Model: t.config.Model,
|
|
Messages: messages,
|
|
}
|
|
|
|
jsonData, err := json.Marshal(requestBody)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error marshaling request: %w", err)
|
|
}
|
|
|
|
// Create HTTP request
|
|
req, err := http.NewRequest("POST", t.config.APIBase+"/v1/chat/completions", bytes.NewBuffer(jsonData))
|
|
if err != nil {
|
|
return "", fmt.Errorf("error creating HTTP request: %w", err)
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+t.config.APIKey)
|
|
|
|
// Send the request
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("error sending request to API: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// Parse the response
|
|
var responseBody ChatResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&responseBody); err != nil {
|
|
return "", fmt.Errorf("error parsing API response: %w", err)
|
|
}
|
|
|
|
// Check for API errors
|
|
if responseBody.Error != nil {
|
|
return "", fmt.Errorf("API error: %s", responseBody.Error.Message)
|
|
}
|
|
|
|
// Check if we have any choices
|
|
if len(responseBody.Choices) == 0 {
|
|
return "", fmt.Errorf("API returned no translation")
|
|
}
|
|
|
|
return responseBody.Choices[0].Message.Content, nil
|
|
}
|
|
|
|
// getSystemPrompt constructs the system prompt for translation
|
|
func (t *Translator) getSystemPrompt(sourceLanguage, targetLanguage string) string {
|
|
basePrompt := t.config.SystemPrompt
|
|
|
|
// Replace variables in the system prompt
|
|
basePrompt = strings.ReplaceAll(basePrompt, "$SOURCE_LANG", sourceLanguage)
|
|
basePrompt = strings.ReplaceAll(basePrompt, "$TARGET_LANG", targetLanguage)
|
|
|
|
// If the source language is specified, include it in the prompt
|
|
sourceLangStr := ""
|
|
if sourceLanguage != "" {
|
|
sourceLangStr = fmt.Sprintf(" from %s", sourceLanguage)
|
|
}
|
|
|
|
// Append the translation instruction to the base prompt
|
|
translationInstruction := fmt.Sprintf("\nTranslate the following text%s to %s. Only output the translated text, without any explanations or additional content.", sourceLangStr, targetLanguage)
|
|
|
|
return basePrompt + translationInstruction
|
|
}
|