|
|
import asyncio
|
|
|
import sys
|
|
|
|
|
|
import loguru
|
|
|
from raganything import RAGAnything, RAGAnythingConfig
|
|
|
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
|
|
|
from lightrag.utils import EmbeddingFunc
|
|
|
|
|
|
|
|
|
async def main():
|
|
|
# 在main函数开头添加日志配置
|
|
|
loguru.logger.remove() # 移除默认配置
|
|
|
loguru.logger.add(sys.stderr, level="INFO") # 输出INFO及以上级别到控制台
|
|
|
|
|
|
# Set up API configuration
|
|
|
api_key = "sk-44ae895eeb614aa1a9c6460579e322f1"
|
|
|
base_url = "https://api.deepseek.com" # Optional
|
|
|
|
|
|
# Create RAGAnything configuration
|
|
|
config = RAGAnythingConfig(
|
|
|
working_dir="./rag_storage",
|
|
|
mineru_parse_method="auto",
|
|
|
enable_image_processing=True,
|
|
|
enable_table_processing=True,
|
|
|
enable_equation_processing=True,
|
|
|
)
|
|
|
|
|
|
# Define LLM model function
|
|
|
def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
|
|
|
return openai_complete_if_cache(
|
|
|
"deepseek-chat",
|
|
|
prompt,
|
|
|
system_prompt=system_prompt,
|
|
|
history_messages=history_messages,
|
|
|
api_key=api_key,
|
|
|
base_url=base_url,
|
|
|
**kwargs,
|
|
|
)
|
|
|
|
|
|
# Define vision model function for image processing
|
|
|
def vision_model_func(
|
|
|
prompt, system_prompt=None, history_messages=[], image_data=None, **kwargs
|
|
|
):
|
|
|
if image_data:
|
|
|
return openai_complete_if_cache(
|
|
|
"GLM-4.1V-9B-Thinking", # 这里设置为了硅基流动中免费视觉模型的模型名称
|
|
|
"",
|
|
|
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="sk-pbqibyjwhrgmnlsmdygplahextfaclgnedetybccknxojlyl",
|
|
|
base_url="https://api.siliconflow.cn/v1/chat/completions",
|
|
|
**kwargs,
|
|
|
)
|
|
|
else:
|
|
|
return llm_model_func(prompt, system_prompt, history_messages, **kwargs)
|
|
|
|
|
|
# Define embedding function
|
|
|
embedding_func = EmbeddingFunc(
|
|
|
embedding_dim=1024, # 这里设置为了硅基流动中免费模型的嵌入维度
|
|
|
max_token_size=8192,
|
|
|
func=lambda texts: openai_embed(
|
|
|
texts,
|
|
|
model="BAAI/bge-m3",
|
|
|
api_key="sk-pbqibyjwhrgmnlsmdygplahextfaclgnedetybccknxojlyl",
|
|
|
base_url="https://api.siliconflow.cn/v1",
|
|
|
),
|
|
|
)
|
|
|
|
|
|
# Initialize RAGAnything
|
|
|
rag = RAGAnything(
|
|
|
config=config,
|
|
|
llm_model_func=llm_model_func,
|
|
|
vision_model_func=vision_model_func,
|
|
|
embedding_func=embedding_func,
|
|
|
)
|
|
|
|
|
|
# Process a document
|
|
|
await rag.process_document_complete(
|
|
|
file_path=r"D:\python\小乔证件\黄琬乔2023蓝桥杯省赛准考证.pdf",
|
|
|
output_dir="./output",
|
|
|
parse_method="auto"
|
|
|
)
|
|
|
|
|
|
print("Processing complete.")
|
|
|
|
|
|
# 调用LightRAG的代码,包括:1、创建索引,2、查询知识库,不应该放在这个模块中完成,这里只需要完成到文本解析完成即可。
|
|
|
# Query the processed content
|
|
|
# Pure text query - for basic knowledge base search
|
|
|
# text_result = await rag.aquery(
|
|
|
# "这个文档是做什么用的?",
|
|
|
# mode="hybrid"
|
|
|
# )
|
|
|
# print("Text query result:", text_result)
|
|
|
|
|
|
# # Multimodal query with specific multimodal content
|
|
|
# multimodal_result = await rag.aquery_with_multimodal(
|
|
|
# "Explain this formula and its relevance to the document content",
|
|
|
# multimodal_content=[{
|
|
|
# "type": "equation",
|
|
|
# "latex": "P(d|q) = \\frac{P(q|d) \\cdot P(d)}{P(q)}",
|
|
|
# "equation_caption": "Document relevance probability"
|
|
|
# }],
|
|
|
# mode="hybrid"
|
|
|
# )
|
|
|
# print("Multimodal query result:", multimodal_result)
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
asyncio.run(main())
|