The init commit:
- simple file command-line translate tools. - support config and command-line args
This commit is contained in:
261
main.go
Normal file
261
main.go
Normal file
@@ -0,0 +1,261 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/alecthomas/kong"
|
||||
kongtoml "github.com/alecthomas/kong-toml"
|
||||
)
|
||||
|
||||
type ChatMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
type chatCompletionReq struct {
|
||||
Model string `json:"model"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
ReasoningEffort string `json:"reasoning_effort,omitempty"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
|
||||
type streamEvent struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
} `json:"choices"`
|
||||
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
func streamChatText(
|
||||
ctx context.Context,
|
||||
client *http.Client,
|
||||
baseURL string,
|
||||
apiKey string,
|
||||
model string,
|
||||
reasoningEffort string,
|
||||
temperature *float64,
|
||||
msgs []ChatMessage,
|
||||
) (io.ReadCloser, error) {
|
||||
if msgs == nil || len(msgs) == 0 {
|
||||
return nil, errors.New("missing messages")
|
||||
}
|
||||
endpoint := strings.TrimRight(baseURL, "/") + "/chat/completions"
|
||||
|
||||
body := chatCompletionReq{
|
||||
Model: model,
|
||||
Messages: msgs,
|
||||
Temperature: temperature,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: true,
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
|
||||
go func() {
|
||||
defer func(pw *io.PipeWriter) {
|
||||
err := pw.Close()
|
||||
if err != nil {
|
||||
log.Println(err)
|
||||
}
|
||||
}(pw)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
_ = pw.CloseWithError(err)
|
||||
return
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
_ = pw.CloseWithError(err)
|
||||
return
|
||||
}
|
||||
defer func(Body io.ReadCloser) {
|
||||
err := Body.Close()
|
||||
if err != nil {
|
||||
log.Println(err)
|
||||
}
|
||||
}(resp.Body)
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode > 299 {
|
||||
b, _ := io.ReadAll(resp.Body)
|
||||
_ = pw.CloseWithError(fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b))))
|
||||
return
|
||||
}
|
||||
|
||||
sc := bufio.NewScanner(resp.Body)
|
||||
sc.Buffer(make([]byte, 0, 64*1024), 2*1024*1024)
|
||||
|
||||
var dataLines []string
|
||||
flushEvent := func() bool {
|
||||
if len(dataLines) == 0 {
|
||||
_ = pw.Close()
|
||||
return true
|
||||
}
|
||||
data := strings.Join(dataLines, "\n")
|
||||
dataLines = dataLines[:0]
|
||||
if data == "[DONE]" {
|
||||
_ = pw.Close()
|
||||
return false
|
||||
}
|
||||
|
||||
var evt streamEvent
|
||||
if err := json.Unmarshal([]byte(data), &evt); err != nil {
|
||||
_ = pw.CloseWithError(fmt.Errorf("failed to unmarshal event: %w, data=%q", err, data))
|
||||
return false
|
||||
}
|
||||
if evt.Error != nil {
|
||||
_ = pw.CloseWithError(fmt.Errorf("api error: %+v", evt.Error))
|
||||
return false
|
||||
}
|
||||
if len(evt.Choices) > 0 {
|
||||
chunk := evt.Choices[0].Delta.Content
|
||||
if chunk != "" {
|
||||
if _, err := io.WriteString(pw, chunk); err != nil {
|
||||
_ = pw.CloseWithError(err)
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
for sc.Scan() {
|
||||
line := sc.Text()
|
||||
|
||||
if line == "" {
|
||||
if ok := flushEvent(); !ok {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
v := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
dataLines = append(dataLines, v)
|
||||
}
|
||||
}
|
||||
|
||||
if len(dataLines) > 0 {
|
||||
_ = flushEvent()
|
||||
return
|
||||
}
|
||||
|
||||
if err := sc.Err(); err != nil {
|
||||
_ = pw.CloseWithError(err)
|
||||
return
|
||||
}
|
||||
}()
|
||||
|
||||
return pr, nil
|
||||
}
|
||||
|
||||
// Translate src to lang
|
||||
func Translate(
|
||||
ctx context.Context,
|
||||
client *http.Client,
|
||||
baseURL string,
|
||||
apiKey string,
|
||||
model string,
|
||||
reasoningEffort string,
|
||||
temperature *float64,
|
||||
src string,
|
||||
lang string,
|
||||
) (io.ReadCloser, error) {
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
msgs := []ChatMessage{
|
||||
{
|
||||
Role: "system",
|
||||
Content: "You are a translation API. Please fully translate the user input to '" + lang + "'. Do not include any formatting changes/non-translation output. No explanation of the translation content. Do not translate proper nouns.",
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: src,
|
||||
},
|
||||
}
|
||||
return streamChatText(ctx, client, baseURL, apiKey, model, reasoningEffort, temperature, msgs)
|
||||
}
|
||||
|
||||
var cli struct {
|
||||
Config string `short:"c" default:"~/.config/translate.toml" help:"Path to config file"`
|
||||
BaseURL string `short:"b" help:"LLM API BASE URL" default:"https://api.openai.com/v1"`
|
||||
ApiKey string `short:"k" help:"LLM API Key"`
|
||||
Model string `short:"m" help:"LLM Model" default:"gpt-5-nano"`
|
||||
ReasoningEffort string `help:"LLM Reasoning Effort (remember some LLM maybe un-support some settings)" default:"minimal"`
|
||||
Temperature *float64 `help:"LLM Temperature"`
|
||||
Lang string `short:"l" help:"The target language" default:"zh-CN"`
|
||||
ReverseLang string `help:"The target language when use -r" default:"en-US"`
|
||||
Reverse bool `short:"r" help:"Set this flag to true to enable reverse translation"`
|
||||
Src []string `arg:"" optional:"" name:"src" help:"Text to translate or leave empty to use stdin"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
_ = kong.Parse(&cli)
|
||||
_ = kong.Parse(
|
||||
&cli,
|
||||
kong.Configuration(kongtoml.Loader, cli.Config),
|
||||
)
|
||||
|
||||
lang := cli.Lang
|
||||
if cli.Reverse {
|
||||
lang = cli.ReverseLang
|
||||
}
|
||||
|
||||
src := strings.Join(cli.Src, " ")
|
||||
|
||||
if src == "" {
|
||||
stdin, err := io.ReadAll(os.Stdin)
|
||||
src = string(stdin)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
stream, err := Translate(context.Background(), http.DefaultClient, cli.BaseURL, cli.ApiKey, cli.Model,
|
||||
cli.ReasoningEffort, cli.Temperature, src, lang)
|
||||
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
_, err = io.Copy(os.Stdout, stream)
|
||||
|
||||
_, _ = os.Stdout.Write([]byte{'\n'})
|
||||
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
err = stream.Close()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user