diff --git a/PathRAG/base.py b/PathRAG/base.py index ba6c988..6cd9ec1 100644 --- a/PathRAG/base.py +++ b/PathRAG/base.py @@ -15,7 +15,7 @@ @dataclass class QueryParam: - mode: Literal["hybrid"] = "global" + mode: Literal["hybrid"] = "hybrid" only_need_context: bool = False only_need_prompt: bool = False response_type: str = "Multiple Paragraphs" diff --git a/PathRAG/llm.py b/PathRAG/llm.py index b814e65..3bfa190 100644 --- a/PathRAG/llm.py +++ b/PathRAG/llm.py @@ -45,6 +45,11 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false" +class BedrockError(Exception): + """Exception raised for AWS Bedrock API errors.""" + pass + + @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10), @@ -62,7 +67,6 @@ async def openai_complete_if_cache( ) -> str: if api_key: os.environ["OPENAI_API_KEY"] = api_key - time.sleep(2) openai_async_client = ( AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url) ) diff --git a/PathRAG/operate.py b/PathRAG/operate.py index ab35567..366e4d2 100644 --- a/PathRAG/operate.py +++ b/PathRAG/operate.py @@ -255,7 +255,6 @@ async def extract_entities( relationships_vdb: BaseVectorStorage, global_config: dict, ) -> Union[BaseGraphStorage, None]: - time.sleep(20) use_llm_func: callable = global_config["llm_model_func"] entity_extract_max_gleaning = global_config["entity_extract_max_gleaning"] @@ -1179,13 +1178,13 @@ async def _find_most_related_edges_from_entities3( edge1 = await knowledge_graph_inst.get_edge(path[1],path[2]) or await knowledge_graph_inst.get_edge(path[2], path[1]) edge2 = await knowledge_graph_inst.get_edge(path[2],path[3]) or await knowledge_graph_inst.get_edge(path[3], path[2]) if edge0==None or edge1==None or edge2==None: - print(path,"边丢失") + print(path, "edge missing") if edge0==None: - print("edge0丢失") + print("edge0 missing") if edge1==None: - print("edge1丢失") + print("edge1 missing") if edge2==None: - print("edge2丢失") + print("edge2 missing") continue e1 = "through edge ("+edge0["keywords"]+") to connect to "+s_name+" and "+b1_name+"." e2 = "through edge ("+edge1["keywords"]+") to connect to "+b1_name+" and "+b2_name+"." @@ -1204,7 +1203,7 @@ async def _find_most_related_edges_from_entities3( edge0 = await knowledge_graph_inst.get_edge(path[0], path[1]) or await knowledge_graph_inst.get_edge(path[1], path[0]) edge1 = await knowledge_graph_inst.get_edge(path[1],path[2]) or await knowledge_graph_inst.get_edge(path[2], path[1]) if edge0==None or edge1==None: - print(path,"边丢失") + print(path, "edge missing") continue e1 = "through edge("+edge0["keywords"]+") to connect to "+s_name+" and "+b_name+"." e2 = "through edge("+edge1["keywords"]+") to connect to "+b_name+" and "+t_name+"." @@ -1219,7 +1218,7 @@ async def _find_most_related_edges_from_entities3( s_name,t_name = path[0],path[1] edge0 = await knowledge_graph_inst.get_edge(path[0], path[1]) or await knowledge_graph_inst.get_edge(path[1], path[0]) if edge0==None: - print(path,"边丢失") + print(path, "edge missing") continue e = "through edge("+edge0["keywords"]+") to connect to "+s_name+" and "+t_name+"." s = await knowledge_graph_inst.get_node(s_name) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..632c02d --- /dev/null +++ b/setup.py @@ -0,0 +1,28 @@ +from setuptools import setup, find_packages + +setup( + name="PathRAG", + version="0.1.0", + description="PathRAG: Pruning Graph-based Retrieval Augmented Generation with Relational Paths", + author="PathRAG Authors", + packages=find_packages(), + python_requires=">=3.9", + install_requires=[ + "accelerate", + "aioboto3", + "aiohttp", + "graspologic", + "hnswlib", + "nano-vectordb", + "networkx", + "numpy", + "ollama", + "openai", + "pydantic", + "tenacity", + "tiktoken", + "torch", + "tqdm", + "transformers", + ], +) diff --git a/v1_test.py b/v1_test.py index 7489ce3..396365b 100644 --- a/v1_test.py +++ b/v1_test.py @@ -2,9 +2,9 @@ from PathRAG import PathRAG, QueryParam from PathRAG.llm import gpt_4o_mini_complete -WORKING_DIR = "" +WORKING_DIR = "./pathrag_cache" # Set your working directory -api_key="" +api_key="your-openai-api-key-here" # Set your OpenAI API key os.environ["OPENAI_API_KEY"] = api_key base_url="https://api.openai.com/v1" os.environ["OPENAI_API_BASE"]=base_url @@ -18,8 +18,8 @@ llm_model_func=gpt_4o_mini_complete, ) -data_file="" -question="" +data_file="./text.txt" # Path to your input document +question="What is the main topic of this document?" # Your query with open(data_file) as f: rag.insert(f.read())