Add files via upload

PathRAG/PathRAG.py

@@ -0,0 +1,562 @@
+import asyncio
+import os
+from tqdm.asyncio import tqdm as tqdm_async
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from functools import partial
+from typing import Type, cast
+
+
+from .llm import (
+    gpt_4o_mini_complete,
+    openai_embedding,
+)
+from .operate import (
+    chunking_by_token_size,
+    extract_entities,
+    kg_query,
+)
+
+from .utils import (
+    EmbeddingFunc,
+    compute_mdhash_id,
+    limit_async_func_call,
+    convert_response_to_json,
+    logger,
+    set_logger,
+)
+from .base import (
+    BaseGraphStorage,
+    BaseKVStorage,
+    BaseVectorStorage,
+    StorageNameSpace,
+    QueryParam,
+)
+
+from .storage import (
+    JsonKVStorage,
+    NanoVectorDBStorage,
+    NetworkXStorage,
+)
+
+
+
+
+def lazy_external_import(module_name: str, class_name: str):
+    """Lazily import a class from an external module based on the package of the caller."""
+
+
+    import inspect
+
+    caller_frame = inspect.currentframe().f_back
+    module = inspect.getmodule(caller_frame)
+    package = module.__package__ if module else None
+
+    def import_class(*args, **kwargs):
+        import importlib
+
+  
+        module = importlib.import_module(module_name, package=package)
+
+
+        cls = getattr(module, class_name)
+        return cls(*args, **kwargs)
+
+    return import_class
+
+
+Neo4JStorage = lazy_external_import(".kg.neo4j_impl", "Neo4JStorage")
+OracleKVStorage = lazy_external_import(".kg.oracle_impl", "OracleKVStorage")
+OracleGraphStorage = lazy_external_import(".kg.oracle_impl", "OracleGraphStorage")
+OracleVectorDBStorage = lazy_external_import(".kg.oracle_impl", "OracleVectorDBStorage")
+MilvusVectorDBStorge = lazy_external_import(".kg.milvus_impl", "MilvusVectorDBStorge")
+MongoKVStorage = lazy_external_import(".kg.mongo_impl", "MongoKVStorage")
+ChromaVectorDBStorage = lazy_external_import(".kg.chroma_impl", "ChromaVectorDBStorage")
+TiDBKVStorage = lazy_external_import(".kg.tidb_impl", "TiDBKVStorage")
+TiDBVectorDBStorage = lazy_external_import(".kg.tidb_impl", "TiDBVectorDBStorage")
+AGEStorage = lazy_external_import(".kg.age_impl", "AGEStorage")
+
+
+def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
+    """
+    Ensure that there is always an event loop available.
+
+    This function tries to get the current event loop. If the current event loop is closed or does not exist,
+    it creates a new event loop and sets it as the current event loop.
+
+    Returns:
+        asyncio.AbstractEventLoop: The current or newly created event loop.
+    """
+    try:
+
+        current_loop = asyncio.get_event_loop()
+        if current_loop.is_closed():
+            raise RuntimeError("Event loop is closed.")
+        return current_loop
+
+    except RuntimeError:
+
+        logger.info("Creating a new event loop in main thread.")
+        new_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(new_loop)
+        return new_loop
+
+
+@dataclass
+class PathRAG:
+    working_dir: str = field(
+        default_factory=lambda: f"./PathRAG_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
+    )
+
+    embedding_cache_config: dict = field(
+        default_factory=lambda: {
+            "enabled": False,
+            "similarity_threshold": 0.95,
+            "use_llm_check": False,
+        }
+    )
+    kv_storage: str = field(default="JsonKVStorage")
+    vector_storage: str = field(default="NanoVectorDBStorage")
+    graph_storage: str = field(default="NetworkXStorage")
+
+    current_log_level = logger.level
+    log_level: str = field(default=current_log_level)
+
+
+    chunk_token_size: int = 1200
+    chunk_overlap_token_size: int = 100
+    tiktoken_model_name: str = "gpt-4o-mini"
+
+
+    entity_extract_max_gleaning: int = 1
+    entity_summary_to_max_tokens: int = 500
+
+
+    node_embedding_algorithm: str = "node2vec"
+    node2vec_params: dict = field(
+        default_factory=lambda: {
+            "dimensions": 1536,
+            "num_walks": 10,
+            "walk_length": 40,
+            "window_size": 2,
+            "iterations": 3,
+            "random_seed": 3,
+        }
+    )
+
+
+    embedding_func: EmbeddingFunc = field(default_factory=lambda: openai_embedding)
+    embedding_batch_num: int = 32
+    embedding_func_max_async: int = 16
+
+
+    llm_model_func: callable = gpt_4o_mini_complete  
+    llm_model_name: str = "meta-llama/Llama-3.2-1B-Instruct"  
+    llm_model_max_token_size: int = 32768
+    llm_model_max_async: int = 16
+    llm_model_kwargs: dict = field(default_factory=dict)
+
+
+    vector_db_storage_cls_kwargs: dict = field(default_factory=dict)
+
+    enable_llm_cache: bool = True
+
+
+    addon_params: dict = field(default_factory=dict)
+    convert_response_to_json_func: callable = convert_response_to_json
+
+    def __post_init__(self):
+        log_file = os.path.join("PathRAG.log")
+        set_logger(log_file)
+        logger.setLevel(self.log_level)
+
+        logger.info(f"Logger initialized for working directory: {self.working_dir}")
+
+
+        self.key_string_value_json_storage_cls: Type[BaseKVStorage] = (
+            self._get_storage_class()[self.kv_storage]
+        )
+        self.vector_db_storage_cls: Type[BaseVectorStorage] = self._get_storage_class()[
+            self.vector_storage
+        ]
+        self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[
+            self.graph_storage
+        ]
+
+        if not os.path.exists(self.working_dir):
+            logger.info(f"Creating working directory {self.working_dir}")
+            os.makedirs(self.working_dir)
+
+        self.llm_response_cache = (
+            self.key_string_value_json_storage_cls(
+                namespace="llm_response_cache",
+                global_config=asdict(self),
+                embedding_func=None,
+            )
+            if self.enable_llm_cache
+            else None
+        )
+        self.embedding_func = limit_async_func_call(self.embedding_func_max_async)(
+            self.embedding_func
+        )
+
+
+        self.full_docs = self.key_string_value_json_storage_cls(
+            namespace="full_docs",
+            global_config=asdict(self),
+            embedding_func=self.embedding_func,
+        )
+        self.text_chunks = self.key_string_value_json_storage_cls(
+            namespace="text_chunks",
+            global_config=asdict(self),
+            embedding_func=self.embedding_func,
+        )
+        self.chunk_entity_relation_graph = self.graph_storage_cls(
+            namespace="chunk_entity_relation",
+            global_config=asdict(self),
+            embedding_func=self.embedding_func,
+        )
+
+
+        self.entities_vdb = self.vector_db_storage_cls(
+            namespace="entities",
+            global_config=asdict(self),
+            embedding_func=self.embedding_func,
+            meta_fields={"entity_name"},
+        )
+        self.relationships_vdb = self.vector_db_storage_cls(
+            namespace="relationships",
+            global_config=asdict(self),
+            embedding_func=self.embedding_func,
+            meta_fields={"src_id", "tgt_id"},
+        )
+        self.chunks_vdb = self.vector_db_storage_cls(
+            namespace="chunks",
+            global_config=asdict(self),
+            embedding_func=self.embedding_func,
+        )
+
+        self.llm_model_func = limit_async_func_call(self.llm_model_max_async)(
+            partial(
+                self.llm_model_func,
+                hashing_kv=self.llm_response_cache
+                if self.llm_response_cache
+                and hasattr(self.llm_response_cache, "global_config")
+                else self.key_string_value_json_storage_cls(
+                    global_config=asdict(self),
+                ),
+                **self.llm_model_kwargs,
+            )
+        )
+
+    def _get_storage_class(self) -> Type[BaseGraphStorage]:
+        return {
+
+            "JsonKVStorage": JsonKVStorage,
+            "OracleKVStorage": OracleKVStorage,
+            "MongoKVStorage": MongoKVStorage,
+            "TiDBKVStorage": TiDBKVStorage,
+
+            "NanoVectorDBStorage": NanoVectorDBStorage,
+            "OracleVectorDBStorage": OracleVectorDBStorage,
+            "MilvusVectorDBStorge": MilvusVectorDBStorge,
+            "ChromaVectorDBStorage": ChromaVectorDBStorage,
+            "TiDBVectorDBStorage": TiDBVectorDBStorage,
+
+            "NetworkXStorage": NetworkXStorage,
+            "Neo4JStorage": Neo4JStorage,
+            "OracleGraphStorage": OracleGraphStorage,
+            "AGEStorage": AGEStorage,
+
+        }
+
+    def insert(self, string_or_strings):
+        
+        loop = always_get_an_event_loop()
+        return loop.run_until_complete(self.ainsert(string_or_strings))
+
+    async def ainsert(self, string_or_strings):
+        update_storage = False
+        try:
+            if isinstance(string_or_strings, str):
+                string_or_strings = [string_or_strings]
+
+            new_docs = {
+                compute_mdhash_id(c.strip(), prefix="doc-"): {"content": c.strip()}
+                for c in string_or_strings
+            }
+            _add_doc_keys = await self.full_docs.filter_keys(list(new_docs.keys()))
+            new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
+            if not len(new_docs):
+                logger.warning("All docs are already in the storage")
+                return
+            update_storage = True
+            logger.info(f"[New Docs] inserting {len(new_docs)} docs")
+
+            inserting_chunks = {}
+            for doc_key, doc in tqdm_async(
+                new_docs.items(), desc="Chunking documents", unit="doc"
+            ):
+                chunks = {
+                    compute_mdhash_id(dp["content"], prefix="chunk-"): {
+                        **dp,
+                        "full_doc_id": doc_key,
+                    }
+                    for dp in chunking_by_token_size(
+                        doc["content"],
+                        overlap_token_size=self.chunk_overlap_token_size,
+                        max_token_size=self.chunk_token_size,
+                        tiktoken_model=self.tiktoken_model_name,
+                    )
+                }
+                inserting_chunks.update(chunks)
+            _add_chunk_keys = await self.text_chunks.filter_keys(
+                list(inserting_chunks.keys())
+            )
+            inserting_chunks = {
+                k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
+            }
+            if not len(inserting_chunks):
+                logger.warning("All chunks are already in the storage")
+                return
+            logger.info(f"[New Chunks] inserting {len(inserting_chunks)} chunks")
+
+            await self.chunks_vdb.upsert(inserting_chunks)
+
+            logger.info("[Entity Extraction]...")
+            maybe_new_kg = await extract_entities(
+                inserting_chunks,
+                knowledge_graph_inst=self.chunk_entity_relation_graph,
+                entity_vdb=self.entities_vdb,
+                relationships_vdb=self.relationships_vdb,
+                global_config=asdict(self),
+            )
+            if maybe_new_kg is None:
+                logger.warning("No new entities and relationships found")
+                return
+            self.chunk_entity_relation_graph = maybe_new_kg
+
+            await self.full_docs.upsert(new_docs)
+            await self.text_chunks.upsert(inserting_chunks)
+        finally:
+            if update_storage:
+                await self._insert_done()
+
+    async def _insert_done(self):
+        tasks = []
+        for storage_inst in [
+            self.full_docs,
+            self.text_chunks,
+            self.llm_response_cache,
+            self.entities_vdb,
+            self.relationships_vdb,
+            self.chunks_vdb,
+            self.chunk_entity_relation_graph,
+        ]:
+            if storage_inst is None:
+                continue
+            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
+        await asyncio.gather(*tasks)
+
+    def insert_custom_kg(self, custom_kg: dict):
+        loop = always_get_an_event_loop()
+        return loop.run_until_complete(self.ainsert_custom_kg(custom_kg))
+
+    async def ainsert_custom_kg(self, custom_kg: dict):
+        update_storage = False
+        try:
+
+            all_chunks_data = {}
+            chunk_to_source_map = {}
+            for chunk_data in custom_kg.get("chunks", []):
+                chunk_content = chunk_data["content"]
+                source_id = chunk_data["source_id"]
+                chunk_id = compute_mdhash_id(chunk_content.strip(), prefix="chunk-")
+
+                chunk_entry = {"content": chunk_content.strip(), "source_id": source_id}
+                all_chunks_data[chunk_id] = chunk_entry
+                chunk_to_source_map[source_id] = chunk_id
+                update_storage = True
+
+            if self.chunks_vdb is not None and all_chunks_data:
+                await self.chunks_vdb.upsert(all_chunks_data)
+            if self.text_chunks is not None and all_chunks_data:
+                await self.text_chunks.upsert(all_chunks_data)
+
+ 
+            all_entities_data = []
+            for entity_data in custom_kg.get("entities", []):
+                entity_name = f'"{entity_data["entity_name"].upper()}"'
+                entity_type = entity_data.get("entity_type", "UNKNOWN")
+                description = entity_data.get("description", "No description provided")
+
+                source_chunk_id = entity_data.get("source_id", "UNKNOWN")
+                source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
+
+
+                if source_id == "UNKNOWN":
+                    logger.warning(
+                        f"Entity '{entity_name}' has an UNKNOWN source_id. Please check the source mapping."
+                    )
+
+
+                node_data = {
+                    "entity_type": entity_type,
+                    "description": description,
+                    "source_id": source_id,
+                }
+
+                await self.chunk_entity_relation_graph.upsert_node(
+                    entity_name, node_data=node_data
+                )
+                node_data["entity_name"] = entity_name
+                all_entities_data.append(node_data)
+                update_storage = True
+
+
+            all_relationships_data = []
+            for relationship_data in custom_kg.get("relationships", []):
+                src_id = f'"{relationship_data["src_id"].upper()}"'
+                tgt_id = f'"{relationship_data["tgt_id"].upper()}"'
+                description = relationship_data["description"]
+                keywords = relationship_data["keywords"]
+                weight = relationship_data.get("weight", 1.0)
+
+                source_chunk_id = relationship_data.get("source_id", "UNKNOWN")
+                source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
+
+
+                if source_id == "UNKNOWN":
+                    logger.warning(
+                        f"Relationship from '{src_id}' to '{tgt_id}' has an UNKNOWN source_id. Please check the source mapping."
+                    )
+
+
+                for need_insert_id in [src_id, tgt_id]:
+                    if not (
+                        await self.chunk_entity_relation_graph.has_node(need_insert_id)
+                    ):
+                        await self.chunk_entity_relation_graph.upsert_node(
+                            need_insert_id,
+                            node_data={
+                                "source_id": source_id,
+                                "description": "UNKNOWN",
+                                "entity_type": "UNKNOWN",
+                            },
+                        )
+
+
+                await self.chunk_entity_relation_graph.upsert_edge(
+                    src_id,
+                    tgt_id,
+                    edge_data={
+                        "weight": weight,
+                        "description": description,
+                        "keywords": keywords,
+                        "source_id": source_id,
+                    },
+                )
+                edge_data = {
+                    "src_id": src_id,
+                    "tgt_id": tgt_id,
+                    "description": description,
+                    "keywords": keywords,
+                }
+                all_relationships_data.append(edge_data)
+                update_storage = True
+
+
+            if self.entities_vdb is not None:
+                data_for_vdb = {
+                    compute_mdhash_id(dp["entity_name"], prefix="ent-"): {
+                        "content": dp["entity_name"] + dp["description"],
+                        "entity_name": dp["entity_name"],
+                    }
+                    for dp in all_entities_data
+                }
+                await self.entities_vdb.upsert(data_for_vdb)
+
+
+            if self.relationships_vdb is not None:
+                data_for_vdb = {
+                    compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): {
+                        "src_id": dp["src_id"],
+                        "tgt_id": dp["tgt_id"],
+                        "content": dp["keywords"]
+                        + dp["src_id"]
+                        + dp["tgt_id"]
+                        + dp["description"],
+                    }
+                    for dp in all_relationships_data
+                }
+                await self.relationships_vdb.upsert(data_for_vdb)
+        finally:
+            if update_storage:
+                await self._insert_done()
+    
+    def query(self, query: str, param: QueryParam = QueryParam()):
+        loop = always_get_an_event_loop()
+        return loop.run_until_complete(self.aquery(query, param))
+    
+    async def aquery(self, query: str, param: QueryParam = QueryParam()):
+        if param.mode in ["hybrid"]:
+            response= await kg_query(
+                query,
+                self.chunk_entity_relation_graph,
+                self.entities_vdb,
+                self.relationships_vdb,
+                self.text_chunks,
+                param,
+                asdict(self),
+                hashing_kv=self.llm_response_cache
+                if self.llm_response_cache
+                and hasattr(self.llm_response_cache, "global_config")
+                else self.key_string_value_json_storage_cls(
+                    global_config=asdict(self),
+                ),
+            )
+            print("response all ready")
+        else:
+            raise ValueError(f"Unknown mode {param.mode}")
+        await self._query_done()
+        return response
+
+        
+    async def _query_done(self):
+        tasks = []
+        for storage_inst in [self.llm_response_cache]:
+            if storage_inst is None:
+                continue
+            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
+        await asyncio.gather(*tasks)
+
+    def delete_by_entity(self, entity_name: str):
+        loop = always_get_an_event_loop()
+        return loop.run_until_complete(self.adelete_by_entity(entity_name))
+
+    async def adelete_by_entity(self, entity_name: str):
+        entity_name = f'"{entity_name.upper()}"'
+
+        try:
+            await self.entities_vdb.delete_entity(entity_name)
+            await self.relationships_vdb.delete_relation(entity_name)
+            await self.chunk_entity_relation_graph.delete_node(entity_name)
+
+            logger.info(
+                f"Entity '{entity_name}' and its relationships have been deleted."
+            )
+            await self._delete_by_entity_done()
+        except Exception as e:
+            logger.error(f"Error while deleting entity '{entity_name}': {e}")
+
+    async def _delete_by_entity_done(self):
+        tasks = []
+        for storage_inst in [
+            self.entities_vdb,
+            self.relationships_vdb,
+            self.chunk_entity_relation_graph,
+        ]:
+            if storage_inst is None:
+                continue
+            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
+        await asyncio.gather(*tasks)

