You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
73 lines
2.3 KiB
73 lines
2.3 KiB
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,
|
|
),
|
|
)
|
|
|