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.

129 lines
4.6 KiB

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

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())