From fb241a467c9b9d886c40e1f4000fc63d56ab29c9 Mon Sep 17 00:00:00 2001 From: Tadashi Date: Thu, 31 Oct 2024 09:22:55 +0700 Subject: [PATCH] fix: nano-graphrag import when not installed --- .../ktem/index/file/graph/nano_pipelines.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/libs/ktem/ktem/index/file/graph/nano_pipelines.py b/libs/ktem/ktem/index/file/graph/nano_pipelines.py index de3bce4..b3f6b95 100644 --- a/libs/ktem/ktem/index/file/graph/nano_pipelines.py +++ b/libs/ktem/ktem/index/file/graph/nano_pipelines.py @@ -29,13 +29,6 @@ try: _find_most_related_text_unit_from_entities, ) from nano_graphrag._utils import EmbeddingFunc - from nano_graphrag.base import ( - BaseGraphStorage, - BaseKVStorage, - BaseVectorStorage, - CommunitySchema, - TextChunkSchema, - ) except ImportError: print( @@ -124,12 +117,12 @@ def clean_quote(input: str) -> str: async def nano_graph_rag_build_local_query_context( graph_func, query, - query_param: QueryParam, + query_param, ): - knowledge_graph_inst: BaseGraphStorage = graph_func.chunk_entity_relation_graph - entities_vdb: BaseVectorStorage = graph_func.entities_vdb - community_reports: BaseKVStorage[CommunitySchema] = graph_func.community_reports - text_chunks_db: BaseKVStorage[TextChunkSchema] = graph_func.text_chunks + knowledge_graph_inst = graph_func.chunk_entity_relation_graph + entities_vdb = graph_func.entities_vdb + community_reports = graph_func.community_reports + text_chunks_db = graph_func.text_chunks results = await entities_vdb.query(query, top_k=query_param.top_k) if not len(results):