Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add OpenAiX translation engine #143

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 119 additions & 0 deletions translate/openai/openai_custom.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
package openai

import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"

"github.com/metatube-community/metatube-sdk-go/common/fetch"
"github.com/metatube-community/metatube-sdk-go/translate"
)

var _ translate.Translator = (*OpenAIX)(nil)

type OpenAIX struct {
APIKey string `json:"openaix-api-key"`
BaseURL string `json:"openaix-base-url"`
Model string `json:"openaix-model"`
SystemPrompt string `json:"openaix-system-prompt"`
}

func (oax *OpenAIX) Translate(q, source, target string) (result string, err error) {
if oax.BaseURL == "" {
return "", translate.ErrTranslator
}

// Prepare the chat message
systemPrompt := oax.SystemPrompt
if systemPrompt == "" {
systemPrompt = `You are a professional translator for adult video content.
Rules:
1. Use official translations for actor/actress names if available, otherwise keep them unchanged
2. Do not invent translations for names without official versions
3. Maintain any numbers, dates, and measurements in their original format
4. Translate naturally and fluently, avoiding word-for-word translation
5. Do not add any explanations or notes
6. Only output the translation`
}

userPrompt := fmt.Sprintf("Please translate the following text from %s to %s:\n\n%s", source, target, q)

model := oax.Model
if model == "" {
model = "gpt-3.5-turbo"
}

reqBody := map[string]interface{}{
"model": model,
"messages": []map[string]string{
{
"role": "system",
"content": systemPrompt,
},
{
"role": "user",
"content": userPrompt,
},
},
"temperature": 0.3,
"max_tokens": 1000,
}

opts := []fetch.Option{
fetch.WithRaiseForStatus(true),
fetch.WithHeader("Content-Type", "application/json"),
fetch.WithHeader("Accept", "application/json"),
}

if oax.APIKey != "" {
opts = append(opts,
fetch.WithHeader("Authorization", "Bearer "+oax.APIKey),
)
}

var resp *http.Response
if resp, err = fetch.Post(oax.BaseURL, fetch.WithJSONBody(reqBody), opts...); err != nil {
return "", fmt.Errorf("request failed: %v", err)
}
defer resp.Body.Close()

respBody, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("failed to read response body: %v", err)
}

var data struct {
Error *struct {
Message string `json:"message"`
Type string `json:"type"`
Code string `json:"code"`
} `json:"error,omitempty"`
Choices []struct {
Message struct {
Content string `json:"content"`
} `json:"message"`
} `json:"choices"`
}

if err = json.NewDecoder(bytes.NewReader(respBody)).Decode(&data); err != nil {
return "", fmt.Errorf("failed to decode response: %v, body: %s", err, string(respBody))
}

if data.Error != nil {
return "", fmt.Errorf("API error: %s (%s)", data.Error.Message, data.Error.Type)
}

if len(data.Choices) > 0 {
result = strings.TrimSpace(data.Choices[0].Message.Content)
} else {
err = fmt.Errorf("no translation result")
}
return
}

func init() {
translate.Register(&OpenAIX{})
}
39 changes: 39 additions & 0 deletions translate/openai/openai_custom_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
package openai

import (
"os"
"testing"
)

func TestOpenAiXTranslate(t *testing.T) {
baseURL := os.Getenv("OPENAI_X_BASE_URL")
if baseURL == "" {
t.Skip("OPENAI_X_BASE_URL not set")
}

apiKey := os.Getenv("OPENAI_X_API_KEY")
model := os.Getenv("OPENAI_X_MODEL")

translator := &OpenAIX{
APIKey: apiKey,
BaseURL: baseURL,
Model: model,
}

for _, unit := range []struct {
text, from, to string
}{
{`Oh yeah! I'm a translator!`, "EN", "ZH"},
{`私は翻訳者です`, "JA", "ZH"},
{`我是一个翻译器`, "ZH", "EN"},
{`PPPE-001 深田えいみ 親友からこっそり彼氏を寝取る巨乳でエッチな痴女お姉さん`, "JA", "ZH"},
{`A Busty and Naughty Sister Who Secretly Takes Her Best Friend's Boyfriend`, "EN", "ZH"},
} {
result, err := translator.Translate(unit.text, unit.from, unit.to)
if err != nil {
t.Errorf("Failed to translate text %q from %q to %q: %v", unit.text, unit.from, unit.to, err)
continue
}
t.Logf("Translated %q (%s->%s): %q", unit.text, unit.from, unit.to, result)
}
}