import lazyllm
from lazyllm import (
OnlineEmbeddingModule, OnlineChatModule, Document, SentenceSplitter,
Retriever, Reranker, ChatPrompter, pipeline
)
# Initialize api key and prompt
api_key = 'sk-'
prompt = """
You will play the role of an AI Q&A assistant and complete a dialogue task.
In this task, you need to provide your answer based on the given context and question.
"""
# Initialize models
embed_model = OnlineEmbeddingModule(source="siliconflow", api_key=api_key)
rerank_model = OnlineEmbeddingModule(source="siliconflow", api_key=api_key, type="rerank")
llm = OnlineChatModule(source="siliconflow", api_key=api_key)
# Define the document management module and create node groups
doc = Document(dataset_path="/home/xxx/database", manager=False, embed=embed_model)
doc.create_node_group(name="block", transform=SentenceSplitter, chunk_size=1024, chunk_overlap=100)
doc.create_node_group(name="line", transform=SentenceSplitter, chunk_size=128, chunk_overlap=20, parent="block")
# Build the RAG pipeline (multi-route retrieval → rerank → prompt formatting → LLM answering)
with pipeline() as ppl:
with lazyllm.parallel().sum as ppl.prl:
prl.r1 = Retriever(doc, group_name='line', similarity="cosine", topk=6, target='block')
prl.r2 = Retriever(doc, group_name='block', similarity="cosine", topk=6)
ppl.reranker = Reranker('ModuleReranker', model=rerank_model, output_format='content',
join=True) | bind(query=ppl.input)
ppl.formatter = (lambda context, query: dict(context_str=str(context), query=query)) | bind(query=ppl.input)
ppl.llm = llm.prompt(lazyllm.ChatPrompter(prompt, extra_keys=["context_str"]))
ppl.start()
query = "What is the Way of Heaven?"
print(ppl(query))