PathRAG/__init__.py

@@ -0,0 +1,3 @@
+from .PathRAG import PathRAG as PathRAG, QueryParam as QueryParam
+
+

PathRAG/base.py

@@ -0,0 +1,135 @@
+from dataclasses import dataclass, field
+from typing import TypedDict, Union, Literal, Generic, TypeVar
+
+import numpy as np
+
+from .utils import EmbeddingFunc
+
+TextChunkSchema = TypedDict(
+    "TextChunkSchema",
+    {"tokens": int, "content": str, "full_doc_id": str, "chunk_order_index": int},
+)
+
+T = TypeVar("T")
+
+
+@dataclass
+class QueryParam:
+    mode: Literal["hybrid"] = "global"
+    only_need_context: bool = False
+    only_need_prompt: bool = False
+    response_type: str = "Multiple Paragraphs"
+    stream: bool = False
+    top_k: int =40
+    max_token_for_text_unit: int = 4000
+    max_token_for_global_context: int = 3000
+    max_token_for_local_context: int = 5000
+
+
+@dataclass
+class StorageNameSpace:
+    namespace: str
+    global_config: dict
+
+    async def index_done_callback(self):
+       
+        pass
+
+    async def query_done_callback(self):
+        
+        pass
+
+
+@dataclass
+class BaseVectorStorage(StorageNameSpace):
+    embedding_func: EmbeddingFunc
+    meta_fields: set = field(default_factory=set)
+
+    async def query(self, query: str, top_k: int) -> list[dict]:
+        raise NotImplementedError
+
+    async def upsert(self, data: dict[str, dict]):
+
+        raise NotImplementedError
+
+
+@dataclass
+class BaseKVStorage(Generic[T], StorageNameSpace):
+    embedding_func: EmbeddingFunc
+
+    async def all_keys(self) -> list[str]:
+        raise NotImplementedError
+
+    async def get_by_id(self, id: str) -> Union[T, None]:
+        raise NotImplementedError
+
+    async def get_by_ids(
+        self, ids: list[str], fields: Union[set[str], None] = None
+    ) -> list[Union[T, None]]:
+        raise NotImplementedError
+
+    async def filter_keys(self, data: list[str]) -> set[str]:
+        
+        raise NotImplementedError
+
+    async def upsert(self, data: dict[str, T]):
+        raise NotImplementedError
+
+    async def drop(self):
+        raise NotImplementedError
+
+
+@dataclass
+class BaseGraphStorage(StorageNameSpace):
+    embedding_func: EmbeddingFunc = None
+
+    async def has_node(self, node_id: str) -> bool:
+        raise NotImplementedError
+
+    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
+        raise NotImplementedError
+
+    async def node_degree(self, node_id: str) -> int:
+        raise NotImplementedError
+
+    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        raise NotImplementedError
+    
+    async def get_pagerank(self,node_id:str) -> float:
+        raise NotImplementedError
+
+    async def get_node(self, node_id: str) -> Union[dict, None]:
+        raise NotImplementedError
+
+    async def get_edge(
+        self, source_node_id: str, target_node_id: str
+    ) -> Union[dict, None]:
+        raise NotImplementedError
+
+    async def get_node_edges(
+        self, source_node_id: str
+    ) -> Union[list[tuple[str, str]], None]:
+        raise NotImplementedError
+    
+    async def get_node_in_edges(
+        self,source_node_id:str
+    ) -> Union[list[tuple[str,str]],None]:
+        raise NotImplementedError
+    async def get_node_out_edges(
+        self,source_node_id:str
+    ) -> Union[list[tuple[str,str]],None]:
+        raise NotImplementedError
+
+    async def upsert_node(self, node_id: str, node_data: dict[str, str]):
+        raise NotImplementedError
+
+    async def upsert_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ):
+        raise NotImplementedError
+
+    async def delete_node(self, node_id: str):
+        raise NotImplementedError
+
+    async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]:
+        raise NotImplementedError("Node embedding is not used in PathRag.")

PathRAG/llm.py

