From 9c9d6614fad9bac87d762dacce90835d7cbad601 Mon Sep 17 00:00:00 2001
From: KingSkyLi <15566300566@163.com>
Date: Fri, 30 Aug 2024 14:40:17 +0800
Subject: [PATCH] fixed;

---
 examples/rag/graph_rag_summary_example.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/examples/rag/graph_rag_summary_example.py b/examples/rag/graph_rag_summary_example.py
index f054acaa9..2ddb5d70d 100644
--- a/examples/rag/graph_rag_summary_example.py
+++ b/examples/rag/graph_rag_summary_example.py
@@ -1,11 +1,12 @@
 import asyncio
 import os
 
-from dbgpt.configs.model_config import ROOT_PATH
+from dbgpt.configs.model_config import ROOT_PATH, MODEL_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.rag.embedding import DefaultEmbeddingFactory
 from dbgpt.rag.retriever import RetrieverStrategy
 from dbgpt.storage.knowledge_graph.community_summary import (
     CommunitySummaryKnowledgeGraph,
@@ -13,6 +14,8 @@
 )
 from dbgpt.storage.knowledge_graph.knowledge_graph import BuiltinKnowledgeGraphConfig
 
+
+
 """GraphRAG example.
     pre-requirements:
     * Set LLM config (url/sk) in `.env`.
@@ -25,6 +28,7 @@
     TUGRAPH_USERNAME=admin
     TUGRAPH_PASSWORD=73@TuGraph
     GRAPH_COMMUNITY_SUMMARY_ENABLED=True
+    TUGRAPH_PLUGIN_NAMES=leiden
     ```
     Examples:
         ..code-block:: shell
@@ -36,9 +40,14 @@ def _create_kg_connector():
     """Create knowledge graph connector."""
     return CommunitySummaryKnowledgeGraph(
         config=CommunitySummaryKnowledgeGraphConfig(
-            name="graph_rag_test",
-            embedding_fn=None,
-            llm_client=OpenAILLMClient(),
+            name="graph_rag_summary_test",
+            embedding_fn=DefaultEmbeddingFactory(
+            default_model_name=os.path.join(MODEL_PATH, "text2vec-large-chinese"),
+        ).create(),
+            llm_client=OpenAILLMClient(
+                api_base=os.getenv('PROXY_SERVER_URL'),
+                api_key=os.getenv('PROXY_API_KEY'),
+            ),
             model_name="gpt-3.5-turbo",
         ),
     )
@@ -63,7 +72,7 @@ async def main():
         "What actions has Megatron taken ?", score_threshold=0.3
     )
     print(f"embedding rag example results:{chunks}")
-    graph_store.delete_vector_name("graph_rag_test")
+    graph_store.delete_vector_name("graph_rag_summary_test")
 
 
 if __name__ == "__main__":