feat: add GraphRAG framework and integrate TuGraph (#1506)

Co-authored-by: KingSkyLi <15566300566@163.com>
Co-authored-by: aries_ckt <916701291@qq.com>
Co-authored-by: Fangyin Cheng <staneyffer@gmail.com>
This commit is contained in:
Florian
2024-05-16 15:39:50 +08:00
committed by GitHub
parent 593e974405
commit a9087c3853
133 changed files with 10139 additions and 6631 deletions

View File

@@ -0,0 +1,66 @@
import asyncio
import os
from dbgpt.configs.model_config import ROOT_PATH
from dbgpt.model.proxy.llms.chatgpt import OpenAILLMClient
from dbgpt.rag import ChunkParameters
from dbgpt.rag.assembler import EmbeddingAssembler
from dbgpt.rag.knowledge import KnowledgeFactory
from dbgpt.storage.vector_store.base import VectorStoreConfig
from dbgpt.storage.vector_store.connector import VectorStoreConnector
"""GraphRAG example.
pre-requirements:
* Set LLM config (url/sk) in `.env`.
* Setup/startup TuGraph from: https://github.com/TuGraph-family/tugraph-db
* Config TuGraph following the format below.
```
GRAPH_STORE_TYPE=TuGraph
TUGRAPH_HOST=127.0.0.1
TUGRAPH_PORT=7687
TUGRAPH_USERNAME=admin
TUGRAPH_PASSWORD=73@TuGraph
```
Examples:
..code-block:: shell
python examples/rag/graph_rag_example.py
"""
def _create_kg_connector():
"""Create knowledge graph connector."""
return VectorStoreConnector(
vector_store_type="KnowledgeGraph",
vector_store_config=VectorStoreConfig(
name="graph_rag_test",
embedding_fn=None,
llm_client=OpenAILLMClient(),
model_name="gpt-3.5-turbo",
),
)
async def main():
file_path = os.path.join(ROOT_PATH, "examples/test_files/tranformers_story.md")
knowledge = KnowledgeFactory.from_file_path(file_path)
vector_connector = _create_kg_connector()
chunk_parameters = ChunkParameters(chunk_strategy="CHUNK_BY_SIZE")
# get embedding assembler
assembler = EmbeddingAssembler.load_from_knowledge(
knowledge=knowledge,
chunk_parameters=chunk_parameters,
vector_store_connector=vector_connector,
)
assembler.persist()
# get embeddings retriever
retriever = assembler.as_retriever(3)
chunks = await retriever.aretrieve_with_scores(
"What actions has Megatron taken ?", score_threshold=0.3
)
print(f"embedding rag example results:{chunks}")
vector_connector.delete_vector_name("graph_rag_test")
if __name__ == "__main__":
asyncio.run(main())