106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
import argparse
|
|
import mimetypes
|
|
from pathlib import Path
|
|
from .common import invert_input_tag_args, read_text_file
|
|
from ..configuration import Config
|
|
from ..message import MessageFilter, Message, Question
|
|
from ..chat import ChatDB, msg_location
|
|
|
|
|
|
class TranslationCmdError(Exception):
|
|
pass
|
|
|
|
|
|
text_separator: str = 'TEXT:'
|
|
|
|
|
|
def assert_document_type_supported_openai(document_file: Path) -> None:
|
|
doctype = mimetypes.guess_type(document_file)
|
|
if doctype != 'text/plain':
|
|
raise TranslationCmdError("AI 'OpenAI' only supports document type 'text/plain''")
|
|
|
|
|
|
def translation_prompt_openai(source_lang: str, target_lang: str) -> str:
|
|
"""
|
|
Return the prompt for GPT that tells it to do the translation.
|
|
"""
|
|
return f"Translate the text below the line {text_separator} from {source_lang} to {target_lang}."
|
|
|
|
|
|
def create_message_openai(chat: ChatDB, args: argparse.Namespace) -> Message:
|
|
"""
|
|
Create a new message from the given arguments and write it to the cache directory.
|
|
|
|
Message format
|
|
1. Translation prompt (tells GPT to do a translation)
|
|
2. Glossary (if specified as an argument)
|
|
3. User provided prompt enhancements
|
|
4. Translation separator
|
|
5. User provided text to be translated
|
|
|
|
The text to be translated is determined as a follows:
|
|
- if a document is provided in the arguments, translate its content
|
|
- if no document is provided, translate the last text argument
|
|
|
|
The other text arguments will be put into the "header" and can be used
|
|
to improve the translation prompt.
|
|
"""
|
|
text_args: list[str] = []
|
|
if args.create is not None:
|
|
text_args = args.create
|
|
elif args.ask is not None:
|
|
text_args = args.ask
|
|
else:
|
|
raise TranslationCmdError("No input text found")
|
|
|
|
# extract user prompt and user text to be translated
|
|
user_text: str
|
|
user_prompt: str
|
|
if args.input_document is not None:
|
|
assert_document_type_supported_openai(Path(args.input_document))
|
|
user_text = read_text_file(Path(args.input_document))
|
|
user_prompt = '\n\n'.join([str(s) for s in text_args])
|
|
else:
|
|
user_text = text_args[-1]
|
|
user_prompt = '\n\n'.join([str(s) for s in text_args[:-1]])
|
|
|
|
# build full question string
|
|
# FIXME: add glossaries if given
|
|
question_text: str = '\n\n'.join([translation_prompt_openai(args.source_lang, args.target_lang),
|
|
user_prompt,
|
|
text_separator,
|
|
user_text])
|
|
# create and write the message
|
|
message = Message(question=Question(question_text),
|
|
tags=args.output_tags,
|
|
ai=args.AI,
|
|
model=args.model)
|
|
# only write the new message to the cache,
|
|
# don't add it to the internal list
|
|
chat.cache_write([message])
|
|
return message
|
|
|
|
|
|
def translation_cmd(args: argparse.Namespace, config: Config) -> None:
|
|
"""
|
|
Handler for the 'translation' command. Creates and executes translation
|
|
requests based on the input and selected AI. Depending on the AI, the
|
|
whole process may be significantly different (e.g. DeepL vs OpenAI).
|
|
"""
|
|
invert_input_tag_args(args)
|
|
mfilter = MessageFilter(tags_or=args.or_tags,
|
|
tags_and=args.and_tags,
|
|
tags_not=args.exclude_tags)
|
|
chat = ChatDB.from_dir(cache_path=Path(config.cache),
|
|
db_path=Path(config.db),
|
|
mfilter=mfilter,
|
|
glob=args.glob,
|
|
loc=msg_location(args.location))
|
|
# if it's a new translation, create and store it immediately
|
|
# FIXME: check AI type
|
|
if args.ask or args.create:
|
|
# message = create_message(chat, args)
|
|
create_message_openai(chat, args)
|
|
if args.create:
|
|
return
|