mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-07-27 13:57:46 +00:00
37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
import asyncio
|
|
|
|
from dbgpt.model.proxy import DeepseekLLMClient
|
|
from dbgpt.rag.knowledge.base import ChunkStrategy
|
|
from dbgpt.rag.transformer.keyword_extractor import KeywordExtractor
|
|
from dbgpt_ext.rag import ChunkParameters
|
|
from dbgpt_ext.rag.knowledge import KnowledgeFactory
|
|
from dbgpt_ext.rag.retriever.doc_tree import DocTreeRetriever
|
|
|
|
|
|
async def main():
|
|
knowledge = KnowledgeFactory.from_file_path("../../docs/docs/awel/awel.md")
|
|
chunk_parameters = ChunkParameters(
|
|
chunk_strategy=ChunkStrategy.CHUNK_BY_MARKDOWN_HEADER.name
|
|
)
|
|
docs = knowledge.load()
|
|
docs = knowledge.extract(docs, chunk_parameters)
|
|
llm_client = DeepseekLLMClient(api_key="your_api_key")
|
|
keyword_extractor = KeywordExtractor(
|
|
llm_client=llm_client, model_name="deepseek-chat"
|
|
)
|
|
# doc tree retriever retriever
|
|
retriever = DocTreeRetriever(
|
|
docs=docs,
|
|
top_k=10,
|
|
keywords_extractor=keyword_extractor,
|
|
with_content=False,
|
|
)
|
|
tree_index = retriever._tree_indexes[0]
|
|
nodes = await retriever.aretrieve("Introduce awel Operators")
|
|
for node in nodes:
|
|
tree_index.display_tree(node)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|