135 lines
4.5 KiB
Python
135 lines
4.5 KiB
Python
import asyncio
|
||
|
||
from raganything import RAGAnything, RAGAnythingConfig
|
||
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
|
||
from lightrag.utils import EmbeddingFunc
|
||
|
||
import Config.Config
|
||
|
||
|
||
async def main():
|
||
# 设置 API 配置
|
||
api_key = Config.Config.ALY_LLM_API_KEY
|
||
base_url = Config.Config.ALY_LLM_BASE_URL
|
||
|
||
# 创建 RAGAnything 配置
|
||
config = RAGAnythingConfig(
|
||
working_dir="./rag_storage",
|
||
parser="mineru", # 选择解析器：mineru 或 docling
|
||
parse_method="auto", # 解析方法：auto, ocr 或 txt
|
||
enable_image_processing=True,
|
||
enable_table_processing=True,
|
||
enable_equation_processing=True,
|
||
)
|
||
|
||
# 定义 LLM 模型函数
|
||
def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
|
||
return openai_complete_if_cache(
|
||
Config.Config.ALY_LLM_MODEL_NAME,
|
||
prompt,
|
||
system_prompt=system_prompt,
|
||
history_messages=history_messages,
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
**kwargs,
|
||
)
|
||
|
||
# 定义视觉模型函数用于图像处理
|
||
def vision_model_func(
|
||
prompt, system_prompt=None, history_messages=[], image_data=None, messages=None, **kwargs
|
||
):
|
||
# 如果提供了messages格式（用于多模态VLM增强查询），直接使用
|
||
if messages:
|
||
return openai_complete_if_cache(
|
||
Config.Config.GLM_MODEL_NAME,
|
||
"",
|
||
system_prompt=None,
|
||
history_messages=[],
|
||
messages=messages,
|
||
api_key=Config.Config.GLM_API_KEY,
|
||
base_url=Config.Config.GLM_BASE_URL,
|
||
**kwargs,
|
||
)
|
||
# 传统单图片格式
|
||
elif image_data:
|
||
return openai_complete_if_cache(
|
||
Config.Config.GLM_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=Config.Config.GLM_API_KEY,
|
||
base_url=Config.Config.GLM_BASE_URL,
|
||
**kwargs,
|
||
)
|
||
# 纯文本格式
|
||
else:
|
||
return llm_model_func(prompt, system_prompt, history_messages, **kwargs)
|
||
|
||
# 定义嵌入函数
|
||
embedding_func = EmbeddingFunc(
|
||
embedding_dim=Config.Config.EMBED_DIM,
|
||
max_token_size=Config.Config.EMBED_MAX_TOKEN_SIZE,
|
||
func=lambda texts: openai_embed(
|
||
texts,
|
||
model=Config.Config.EMBED_MODEL_NAME,
|
||
api_key=Config.Config.EMBED_API_KEY,
|
||
base_url=Config.Config.EMBED_BASE_URL,
|
||
),
|
||
)
|
||
from functools import partial
|
||
from lightrag.rerank import cohere_rerank
|
||
|
||
rerank_model_func = partial(
|
||
cohere_rerank,
|
||
model=Config.Config.RERANK_MODEL,
|
||
api_key=Config.Config.RERANK_BINDING_API_KEY,
|
||
base_url=Config.Config.RERANK_BASE_URL,
|
||
)
|
||
|
||
# 初始化 RAGAnything
|
||
rag = RAGAnything(
|
||
config=config,
|
||
llm_model_func=llm_model_func,
|
||
vision_model_func=vision_model_func,
|
||
embedding_func=embedding_func,
|
||
rerank_model_func=rerank_model_func,
|
||
)
|
||
|
||
# 处理文档
|
||
await rag.process_document_complete(
|
||
file_path="./Doc/GeoGebra5经典版指令汇编201903061.pdf",
|
||
output_dir="./output",
|
||
parse_method="auto"
|
||
)
|
||
|
||
# 查询处理后的内容
|
||
# 纯文本查询 - 基本知识库搜索
|
||
text_result = await rag.aquery(
|
||
"文档的主要内容是什么？",
|
||
mode="hybrid"
|
||
)
|
||
print("文本查询结果:", text_result)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|