From bc79f6650e57236b572e81e022d3fdbfadb51552 Mon Sep 17 00:00:00 2001 From: Gurjot Singh Date: Tue, 14 Jan 2025 22:23:14 +0530 Subject: [PATCH] Fix linting errors --- lightrag/lightrag.py | 40 ++++++++++++++++++++-------------------- lightrag/operate.py | 23 ++++++++++++++++++----- test.py | 2 +- 3 files changed, 39 insertions(+), 26 deletions(-) diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index e8859071..cacdfc50 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -756,10 +756,7 @@ class LightRAG: return response def query_with_separate_keyword_extraction( - self, - query: str, - prompt: str, - param: QueryParam = QueryParam() + self, query: str, prompt: str, param: QueryParam = QueryParam() ): """ 1. Extract keywords from the 'query' using new function in operate.py. @@ -767,13 +764,12 @@ class LightRAG: """ loop = always_get_an_event_loop() - return loop.run_until_complete(self.aquery_with_separate_keyword_extraction(query, prompt, param)) - + return loop.run_until_complete( + self.aquery_with_separate_keyword_extraction(query, prompt, param) + ) + async def aquery_with_separate_keyword_extraction( - self, - query: str, - prompt: str, - param: QueryParam = QueryParam() + self, query: str, prompt: str, param: QueryParam = QueryParam() ): """ 1. Calls extract_keywords_only to get HL/LL keywords from 'query'. @@ -788,20 +784,21 @@ class LightRAG: text=query, param=param, global_config=asdict(self), - hashing_kv=self.llm_response_cache or self.key_string_value_json_storage_cls( + hashing_kv=self.llm_response_cache + or self.key_string_value_json_storage_cls( namespace="llm_response_cache", global_config=asdict(self), embedding_func=None, - ) + ), ) - - param.hl_keywords=hl_keywords, - param.ll_keywords=ll_keywords, - + + param.hl_keywords = (hl_keywords,) + param.ll_keywords = (ll_keywords,) + # --------------------- # STEP 2: Final Query Logic # --------------------- - + # Create a new string with the prompt and the keywords ll_keywords_str = ", ".join(ll_keywords) hl_keywords_str = ", ".join(hl_keywords) @@ -817,7 +814,8 @@ class LightRAG: param, asdict(self), hashing_kv=self.llm_response_cache - if self.llm_response_cache and hasattr(self.llm_response_cache, "global_config") + if self.llm_response_cache + and hasattr(self.llm_response_cache, "global_config") else self.key_string_value_json_storage_cls( namespace="llm_response_cache", global_config=asdict(self), @@ -832,7 +830,8 @@ class LightRAG: param, asdict(self), hashing_kv=self.llm_response_cache - if self.llm_response_cache and hasattr(self.llm_response_cache, "global_config") + if self.llm_response_cache + and hasattr(self.llm_response_cache, "global_config") else self.key_string_value_json_storage_cls( namespace="llm_response_cache", global_config=asdict(self), @@ -850,7 +849,8 @@ class LightRAG: param, asdict(self), hashing_kv=self.llm_response_cache - if self.llm_response_cache and hasattr(self.llm_response_cache, "global_config") + if self.llm_response_cache + and hasattr(self.llm_response_cache, "global_config") else self.key_string_value_json_storage_cls( namespace="llm_response_cache", global_config=asdict(self), diff --git a/lightrag/operate.py b/lightrag/operate.py index f4993873..7df489b3 100644 --- a/lightrag/operate.py +++ b/lightrag/operate.py @@ -680,6 +680,7 @@ async def kg_query( ) return response + async def kg_query_with_keywords( query: str, knowledge_graph_inst: BaseGraphStorage, @@ -717,8 +718,10 @@ async def kg_query_with_keywords( # If neither has any keywords, you could handle that logic here. if not hl_keywords and not ll_keywords: - logger.warning("No keywords found in query_param. Could default to global mode or fail.") - return PROMPTS["fail_response"] + logger.warning( + "No keywords found in query_param. Could default to global mode or fail." + ) + return PROMPTS["fail_response"] if not ll_keywords and query_param.mode in ["local", "hybrid"]: logger.warning("low_level_keywords is empty, switching to global mode.") query_param.mode = "global" @@ -727,8 +730,16 @@ async def kg_query_with_keywords( query_param.mode = "local" # Flatten low-level and high-level keywords if needed - ll_keywords_flat = [item for sublist in ll_keywords for item in sublist] if any(isinstance(i, list) for i in ll_keywords) else ll_keywords - hl_keywords_flat = [item for sublist in hl_keywords for item in sublist] if any(isinstance(i, list) for i in hl_keywords) else hl_keywords + ll_keywords_flat = ( + [item for sublist in ll_keywords for item in sublist] + if any(isinstance(i, list) for i in ll_keywords) + else ll_keywords + ) + hl_keywords_flat = ( + [item for sublist in hl_keywords for item in sublist] + if any(isinstance(i, list) for i in hl_keywords) + else hl_keywords + ) # Join the flattened lists ll_keywords_str = ", ".join(ll_keywords_flat) if ll_keywords_flat else "" @@ -766,7 +777,7 @@ async def kg_query_with_keywords( if query_param.only_need_prompt: return sys_prompt - + # Now call the LLM with the final system prompt response = await use_model_func( query, @@ -803,6 +814,7 @@ async def kg_query_with_keywords( ) return response + async def extract_keywords_only( text: str, param: QueryParam, @@ -881,6 +893,7 @@ async def extract_keywords_only( ) return hl_keywords, ll_keywords + async def _build_query_context( query: list, knowledge_graph_inst: BaseGraphStorage, diff --git a/test.py b/test.py index 895f0b30..80bcaa6d 100644 --- a/test.py +++ b/test.py @@ -39,4 +39,4 @@ print( # Perform hybrid search print( rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")) -) \ No newline at end of file +)