from lightrag.llm.openai import openai_complete_if_cache, openai_embed
from lightrag.utils import EmbeddingFunc

from Config.Config import *


def create_llm_model_func(history_messages):
    def llm_model_func(prompt, system_prompt=None, **kwargs):
        return openai_complete_if_cache(
            LLM_MODEL_NAME,
            prompt,
            system_prompt=system_prompt,
            history_messages=history_messages,
            api_key=LLM_API_KEY,
            base_url=LLM_BASE_URL,
            **kwargs,
        )

    return llm_model_func


def create_vision_model_func(llm_model_func):
    def vision_model_func(
            prompt, system_prompt=None, history_messages=[], image_data=None, **kwargs
    ):
        if image_data:
            return openai_complete_if_cache(
                VISION_MODEL_NAME,
                "",
                system_prompt=None,
                history_messages=[],
                messages=[
                    {"role": "system", "content": system_prompt}
                    if system_prompt
                    else None,
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_data}"
                                },
                            },
                        ],
                    }
                    if image_data
                    else {"role": "user", "content": prompt},
                ],
                api_key=VISION_API_KEY,
                base_url=VISION_BASE_URL,
                **kwargs,
            )
        else:
            return llm_model_func(prompt, system_prompt, history_messages, **kwargs)

    return vision_model_func


def create_embedding_func():
    return EmbeddingFunc(
        embedding_dim=1024,
        max_token_size=8192,
        func=lambda texts: openai_embed(
            texts,
            model=EMBED_MODEL_NAME,
            api_key=EMBED_API_KEY,
            base_url=EMBED_BASE_URL,
        ),
    )