@@ -0,0 +1,1036 @@
+import base64
+import copy
+import json
+import os
+import re
+import struct
+from functools import lru_cache
+from typing import List, Dict, Callable, Any, Union, Optional
+import aioboto3
+import aiohttp
+import numpy as np
+import ollama
+import torch
+import time
+from openai import (
+    AsyncOpenAI,
+    APIConnectionError,
+    RateLimitError,
+    Timeout,
+    AsyncAzureOpenAI,
+)
+from pydantic import BaseModel, Field
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    retry_if_exception_type,
+)
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from .utils import (
+    wrap_embedding_func_with_attrs,
+    locate_json_string_body_from_string,
+    safe_unicode_decode,
+    logger,
+)
+
+import sys
+
+if sys.version_info < (3, 9):
+    from typing import AsyncIterator
+else:
+    from collections.abc import AsyncIterator
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def openai_complete_if_cache(
+    model,
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    base_url="https://api.openai.com/v1",
+
+    api_key="",
+    **kwargs,
+) -> 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)
+    )
+    kwargs.pop("hashing_kv", None)
+    kwargs.pop("keyword_extraction", None)
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+    logger.debug("===== Query Input to LLM =====")
+    logger.debug(f"Query: {prompt}")
+    logger.debug(f"System prompt: {system_prompt}")
+    logger.debug("Full context:")
+    if "response_format" in kwargs:
+        response = await openai_async_client.beta.chat.completions.parse(
+            model=model, messages=messages, **kwargs
+        )
+    else:
+        response = await openai_async_client.chat.completions.create(
+            model=model, messages=messages, **kwargs
+        )
+
+    if hasattr(response, "__aiter__"):
+
+        async def inner():
+            async for chunk in response:
+                content = chunk.choices[0].delta.content
+                if content is None:
+                    continue
+                if r"\u" in content:
+                    content = safe_unicode_decode(content.encode("utf-8"))
+                yield content
+
+        return inner()
+    else:
+        content = response.choices[0].message.content
+        if r"\u" in content:
+            content = safe_unicode_decode(content.encode("utf-8"))
+        return content
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def azure_openai_complete_if_cache(
+    model,
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    base_url=None,
+    api_key=None,
+    api_version=None,
+    **kwargs,
+):
+    if api_key:
+        os.environ["AZURE_OPENAI_API_KEY"] = api_key
+    if base_url:
+        os.environ["AZURE_OPENAI_ENDPOINT"] = base_url
+    if api_version:
+        os.environ["AZURE_OPENAI_API_VERSION"] = api_version
+
+    openai_async_client = AsyncAzureOpenAI(
+        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
+        api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
+    )
+    kwargs.pop("hashing_kv", None)
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    if prompt is not None:
+        messages.append({"role": "user", "content": prompt})
+
+    response = await openai_async_client.chat.completions.create(
+        model=model, messages=messages, **kwargs
+    )
+    content = response.choices[0].message.content
+
+    return content
+
+
+
+
+
+@retry(
+    stop=stop_after_attempt(5),
+    wait=wait_exponential(multiplier=1, max=60),
+    retry=retry_if_exception_type((BedrockError)),
+)
+async def bedrock_complete_if_cache(
+    model,
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    aws_access_key_id=None,
+    aws_secret_access_key=None,
+    aws_session_token=None,
+    **kwargs,
+) -> str:
+    os.environ["AWS_ACCESS_KEY_ID"] = os.environ.get(
+        "AWS_ACCESS_KEY_ID", aws_access_key_id
+    )
+    os.environ["AWS_SECRET_ACCESS_KEY"] = os.environ.get(
+        "AWS_SECRET_ACCESS_KEY", aws_secret_access_key
+    )
+    os.environ["AWS_SESSION_TOKEN"] = os.environ.get(
+        "AWS_SESSION_TOKEN", aws_session_token
+    )
+    kwargs.pop("hashing_kv", None)
+
+    messages = []
+    for history_message in history_messages:
+        message = copy.copy(history_message)
+        message["content"] = [{"text": message["content"]}]
+        messages.append(message)
+
+
+    messages.append({"role": "user", "content": [{"text": prompt}]})
+
+
+    args = {"modelId": model, "messages": messages}
+
+
+    if system_prompt:
+        args["system"] = [{"text": system_prompt}]
+
+
+    inference_params_map = {
+        "max_tokens": "maxTokens",
+        "top_p": "topP",
+        "stop_sequences": "stopSequences",
+    }
+    if inference_params := list(
+        set(kwargs) & set(["max_tokens", "temperature", "top_p", "stop_sequences"])
+    ):
+        args["inferenceConfig"] = {}
+        for param in inference_params:
+            args["inferenceConfig"][inference_params_map.get(param, param)] = (
+                kwargs.pop(param)
+            )
+
+
+    session = aioboto3.Session()
+    async with session.client("bedrock-runtime") as bedrock_async_client:
+        try:
+            response = await bedrock_async_client.converse(**args, **kwargs)
+        except Exception as e:
+            raise BedrockError(e)
+
+    return response["output"]["message"]["content"][0]["text"]
+
+
+@lru_cache(maxsize=1)
+def initialize_hf_model(model_name):
+    hf_tokenizer = AutoTokenizer.from_pretrained(
+        model_name, device_map="auto", trust_remote_code=True
+    )
+    hf_model = AutoModelForCausalLM.from_pretrained(
+        model_name, device_map="auto", trust_remote_code=True
+    )
+    if hf_tokenizer.pad_token is None:
+        hf_tokenizer.pad_token = hf_tokenizer.eos_token
+
+    return hf_model, hf_tokenizer
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def hf_model_if_cache(
+    model,
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    **kwargs,
+) -> str:
+    model_name = model
+    hf_model, hf_tokenizer = initialize_hf_model(model_name)
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+    kwargs.pop("hashing_kv", None)
+    input_prompt = ""
+    try:
+        input_prompt = hf_tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+    except Exception:
+        try:
+            ori_message = copy.deepcopy(messages)
+            if messages[0]["role"] == "system":
+                messages[1]["content"] = (
+                    "<system>"
+                    + messages[0]["content"]
+                    + "</system>\n"
+                    + messages[1]["content"]
+                )
+                messages = messages[1:]
+                input_prompt = hf_tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True
+                )
+        except Exception:
+            len_message = len(ori_message)
+            for msgid in range(len_message):
+                input_prompt = (
+                    input_prompt
+                    + "<"
+                    + ori_message[msgid]["role"]
+                    + ">"
+                    + ori_message[msgid]["content"]
+                    + "</"
+                    + ori_message[msgid]["role"]
+                    + ">\n"
+                )
+
+    input_ids = hf_tokenizer(
+        input_prompt, return_tensors="pt", padding=True, truncation=True
+    ).to("cuda")
+    inputs = {k: v.to(hf_model.device) for k, v in input_ids.items()}
+    output = hf_model.generate(
+        **input_ids, max_new_tokens=512, num_return_sequences=1, early_stopping=True
+    )
+    response_text = hf_tokenizer.decode(
+        output[0][len(inputs["input_ids"][0]) :], skip_special_tokens=True
+    )
+
+    return response_text
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def ollama_model_if_cache(
+    model,
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    **kwargs,
+) -> Union[str, AsyncIterator[str]]:
+    stream = True if kwargs.get("stream") else False
+    kwargs.pop("max_tokens", None)
+
+    host = kwargs.pop("host", None)
+    timeout = kwargs.pop("timeout", None)
+    kwargs.pop("hashing_kv", None)
+    ollama_client = ollama.AsyncClient(host=host, timeout=timeout)
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+    response = await ollama_client.chat(model=model, messages=messages, **kwargs)
+    if stream:
+
+
+        async def inner():
+            async for chunk in response:
+                yield chunk["message"]["content"]
+
+        return inner()
+    else:
+        return response["message"]["content"]
+
+
+@lru_cache(maxsize=1)
+def initialize_lmdeploy_pipeline(
+    model,
+    tp=1,
+    chat_template=None,
+    log_level="WARNING",
+    model_format="hf",
+    quant_policy=0,
+):
+    from lmdeploy import pipeline, ChatTemplateConfig, TurbomindEngineConfig
+
+    lmdeploy_pipe = pipeline(
+        model_path=model,
+        backend_config=TurbomindEngineConfig(
+            tp=tp, model_format=model_format, quant_policy=quant_policy
+        ),
+        chat_template_config=(
+            ChatTemplateConfig(model_name=chat_template) if chat_template else None
+        ),
+        log_level="WARNING",
+    )
+    return lmdeploy_pipe
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def lmdeploy_model_if_cache(
+    model,
+    prompt,
+    system_prompt=None,
+    history_messages=[],
+    chat_template=None,
+    model_format="hf",
+    quant_policy=0,
+    **kwargs,
+) -> str:
+
+    try:
+        import lmdeploy
+        from lmdeploy import version_info, GenerationConfig
+    except Exception:
+        raise ImportError("Please install lmdeploy before initialize lmdeploy backend.")
+    kwargs.pop("hashing_kv", None)
+    kwargs.pop("response_format", None)
+    max_new_tokens = kwargs.pop("max_tokens", 512)
+    tp = kwargs.pop("tp", 1)
+    skip_special_tokens = kwargs.pop("skip_special_tokens", True)
+    do_preprocess = kwargs.pop("do_preprocess", True)
+    do_sample = kwargs.pop("do_sample", False)
+    gen_params = kwargs
+
+    version = version_info
+    if do_sample is not None and version < (0, 6, 0):
+        raise RuntimeError(
+            "`do_sample` parameter is not supported by lmdeploy until "
+            f"v0.6.0, but currently using lmdeloy {lmdeploy.__version__}"
+        )
+    else:
+        do_sample = True
+        gen_params.update(do_sample=do_sample)
+
+    lmdeploy_pipe = initialize_lmdeploy_pipeline(
+        model=model,
+        tp=tp,
+        chat_template=chat_template,
+        model_format=model_format,
+        quant_policy=quant_policy,
+        log_level="WARNING",
+    )
+
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+    gen_config = GenerationConfig(
+        skip_special_tokens=skip_special_tokens,
+        max_new_tokens=max_new_tokens,
+        **gen_params,
+    )
+
+    response = ""
+    async for res in lmdeploy_pipe.generate(
+        messages,
+        gen_config=gen_config,
+        do_preprocess=do_preprocess,
+        stream_response=False,
+        session_id=1,
+    ):
+        response += res.response
+    return response
+
+
+class GPTKeywordExtractionFormat(BaseModel):
+    high_level_keywords: List[str]
+    low_level_keywords: List[str]
+
+
+async def openai_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> Union[str, AsyncIterator[str]]:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    if keyword_extraction:
+        kwargs["response_format"] = "json"
+    model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
+    return await openai_complete_if_cache(
+        model_name,
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+
+
+async def gpt_4o_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    if keyword_extraction:
+        kwargs["response_format"] = GPTKeywordExtractionFormat
+    return await openai_complete_if_cache(
+        "gpt-4o",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+
+
+async def gpt_4o_mini_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    if keyword_extraction:
+        kwargs["response_format"] = GPTKeywordExtractionFormat
+    return await openai_complete_if_cache(
+        "gpt-4o-mini",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+
+
+async def nvidia_openai_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    result = await openai_complete_if_cache(
+        "nvidia/llama-3.1-nemotron-70b-instruct", 
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        base_url="https://integrate.api.nvidia.com/v1",
+        **kwargs,
+    )
+    if keyword_extraction:  # TODO: use JSON API
+        return locate_json_string_body_from_string(result)
+    return result
+
+
+async def azure_openai_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    result = await azure_openai_complete_if_cache(
+        "conversation-4o-mini",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+    if keyword_extraction:  # TODO: use JSON API
+        return locate_json_string_body_from_string(result)
+    return result
+
+
+async def bedrock_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    result = await bedrock_complete_if_cache(
+        "anthropic.claude-3-haiku-20240307-v1:0",
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+    if keyword_extraction:  # TODO: use JSON API
+        return locate_json_string_body_from_string(result)
+    return result
+
+
+async def hf_model_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> str:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
+    result = await hf_model_if_cache(
+        model_name,
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+    if keyword_extraction:  # TODO: use JSON API
+        return locate_json_string_body_from_string(result)
+    return result
+
+
+async def ollama_model_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+) -> Union[str, AsyncIterator[str]]:
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+    if keyword_extraction:
+        kwargs["format"] = "json"
+    model_name = kwargs["hashing_kv"].global_config["llm_model_name"]
+    return await ollama_model_if_cache(
+        model_name,
+        prompt,
+        system_prompt=system_prompt,
+        history_messages=history_messages,
+        **kwargs,
+    )
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def zhipu_complete_if_cache(
+    prompt: Union[str, List[Dict[str, str]]],
+    model: str = "glm-4-flashx", 
+    api_key: Optional[str] = None,
+    system_prompt: Optional[str] = None,
+    history_messages: List[Dict[str, str]] = [],
+    **kwargs,
+) -> str:
+
+    try:
+        from zhipuai import ZhipuAI
+    except ImportError:
+        raise ImportError("Please install zhipuai before initialize zhipuai backend.")
+
+    if api_key:
+        client = ZhipuAI(api_key=api_key)
+    else:
+
+        client = ZhipuAI()
+
+    messages = []
+
+    if not system_prompt:
+        system_prompt = "You are a helpful assistant. 注意内容里的敏感词用***替换。"
+
+
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+
+
+    logger.debug("===== Query Input to LLM =====")
+    logger.debug(f"Query: {prompt}")
+    logger.debug(f"System prompt: {system_prompt}")
+
+
+    kwargs = {
+        k: v for k, v in kwargs.items() if k not in ["hashing_kv", "keyword_extraction"]
+    }
+
+    response = client.chat.completions.create(model=model, messages=messages, **kwargs)
+
+    return response.choices[0].message.content
+
+
+async def zhipu_complete(
+    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
+):
+
+    keyword_extraction = kwargs.pop("keyword_extraction", None)
+
+    if keyword_extraction:
+
+        extraction_prompt = """You are a helpful assistant that extracts keywords from text.
+        Please analyze the content and extract two types of keywords:
+        1. High-level keywords: Important concepts and main themes
+        2. Low-level keywords: Specific details and supporting elements
+
+        Return your response in this exact JSON format:
+        {
+            "high_level_keywords": ["keyword1", "keyword2"],
+            "low_level_keywords": ["keyword1", "keyword2", "keyword3"]
+        }
+
+        Only return the JSON, no other text."""
+
+        if system_prompt:
+            system_prompt = f"{system_prompt}\n\n{extraction_prompt}"
+        else:
+            system_prompt = extraction_prompt
+
+        try:
+            response = await zhipu_complete_if_cache(
+                prompt=prompt,
+                system_prompt=system_prompt,
+                history_messages=history_messages,
+                **kwargs,
+            )
+
+
+            try:
+                data = json.loads(response)
+                return GPTKeywordExtractionFormat(
+                    high_level_keywords=data.get("high_level_keywords", []),
+                    low_level_keywords=data.get("low_level_keywords", []),
+                )
+            except json.JSONDecodeError:
+                match = re.search(r"\{[\s\S]*\}", response)
+                if match:
+                    try:
+                        data = json.loads(match.group())
+                        return GPTKeywordExtractionFormat(
+                            high_level_keywords=data.get("high_level_keywords", []),
+                            low_level_keywords=data.get("low_level_keywords", []),
+                        )
+                    except json.JSONDecodeError:
+                        pass
+
+                logger.warning(
+                    f"Failed to parse keyword extraction response: {response}"
+                )
+                return GPTKeywordExtractionFormat(
+                    high_level_keywords=[], low_level_keywords=[]
+                )
+        except Exception as e:
+            logger.error(f"Error during keyword extraction: {str(e)}")
+            return GPTKeywordExtractionFormat(
+                high_level_keywords=[], low_level_keywords=[]
+            )
+    else:
+
+        return await zhipu_complete_if_cache(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+            **kwargs,
+        )
+
+
+@wrap_embedding_func_with_attrs(embedding_dim=1024, max_token_size=8192)
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=60),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def zhipu_embedding(
+    texts: list[str], model: str = "embedding-3", api_key: str = None, **kwargs
+) -> np.ndarray:
+
+    try:
+        from zhipuai import ZhipuAI
+    except ImportError:
+        raise ImportError("Please install zhipuai before initialize zhipuai backend.")
+    if api_key:
+        client = ZhipuAI(api_key=api_key)
+    else:
+
+        client = ZhipuAI()
+
+    if isinstance(texts, str):
+        texts = [texts]
+
+    embeddings = []
+    for text in texts:
+        try:
+            response = client.embeddings.create(model=model, input=[text], **kwargs)
+            embeddings.append(response.data[0].embedding)
+        except Exception as e:
+            raise Exception(f"Error calling ChatGLM Embedding API: {str(e)}")
+
+    return np.array(embeddings)
+
+
+@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8192)
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=60),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def openai_embedding(
+    texts: list[str],
+    model: str = "text-embedding-3-small",
+    base_url="https://api.openai.com/v1",
+    api_key="",
+) -> np.ndarray:
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+
+    openai_async_client = (
+        AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
+    )
+    response = await openai_async_client.embeddings.create(
+        model=model, input=texts, encoding_format="float"
+    )
+    return np.array([dp.embedding for dp in response.data])
+
+
+async def fetch_data(url, headers, data):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, headers=headers, json=data) as response:
+            response_json = await response.json()
+            data_list = response_json.get("data", [])
+            return data_list
+
+
+async def jina_embedding(
+    texts: list[str],
+    dimensions: int = 1024,
+    late_chunking: bool = False,
+    base_url: str = None,
+    api_key: str = None,
+) -> np.ndarray:
+    if api_key:
+        os.environ["JINA_API_KEY"] = api_key
+    url = "https://api.jina.ai/v1/embeddings" if not base_url else base_url
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {os.environ['JINA_API_KEY']}",
+    }
+    data = {
+        "model": "jina-embeddings-v3",
+        "normalized": True,
+        "embedding_type": "float",
+        "dimensions": f"{dimensions}",
+        "late_chunking": late_chunking,
+        "input": texts,
+    }
+    data_list = await fetch_data(url, headers, data)
+    return np.array([dp["embedding"] for dp in data_list])
+
+
+@wrap_embedding_func_with_attrs(embedding_dim=2048, max_token_size=512)
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=60),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def nvidia_openai_embedding(
+    texts: list[str],
+    model: str = "nvidia/llama-3.2-nv-embedqa-1b-v1",
+    base_url: str = "https://integrate.api.nvidia.com/v1",
+    api_key: str = None,
+    input_type: str = "passage",  
+    trunc: str = "NONE",  
+    encode: str = "float",  
+) -> np.ndarray:
+    if api_key:
+        os.environ["OPENAI_API_KEY"] = api_key
+
+    openai_async_client = (
+        AsyncOpenAI() if base_url is None else AsyncOpenAI(base_url=base_url)
+    )
+    response = await openai_async_client.embeddings.create(
+        model=model,
+        input=texts,
+        encoding_format=encode,
+        extra_body={"input_type": input_type, "truncate": trunc},
+    )
+    return np.array([dp.embedding for dp in response.data])
+
+
+@wrap_embedding_func_with_attrs(embedding_dim=1536, max_token_size=8191)
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=10),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def azure_openai_embedding(
+    texts: list[str],
+    model: str = "text-embedding-3-small",
+    base_url: str = None,
+    api_key: str = None,
+    api_version: str = None,
+) -> np.ndarray:
+    if api_key:
+        os.environ["AZURE_OPENAI_API_KEY"] = api_key
+    if base_url:
+        os.environ["AZURE_OPENAI_ENDPOINT"] = base_url
+    if api_version:
+        os.environ["AZURE_OPENAI_API_VERSION"] = api_version
+
+    openai_async_client = AsyncAzureOpenAI(
+        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
+        api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
+    )
+
+    response = await openai_async_client.embeddings.create(
+        model=model, input=texts, encoding_format="float"
+    )
+    return np.array([dp.embedding for dp in response.data])
+
+
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=4, max=60),
+    retry=retry_if_exception_type((RateLimitError, APIConnectionError, Timeout)),
+)
+async def siliconcloud_embedding(
+    texts: list[str],
+    model: str = "netease-youdao/bce-embedding-base_v1",
+    base_url: str = "https://api.siliconflow.cn/v1/embeddings",
+    max_token_size: int = 512,
+    api_key: str = None,
+) -> np.ndarray:
+    if api_key and not api_key.startswith("Bearer "):
+        api_key = "Bearer " + api_key
+
+    headers = {"Authorization": api_key, "Content-Type": "application/json"}
+
+    truncate_texts = [text[0:max_token_size] for text in texts]
+
+    payload = {"model": model, "input": truncate_texts, "encoding_format": "base64"}
+
+    base64_strings = []
+    async with aiohttp.ClientSession() as session:
+        async with session.post(base_url, headers=headers, json=payload) as response:
+            content = await response.json()
+            if "code" in content:
+                raise ValueError(content)
+            base64_strings = [item["embedding"] for item in content["data"]]
+
+    embeddings = []
+    for string in base64_strings:
+        decode_bytes = base64.b64decode(string)
+        n = len(decode_bytes) // 4
+        float_array = struct.unpack("<" + "f" * n, decode_bytes)
+        embeddings.append(float_array)
+    return np.array(embeddings)
+
+
+async def bedrock_embedding(
+    texts: list[str],
+    model: str = "amazon.titan-embed-text-v2:0",
+    aws_access_key_id=None,
+    aws_secret_access_key=None,
+    aws_session_token=None,
+) -> np.ndarray:
+    os.environ["AWS_ACCESS_KEY_ID"] = os.environ.get(
+        "AWS_ACCESS_KEY_ID", aws_access_key_id
+    )
+    os.environ["AWS_SECRET_ACCESS_KEY"] = os.environ.get(
+        "AWS_SECRET_ACCESS_KEY", aws_secret_access_key
+    )
+    os.environ["AWS_SESSION_TOKEN"] = os.environ.get(
+        "AWS_SESSION_TOKEN", aws_session_token
+    )
+
+    session = aioboto3.Session()
+    async with session.client("bedrock-runtime") as bedrock_async_client:
+        if (model_provider := model.split(".")[0]) == "amazon":
+            embed_texts = []
+            for text in texts:
+                if "v2" in model:
+                    body = json.dumps(
+                        {
+                            "inputText": text,
+                            
+                            "embeddingTypes": ["float"],
+                        }
+                    )
+                elif "v1" in model:
+                    body = json.dumps({"inputText": text})
+                else:
+                    raise ValueError(f"Model {model} is not supported!")
+
+                response = await bedrock_async_client.invoke_model(
+                    modelId=model,
+                    body=body,
+                    accept="application/json",
+                    contentType="application/json",
+                )
+
+                response_body = await response.get("body").json()
+
+                embed_texts.append(response_body["embedding"])
+        elif model_provider == "cohere":
+            body = json.dumps(
+                {"texts": texts, "input_type": "search_document", "truncate": "NONE"}
+            )
+
+            response = await bedrock_async_client.invoke_model(
+                model=model,
+                body=body,
+                accept="application/json",
+                contentType="application/json",
+            )
+
+            response_body = json.loads(response.get("body").read())
+
+            embed_texts = response_body["embeddings"]
+        else:
+            raise ValueError(f"Model provider '{model_provider}' is not supported!")
+
+        return np.array(embed_texts)
+
+
+async def hf_embedding(texts: list[str], tokenizer, embed_model) -> np.ndarray:
+    device = next(embed_model.parameters()).device
+    input_ids = tokenizer(
+        texts, return_tensors="pt", padding=True, truncation=True
+    ).input_ids.to(device)
+    with torch.no_grad():
+        outputs = embed_model(input_ids)
+        embeddings = outputs.last_hidden_state.mean(dim=1)
+    if embeddings.dtype == torch.bfloat16:
+        return embeddings.detach().to(torch.float32).cpu().numpy()
+    else:
+        return embeddings.detach().cpu().numpy()
+
+
+async def ollama_embedding(texts: list[str], embed_model, **kwargs) -> np.ndarray:
+
+    embed_text = []
+    ollama_client = ollama.Client(**kwargs)
+    for text in texts:
+        data = ollama_client.embeddings(model=embed_model, prompt=text)
+        embed_text.append(data["embedding"])
+
+    return embed_text
+
+
+async def ollama_embed(texts: list[str], embed_model, **kwargs) -> np.ndarray:
+    ollama_client = ollama.Client(**kwargs)
+    data = ollama_client.embed(model=embed_model, input=texts)
+    return data["embeddings"]
+
+
+class Model(BaseModel):
+
+    gen_func: Callable[[Any], str] = Field(
+        ...,
+        description="A function that generates the response from the llm. The response must be a string",
+    )
+    kwargs: Dict[str, Any] = Field(
+        ...,
+        description="The arguments to pass to the callable function. Eg. the api key, model name, etc",
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class MultiModel:
+    def __init__(self, models: List[Model]):
+        self._models = models
+        self._current_model = 0
+
+    def _next_model(self):
+        self._current_model = (self._current_model + 1) % len(self._models)
+        return self._models[self._current_model]
+
+    async def llm_model_func(
+        self, prompt, system_prompt=None, history_messages=[], **kwargs
+    ) -> str:
+        kwargs.pop("model", None)  
+        kwargs.pop("keyword_extraction", None)
+        kwargs.pop("mode", None)
+        next_model = self._next_model()
+        args = dict(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+            **kwargs,
+            **next_model.kwargs,
+        )
+
+        return await next_model.gen_func(**args)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    async def main():
+        result = await gpt_4o_mini_complete("How are you?")
+        print(result)
+
+    asyncio.run(main())

PathRAG/operate.py

@@ -0,0 +1,1239 @@
+import asyncio
+import json
+import re
+from tqdm.asyncio import tqdm as tqdm_async
+from typing import Union
+from collections import Counter, defaultdict
+import warnings
+import tiktoken
+import time
+import csv
+from .utils import (
+    logger,
+    clean_str,
+    compute_mdhash_id,
+    decode_tokens_by_tiktoken,
+    encode_string_by_tiktoken,
+    is_float_regex,
+    list_of_list_to_csv,
+    pack_user_ass_to_openai_messages,
+    split_string_by_multi_markers,
+    truncate_list_by_token_size,
+    process_combine_contexts,
+    compute_args_hash,
+    handle_cache,
+    save_to_cache,
+    CacheData,
+)
+from .base import (
+    BaseGraphStorage,
+    BaseKVStorage,
+    BaseVectorStorage,
+    TextChunkSchema,
+    QueryParam,
+)
+from .prompt import GRAPH_FIELD_SEP, PROMPTS
+
+
+def chunking_by_token_size(
+    content: str, overlap_token_size=128, max_token_size=1024, tiktoken_model="gpt-4o"
+):
+    tokens = encode_string_by_tiktoken(content, model_name=tiktoken_model)
+    results = []
+    for index, start in enumerate(
+        range(0, len(tokens), max_token_size - overlap_token_size)
+    ):
+        chunk_content = decode_tokens_by_tiktoken(
+            tokens[start : start + max_token_size], model_name=tiktoken_model
+        )
+        results.append(
+            {
+                "tokens": min(max_token_size, len(tokens) - start),
+                "content": chunk_content.strip(),
+                "chunk_order_index": index,
+            }
+        )
+    return results
+
+
+async def _handle_entity_relation_summary(
+    entity_or_relation_name: str,
+    description: str,
+    global_config: dict,
+) -> str:
+    use_llm_func: callable = global_config["llm_model_func"]
+    llm_max_tokens = global_config["llm_model_max_token_size"]
+    tiktoken_model_name = global_config["tiktoken_model_name"]
+    summary_max_tokens = global_config["entity_summary_to_max_tokens"]
+    language = global_config["addon_params"].get(
+        "language", PROMPTS["DEFAULT_LANGUAGE"]
+    )
+
+    tokens = encode_string_by_tiktoken(description, model_name=tiktoken_model_name)
+    if len(tokens) < summary_max_tokens: 
+        return description
+    prompt_template = PROMPTS["summarize_entity_descriptions"]
+    use_description = decode_tokens_by_tiktoken(
+        tokens[:llm_max_tokens], model_name=tiktoken_model_name
+    )
+    context_base = dict(
+        entity_name=entity_or_relation_name,
+        description_list=use_description.split(GRAPH_FIELD_SEP),
+        language=language,
+    )
+    use_prompt = prompt_template.format(**context_base)
+    logger.debug(f"Trigger summary: {entity_or_relation_name}")
+    summary = await use_llm_func(use_prompt, max_tokens=summary_max_tokens)
+    return summary
+
+
+async def _handle_single_entity_extraction(
+    record_attributes: list[str],
+    chunk_key: str,
+):
+    if len(record_attributes) < 4 or record_attributes[0] != '"entity"':
+        return None
+   
+    entity_name = clean_str(record_attributes[1].upper())
+    if not entity_name.strip():
+        return None
+    entity_type = clean_str(record_attributes[2].upper())
+    entity_description = clean_str(record_attributes[3])
+    entity_source_id = chunk_key
+    return dict(
+        entity_name=entity_name,
+        entity_type=entity_type,
+        description=entity_description,
+        source_id=entity_source_id,
+    )
+
+
+async def _handle_single_relationship_extraction(
+    record_attributes: list[str],
+    chunk_key: str,
+):
+    if len(record_attributes) < 5 or record_attributes[0] != '"relationship"':
+        return None
+   
+    source = clean_str(record_attributes[1].upper())
+    target = clean_str(record_attributes[2].upper())
+    edge_description = clean_str(record_attributes[3])
+
+    edge_keywords = clean_str(record_attributes[4])
+    edge_source_id = chunk_key
+    weight = (
+        float(record_attributes[-1]) if is_float_regex(record_attributes[-1]) else 1.0
+    )
+    return dict(
+        src_id=source,
+        tgt_id=target,
+        weight=weight,
+        description=edge_description,
+        keywords=edge_keywords,
+        source_id=edge_source_id,
+    )
+
+
+async def _merge_nodes_then_upsert(
+    entity_name: str,
+    nodes_data: list[dict],
+    knowledge_graph_inst: BaseGraphStorage,
+    global_config: dict,
+):
+    already_entity_types = []
+    already_source_ids = []
+    already_description = []
+
+    already_node = await knowledge_graph_inst.get_node(entity_name)
+    if already_node is not None:
+        already_entity_types.append(already_node["entity_type"])
+        already_source_ids.extend(
+            split_string_by_multi_markers(already_node["source_id"], [GRAPH_FIELD_SEP])
+        )
+        already_description.append(already_node["description"])
+
+    entity_type = sorted(
+        Counter(
+            [dp["entity_type"] for dp in nodes_data] + already_entity_types
+        ).items(),
+        key=lambda x: x[1],
+        reverse=True,
+    )[0][0]
+    description = GRAPH_FIELD_SEP.join(
+        sorted(set([dp["description"] for dp in nodes_data] + already_description))
+    )
+    source_id = GRAPH_FIELD_SEP.join(
+        set([dp["source_id"] for dp in nodes_data] + already_source_ids)
+    )
+    description = await _handle_entity_relation_summary(
+        entity_name, description, global_config
+    )
+    node_data = dict(
+        entity_type=entity_type,
+        description=description,
+        source_id=source_id,
+    )
+    await knowledge_graph_inst.upsert_node(
+        entity_name,
+        node_data=node_data,
+    )
+    node_data["entity_name"] = entity_name
+    return node_data
+
+
+async def _merge_edges_then_upsert(
+    src_id: str,
+    tgt_id: str,
+    edges_data: list[dict],
+    knowledge_graph_inst: BaseGraphStorage,
+    global_config: dict,
+):
+    already_weights = []
+    already_source_ids = []
+    already_description = []
+    already_keywords = []
+
+    if await knowledge_graph_inst.has_edge(src_id, tgt_id):
+        already_edge = await knowledge_graph_inst.get_edge(src_id, tgt_id)
+        already_weights.append(already_edge["weight"])
+        already_source_ids.extend(
+            split_string_by_multi_markers(already_edge["source_id"], [GRAPH_FIELD_SEP])
+        )
+        already_description.append(already_edge["description"])
+        already_keywords.extend(
+            split_string_by_multi_markers(already_edge["keywords"], [GRAPH_FIELD_SEP])
+        )
+
+    weight = sum([dp["weight"] for dp in edges_data] + already_weights)
+    description = GRAPH_FIELD_SEP.join(
+        sorted(set([dp["description"] for dp in edges_data] + already_description))
+    )
+    keywords = GRAPH_FIELD_SEP.join(
+        sorted(set([dp["keywords"] for dp in edges_data] + already_keywords))
+    )
+    source_id = GRAPH_FIELD_SEP.join(
+        set([dp["source_id"] for dp in edges_data] + already_source_ids)
+    )
+    for need_insert_id in [src_id, tgt_id]:
+        if not (await knowledge_graph_inst.has_node(need_insert_id)):
+            await knowledge_graph_inst.upsert_node(
+                need_insert_id,
+                node_data={
+                    "source_id": source_id,
+                    "description": description,
+                    "entity_type": '"UNKNOWN"',
+                },
+            )
+    description = await _handle_entity_relation_summary(
+        f"({src_id}, {tgt_id})", description, global_config
+    )
+    await knowledge_graph_inst.upsert_edge(
+        src_id,
+        tgt_id,
+        edge_data=dict(
+            weight=weight,
+            description=description,
+            keywords=keywords,
+            source_id=source_id,
+        ),
+    )
+
+    edge_data = dict(
+        src_id=src_id,
+        tgt_id=tgt_id,
+        description=description,
+        keywords=keywords,
+    )
+
+    return edge_data
+
+
+async def extract_entities(
+    chunks: dict[str, TextChunkSchema],
+    knowledge_graph_inst: BaseGraphStorage,
+    entity_vdb: BaseVectorStorage,
+    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"]
+
+    ordered_chunks = list(chunks.items())
+  
+    language = global_config["addon_params"].get(
+        "language", PROMPTS["DEFAULT_LANGUAGE"]
+    )
+    entity_types = global_config["addon_params"].get(
+        "entity_types", PROMPTS["DEFAULT_ENTITY_TYPES"]
+    )
+    example_number = global_config["addon_params"].get("example_number", None)
+    if example_number and example_number < len(PROMPTS["entity_extraction_examples"]):
+        examples = "\n".join(
+            PROMPTS["entity_extraction_examples"][: int(example_number)]
+        )
+    else:
+        examples = "\n".join(PROMPTS["entity_extraction_examples"])
+
+    example_context_base = dict(
+        tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
+        record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
+        completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
+        entity_types=",".join(entity_types),
+        language=language,
+    )
+  
+    examples = examples.format(**example_context_base)
+
+    entity_extract_prompt = PROMPTS["entity_extraction"]
+    context_base = dict(
+        tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
+        record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"],
+        completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
+        entity_types=",".join(entity_types),
+        examples=examples,
+        language=language,
+    )
+
+    continue_prompt = PROMPTS["entiti_continue_extraction"]
+    if_loop_prompt = PROMPTS["entiti_if_loop_extraction"]
+
+    already_processed = 0
+    already_entities = 0
+    already_relations = 0
+
+    async def _process_single_content(chunk_key_dp: tuple[str, TextChunkSchema]):
+        nonlocal already_processed, already_entities, already_relations
+        chunk_key = chunk_key_dp[0]
+        chunk_dp = chunk_key_dp[1]
+        content = chunk_dp["content"]
+        hint_prompt = entity_extract_prompt.format(
+            **context_base, input_text="{input_text}"
+        ).format(**context_base, input_text=content)
+
+        final_result = await use_llm_func(hint_prompt)
+        history = pack_user_ass_to_openai_messages(hint_prompt, final_result)
+        for now_glean_index in range(entity_extract_max_gleaning):
+            glean_result = await use_llm_func(continue_prompt, history_messages=history)
+
+            history += pack_user_ass_to_openai_messages(continue_prompt, glean_result)
+            final_result += glean_result
+            if now_glean_index == entity_extract_max_gleaning - 1:
+                break
+
+            if_loop_result: str = await use_llm_func(
+                if_loop_prompt, history_messages=history
+            )
+            if_loop_result = if_loop_result.strip().strip('"').strip("'").lower()
+            if if_loop_result != "yes":
+                break
+
+        records = split_string_by_multi_markers(
+            final_result,
+            [context_base["record_delimiter"], context_base["completion_delimiter"]],
+        )
+
+        maybe_nodes = defaultdict(list)
+        maybe_edges = defaultdict(list)
+        for record in records:
+            record = re.search(r"\((.*)\)", record)
+            if record is None:
+                continue
+            record = record.group(1)
+            record_attributes = split_string_by_multi_markers(
+                record, [context_base["tuple_delimiter"]]
+            )
+            if_entities = await _handle_single_entity_extraction(
+                record_attributes, chunk_key
+            )
+            if if_entities is not None:
+                maybe_nodes[if_entities["entity_name"]].append(if_entities)
+                continue
+
+            if_relation = await _handle_single_relationship_extraction(
+                record_attributes, chunk_key
+            )
+            if if_relation is not None:
+                maybe_edges[(if_relation["src_id"], if_relation["tgt_id"])].append(
+                    if_relation
+                )
+        already_processed += 1
+        already_entities += len(maybe_nodes)
+        already_relations += len(maybe_edges)
+        now_ticks = PROMPTS["process_tickers"][
+            already_processed % len(PROMPTS["process_tickers"])
+        ]
+        print(
+            f"{now_ticks} Processed {already_processed} chunks, {already_entities} entities(duplicated), {already_relations} relations(duplicated)\r",
+            end="",
+            flush=True,
+        )
+        return dict(maybe_nodes), dict(maybe_edges)
+
+    results = []
+    for result in tqdm_async(
+        asyncio.as_completed([_process_single_content(c) for c in ordered_chunks]),
+        total=len(ordered_chunks),
+        desc="Extracting entities from chunks",
+        unit="chunk",
+    ):
+        results.append(await result)
+
+    maybe_nodes = defaultdict(list)
+    maybe_edges = defaultdict(list)
+    for m_nodes, m_edges in results:
+        for k, v in m_nodes.items():
+            maybe_nodes[k].extend(v)
+        for k, v in m_edges.items():
+            maybe_edges[k].extend(v)
+    logger.info("Inserting entities into storage...")
+    all_entities_data = []
+    for result in tqdm_async(
+        asyncio.as_completed(
+            [
+                _merge_nodes_then_upsert(k, v, knowledge_graph_inst, global_config)
+                for k, v in maybe_nodes.items()
+            ]
+        ),
+        total=len(maybe_nodes),
+        desc="Inserting entities",
+        unit="entity",
+    ):
+        all_entities_data.append(await result)
+
+    logger.info("Inserting relationships into storage...")
+    all_relationships_data = []
+    for result in tqdm_async(
+        asyncio.as_completed(
+            [
+                _merge_edges_then_upsert(
+                    k[0], k[1], v, knowledge_graph_inst, global_config
+                )
+                for k, v in maybe_edges.items()
+            ]
+        ),
+        total=len(maybe_edges),
+        desc="Inserting relationships",
+        unit="relationship",
+    ):
+        all_relationships_data.append(await result)
+
+    if not len(all_entities_data) and not len(all_relationships_data):
+        logger.warning(
+            "Didn't extract any entities and relationships, maybe your LLM is not working"
+        )
+        return None
+
+    if not len(all_entities_data):
+        logger.warning("Didn't extract any entities")
+    if not len(all_relationships_data):
+        logger.warning("Didn't extract any relationships")
+
+    if entity_vdb is not None:
+        data_for_vdb = {
+            compute_mdhash_id(dp["entity_name"], prefix="ent-"): {
+                "content": dp["entity_name"] + dp["description"],
+                "entity_name": dp["entity_name"],
+            }
+            for dp in all_entities_data
+        }
+        await entity_vdb.upsert(data_for_vdb)
+
+    if relationships_vdb is not None:
+        data_for_vdb = {
+            compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): {
+                "src_id": dp["src_id"],
+                "tgt_id": dp["tgt_id"],
+                "content": dp["keywords"]
+                + dp["src_id"]
+                + dp["tgt_id"]
+                + dp["description"],
+            }
+            for dp in all_relationships_data
+        }
+        await relationships_vdb.upsert(data_for_vdb)
+
+    return knowledge_graph_inst
+
+
+
+async def kg_query(
+    query,
+    knowledge_graph_inst: BaseGraphStorage,
+    entities_vdb: BaseVectorStorage,
+    relationships_vdb: BaseVectorStorage,
+    text_chunks_db: BaseKVStorage[TextChunkSchema],
+    query_param: QueryParam,
+    global_config: dict,
+    hashing_kv: BaseKVStorage = None,
+) -> str:
+
+    use_model_func = global_config["llm_model_func"]
+    args_hash = compute_args_hash(query_param.mode, query)
+    cached_response, quantized, min_val, max_val = await handle_cache(
+        hashing_kv, args_hash, query, query_param.mode
+    )
+    if cached_response is not None:
+        return cached_response
+
+    example_number = global_config["addon_params"].get("example_number", None)
+    if example_number and example_number < len(PROMPTS["keywords_extraction_examples"]):
+        examples = "\n".join(
+            PROMPTS["keywords_extraction_examples"][: int(example_number)]
+        )
+    else:
+        examples = "\n".join(PROMPTS["keywords_extraction_examples"])
+    language = global_config["addon_params"].get(
+        "language", PROMPTS["DEFAULT_LANGUAGE"]
+    )
+
+    if query_param.mode not in ["hybrid"]:
+        logger.error(f"Unknown mode {query_param.mode} in kg_query")
+        return PROMPTS["fail_response"]
+
+
+    kw_prompt_temp = PROMPTS["keywords_extraction"]
+    kw_prompt = kw_prompt_temp.format(query=query, examples=examples, language=language)
+    result = await use_model_func(kw_prompt, keyword_extraction=True)
+    logger.info("kw_prompt result:")
+    print(result)
+    try:
+
+        match = re.search(r"\{.*\}", result, re.DOTALL)
+        if match:
+            result = match.group(0)
+            keywords_data = json.loads(result)
+
+            hl_keywords = keywords_data.get("high_level_keywords", [])
+            ll_keywords = keywords_data.get("low_level_keywords", [])
+        else:
+            logger.error("No JSON-like structure found in the result.")
+            return PROMPTS["fail_response"]
+
+
+    except json.JSONDecodeError as e:
+        print(f"JSON parsing error: {e} {result}")
+        return PROMPTS["fail_response"]
+
+
+    if hl_keywords == [] and ll_keywords == []:
+        logger.warning("low_level_keywords and high_level_keywords is empty")
+        return PROMPTS["fail_response"]
+    if ll_keywords == [] and query_param.mode in ["hybrid"]:
+        logger.warning("low_level_keywords is empty")
+        return PROMPTS["fail_response"]
+    else:
+        ll_keywords = ", ".join(ll_keywords)
+    if hl_keywords == [] and query_param.mode in ["hybrid"]:
+        logger.warning("high_level_keywords is empty")
+        return PROMPTS["fail_response"]
+    else:
+        hl_keywords = ", ".join(hl_keywords)
+
+
+    keywords = [ll_keywords, hl_keywords]
+    context= await _build_query_context(
+        keywords,
+        knowledge_graph_inst,
+        entities_vdb,
+        relationships_vdb,
+        text_chunks_db,
+        query_param,
+    )
+
+    
+
+    if query_param.only_need_context:
+        return context
+    if context is None:
+        return PROMPTS["fail_response"]
+    sys_prompt_temp = PROMPTS["rag_response"]
+    sys_prompt = sys_prompt_temp.format(
+        context_data=context, response_type=query_param.response_type
+    )
+    if query_param.only_need_prompt:
+        return sys_prompt
+    response = await use_model_func(
+        query,
+        system_prompt=sys_prompt,
+        stream=query_param.stream,
+    )
+    if isinstance(response, str) and len(response) > len(sys_prompt):
+        response = (
+            response.replace(sys_prompt, "")
+            .replace("user", "")
+            .replace("model", "")
+            .replace(query, "")
+            .replace("<system>", "")
+            .replace("</system>", "")
+            .strip()
+        )
+
+
+    await save_to_cache(
+        hashing_kv,
+        CacheData(
+            args_hash=args_hash,
+            content=response,
+            prompt=query,
+            quantized=quantized,
+            min_val=min_val,
+            max_val=max_val,
+            mode=query_param.mode,
+        ),
+    )
+    return response
+
+
+async def _build_query_context(
+    query: list,
+    knowledge_graph_inst: BaseGraphStorage,
+    entities_vdb: BaseVectorStorage,
+    relationships_vdb: BaseVectorStorage,
+    text_chunks_db: BaseKVStorage[TextChunkSchema],
+    query_param: QueryParam,
+):
+    ll_entities_context, ll_relations_context, ll_text_units_context = "", "", ""
+    hl_entities_context, hl_relations_context, hl_text_units_context = "", "", ""
+
+    ll_kewwords, hl_keywrds = query[0], query[1]
+    if query_param.mode in ["local", "hybrid"]:
+        if ll_kewwords == "":
+            ll_entities_context, ll_relations_context, ll_text_units_context = (
+                "",
+                "",
+                "",
+            )
+            warnings.warn(
+                "Low Level context is None. Return empty Low entity/relationship/source"
+            )
+            query_param.mode = "global"
+        else:
+            (
+                ll_entities_context,
+                ll_relations_context,
+                ll_text_units_context,
+            ) = await _get_node_data(
+                ll_kewwords,
+                knowledge_graph_inst,
+                entities_vdb,
+                text_chunks_db,
+                query_param,
+            )
+    if query_param.mode in ["hybrid"]:
+        if hl_keywrds == "":
+            hl_entities_context, hl_relations_context, hl_text_units_context = (
+                "",
+                "",
+                "",
+            )
+            warnings.warn(
+                "High Level context is None. Return empty High entity/relationship/source"
+            )
+            query_param.mode = "local"
+        else:
+            (
+                hl_entities_context,
+                hl_relations_context,
+                hl_text_units_context,
+            ) = await _get_edge_data(
+                hl_keywrds,
+                knowledge_graph_inst,
+                relationships_vdb,
+                text_chunks_db,
+                query_param,
+            )
+            if (
+                hl_entities_context == ""
+                and hl_relations_context == ""
+                and hl_text_units_context == ""
+            ):
+                logger.warn("No high level context found. Switching to local mode.")
+                query_param.mode = "local"
+    if query_param.mode == "hybrid":
+        entities_context, relations_context, text_units_context = combine_contexts(
+            [hl_entities_context, hl_relations_context],
+            [ll_entities_context, ll_relations_context],
+            [hl_text_units_context, ll_text_units_context],
+        )
+
+
+    return f"""
+-----global-information-----
+-----high-level entity information-----
+```csv
+{hl_entities_context}
+```
+-----high-level relationship information-----
+```csv
+{hl_relations_context}
+```
+-----Sources-----
+```csv
+{text_units_context}
+```
+-----local-information-----
+-----low-level entity information-----
+```csv
+{ll_entities_context}
+```
+-----low-level relationship information-----
+```csv
+{ll_relations_context}
+```
+"""
+
+async def _get_node_data(
+    query,
+    knowledge_graph_inst: BaseGraphStorage,
+    entities_vdb: BaseVectorStorage,
+    text_chunks_db: BaseKVStorage[TextChunkSchema],
+    query_param: QueryParam,
+):
+
+    results = await entities_vdb.query(query, top_k=query_param.top_k)
+    if not len(results):
+        return "", "", ""
+
+    node_datas = await asyncio.gather(
+        *[knowledge_graph_inst.get_node(r["entity_name"]) for r in results]
+    )
+    if not all([n is not None for n in node_datas]):
+        logger.warning("Some nodes are missing, maybe the storage is damaged")
+
+
+    node_degrees = await asyncio.gather(
+        *[knowledge_graph_inst.node_degree(r["entity_name"]) for r in results]
+    )
+    node_datas = [
+        {**n, "entity_name": k["entity_name"], "rank": d}
+        for k, n, d in zip(results, node_datas, node_degrees)
+        if n is not None
+    ]  
+    use_text_units = await _find_most_related_text_unit_from_entities(
+        node_datas, query_param, text_chunks_db, knowledge_graph_inst
+    )
+
+
+    use_relations= await _find_most_related_edges_from_entities3(
+        node_datas, query_param, knowledge_graph_inst
+    )
+
+    logger.info(
+        f"Local query uses {len(node_datas)} entites, {len(use_relations)} relations, {len(use_text_units)} text units"
+    )
+
+
+    entites_section_list = [["id", "entity", "type", "description", "rank"]]
+    for i, n in enumerate(node_datas):
+        entites_section_list.append(
+            [
+                i,
+                n["entity_name"],
+                n.get("entity_type", "UNKNOWN"),
+                n.get("description", "UNKNOWN"),
+                n["rank"],
+            ]
+        )
+    entities_context = list_of_list_to_csv(entites_section_list)
+
+    relations_section_list=[["id","context"]]
+    for i,e in enumerate(use_relations):
+        relations_section_list.append([i,e])
+    relations_context=list_of_list_to_csv(relations_section_list)
+
+    text_units_section_list = [["id", "content"]]
+    for i, t in enumerate(use_text_units):
+        text_units_section_list.append([i, t["content"]])
+    text_units_context = list_of_list_to_csv(text_units_section_list)
+    
+    return entities_context,relations_context,text_units_context
+
+
+async def _find_most_related_text_unit_from_entities(
+    node_datas: list[dict],
+    query_param: QueryParam,
+    text_chunks_db: BaseKVStorage[TextChunkSchema],
+    knowledge_graph_inst: BaseGraphStorage,
+):
+    text_units = [
+        split_string_by_multi_markers(dp["source_id"], [GRAPH_FIELD_SEP])
+        for dp in node_datas
+    ]
+    edges = await asyncio.gather(
+        *[knowledge_graph_inst.get_node_edges(dp["entity_name"]) for dp in node_datas]
+    )
+    all_one_hop_nodes = set()
+    for this_edges in edges:
+        if not this_edges:
+            continue
+        all_one_hop_nodes.update([e[1] for e in this_edges])
+
+    all_one_hop_nodes = list(all_one_hop_nodes)
+    all_one_hop_nodes_data = await asyncio.gather(
+        *[knowledge_graph_inst.get_node(e) for e in all_one_hop_nodes]
+    )
+
+
+    all_one_hop_text_units_lookup = {
+        k: set(split_string_by_multi_markers(v["source_id"], [GRAPH_FIELD_SEP]))
+        for k, v in zip(all_one_hop_nodes, all_one_hop_nodes_data)
+        if v is not None and "source_id" in v  
+    }
+
+    all_text_units_lookup = {}
+    for index, (this_text_units, this_edges) in enumerate(zip(text_units, edges)):
+        for c_id in this_text_units:
+            if c_id not in all_text_units_lookup:
+                all_text_units_lookup[c_id] = {
+                    "data": await text_chunks_db.get_by_id(c_id),
+                    "order": index,
+                    "relation_counts": 0,
+                }
+
+            if this_edges:
+                for e in this_edges:
+                    if (
+                        e[1] in all_one_hop_text_units_lookup
+                        and c_id in all_one_hop_text_units_lookup[e[1]]
+                    ):
+                        all_text_units_lookup[c_id]["relation_counts"] += 1
+
+
+    all_text_units = [
+        {"id": k, **v}
+        for k, v in all_text_units_lookup.items()
+        if v is not None and v.get("data") is not None and "content" in v["data"]
+    ]
+
+    if not all_text_units:
+        logger.warning("No valid text units found")
+        return []
+
+    all_text_units = sorted(
+        all_text_units, key=lambda x: (x["order"], -x["relation_counts"])
+    )
+
+    all_text_units = truncate_list_by_token_size(
+        all_text_units,
+        key=lambda x: x["data"]["content"],
+        max_token_size=query_param.max_token_for_text_unit,
+    )
+
+    all_text_units = [t["data"] for t in all_text_units]
+    return all_text_units
+
+async def _get_edge_data(
+    keywords,
+    knowledge_graph_inst: BaseGraphStorage,
+    relationships_vdb: BaseVectorStorage,
+    text_chunks_db: BaseKVStorage[TextChunkSchema],
+    query_param: QueryParam,
+):
+    results = await relationships_vdb.query(keywords, top_k=query_param.top_k)
+
+    if not len(results):
+        return "", "", ""
+
+    edge_datas = await asyncio.gather(
+        *[knowledge_graph_inst.get_edge(r["src_id"], r["tgt_id"]) for r in results]
+    )
+
+    if not all([n is not None for n in edge_datas]):
+        logger.warning("Some edges are missing, maybe the storage is damaged")
+    edge_degree = await asyncio.gather(
+        *[knowledge_graph_inst.edge_degree(r["src_id"], r["tgt_id"]) for r in results]
+    )
+    edge_datas = [
+        {"src_id": k["src_id"], "tgt_id": k["tgt_id"], "rank": d, **v}
+        for k, v, d in zip(results, edge_datas, edge_degree)
+        if v is not None
+    ]
+    edge_datas = sorted(
+        edge_datas, key=lambda x: (x["rank"], x["weight"]), reverse=True
+    )
+    edge_datas = truncate_list_by_token_size(
+        edge_datas,
+        key=lambda x: x["description"],
+        max_token_size=query_param.max_token_for_global_context,
+    )
+
+    use_entities = await _find_most_related_entities_from_relationships(
+        edge_datas, query_param, knowledge_graph_inst
+    )
+    use_text_units = await _find_related_text_unit_from_relationships(
+        edge_datas, query_param, text_chunks_db, knowledge_graph_inst
+    )
+    logger.info(
+        f"Global query uses {len(use_entities)} entites, {len(edge_datas)} relations, {len(use_text_units)} text units"
+    )
+
+    relations_section_list = [
+        ["id", "source", "target", "description", "keywords", "weight", "rank"]
+    ]
+    for i, e in enumerate(edge_datas):
+        relations_section_list.append(
+            [
+                i,
+                e["src_id"],
+                e["tgt_id"],
+                e["description"],
+                e["keywords"],
+                e["weight"],
+                e["rank"],
+            ]
+        )
+    relations_context = list_of_list_to_csv(relations_section_list)
+
+    entites_section_list = [["id", "entity", "type", "description", "rank"]]
+    for i, n in enumerate(use_entities):
+        entites_section_list.append(
+            [
+                i,
+                n["entity_name"],
+                n.get("entity_type", "UNKNOWN"),
+                n.get("description", "UNKNOWN"),
+                n["rank"],
+            ]
+        )
+    entities_context = list_of_list_to_csv(entites_section_list)
+
+    text_units_section_list = [["id", "content"]]
+    for i, t in enumerate(use_text_units):
+        text_units_section_list.append([i, t["content"]])
+    text_units_context = list_of_list_to_csv(text_units_section_list)
+    return entities_context, relations_context, text_units_context
+
+
+async def _find_most_related_entities_from_relationships(
+    edge_datas: list[dict],
+    query_param: QueryParam,
+    knowledge_graph_inst: BaseGraphStorage,
+):
+    entity_names = []
+    seen = set()
+
+    for e in edge_datas:
+        if e["src_id"] not in seen:
+            entity_names.append(e["src_id"])
+            seen.add(e["src_id"])
+        if e["tgt_id"] not in seen:
+            entity_names.append(e["tgt_id"])
+            seen.add(e["tgt_id"])
+
+    node_datas = await asyncio.gather(
+        *[knowledge_graph_inst.get_node(entity_name) for entity_name in entity_names]
+    )
+
+    node_degrees = await asyncio.gather(
+        *[knowledge_graph_inst.node_degree(entity_name) for entity_name in entity_names]
+    )
+    node_datas = [
+        {**n, "entity_name": k, "rank": d}
+        for k, n, d in zip(entity_names, node_datas, node_degrees)
+    ]
+
+    node_datas = truncate_list_by_token_size(
+        node_datas,
+        key=lambda x: x["description"],
+        max_token_size=query_param.max_token_for_local_context,
+    )
+
+    return node_datas
+
+
+async def _find_related_text_unit_from_relationships(
+    edge_datas: list[dict],
+    query_param: QueryParam,
+    text_chunks_db: BaseKVStorage[TextChunkSchema],
+    knowledge_graph_inst: BaseGraphStorage,
+):
+    text_units = [
+        split_string_by_multi_markers(dp["source_id"], [GRAPH_FIELD_SEP])
+        for dp in edge_datas
+    ]
+    all_text_units_lookup = {}
+
+    for index, unit_list in enumerate(text_units):
+        for c_id in unit_list:
+            if c_id not in all_text_units_lookup:
+                chunk_data = await text_chunks_db.get_by_id(c_id)
+
+                if chunk_data is not None and "content" in chunk_data:
+                    all_text_units_lookup[c_id] = {
+                        "data": chunk_data,
+                        "order": index,
+                    }
+
+    if not all_text_units_lookup:
+        logger.warning("No valid text chunks found")
+        return []
+
+    all_text_units = [{"id": k, **v} for k, v in all_text_units_lookup.items()]
+    all_text_units = sorted(all_text_units, key=lambda x: x["order"])
+
+
+    valid_text_units = [
+        t for t in all_text_units if t["data"] is not None and "content" in t["data"]
+    ]
+
+    if not valid_text_units:
+        logger.warning("No valid text chunks after filtering")
+        return []
+
+    truncated_text_units = truncate_list_by_token_size(
+        valid_text_units,
+        key=lambda x: x["data"]["content"],
+        max_token_size=query_param.max_token_for_text_unit,
+    )
+
+    all_text_units: list[TextChunkSchema] = [t["data"] for t in truncated_text_units]
+
+    return all_text_units
+
+
+def combine_contexts(entities, relationships, sources):
+
+    hl_entities, ll_entities = entities[0], entities[1]
+    hl_relationships, ll_relationships = relationships[0], relationships[1]
+    hl_sources, ll_sources = sources[0], sources[1]
+
+    combined_entities = process_combine_contexts(hl_entities, ll_entities)
+
+    combined_relationships = process_combine_contexts(
+        hl_relationships, ll_relationships
+    )
+
+    combined_sources = process_combine_contexts(hl_sources, ll_sources)
+
+    return combined_entities, combined_relationships, combined_sources
+
+
+import networkx as nx
+from collections import defaultdict
+async def find_paths_and_edges_with_stats(graph, target_nodes):
+
+    result = defaultdict(lambda: {"paths": [], "edges": set()})
+    path_stats = {"1-hop": 0, "2-hop": 0, "3-hop": 0}   
+    one_hop_paths = []
+    two_hop_paths = []
+    three_hop_paths = []
+
+    async def dfs(current, target, path, depth):
+
+        if depth > 3: 
+            return
+        if current == target: 
+            result[(path[0], target)]["paths"].append(list(path))
+            for u, v in zip(path[:-1], path[1:]):
+                result[(path[0], target)]["edges"].add(tuple(sorted((u, v))))
+            if depth == 1:
+                path_stats["1-hop"] += 1
+                one_hop_paths.append(list(path))
+            elif depth == 2:
+                path_stats["2-hop"] += 1
+                two_hop_paths.append(list(path))
+            elif depth == 3:
+                path_stats["3-hop"] += 1
+                three_hop_paths.append(list(path))
+            return
+        neighbors = graph.neighbors(current) 
+        for neighbor in neighbors:
+            if neighbor not in path:  
+                await dfs(neighbor, target, path + [neighbor], depth + 1)
+
+    for node1 in target_nodes:
+        for node2 in target_nodes:
+            if node1 != node2:
+                await dfs(node1, node2, [node1], 0)
+
+    for key in result:
+        result[key]["edges"] = list(result[key]["edges"])
+
+    return dict(result), path_stats , one_hop_paths, two_hop_paths, three_hop_paths
+def bfs_weighted_paths(G, path, source, target, threshold, alpha):
+    results = [] 
+    edge_weights = defaultdict(float)  
+    node = source
+    follow_dict = {}
+
+    for p in path:
+        for i in range(len(p) - 1):  
+            current = p[i]
+            next_num = p[i + 1]
+
+            if current in follow_dict:
+                follow_dict[current].add(next_num)
+            else:
+                follow_dict[current] = {next_num}
+
+    for neighbor in follow_dict[node]:
+        edge_weights[(node, neighbor)] += 1/len(follow_dict[node])
+
+        if neighbor == target:
+            results.append(([node, neighbor]))
+            continue
+        
+        if edge_weights[(node, neighbor)] > threshold:
+
+            for second_neighbor in follow_dict[neighbor]:
+                weight = edge_weights[(node, neighbor)] * alpha / len(follow_dict[neighbor])
+                edge_weights[(neighbor, second_neighbor)] += weight
+
+                if second_neighbor == target:
+                    results.append(([node, neighbor, second_neighbor]))
+                    continue
+
+                if edge_weights[(neighbor, second_neighbor)] > threshold:    
+
+                    for third_neighbor in follow_dict[second_neighbor]:
+                        weight = edge_weights[(neighbor, second_neighbor)] * alpha / len(follow_dict[second_neighbor]) 
+                        edge_weights[(second_neighbor, third_neighbor)] += weight
+
+                        if third_neighbor == target :
+                            results.append(([node, neighbor, second_neighbor, third_neighbor]))
+                            continue
+    path_weights = []
+    for p in path:
+        path_weight = 0
+        for i in range(len(p) - 1):
+            edge = (p[i], p[i + 1])
+            path_weight += edge_weights.get(edge, 0)  
+        path_weights.append(path_weight/(len(p)-1))
+
+    combined = [(p, w) for p, w in zip(path, path_weights)]
+
+    return combined
+async def _find_most_related_edges_from_entities3(
+    node_datas: list[dict],
+    query_param: QueryParam,
+    knowledge_graph_inst: BaseGraphStorage,
+):  
+
+    G = nx.Graph()
+    edges = await knowledge_graph_inst.edges()
+    nodes = await knowledge_graph_inst.nodes()
+
+    for u, v in edges:
+        G.add_edge(u, v) 
+    G.add_nodes_from(nodes)
+    source_nodes = [dp["entity_name"] for dp in node_datas]
+    result, path_stats, one_hop_paths, two_hop_paths, three_hop_paths = await find_paths_and_edges_with_stats(G, source_nodes)
+
+
+    threshold = 0.3
+    alpha = 0.8 
+    all_results = []
+    
+    for node1 in source_nodes: 
+        for node2 in source_nodes: 
+            if node1 != node2: 
+                if (node1, node2) in result:
+                    sub_G = nx.Graph()
+                    paths = result[(node1,node2)]['paths']
+                    edges = result[(node1,node2)]['edges']
+                    sub_G.add_edges_from(edges)
+                    results = bfs_weighted_paths(G, paths, node1, node2, threshold, alpha)
+                    all_results+= results
+    all_results = sorted(all_results, key=lambda x: x[1], reverse=True)
+    seen = set()
+    result_edge = []
+    for edge, weight in all_results:
+        sorted_edge = tuple(sorted(edge))
+        if sorted_edge not in seen:
+            seen.add(sorted_edge)  
+            result_edge.append((edge, weight))  
+
+    
+    length_1 = int(len(one_hop_paths)/2)
+    length_2 = int(len(two_hop_paths)/2) 
+    length_3 = int(len(three_hop_paths)/2) 
+    results = []
+    if one_hop_paths!=[]:
+        results = one_hop_paths[0:length_1]
+    if two_hop_paths!=[]:
+        results = results + two_hop_paths[0:length_2]
+    if three_hop_paths!=[]:
+        results  =results + three_hop_paths[0:length_3]
+
+    length = len(results)
+    total_edges = 15
+    if length < total_edges:
+        total_edges = length
+    sort_result = []
+    if result_edge:
+        if len(result_edge)>total_edges:
+            sort_result = result_edge[0:total_edges]
+        else : 
+            sort_result = result_edge
+    final_result = []
+    for edge, weight in sort_result:
+        final_result.append(edge)
+
+    relationship = []
+
+    for path in final_result:
+        if len(path) == 4:
+            s_name,b1_name,b2_name,t_name = path[0],path[1],path[2],path[3]
+            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])
+            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,"边丢失")
+                if edge0==None:
+                    print("edge0丢失")
+                if edge1==None:
+                    print("edge1丢失")
+                if edge2==None:
+                    print("edge2丢失")
+                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+"."
+            e3 = "through edge ("+edge2["keywords"]+") to connect to "+b2_name+" and "+t_name+"."
+            s = await knowledge_graph_inst.get_node(s_name)
+            s = "The entity "+s_name+" is a "+s["entity_type"]+" with the description("+s["description"]+")"
+            b1 = await knowledge_graph_inst.get_node(b1_name)
+            b1 = "The entity "+b1_name+" is a "+b1["entity_type"]+" with the description("+b1["description"]+")"
+            b2 = await knowledge_graph_inst.get_node(b2_name)
+            b2 = "The entity "+b2_name+" is a "+b2["entity_type"]+" with the description("+b2["description"]+")"
+            t = await knowledge_graph_inst.get_node(t_name)
+            t = "The entity "+t_name+" is a "+t["entity_type"]+" with the description("+t["description"]+")"
+            relationship.append([s+e1+b1+"and"+b1+e2+b2+"and"+b2+e3+t])
+        elif len(path) == 3:
+            s_name,b_name,t_name = path[0],path[1],path[2]
+            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,"边丢失")
+                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+"."
+            s = await knowledge_graph_inst.get_node(s_name)
+            s = "The entity "+s_name+" is a "+s["entity_type"]+" with the description("+s["description"]+")"
+            b = await knowledge_graph_inst.get_node(b_name)
+            b = "The entity "+b_name+" is a "+b["entity_type"]+" with the description("+b["description"]+")"
+            t = await knowledge_graph_inst.get_node(t_name)
+            t = "The entity "+t_name+" is a "+t["entity_type"]+" with the description("+t["description"]+")"
+            relationship.append([s+e1+b+"and"+b+e2+t])
+        elif len(path) == 2:
+            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,"边丢失")
+                continue
+            e = "through edge("+edge0["keywords"]+") to connect to "+s_name+" and "+t_name+"."
+            s = await knowledge_graph_inst.get_node(s_name)
+            s = "The entity "+s_name+" is a "+s["entity_type"]+" with the description("+s["description"]+")"
+            t = await knowledge_graph_inst.get_node(t_name)
+            t = "The entity "+t_name+" is a "+t["entity_type"]+" with the description("+t["description"]+")"
+            relationship.append([s+e+t])
+
+
+    relationship = truncate_list_by_token_size(
+          relationship, 
+          key=lambda x: x[0],
+          max_token_size=query_param.max_token_for_local_context,
+    )
+
+    reversed_relationship = relationship[::-1]
+    return reversed_relationship

