taskAI.py

from llama_index.llms.openai_like import OpenAILike
from llama_index.core.llms import ChatMessage  # , MessageRole
from llama_index.core import ChatPromptTemplate

from util import mylogger

logger = mylogger(__name__,'%(asctime)s:%(levelname)s:%(message)s')
info = logger.info


## define templates

### topic,input
JD_PREPROCESS = ChatPromptTemplate(
    [
        ChatMessage(
            role="system",
            content="You are a content extractor. You never paraphrase; you only reduce content at the sentence level. Your mission is to extract information directly related to {topic} from user input. Make sure output contains complete information.",
        ),
        ChatMessage(role="user", content="{input}"),
    ]
)

### input
CV_PREPROCESS = ChatPromptTemplate(
    [
        ChatMessage(
            role="system",
            content="You are an AI text converter alternative to pandoc. Your mission is to convert the input content into markdown. Regarding styles, only keep headers, lists and links, and remove other styles.",
        ),
        ChatMessage(role="user", content="{input}"),
    ]
)

## basic func


def oai(base: str, key: str, model: str, **kwargs) -> OpenAILike:
    return OpenAILike(
        api_base=base,
        api_key=key,
        model=model,
        is_chat_model=True,
        context_window=window_size,
        **kwargs,
    )


## tasks
class TaskAI(OpenAILike):
    def __init__(self, api: dict[str, str], **kwargs):
        def guess_window_size(model=api["model"]):
            _mid = model.lower()
            windows: dict = {
                8000: ["gemma", "8k"],
                16000: ["16k"],
                32000: ["mistral", "mixtral", "32k"],
            }
            window_size = 3900
            for ws, names in windows.items():
                if any([n in _mid for n in names]):
                    window_size = ws
            info(f"use context window size: {window_size} for {model}")
            return window_size

        super().__init__(
            api_base=api["base"], api_key=api["key"], model=api["model"], is_chat_model=True, context_window=guess_window_size(), **kwargs
        )

    def jd_preprocess(self, topic: str, input: str):
        return self.stream_chat(JD_PREPROCESS.format_messages(topic=topic, input=input))

    def cv_preprocess(self, input: str):
        return self.stream_chat(CV_PREPROCESS.format_messages(input=input))