Add graphrag (#1793)

### What problem does this PR solve? #1594 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
infiniflow · Aug 2, 2024 · 152072f · 152072f
1 parent 80032b1
commit 152072f
Show file tree

Hide file tree

Showing 74 changed files with 2,522 additions and 105 deletions.
diff --git a/graph/README.md → agent/README.md b/graph/README.md → agent/README.md
diff --git a/graph/README_zh.md → agent/README_zh.md b/graph/README_zh.md → agent/README_zh.md
diff --git a/graph/__init__.py → agent/__init__.py b/graph/__init__.py → agent/__init__.py
diff --git a/graph/canvas.py → agent/canvas.py b/graph/canvas.py → agent/canvas.py
@@ -22,9 +22,9 @@
 
 import pandas as pd
 
-from graph.component import component_class
-from graph.component.base import ComponentBase
-from graph.settings import flow_logger, DEBUG
+from agent.component import component_class
+from agent.component.base import ComponentBase
+from agent.settings import flow_logger, DEBUG
 
 
 class Canvas(ABC):

diff --git a/graph/component/__init__.py → agent/component/__init__.py b/graph/component/__init__.py → agent/component/__init__.py
diff --git a/graph/component/answer.py → agent/component/answer.py b/graph/component/answer.py → agent/component/answer.py
@@ -19,7 +19,7 @@
 
 import pandas as pd
 
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class AnswerParam(ComponentParamBase):

diff --git a/graph/component/arxiv.py → agent/component/arxiv.py b/graph/component/arxiv.py → agent/component/arxiv.py
@@ -13,13 +13,11 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import random
 from abc import ABC
-from functools import partial
 import arxiv
 import pandas as pd
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class ArXivParam(ComponentParamBase):

diff --git a/graph/component/baidu.py → agent/component/baidu.py b/graph/component/baidu.py → agent/component/baidu.py
@@ -19,8 +19,8 @@
 import pandas as pd
 import requests
 import re
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class BaiduParam(ComponentParamBase):

diff --git a/graph/component/base.py → agent/component/base.py b/graph/component/base.py → agent/component/base.py
@@ -23,8 +23,8 @@
 
 import pandas as pd
 
-from graph import settings
-from graph.settings import flow_logger, DEBUG
+from agent import settings
+from agent.settings import flow_logger, DEBUG
 
 _FEEDED_DEPRECATED_PARAMS = "_feeded_deprecated_params"
 _DEPRECATED_PARAMS = "_deprecated_params"

diff --git a/graph/component/begin.py → agent/component/begin.py b/graph/component/begin.py → agent/component/begin.py
@@ -13,11 +13,10 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import json
 from functools import partial
-
 import pandas as pd
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.component.base import ComponentBase, ComponentParamBase
+
 
 class BeginParam(ComponentParamBase):
 

diff --git a/graph/component/bing.py → agent/component/bing.py b/graph/component/bing.py → agent/component/bing.py
@@ -16,8 +16,8 @@
 from abc import ABC
 import requests
 import pandas as pd
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class BingParam(ComponentParamBase):

diff --git a/graph/component/categorize.py → agent/component/categorize.py b/graph/component/categorize.py → agent/component/categorize.py
@@ -14,13 +14,10 @@
 #  limitations under the License.
 #
 from abc import ABC
-
-import pandas as pd
-
 from api.db import LLMType
 from api.db.services.llm_service import LLMBundle
-from graph.component import GenerateParam, Generate
-from graph.settings import DEBUG
+from agent.component import GenerateParam, Generate
+from agent.settings import DEBUG
 
 
 class CategorizeParam(GenerateParam):

diff --git a/graph/component/cite.py → agent/component/cite.py b/graph/component/cite.py → agent/component/cite.py
@@ -21,7 +21,7 @@
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
 from api.settings import retrievaler
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class CiteParam(ComponentParamBase):

diff --git a/graph/component/duckduckgo.py → agent/component/duckduckgo.py b/graph/component/duckduckgo.py → agent/component/duckduckgo.py
@@ -13,13 +13,11 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import random
 from abc import ABC
-from functools import partial
 from duckduckgo_search import DDGS
 import pandas as pd
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class DuckDuckGoParam(ComponentParamBase):

diff --git a/graph/component/generate.py → agent/component/generate.py b/graph/component/generate.py → agent/component/generate.py
@@ -15,13 +15,11 @@
 #
 import re
 from functools import partial
-
 import pandas as pd
-
 from api.db import LLMType
 from api.db.services.llm_service import LLMBundle
 from api.settings import retrievaler
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class GenerateParam(ComponentParamBase):

diff --git a/graph/component/google.py → agent/component/google.py b/graph/component/google.py → agent/component/google.py
@@ -16,8 +16,8 @@
 from abc import ABC
 from serpapi import GoogleSearch
 import pandas as pd
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class GoogleParam(ComponentParamBase):

diff --git a/graph/component/googlescholar.py → agent/component/googlescholar.py b/graph/component/googlescholar.py → agent/component/googlescholar.py
@@ -15,8 +15,8 @@
 #
 from abc import ABC
 import pandas as pd
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 from scholarly import scholarly
 
 

diff --git a/graph/component/keyword.py → agent/component/keyword.py b/graph/component/keyword.py → agent/component/keyword.py
@@ -17,8 +17,8 @@
 from abc import ABC
 from api.db import LLMType
 from api.db.services.llm_service import LLMBundle
-from graph.component import GenerateParam, Generate
-from graph.settings import DEBUG
+from agent.component import GenerateParam, Generate
+from agent.settings import DEBUG
 
 
 class KeywordExtractParam(GenerateParam):

diff --git a/graph/component/message.py → agent/component/message.py b/graph/component/message.py → agent/component/message.py
@@ -16,10 +16,7 @@
 import random
 from abc import ABC
 from functools import partial
-
-import pandas as pd
-
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class MessageParam(ComponentParamBase):

diff --git a/graph/component/pubmed.py → agent/component/pubmed.py b/graph/component/pubmed.py → agent/component/pubmed.py
@@ -13,14 +13,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import random
 from abc import ABC
-from functools import partial
 from Bio import Entrez
 import pandas as pd
 import xml.etree.ElementTree as ET
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class PubMedParam(ComponentParamBase):

diff --git a/graph/component/relevant.py → agent/component/relevant.py b/graph/component/relevant.py → agent/component/relevant.py
@@ -16,7 +16,7 @@
 from abc import ABC
 from api.db import LLMType
 from api.db.services.llm_service import LLMBundle
-from graph.component import GenerateParam, Generate
+from agent.component import GenerateParam, Generate
 from rag.utils import num_tokens_from_string, encoder
 
 

diff --git a/graph/component/retrieval.py → agent/component/retrieval.py b/graph/component/retrieval.py → agent/component/retrieval.py
@@ -21,7 +21,7 @@
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.llm_service import LLMBundle
 from api.settings import retrievaler
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class RetrievalParam(ComponentParamBase):

diff --git a/graph/component/rewrite.py → agent/component/rewrite.py b/graph/component/rewrite.py → agent/component/rewrite.py
@@ -16,7 +16,7 @@
 from abc import ABC
 from api.db import LLMType
 from api.db.services.llm_service import LLMBundle
-from graph.component import GenerateParam, Generate
+from agent.component import GenerateParam, Generate
 
 
 class RewriteQuestionParam(GenerateParam):

diff --git a/graph/component/switch.py → agent/component/switch.py b/graph/component/switch.py → agent/component/switch.py
@@ -16,12 +16,7 @@
 from abc import ABC
 
 import pandas as pd
-
-from api.db import LLMType
-from api.db.services.knowledgebase_service import KnowledgebaseService
-from api.db.services.llm_service import LLMBundle
-from api.settings import retrievaler
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class SwitchParam(ComponentParamBase):

diff --git a/graph/component/wikipedia.py → agent/component/wikipedia.py b/graph/component/wikipedia.py → agent/component/wikipedia.py
@@ -18,8 +18,8 @@
 from functools import partial
 import wikipedia
 import pandas as pd
-from graph.settings import DEBUG
-from graph.component.base import ComponentBase, ComponentParamBase
+from agent.settings import DEBUG
+from agent.component.base import ComponentBase, ComponentParamBase
 
 
 class WikipediaParam(ComponentParamBase):

diff --git a/graph/settings.py → agent/settings.py b/graph/settings.py → agent/settings.py
diff --git a/graph/templates/HR_callout_zh.json → agent/templates/HR_callout_zh.json b/graph/templates/HR_callout_zh.json → agent/templates/HR_callout_zh.json
diff --git a/graph/templates/customer_service.json → agent/templates/customer_service.json b/graph/templates/customer_service.json → agent/templates/customer_service.json
diff --git a/graph/templates/general_chat_bot.json → agent/templates/general_chat_bot.json b/graph/templates/general_chat_bot.json → agent/templates/general_chat_bot.json
diff --git a/graph/templates/interpreter.json → agent/templates/interpreter.json b/graph/templates/interpreter.json → agent/templates/interpreter.json
diff --git a/graph/templates/websearch_assistant.json → agent/templates/websearch_assistant.json b/graph/templates/websearch_assistant.json → agent/templates/websearch_assistant.json
diff --git a/graph/test/client.py → agent/test/client.py b/graph/test/client.py → agent/test/client.py
@@ -16,9 +16,8 @@
 import argparse
 import os
 from functools import partial
-import readline
-from graph.canvas import Canvas
-from graph.settings import DEBUG
+from agent.canvas import Canvas
+from agent.settings import DEBUG
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()

diff --git a/graph/test/dsl_examples/categorize.json → agent/test/dsl_examples/categorize.json b/graph/test/dsl_examples/categorize.json → agent/test/dsl_examples/categorize.json
diff --git a/...h/test/dsl_examples/customer_service.json → ...t/test/dsl_examples/customer_service.json b/...h/test/dsl_examples/customer_service.json → ...t/test/dsl_examples/customer_service.json
diff --git a/graph/test/dsl_examples/headhunter_zh.json → agent/test/dsl_examples/headhunter_zh.json b/graph/test/dsl_examples/headhunter_zh.json → agent/test/dsl_examples/headhunter_zh.json
diff --git a/graph/test/dsl_examples/intergreper.json → agent/test/dsl_examples/intergreper.json b/graph/test/dsl_examples/intergreper.json → agent/test/dsl_examples/intergreper.json
diff --git a/graph/test/dsl_examples/interpreter.json → agent/test/dsl_examples/interpreter.json b/graph/test/dsl_examples/interpreter.json → agent/test/dsl_examples/interpreter.json
diff --git a/...mples/keyword_wikipedia_and_generate.json → ...mples/keyword_wikipedia_and_generate.json b/...mples/keyword_wikipedia_and_generate.json → ...mples/keyword_wikipedia_and_generate.json
diff --git a/.../dsl_examples/retrieval_and_generate.json → .../dsl_examples/retrieval_and_generate.json b/.../dsl_examples/retrieval_and_generate.json → .../dsl_examples/retrieval_and_generate.json
diff --git a/...es/retrieval_categorize_and_generate.json → ...es/retrieval_categorize_and_generate.json b/...es/retrieval_categorize_and_generate.json → ...es/retrieval_categorize_and_generate.json
diff --git a/...ples/retrieval_relevant_and_generate.json → ...ples/retrieval_relevant_and_generate.json b/...ples/retrieval_relevant_and_generate.json → ...ples/retrieval_relevant_and_generate.json
diff --git a/..._relevant_keyword_baidu_and_generate.json → ..._relevant_keyword_baidu_and_generate.json b/..._relevant_keyword_baidu_and_generate.json → ..._relevant_keyword_baidu_and_generate.json
diff --git a/...rieval_relevant_rewrite_and_generate.json → ...rieval_relevant_rewrite_and_generate.json b/...rieval_relevant_rewrite_and_generate.json → ...rieval_relevant_rewrite_and_generate.json
diff --git a/api/apps/api_app.py b/api/apps/api_app.py
@@ -20,7 +20,7 @@
 from flask import request, Response
 from flask_login import login_required, current_user
 
-from api.db import FileType, ParserType, FileSource, LLMType
+from api.db import FileType, ParserType, FileSource
 from api.db.db_models import APIToken, API4Conversation, Task, File
 from api.db.services import duplicate_name
 from api.db.services.api_service import APITokenService, API4ConversationService
@@ -29,7 +29,6 @@
 from api.db.services.file2document_service import File2DocumentService
 from api.db.services.file_service import FileService
 from api.db.services.knowledgebase_service import KnowledgebaseService
-from api.db.services.llm_service import TenantLLMService
 from api.db.services.task_service import queue_tasks, TaskService
 from api.db.services.user_service import UserTenantService
 from api.settings import RetCode, retrievaler
@@ -38,7 +37,6 @@
 from itsdangerous import URLSafeTimedSerializer
 
 from api.utils.file_utils import filename_type, thumbnail
-from rag.nlp import keyword_extraction
 from rag.utils.minio_conn import MINIO
 
 

diff --git a/api/apps/canvas_app.py b/api/apps/canvas_app.py
@@ -15,15 +15,12 @@
 #
 import json
 from functools import partial
-
 from flask import request, Response
 from flask_login import login_required, current_user
-
-from api.db.db_models import UserCanvas
 from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
 from api.utils import get_uuid
 from api.utils.api_utils import get_json_result, server_error_response, validate_request
-from graph.canvas import Canvas
+from agent.canvas import Canvas
 
 
 @manager.route('/templates', methods=['GET'])