PathRAG/prompt.py

@@ -0,0 +1,286 @@
+GRAPH_FIELD_SEP = "<SEP>"
+
+PROMPTS = {}
+
+PROMPTS["DEFAULT_LANGUAGE"] = "English"
+PROMPTS["DEFAULT_TUPLE_DELIMITER"] = "<|>"
+PROMPTS["DEFAULT_RECORD_DELIMITER"] = "##"
+PROMPTS["DEFAULT_COMPLETION_DELIMITER"] = "<|COMPLETE|>"
+PROMPTS["process_tickers"] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+
+PROMPTS["DEFAULT_ENTITY_TYPES"] = ["organization", "person", "geo", "event", "category"]
+
+PROMPTS["entity_extraction"] = """-Goal-
+Given a text document that is potentially relevant to this activity and a list of entity types, identify all entities of those types from the text and all relationships among the identified entities.
+Use {language} as output language.
+
+-Steps-
+1. Identify all entities. For each identified entity, extract the following information:
+- entity_name: Name of the entity, use same language as input text. If English, capitalized the name.
+- entity_type: One of the following types: [{entity_types}]
+- entity_description: Comprehensive description of the entity's attributes and activities
+Format each entity as ("entity"{tuple_delimiter}<entity_name>{tuple_delimiter}<entity_type>{tuple_delimiter}<entity_description>)
+
+2. From the entities identified in step 1, identify all pairs of (source_entity, target_entity) that are *clearly related* to each other.
+For each pair of related entities, extract the following information:
+- source_entity: name of the source entity, as identified in step 1
+- target_entity: name of the target entity, as identified in step 1
+- relationship_description: explanation as to why you think the source entity and the target entity are related to each other
+- relationship_strength: a numeric score indicating strength of the relationship between the source entity and target entity
+- relationship_keywords: one or more high-level key words that summarize the overarching nature of the relationship, focusing on concepts or themes rather than specific details
+Format each relationship as ("relationship"{tuple_delimiter}<source_entity>{tuple_delimiter}<target_entity>{tuple_delimiter}<relationship_description>{tuple_delimiter}<relationship_keywords>{tuple_delimiter}<relationship_strength>)
+
+3. Identify high-level key words that summarize the main concepts, themes, or topics of the entire text. These should capture the overarching ideas present in the document.
+Format the content-level key words as ("content_keywords"{tuple_delimiter}<high_level_keywords>)
+
+4. Return output in {language} as a single list of all the entities and relationships identified in steps 1 and 2. Use **{record_delimiter}** as the list delimiter.
+
+5. When finished, output {completion_delimiter}
+
+######################
+-Examples-
+######################
+{examples}
+
+#############################
+-Real Data-
+######################
+Entity_types: {entity_types}
+Text: {input_text}
+######################
+Output:
+"""
+
+PROMPTS["entity_extraction_examples"] = [
+    """Example 1:
+
+Entity_types: [person, technology, mission, organization, location]
+Text:
+while Alex clenched his jaw, the buzz of frustration dull against the backdrop of Taylor's authoritarian certainty. It was this competitive undercurrent that kept him alert, the sense that his and Jordan's shared commitment to discovery was an unspoken rebellion against Cruz's narrowing vision of control and order.
+
+Then Taylor did something unexpected. They paused beside Jordan and, for a moment, observed the device with something akin to reverence. “If this tech can be understood..." Taylor said, their voice quieter, "It could change the game for us. For all of us.”
+
+The underlying dismissal earlier seemed to falter, replaced by a glimpse of reluctant respect for the gravity of what lay in their hands. Jordan looked up, and for a fleeting heartbeat, their eyes locked with Taylor's, a wordless clash of wills softening into an uneasy truce.
+
+It was a small transformation, barely perceptible, but one that Alex noted with an inward nod. They had all been brought here by different paths
+################
+Output:
+("entity"{tuple_delimiter}"Alex"{tuple_delimiter}"person"{tuple_delimiter}"Alex is a character who experiences frustration and is observant of the dynamics among other characters."){record_delimiter}
+("entity"{tuple_delimiter}"Taylor"{tuple_delimiter}"person"{tuple_delimiter}"Taylor is portrayed with authoritarian certainty and shows a moment of reverence towards a device, indicating a change in perspective."){record_delimiter}
+("entity"{tuple_delimiter}"Jordan"{tuple_delimiter}"person"{tuple_delimiter}"Jordan shares a commitment to discovery and has a significant interaction with Taylor regarding a device."){record_delimiter}
+("entity"{tuple_delimiter}"Cruz"{tuple_delimiter}"person"{tuple_delimiter}"Cruz is associated with a vision of control and order, influencing the dynamics among other characters."){record_delimiter}
+("entity"{tuple_delimiter}"The Device"{tuple_delimiter}"technology"{tuple_delimiter}"The Device is central to the story, with potential game-changing implications, and is revered by Taylor."){record_delimiter}
+("relationship"{tuple_delimiter}"Alex"{tuple_delimiter}"Taylor"{tuple_delimiter}"Alex is affected by Taylor's authoritarian certainty and observes changes in Taylor's attitude towards the device."{tuple_delimiter}"power dynamics, perspective shift"{tuple_delimiter}7){record_delimiter}
+("relationship"{tuple_delimiter}"Alex"{tuple_delimiter}"Jordan"{tuple_delimiter}"Alex and Jordan share a commitment to discovery, which contrasts with Cruz's vision."{tuple_delimiter}"shared goals, rebellion"{tuple_delimiter}6){record_delimiter}
+("relationship"{tuple_delimiter}"Taylor"{tuple_delimiter}"Jordan"{tuple_delimiter}"Taylor and Jordan interact directly regarding the device, leading to a moment of mutual respect and an uneasy truce."{tuple_delimiter}"conflict resolution, mutual respect"{tuple_delimiter}8){record_delimiter}
+("relationship"{tuple_delimiter}"Jordan"{tuple_delimiter}"Cruz"{tuple_delimiter}"Jordan's commitment to discovery is in rebellion against Cruz's vision of control and order."{tuple_delimiter}"ideological conflict, rebellion"{tuple_delimiter}5){record_delimiter}
+("relationship"{tuple_delimiter}"Taylor"{tuple_delimiter}"The Device"{tuple_delimiter}"Taylor shows reverence towards the device, indicating its importance and potential impact."{tuple_delimiter}"reverence, technological significance"{tuple_delimiter}9){record_delimiter}
+("content_keywords"{tuple_delimiter}"power dynamics, ideological conflict, discovery, rebellion"){completion_delimiter}
+#############################""",
+    """Example 2:
+
+Entity_types: [person, technology, mission, organization, location]
+Text:
+They were no longer mere operatives; they had become guardians of a threshold, keepers of a message from a realm beyond stars and stripes. This elevation in their mission could not be shackled by regulations and established protocols—it demanded a new perspective, a new resolve.
+
+Tension threaded through the dialogue of beeps and static as communications with Washington buzzed in the background. The team stood, a portentous air enveloping them. It was clear that the decisions they made in the ensuing hours could redefine humanity's place in the cosmos or condemn them to ignorance and potential peril.
+
+Their connection to the stars solidified, the group moved to address the crystallizing warning, shifting from passive recipients to active participants. Mercer's latter instincts gained precedence— the team's mandate had evolved, no longer solely to observe and report but to interact and prepare. A metamorphosis had begun, and Operation: Dulce hummed with the newfound frequency of their daring, a tone set not by the earthly
+#############
+Output:
+("entity"{tuple_delimiter}"Washington"{tuple_delimiter}"location"{tuple_delimiter}"Washington is a location where communications are being received, indicating its importance in the decision-making process."){record_delimiter}
+("entity"{tuple_delimiter}"Operation: Dulce"{tuple_delimiter}"mission"{tuple_delimiter}"Operation: Dulce is described as a mission that has evolved to interact and prepare, indicating a significant shift in objectives and activities."){record_delimiter}
+("entity"{tuple_delimiter}"The team"{tuple_delimiter}"organization"{tuple_delimiter}"The team is portrayed as a group of individuals who have transitioned from passive observers to active participants in a mission, showing a dynamic change in their role."){record_delimiter}
+("relationship"{tuple_delimiter}"The team"{tuple_delimiter}"Washington"{tuple_delimiter}"The team receives communications from Washington, which influences their decision-making process."{tuple_delimiter}"decision-making, external influence"{tuple_delimiter}7){record_delimiter}
+("relationship"{tuple_delimiter}"The team"{tuple_delimiter}"Operation: Dulce"{tuple_delimiter}"The team is directly involved in Operation: Dulce, executing its evolved objectives and activities."{tuple_delimiter}"mission evolution, active participation"{tuple_delimiter}9){completion_delimiter}
+("content_keywords"{tuple_delimiter}"mission evolution, decision-making, active participation, cosmic significance"){completion_delimiter}
+#############################""",
+    """Example 3:
+
+Entity_types: [person, role, technology, organization, event, location, concept]
+Text:
+their voice slicing through the buzz of activity. "Control may be an illusion when facing an intelligence that literally writes its own rules," they stated stoically, casting a watchful eye over the flurry of data.
+
+"It's like it's learning to communicate," offered Sam Rivera from a nearby interface, their youthful energy boding a mix of awe and anxiety. "This gives talking to strangers' a whole new meaning."
+
+Alex surveyed his team—each face a study in concentration, determination, and not a small measure of trepidation. "This might well be our first contact," he acknowledged, "And we need to be ready for whatever answers back."
+
+Together, they stood on the edge of the unknown, forging humanity's response to a message from the heavens. The ensuing silence was palpable—a collective introspection about their role in this grand cosmic play, one that could rewrite human history.
+
+The encrypted dialogue continued to unfold, its intricate patterns showing an almost uncanny anticipation
+#############
+Output:
+("entity"{tuple_delimiter}"Sam Rivera"{tuple_delimiter}"person"{tuple_delimiter}"Sam Rivera is a member of a team working on communicating with an unknown intelligence, showing a mix of awe and anxiety."){record_delimiter}
+("entity"{tuple_delimiter}"Alex"{tuple_delimiter}"person"{tuple_delimiter}"Alex is the leader of a team attempting first contact with an unknown intelligence, acknowledging the significance of their task."){record_delimiter}
+("entity"{tuple_delimiter}"Control"{tuple_delimiter}"concept"{tuple_delimiter}"Control refers to the ability to manage or govern, which is challenged by an intelligence that writes its own rules."){record_delimiter}
+("entity"{tuple_delimiter}"Intelligence"{tuple_delimiter}"concept"{tuple_delimiter}"Intelligence here refers to an unknown entity capable of writing its own rules and learning to communicate."){record_delimiter}
+("entity"{tuple_delimiter}"First Contact"{tuple_delimiter}"event"{tuple_delimiter}"First Contact is the potential initial communication between humanity and an unknown intelligence."){record_delimiter}
+("entity"{tuple_delimiter}"Humanity's Response"{tuple_delimiter}"event"{tuple_delimiter}"Humanity's Response is the collective action taken by Alex's team in response to a message from an unknown intelligence."){record_delimiter}
+("relationship"{tuple_delimiter}"Sam Rivera"{tuple_delimiter}"Intelligence"{tuple_delimiter}"Sam Rivera is directly involved in the process of learning to communicate with the unknown intelligence."{tuple_delimiter}"communication, learning process"{tuple_delimiter}9){record_delimiter}
+("relationship"{tuple_delimiter}"Alex"{tuple_delimiter}"First Contact"{tuple_delimiter}"Alex leads the team that might be making the First Contact with the unknown intelligence."{tuple_delimiter}"leadership, exploration"{tuple_delimiter}10){record_delimiter}
+("relationship"{tuple_delimiter}"Alex"{tuple_delimiter}"Humanity's Response"{tuple_delimiter}"Alex and his team are the key figures in Humanity's Response to the unknown intelligence."{tuple_delimiter}"collective action, cosmic significance"{tuple_delimiter}8){record_delimiter}
+("relationship"{tuple_delimiter}"Control"{tuple_delimiter}"Intelligence"{tuple_delimiter}"The concept of Control is challenged by the Intelligence that writes its own rules."{tuple_delimiter}"power dynamics, autonomy"{tuple_delimiter}7){record_delimiter}
+("content_keywords"{tuple_delimiter}"first contact, control, communication, cosmic significance"){completion_delimiter}
+#############################""",
+]
+
+PROMPTS[
+    "summarize_entity_descriptions"
+] = """You are a helpful assistant responsible for generating a comprehensive summary of the data provided below.
+Given one or two entities, and a list of descriptions, all related to the same entity or group of entities.
+Please concatenate all of these into a single, comprehensive description. Make sure to include information collected from all the descriptions.
+If the provided descriptions are contradictory, please resolve the contradictions and provide a single, coherent summary.
+Make sure it is written in third person, and include the entity names so we the have full context.
+Use {language} as output language.
+
+#######
+-Data-
+Entities: {entity_name}
+Description List: {description_list}
+#######
+Output:
+"""
+
+PROMPTS[
+    "entiti_continue_extraction"
+] = """MANY entities were missed in the last extraction.  Add them below using the same format:
+"""
+
+PROMPTS[
+    "entiti_if_loop_extraction"
+] = """It appears some entities may have still been missed.  Answer YES | NO if there are still entities that need to be added.
+"""
+
+PROMPTS["fail_response"] = "Sorry, I'm not able to provide an answer to that question."
+
+PROMPTS["rag_response"] = """---Role---
+
+You are a helpful assistant responding to questions about data in the tables provided.
+
+
+---Goal---
+
+Generate a response of the target length and format that responds to the user's question, summarizing all information in the input data tables appropriate for the response length and format, and incorporating any relevant general knowledge.
+If you don't know the answer, just say so. Do not make anything up.
+Do not include information where the supporting evidence for it is not provided.
+
+---Target response length and format---
+
+{response_type}
+
+---Data tables---
+
+{context_data}
+
+Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+"""
+
+PROMPTS["keywords_extraction"] = """---Role---
+
+You are a helpful assistant tasked with identifying both high-level and low-level keywords in the user's query.
+
+---Goal---
+
+Given the query, list both high-level and low-level keywords. High-level keywords focus on overarching concepts or themes, while low-level keywords focus on specific entities, details, or concrete terms.
+
+---Instructions---
+
+- Output the keywords in JSON format.
+- The JSON should have two keys:
+  - "high_level_keywords" for overarching concepts or themes.
+  - "low_level_keywords" for specific entities or details.
+
+######################
+-Examples-
+######################
+{examples}
+
+#############################
+-Real Data-
+######################
+Query: {query}
+######################
+The `Output` should be human text, not unicode characters. Keep the same language as `Query`.
+Output:
+
+"""
+
+PROMPTS["keywords_extraction_examples"] = [
+    """Example 1:
+
+Query: "How does international trade influence global economic stability?"
+################
+Output:
+{{
+  "high_level_keywords": ["International trade", "Global economic stability", "Economic impact"],
+  "low_level_keywords": ["Trade agreements", "Tariffs", "Currency exchange", "Imports", "Exports"]
+}}
+#############################""",
+    """Example 2:
+
+Query: "What are the environmental consequences of deforestation on biodiversity?"
+################
+Output:
+{{
+  "high_level_keywords": ["Environmental consequences", "Deforestation", "Biodiversity loss"],
+  "low_level_keywords": ["Species extinction", "Habitat destruction", "Carbon emissions", "Rainforest", "Ecosystem"]
+}}
+#############################""",
+    """Example 3:
+
+Query: "What is the role of education in reducing poverty?"
+################
+Output:
+{{
+  "high_level_keywords": ["Education", "Poverty reduction", "Socioeconomic development"],
+  "low_level_keywords": ["School access", "Literacy rates", "Job training", "Income inequality"]
+}}
+#############################""",
+]
+
+
+PROMPTS["naive_rag_response"] = """---Role---
+
+You are a helpful assistant responding to questions about documents provided.
+
+
+---Goal---
+
+Generate a response of the target length and format that responds to the user's question, summarizing all information in the input data tables appropriate for the response length and format, and incorporating any relevant general knowledge.
+If you don't know the answer, just say so. Do not make anything up.
+Do not include information where the supporting evidence for it is not provided.
+
+---Target response length and format---
+
+{response_type}
+
+---Documents---
+
+{content_data}
+
+Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+"""
+
+PROMPTS[
+    "similarity_check"
+] = """Please analyze the similarity between these two questions:
+
+Question 1: {original_prompt}
+Question 2: {cached_prompt}
+
+Please evaluate the following two points and provide a similarity score between 0 and 1 directly:
+1. Whether these two questions are semantically similar
+2. Whether the answer to Question 2 can be used to answer Question 1
+Similarity score criteria:
+0: Completely unrelated or answer cannot be reused, including but not limited to:
+   - The questions have different topics
+   - The locations mentioned in the questions are different
+   - The times mentioned in the questions are different
+   - The specific individuals mentioned in the questions are different
+   - The specific events mentioned in the questions are different
+   - The background information in the questions is different
+   - The key conditions in the questions are different
+1: Identical and answer can be directly reused
+0.5: Partially related and answer needs modification to be used
+Return only a number between 0-1, without any additional content.
+"""

PathRAG/storage.py

@@ -0,0 +1,341 @@
+import asyncio
+import html
+import os
+from tqdm.asyncio import tqdm as tqdm_async
+from dataclasses import dataclass
+from typing import Any, Union, cast
+import networkx as nx
+import numpy as np
+from nano_vectordb import NanoVectorDB
+
+from .utils import (
+    logger,
+    load_json,
+    write_json,
+    compute_mdhash_id,
+)
+
+from .base import (
+    BaseGraphStorage,
+    BaseKVStorage,
+    BaseVectorStorage,
+)
+
+
+@dataclass
+class JsonKVStorage(BaseKVStorage):
+    def __post_init__(self):
+        working_dir = self.global_config["working_dir"]
+        self._file_name = os.path.join(working_dir, f"kv_store_{self.namespace}.json")
+        self._data = load_json(self._file_name) or {}
+        logger.info(f"Load KV {self.namespace} with {len(self._data)} data")
+
+    async def all_keys(self) -> list[str]:
+        return list(self._data.keys())
+
+    async def index_done_callback(self):
+        write_json(self._data, self._file_name)
+
+    async def get_by_id(self, id):
+        return self._data.get(id, None)
+
+    async def get_by_ids(self, ids, fields=None):
+        if fields is None:
+            return [self._data.get(id, None) for id in ids]
+        return [
+            (
+                {k: v for k, v in self._data[id].items() if k in fields}
+                if self._data.get(id, None)
+                else None
+            )
+            for id in ids
+        ]
+
+    async def filter_keys(self, data: list[str]) -> set[str]:
+        return set([s for s in data if s not in self._data])
+
+    async def upsert(self, data: dict[str, dict]):
+        left_data = {k: v for k, v in data.items() if k not in self._data}
+        self._data.update(left_data)
+        return left_data
+
+    async def drop(self):
+        self._data = {}
+
+
+@dataclass
+class NanoVectorDBStorage(BaseVectorStorage):
+    cosine_better_than_threshold: float = 0.2
+
+    def __post_init__(self):
+        self._client_file_name = os.path.join(
+            self.global_config["working_dir"], f"vdb_{self.namespace}.json"
+        )
+        self._max_batch_size = self.global_config["embedding_batch_num"]
+        self._client = NanoVectorDB(
+            self.embedding_func.embedding_dim, storage_file=self._client_file_name
+        )
+        self.cosine_better_than_threshold = self.global_config.get(
+            "cosine_better_than_threshold", self.cosine_better_than_threshold
+        )
+
+    async def upsert(self, data: dict[str, dict]):
+        logger.info(f"Inserting {len(data)} vectors to {self.namespace}")
+        if not len(data):
+            logger.warning("You insert an empty data to vector DB")
+            return []
+        list_data = [
+            {
+                "__id__": k,
+                **{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
+            }
+            for k, v in data.items()
+        ]
+        contents = [v["content"] for v in data.values()]
+        batches = [
+            contents[i : i + self._max_batch_size]
+            for i in range(0, len(contents), self._max_batch_size)
+        ]
+
+        async def wrapped_task(batch):
+            result = await self.embedding_func(batch)
+            pbar.update(1)
+            return result
+
+        embedding_tasks = [wrapped_task(batch) for batch in batches]
+        pbar = tqdm_async(
+            total=len(embedding_tasks), desc="Generating embeddings", unit="batch"
+        )
+        embeddings_list = await asyncio.gather(*embedding_tasks)
+
+        embeddings = np.concatenate(embeddings_list)
+        if len(embeddings) == len(list_data):
+            for i, d in enumerate(list_data):
+                d["__vector__"] = embeddings[i]
+            results = self._client.upsert(datas=list_data)
+            return results
+        else:
+            # sometimes the embedding is not returned correctly. just log it.
+            logger.error(
+                f"embedding is not 1-1 with data, {len(embeddings)} != {len(list_data)}"
+            )
+
+    async def query(self, query: str, top_k=5):
+        embedding = await self.embedding_func([query])
+        embedding = embedding[0]
+        results = self._client.query(
+            query=embedding,
+            top_k=top_k,
+            better_than_threshold=self.cosine_better_than_threshold,
+        )
+        results = [
+            {**dp, "id": dp["__id__"], "distance": dp["__metrics__"]} for dp in results
+        ]
+        return results
+
+    @property
+    def client_storage(self):
+        return getattr(self._client, "_NanoVectorDB__storage")
+
+    async def delete_entity(self, entity_name: str):
+        try:
+            entity_id = [compute_mdhash_id(entity_name, prefix="ent-")]
+
+            if self._client.get(entity_id):
+                self._client.delete(entity_id)
+                logger.info(f"Entity {entity_name} have been deleted.")
+            else:
+                logger.info(f"No entity found with name {entity_name}.")
+        except Exception as e:
+            logger.error(f"Error while deleting entity {entity_name}: {e}")
+
+    async def delete_relation(self, entity_name: str):
+        try:
+            relations = [
+                dp
+                for dp in self.client_storage["data"]
+                if dp["src_id"] == entity_name or dp["tgt_id"] == entity_name
+            ]
+            ids_to_delete = [relation["__id__"] for relation in relations]
+
+            if ids_to_delete:
+                self._client.delete(ids_to_delete)
+                logger.info(
+                    f"All relations related to entity {entity_name} have been deleted."
+                )
+            else:
+                logger.info(f"No relations found for entity {entity_name}.")
+        except Exception as e:
+            logger.error(
+                f"Error while deleting relations for entity {entity_name}: {e}"
+            )
+
+    async def index_done_callback(self):
+        self._client.save()
+
+
+@dataclass
+class NetworkXStorage(BaseGraphStorage):
+    @staticmethod
+    def load_nx_graph(file_name) -> nx.DiGraph:
+        if os.path.exists(file_name):
+            return nx.read_graphml(file_name)
+        return None
+    # def load_nx_graph(file_name) -> nx.Graph:
+    #     if os.path.exists(file_name):
+    #         return nx.read_graphml(file_name)
+    #     return None
+
+    @staticmethod
+    def write_nx_graph(graph: nx.DiGraph, file_name):
+        logger.info(
+            f"Writing graph with {graph.number_of_nodes()} nodes, {graph.number_of_edges()} edges"
+        )
+        nx.write_graphml(graph, file_name)
+
+    @staticmethod
+    def stable_largest_connected_component(graph: nx.Graph) -> nx.Graph:
+        """Refer to https://github.com/microsoft/graphrag/index/graph/utils/stable_lcc.py
+        Return the largest connected component of the graph, with nodes and edges sorted in a stable way.
+        """
+        from graspologic.utils import largest_connected_component
+
+        graph = graph.copy()
+        graph = cast(nx.Graph, largest_connected_component(graph))
+        node_mapping = {
+            node: html.unescape(node.upper().strip()) for node in graph.nodes()
+        }  # type: ignore
+        graph = nx.relabel_nodes(graph, node_mapping)
+        return NetworkXStorage._stabilize_graph(graph)
+
+    @staticmethod
+    def _stabilize_graph(graph: nx.Graph) -> nx.Graph:
+        """Refer to https://github.com/microsoft/graphrag/index/graph/utils/stable_lcc.py
+        Ensure an undirected graph with the same relationships will always be read the same way.
+        """
+        fixed_graph = nx.DiGraph() if graph.is_directed() else nx.Graph()
+
+        sorted_nodes = graph.nodes(data=True)
+        sorted_nodes = sorted(sorted_nodes, key=lambda x: x[0])
+
+        fixed_graph.add_nodes_from(sorted_nodes)
+        edges = list(graph.edges(data=True))
+
+        if not graph.is_directed():
+
+            def _sort_source_target(edge):
+                source, target, edge_data = edge
+                if source > target:
+                    temp = source
+                    source = target
+                    target = temp
+                return source, target, edge_data
+
+            edges = [_sort_source_target(edge) for edge in edges]
+
+        def _get_edge_key(source: Any, target: Any) -> str:
+            return f"{source} -> {target}"
+
+        edges = sorted(edges, key=lambda x: _get_edge_key(x[0], x[1]))
+
+        fixed_graph.add_edges_from(edges)
+        return fixed_graph
+
+    def __post_init__(self):
+        self._graphml_xml_file = os.path.join(
+            self.global_config["working_dir"], f"graph_{self.namespace}.graphml"
+        )
+        preloaded_graph = NetworkXStorage.load_nx_graph(self._graphml_xml_file)
+        if preloaded_graph is not None:
+            logger.info(
+                f"Loaded graph from {self._graphml_xml_file} with {preloaded_graph.number_of_nodes()} nodes, {preloaded_graph.number_of_edges()} edges"
+            )
+        self._graph = preloaded_graph or nx.DiGraph()
+        self._node_embed_algorithms = {
+            "node2vec": self._node2vec_embed,
+        }
+
+    async def index_done_callback(self):
+        NetworkXStorage.write_nx_graph(self._graph, self._graphml_xml_file)
+
+    async def has_node(self, node_id: str) -> bool:
+        return self._graph.has_node(node_id)
+
+    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
+        return self._graph.has_edge(source_node_id, target_node_id)
+
+    async def get_node(self, node_id: str) -> Union[dict, None]:
+        return self._graph.nodes.get(node_id)
+
+    async def node_degree(self, node_id: str) -> int:
+        return self._graph.degree(node_id)
+
+    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        return self._graph.degree(src_id) + self._graph.degree(tgt_id)
+
+    async def get_edge(
+        self, source_node_id: str, target_node_id: str
+    ) -> Union[dict, None]:
+        return self._graph.edges.get((source_node_id, target_node_id))
+
+    async def get_node_edges(self, source_node_id: str):
+        if self._graph.has_node(source_node_id):
+            return list(self._graph.edges(source_node_id))
+        return None
+    async def get_node_in_edges(self, source_node_id: str):
+        if self._graph.has_node(source_node_id):
+            return list(self._graph.in_edges(source_node_id))
+        return None
+    async def get_node_out_edges(self, source_node_id: str):
+        if self._graph.has_node(source_node_id):
+            return list(self._graph.out_edges(source_node_id))
+        return None
+    
+    async def get_pagerank(self,source_node_id:str):
+        pagerank_list=nx.pagerank(self._graph)
+        if source_node_id in pagerank_list:
+            return pagerank_list[source_node_id]
+        else:
+            print("pagerank failed")
+
+    async def upsert_node(self, node_id: str, node_data: dict[str, str]):
+        self._graph.add_node(node_id, **node_data)
+
+    async def upsert_edge(
+        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
+    ):
+        self._graph.add_edge(source_node_id, target_node_id, **edge_data)
+
+    async def delete_node(self, node_id: str):
+        """
+        Delete a node from the graph based on the specified node_id.
+
+        :param node_id: The node_id to delete
+        """
+        if self._graph.has_node(node_id):
+            self._graph.remove_node(node_id)
+            logger.info(f"Node {node_id} deleted from the graph.")
+        else:
+            logger.warning(f"Node {node_id} not found in the graph for deletion.")
+
+    async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]:
+        if algorithm not in self._node_embed_algorithms:
+            raise ValueError(f"Node embedding algorithm {algorithm} not supported")
+        return await self._node_embed_algorithms[algorithm]()
+
+    # @TODO: NOT USED
+    async def _node2vec_embed(self):
+        from graspologic import embed
+
+        embeddings, nodes = embed.node2vec_embed(
+            self._graph,
+            **self.global_config["node2vec_params"],
+        )
+
+        nodes_ids = [self._graph.nodes[node_id]["id"] for node_id in nodes]
+        return embeddings, nodes_ids
+    
+    async def edges(self):
+        return self._graph.edges()
+    async def nodes(self):
+        return self._graph.nodes()

PathRAG/utils.py

@@ -0,0 +1,527 @@
+import asyncio
+import html
+import io
+import csv
+import json
+import logging
+import os
+import re
+from dataclasses import dataclass
+from functools import wraps
+from hashlib import md5
+from typing import Any, Union, List, Optional
+import xml.etree.ElementTree as ET
+
+import numpy as np
+import tiktoken
+
+from PathRAG.prompt import PROMPTS
+
+
+class UnlimitedSemaphore:
+
+
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
+
+ENCODER = None
+
+logger = logging.getLogger("PathRAG")
+
+
+def set_logger(log_file: str):
+    logger.setLevel(logging.DEBUG)
+
+    file_handler = logging.FileHandler(log_file)
+    file_handler.setLevel(logging.DEBUG)
+
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    file_handler.setFormatter(formatter)
+
+    if not logger.handlers:
+        logger.addHandler(file_handler)
+
+
+@dataclass
+class EmbeddingFunc:
+    embedding_dim: int
+    max_token_size: int
+    func: callable
+    concurrent_limit: int = 16
+
+    def __post_init__(self):
+        if self.concurrent_limit != 0:
+            self._semaphore = asyncio.Semaphore(self.concurrent_limit)
+        else:
+            self._semaphore = UnlimitedSemaphore()
+
+    async def __call__(self, *args, **kwargs) -> np.ndarray:
+        async with self._semaphore:
+            return await self.func(*args, **kwargs)
+
+
+def locate_json_string_body_from_string(content: str) -> Union[str, None]:
+
+    try:
+        maybe_json_str = re.search(r"{.*}", content, re.DOTALL)
+        if maybe_json_str is not None:
+            maybe_json_str = maybe_json_str.group(0)
+            maybe_json_str = maybe_json_str.replace("\\n", "")
+            maybe_json_str = maybe_json_str.replace("\n", "")
+            maybe_json_str = maybe_json_str.replace("'", '"')
+
+            return maybe_json_str
+    except Exception:
+        pass
+
+
+        return None
+
+
+def convert_response_to_json(response: str) -> dict:
+    json_str = locate_json_string_body_from_string(response)
+    assert json_str is not None, f"Unable to parse JSON from response: {response}"
+    try:
+        data = json.loads(json_str)
+        return data
+    except json.JSONDecodeError as e:
+        logger.error(f"Failed to parse JSON: {json_str}")
+        raise e from None
+
+
+def compute_args_hash(*args):
+    return md5(str(args).encode()).hexdigest()
+
+
+def compute_mdhash_id(content, prefix: str = ""):
+    return prefix + md5(content.encode()).hexdigest()
+
+
+def limit_async_func_call(max_size: int, waitting_time: float = 0.0001):
+
+
+    def final_decro(func):
+
+        __current_size = 0
+
+        @wraps(func)
+        async def wait_func(*args, **kwargs):
+            nonlocal __current_size
+            while __current_size >= max_size:
+                await asyncio.sleep(waitting_time)
+            __current_size += 1
+            result = await func(*args, **kwargs)
+            __current_size -= 1
+            return result
+
+        return wait_func
+
+    return final_decro
+
+
+def wrap_embedding_func_with_attrs(**kwargs):
+
+
+    def final_decro(func) -> EmbeddingFunc:
+        new_func = EmbeddingFunc(**kwargs, func=func)
+        return new_func
+
+    return final_decro
+
+
+def load_json(file_name):
+    if not os.path.exists(file_name):
+        return None
+    with open(file_name, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def write_json(json_obj, file_name):
+    with open(file_name, "w", encoding="utf-8") as f:
+        json.dump(json_obj, f, indent=2, ensure_ascii=False)
+
+
+def encode_string_by_tiktoken(content: str, model_name: str = "gpt-4o-mini"):
+    global ENCODER
+    if ENCODER is None:
+        ENCODER = tiktoken.encoding_for_model(model_name)
+    tokens = ENCODER.encode(content)
+    return tokens
+
+
+def decode_tokens_by_tiktoken(tokens: list[int], model_name: str = "gpt-4o-mini"):
+    global ENCODER
+    if ENCODER is None:
+        ENCODER = tiktoken.encoding_for_model(model_name)
+    content = ENCODER.decode(tokens)
+    return content
+
+
+def pack_user_ass_to_openai_messages(*args: str):
+    roles = ["user", "assistant"]
+    return [
+        {"role": roles[i % 2], "content": content} for i, content in enumerate(args)
+    ]
+
+
+def split_string_by_multi_markers(content: str, markers: list[str]) -> list[str]:
+
+    if not markers:
+        return [content]
+    results = re.split("|".join(re.escape(marker) for marker in markers), content)
+    return [r.strip() for r in results if r.strip()]
+
+
+
+def clean_str(input: Any) -> str:
+
+
+    if not isinstance(input, str):
+        return input
+
+    result = html.unescape(input.strip())
+
+    return re.sub(r"[\x00-\x1f\x7f-\x9f]", "", result)
+
+
+def is_float_regex(value):
+    return bool(re.match(r"^[-+]?[0-9]*\.?[0-9]+$", value))
+
+
+def truncate_list_by_token_size(list_data: list, key: callable, max_token_size: int):
+
+    if max_token_size <= 0:
+        return []
+    tokens = 0
+    for i, data in enumerate(list_data):
+        tokens += len(encode_string_by_tiktoken(key(data)))
+        if tokens > max_token_size:
+            return list_data[:i]
+    return list_data
+
+
+def list_of_list_to_csv(data: List[List[str]]) -> str:
+    output = io.StringIO()
+    writer = csv.writer(output)
+    writer.writerows(data)
+    return output.getvalue()
+
+
+def csv_string_to_list(csv_string: str) -> List[List[str]]:
+    output = io.StringIO(csv_string)
+    reader = csv.reader(output)
+    return [row for row in reader]
+
+
+def save_data_to_file(data, file_name):
+    with open(file_name, "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=4)
+
+
+def xml_to_json(xml_file):
+    try:
+        tree = ET.parse(xml_file)
+        root = tree.getroot()
+
+        print(f"Root element: {root.tag}")
+        print(f"Root attributes: {root.attrib}")
+
+        data = {"nodes": [], "edges": []}
+        namespace = {"": "http://graphml.graphdrawing.org/xmlns"}
+
+        for node in root.findall(".//node", namespace):
+            node_data = {
+                "id": node.get("id").strip('"'),
+                "entity_type": node.find("./data[@key='d0']", namespace).text.strip('"')
+                if node.find("./data[@key='d0']", namespace) is not None
+                else "",
+                "description": node.find("./data[@key='d1']", namespace).text
+                if node.find("./data[@key='d1']", namespace) is not None
+                else "",
+                "source_id": node.find("./data[@key='d2']", namespace).text
+                if node.find("./data[@key='d2']", namespace) is not None
+                else "",
+            }
+            data["nodes"].append(node_data)
+
+        for edge in root.findall(".//edge", namespace):
+            edge_data = {
+                "source": edge.get("source").strip('"'),
+                "target": edge.get("target").strip('"'),
+                "weight": float(edge.find("./data[@key='d3']", namespace).text)
+                if edge.find("./data[@key='d3']", namespace) is not None
+                else 0.0,
+                "description": edge.find("./data[@key='d4']", namespace).text
+                if edge.find("./data[@key='d4']", namespace) is not None
+                else "",
+                "keywords": edge.find("./data[@key='d5']", namespace).text
+                if edge.find("./data[@key='d5']", namespace) is not None
+                else "",
+                "source_id": edge.find("./data[@key='d6']", namespace).text
+                if edge.find("./data[@key='d6']", namespace) is not None
+                else "",
+            }
+            data["edges"].append(edge_data)
+
+        print(f"Found {len(data['nodes'])} nodes and {len(data['edges'])} edges")
+
+        return data
+    except ET.ParseError as e:
+        print(f"Error parsing XML file: {e}")
+        return None
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return None
+
+
+def process_combine_contexts(hl, ll):
+    header = None
+    list_hl = csv_string_to_list(hl.strip())
+    list_ll = csv_string_to_list(ll.strip())
+
+    if list_hl:
+        header = list_hl[0]
+        list_hl = list_hl[1:]
+    if list_ll:
+        header = list_ll[0]
+        list_ll = list_ll[1:]
+    if header is None:
+        return ""
+
+    if list_hl:
+        list_hl = [",".join(item[1:]) for item in list_hl if item]
+    if list_ll:
+        list_ll = [",".join(item[1:]) for item in list_ll if item]
+
+    combined_sources = []
+    seen = set()
+
+    for item in list_hl + list_ll:
+        if item and item not in seen:
+            combined_sources.append(item)
+            seen.add(item)
+
+    combined_sources_result = [",\t".join(header)]
+
+    for i, item in enumerate(combined_sources, start=1):
+        combined_sources_result.append(f"{i},\t{item}")
+
+    combined_sources_result = "\n".join(combined_sources_result)
+
+    return combined_sources_result
+
+
+async def get_best_cached_response(
+    hashing_kv,
+    current_embedding,
+    similarity_threshold=0.95,
+    mode="default",
+    use_llm_check=False,
+    llm_func=None,
+    original_prompt=None,
+) -> Union[str, None]:
+
+    mode_cache = await hashing_kv.get_by_id(mode)
+    if not mode_cache:
+        return None
+
+    best_similarity = -1
+    best_response = None
+    best_prompt = None
+    best_cache_id = None
+
+
+    for cache_id, cache_data in mode_cache.items():
+        if cache_data["embedding"] is None:
+            continue
+
+
+        cached_quantized = np.frombuffer(
+            bytes.fromhex(cache_data["embedding"]), dtype=np.uint8
+        ).reshape(cache_data["embedding_shape"])
+        cached_embedding = dequantize_embedding(
+            cached_quantized,
+            cache_data["embedding_min"],
+            cache_data["embedding_max"],
+        )
+
+        similarity = cosine_similarity(current_embedding, cached_embedding)
+        if similarity > best_similarity:
+            best_similarity = similarity
+            best_response = cache_data["return"]
+            best_prompt = cache_data["original_prompt"]
+            best_cache_id = cache_id
+
+    if best_similarity > similarity_threshold:
+
+        if use_llm_check and llm_func and original_prompt and best_prompt:
+            compare_prompt = PROMPTS["similarity_check"].format(
+                original_prompt=original_prompt, cached_prompt=best_prompt
+            )
+
+            try:
+                llm_result = await llm_func(compare_prompt)
+                llm_result = llm_result.strip()
+                llm_similarity = float(llm_result)
+
+
+                best_similarity = llm_similarity
+                if best_similarity < similarity_threshold:
+                    log_data = {
+                        "event": "llm_check_cache_rejected",
+                        "original_question": original_prompt[:100] + "..."
+                        if len(original_prompt) > 100
+                        else original_prompt,
+                        "cached_question": best_prompt[:100] + "..."
+                        if len(best_prompt) > 100
+                        else best_prompt,
+                        "similarity_score": round(best_similarity, 4),
+                        "threshold": similarity_threshold,
+                    }
+                    logger.info(json.dumps(log_data, ensure_ascii=False))
+                    return None
+            except Exception as e:   
+                logger.warning(f"LLM similarity check failed: {e}")
+                return None  
+
+        prompt_display = (
+            best_prompt[:50] + "..." if len(best_prompt) > 50 else best_prompt
+        )
+        log_data = {
+            "event": "cache_hit",
+            "mode": mode,
+            "similarity": round(best_similarity, 4),
+            "cache_id": best_cache_id,
+            "original_prompt": prompt_display,
+        }
+        logger.info(json.dumps(log_data, ensure_ascii=False))
+        return best_response
+    return None
+
+
+def cosine_similarity(v1, v2):
+
+    dot_product = np.dot(v1, v2)
+    norm1 = np.linalg.norm(v1)
+    norm2 = np.linalg.norm(v2)
+    return dot_product / (norm1 * norm2)
+
+
+def quantize_embedding(embedding: np.ndarray, bits=8) -> tuple:
+
+
+    min_val = embedding.min()
+    max_val = embedding.max()
+
+
+    scale = (2**bits - 1) / (max_val - min_val)
+    quantized = np.round((embedding - min_val) * scale).astype(np.uint8)
+
+    return quantized, min_val, max_val
+
+
+def dequantize_embedding(
+    quantized: np.ndarray, min_val: float, max_val: float, bits=8
+) -> np.ndarray:
+
+    scale = (max_val - min_val) / (2**bits - 1)
+    return (quantized * scale + min_val).astype(np.float32)
+
+
+async def handle_cache(hashing_kv, args_hash, prompt, mode="default"):
+
+    if hashing_kv is None:
+        return None, None, None, None
+
+
+    if mode == "naive":
+        mode_cache = await hashing_kv.get_by_id(mode) or {}
+        if args_hash in mode_cache:
+            return mode_cache[args_hash]["return"], None, None, None
+        return None, None, None, None
+
+
+    embedding_cache_config = hashing_kv.global_config.get(
+        "embedding_cache_config",
+        {"enabled": False, "similarity_threshold": 0.95, "use_llm_check": False},
+    )
+    is_embedding_cache_enabled = embedding_cache_config["enabled"]
+    use_llm_check = embedding_cache_config.get("use_llm_check", False)
+
+    quantized = min_val = max_val = None
+    if is_embedding_cache_enabled:
+
+        embedding_model_func = hashing_kv.global_config["embedding_func"]["func"]
+        llm_model_func = hashing_kv.global_config.get("llm_model_func")
+
+        current_embedding = await embedding_model_func([prompt])
+        quantized, min_val, max_val = quantize_embedding(current_embedding[0])
+        best_cached_response = await get_best_cached_response(
+            hashing_kv,
+            current_embedding[0],
+            similarity_threshold=embedding_cache_config["similarity_threshold"],
+            mode=mode,
+            use_llm_check=use_llm_check,
+            llm_func=llm_model_func if use_llm_check else None,
+            original_prompt=prompt if use_llm_check else None,
+        )
+        if best_cached_response is not None:
+            return best_cached_response, None, None, None
+    else:
+
+        mode_cache = await hashing_kv.get_by_id(mode) or {}
+        if args_hash in mode_cache:
+            return mode_cache[args_hash]["return"], None, None, None
+
+    return None, quantized, min_val, max_val
+
+
+@dataclass
+class CacheData:
+    args_hash: str
+    content: str
+    prompt: str
+    quantized: Optional[np.ndarray] = None
+    min_val: Optional[float] = None
+    max_val: Optional[float] = None
+    mode: str = "default"
+
+
+async def save_to_cache(hashing_kv, cache_data: CacheData):
+    if hashing_kv is None or hasattr(cache_data.content, "__aiter__"):
+        return
+
+    mode_cache = await hashing_kv.get_by_id(cache_data.mode) or {}
+
+    mode_cache[cache_data.args_hash] = {
+        "return": cache_data.content,
+        "embedding": cache_data.quantized.tobytes().hex()
+        if cache_data.quantized is not None
+        else None,
+        "embedding_shape": cache_data.quantized.shape
+        if cache_data.quantized is not None
+        else None,
+        "embedding_min": cache_data.min_val,
+        "embedding_max": cache_data.max_val,
+        "original_prompt": cache_data.prompt,
+    }
+
+    await hashing_kv.upsert({cache_data.mode: mode_cache})
+
+
+def safe_unicode_decode(content):
+    unicode_escape_pattern = re.compile(r"\\u([0-9a-fA-F]{4})")
+    def replace_unicode_escape(match):
+        return chr(int(match.group(1), 16))
+
+    decoded_content = unicode_escape_pattern.sub(
+        replace_unicode_escape, content.decode("utf-8")
+    )
+
+    return decoded_content

PathRAG/v1_test.py

@@ -0,0 +1,49 @@
+import os
+from PathRAG import PathRAG, QueryParam
+from PathRAG.llm import gpt_4o_mini_complete
+import time
+import openai
+import json
+import os
+import re
+from openai import OpenAI
+import sys
+import re
+from collections import defaultdict
+
+WORKING_DIR = ""
+
+api_key=""
+os.environ["OPENAI_API_KEY"] = api_key
+base_url="https://api.openai.com/v1"
+os.environ["OPENAI_API_BASE"]=base_url
+
+
+if not os.path.exists(WORKING_DIR):
+    os.mkdir(WORKING_DIR)
+
+rag = PathRAG(
+    working_dir=WORKING_DIR,
+    llm_model_func=gpt_4o_mini_complete,  
+)
+
+data_file=""
+question=""
+with open(data_file) as f:
+    rag.insert(f.read())
+
+print(rag.query(question, param=QueryParam(mode="hybrid")))
+
+
+
+
+
+
+
+
+
+
+
+
+
+

requirements.txt

@@ -0,0 +1,28 @@
+accelerate
+aioboto3
+aiohttp
+
+# database packages
+graspologic
+hnswlib
+nano-vectordb
+neo4j
+networkx
+ollama
+openai
+oracledb
+psycopg[binary,pool]
+pymilvus
+pymongo
+pymysql
+pyvis
+# lmdeploy[all]
+sqlalchemy
+tenacity
+
+
+# LLM packages
+tiktoken
+torch
+transformers
+xxhash