diff --git a/rag/graphrag/utils.py b/rag/graphrag/utils.py index 038646199..563647ea8 100644 --- a/rag/graphrag/utils.py +++ b/rag/graphrag/utils.py @@ -168,6 +168,7 @@ def dict_has_keys_with_types(data: dict, expected_fields: list[tuple[str, type]] def get_llm_cache(llmnm, txt, history, genconf): + """Return a cached LLM completion for the given model/text/history/config, or None on miss.""" hasher = xxhash.xxh64() hasher.update((str(llmnm)+str(txt)+str(history)+str(genconf)).encode("utf-8")) @@ -179,6 +180,7 @@ def get_llm_cache(llmnm, txt, history, genconf): def set_llm_cache(llmnm, txt, v, history, genconf): + """Store an LLM completion *v* in Redis keyed by a hash of model/text/history/config.""" hasher = xxhash.xxh64() hasher.update((str(llmnm)+str(txt)+str(history)+str(genconf)).encode("utf-8")) k = hasher.hexdigest() @@ -186,6 +188,7 @@ def set_llm_cache(llmnm, txt, v, history, genconf): def get_embed_cache(llmnm, txt): + """Return a cached embedding vector (numpy array) for *llmnm*/*txt*, or None on miss.""" hasher = xxhash.xxh64() hasher.update(str(llmnm).encode("utf-8")) hasher.update(str(txt).encode("utf-8")) @@ -198,6 +201,7 @@ def get_embed_cache(llmnm, txt): def set_embed_cache(llmnm, txt, arr): + """Store embedding *arr* in Redis for the given model name and input text.""" hasher = xxhash.xxh64() hasher.update(str(llmnm).encode("utf-8")) hasher.update(str(txt).encode("utf-8")) @@ -207,7 +211,35 @@ def set_embed_cache(llmnm, txt, arr): REDIS_CONN.set(k, arr.encode("utf-8"), 24 * 3600) +def _batch_embed_cache_misses(llmnm: str, keys: list) -> "list[bool]": + """Return a boolean miss-mask for *keys* using a single MGET round-trip. + + Avoids per-item REDIS_CONN.get() calls (which would block the event loop + when called from an async context) by issuing one batched MGET instead. + """ + if not keys: + return [] + hashes = [] + for key in keys: + h = xxhash.xxh64() + h.update(str(llmnm).encode("utf-8")) + h.update(str(key).encode("utf-8")) + hashes.append(h.hexdigest()) + return [v is None for v in REDIS_CONN.mget(hashes)] + + +def _write_embed_cache_batch(llmnm: str, keys: list, embeddings) -> None: + """Write a batch of embeddings to the Redis embed cache synchronously. + + Intended for use with thread_pool_exec so that the synchronous Redis SET + calls do not block the event loop. + """ + for key, ebd in zip(keys, embeddings): + set_embed_cache(llmnm, key, ebd) + + def get_tags_from_cache(kb_ids): + """Return cached tag data for the given kb_ids from Redis, or None on miss.""" hasher = xxhash.xxh64() hasher.update(str(kb_ids).encode("utf-8")) @@ -219,6 +251,7 @@ def get_tags_from_cache(kb_ids): def set_tags_to_cache(kb_ids, tags): + """Persist tag data for *kb_ids* in Redis.""" hasher = xxhash.xxh64() hasher.update(str(kb_ids).encode("utf-8")) @@ -263,6 +296,7 @@ def tidy_graph(graph: nx.Graph, callback, check_attribute: bool = True): def get_from_to(node1, node2): + """Return a canonical (lesser, greater) node pair for consistent undirected edge keying.""" if node1 < node2: return (node1, node2) else: @@ -303,6 +337,7 @@ def graph_merge(g1: nx.Graph, g2: nx.Graph, change: GraphChange): def compute_args_hash(*args): + """Return a hex MD5 digest of the string representation of *args* (used as a cache key).""" return md5(str(args).encode()).hexdigest() @@ -310,6 +345,7 @@ def handle_single_entity_extraction( record_attributes: list[str], chunk_key: str, ): + """Parse one entity record from LLM output and return a node-attribute dict, or None.""" if len(record_attributes) < 4 or record_attributes[0] != '"entity"': return None # add this record as a node in the G @@ -328,6 +364,7 @@ def handle_single_entity_extraction( def handle_single_relationship_extraction(record_attributes: list[str], chunk_key: str): + """Parse one relationship record from LLM output and return an edge-attribute dict, or None.""" if len(record_attributes) < 5 or record_attributes[0] != '"relationship"': return None # add this record as edge @@ -351,6 +388,7 @@ def handle_single_relationship_extraction(record_attributes: list[str], chunk_ke def pack_user_ass_to_openai_messages(*args: str): + """Interleave *args* as alternating user/assistant messages in OpenAI chat format.""" roles = ["user", "assistant"] return [{"role": roles[i % 2], "content": content} for i, content in enumerate(args)] @@ -364,14 +402,17 @@ def split_string_by_multi_markers(content: str, markers: list[str]) -> list[str] def is_float_regex(value): + """Return True if *value* is a string representation of a float or integer.""" return bool(re.match(r"^[-+]?[0-9]*\.?[0-9]+$", value)) def chunk_id(chunk): + """Return a deterministic hex ID for *chunk* derived from its content and kb_id.""" return xxhash.xxh64((chunk["content_with_weight"] + chunk["kb_id"]).encode("utf-8")).hexdigest() async def graph_node_to_chunk(kb_id, embd_mdl, ent_name, meta, chunks, nhop_neighbors=None): + """Convert a graph node (entity) to an embeddable chunk and append it to *chunks*.""" global chat_limiter enable_timeout_assertion = os.environ.get("ENABLE_TIMEOUT_ASSERTION") chunk = { @@ -410,6 +451,7 @@ async def graph_node_to_chunk(kb_id, embd_mdl, ent_name, meta, chunks, nhop_neig @timeout(3, 3) async def get_relation(tenant_id, kb_id, from_ent_name, to_ent_name, size=1): + """Retrieve edge metadata between entity names from the document store.""" ents = from_ent_name if isinstance(ents, str): ents = [from_ent_name] @@ -431,6 +473,7 @@ async def get_relation(tenant_id, kb_id, from_ent_name, to_ent_name, size=1): async def graph_edge_to_chunk(kb_id, embd_mdl, from_ent_name, to_ent_name, meta, chunks): + """Convert a graph edge (relation) to an embeddable chunk and append it to *chunks*.""" enable_timeout_assertion = os.environ.get("ENABLE_TIMEOUT_ASSERTION") chunk = { "id": get_uuid(), @@ -466,6 +509,7 @@ async def graph_edge_to_chunk(kb_id, embd_mdl, from_ent_name, to_ent_name, meta, async def does_graph_contains(tenant_id, kb_id, doc_id): + """Return True if *doc_id* is recorded as a source document in the stored graph for *kb_id*.""" # Get doc_ids of graph fields = ["source_id"] condition = { @@ -485,6 +529,7 @@ async def does_graph_contains(tenant_id, kb_id, doc_id): async def get_graph_doc_ids(tenant_id, kb_id) -> list[str]: + """Return the list of document IDs referenced by the stored graph for *kb_id*.""" conds = {"fields": ["source_id"], "removed_kwd": "N", "size": 1, "knowledge_graph_kwd": ["graph"]} res = await settings.retriever.search(conds, search.index_name(tenant_id), [kb_id]) doc_ids = [] @@ -496,6 +541,7 @@ async def get_graph_doc_ids(tenant_id, kb_id) -> list[str]: async def get_graph(tenant_id, kb_id, exclude_rebuild=None): + """Load the knowledge-graph for *kb_id* from the document store, rebuilding if marked removed.""" conds = {"fields": ["content_with_weight", "removed_kwd", "source_id"], "size": 1, "knowledge_graph_kwd": ["graph"]} res = await settings.retriever.search(conds, search.index_name(tenant_id), [kb_id]) if not res.total == 0: @@ -515,6 +561,12 @@ async def get_graph(tenant_id, kb_id, exclude_rebuild=None): async def set_graph(tenant_id: str, kb_id: str, embd_mdl, graph: nx.Graph, change: GraphChange, callback): + """Persist a knowledge-graph snapshot to the document store. + + Converts *graph* nodes and edges to embedding chunks, pre-warms the Redis + embed cache for all cache-miss entities/relations in bulk before spawning + per-item tasks, then atomically replaces the old graph chunks in the store. + """ global chat_limiter start = asyncio.get_running_loop().time() @@ -552,6 +604,42 @@ async def set_graph(tenant_id: str, kb_id: str, embd_mdl, graph: nx.Graph, chang } ) + # ── batch pre-warm entity embeddings ─────────────────────────────────────── + # Without this, set_graph spawns one asyncio task per entity, each calling + # embd_mdl.encode([single_name]). For 17 k+ nodes that is 17 k round-trips. + # Pre-warming the cache here collapses N calls to ceil(N/_INSERT_BULK_SIZE). + _node_list = list(change.added_updated_nodes) + _node_misses = await thread_pool_exec( + _batch_embed_cache_misses, embd_mdl.llm_name, _node_list + ) + _uncached_node_names = [n for n, miss in zip(_node_list, _node_misses) if miss] + logging.debug( + "set_graph node pre-warm: %d nodes, %d cache misses", + len(_node_list), len(_uncached_node_names), + ) + if _uncached_node_names: + _enable_ta = os.environ.get("ENABLE_TIMEOUT_ASSERTION") + _timeout = 3 if _enable_ta else 30000000 + for _i in range(0, len(_uncached_node_names), _INSERT_BULK_SIZE): + _batch = _uncached_node_names[_i : _i + _INSERT_BULK_SIZE] + async with chat_limiter: + _ebds, _ = await asyncio.wait_for( + thread_pool_exec(embd_mdl.encode, _batch), + timeout=_timeout, + ) + await thread_pool_exec(_write_embed_cache_batch, embd_mdl.llm_name, _batch, _ebds) + logging.debug( + "set_graph node pre-warm: wrote batch %d/%d (%d items)", + _i // _INSERT_BULK_SIZE + 1, + (len(_uncached_node_names) + _INSERT_BULK_SIZE - 1) // _INSERT_BULK_SIZE, + len(_batch), + ) + if callback: + callback(msg=f"Batch-embedded {len(_uncached_node_names)} entity names " + f"({(len(_uncached_node_names) + _INSERT_BULK_SIZE - 1) // _INSERT_BULK_SIZE} " + f"batches of {_INSERT_BULK_SIZE}).") + # ── end batch pre-warm ────────────────────────────────────────────────────── + tasks = [] for ii, node in enumerate(change.added_updated_nodes): node_attrs = graph.nodes[node] @@ -570,6 +658,50 @@ async def set_graph(tenant_id: str, kb_id: str, embd_mdl, graph: nx.Graph, chang await asyncio.gather(*tasks, return_exceptions=True) raise + # ── batch pre-warm edge embeddings ───────────────────────────────────────── + # Mirror of the node pre-warm above for relation chunks. + # Cache key = "A->B" (matches graph_edge_to_chunk lookup key) + # Encoded text = "A->B: " (matches graph_edge_to_chunk encode text) + _all_edge_data = [ + (_fn, _tn, graph.get_edge_data(_fn, _tn)) + for _fn, _tn in change.added_updated_edges + ] + _all_edge_data = [(f, t, a) for f, t, a in _all_edge_data if a] + _edge_lookup_keys = [f"{f}->{t}" for f, t, _ in _all_edge_data] + _edge_misses = await thread_pool_exec( + _batch_embed_cache_misses, embd_mdl.llm_name, _edge_lookup_keys + ) if _all_edge_data else [] + _uncached_edge_items = [item for item, miss in zip(_all_edge_data, _edge_misses) if miss] + logging.debug( + "set_graph edge pre-warm: %d edges, %d cache misses", + len(_all_edge_data), len(_uncached_edge_items), + ) + if _uncached_edge_items: + _edge_keys = [f"{f}->{t}" for f, t, _ in _uncached_edge_items] + _edge_texts = [f"{f}->{t}: {a['description']}" for f, t, a in _uncached_edge_items] + _enable_ta = os.environ.get("ENABLE_TIMEOUT_ASSERTION") + _timeout = 3 if _enable_ta else 30000000 + for _i in range(0, len(_edge_texts), _INSERT_BULK_SIZE): + _btexts = _edge_texts[_i : _i + _INSERT_BULK_SIZE] + _bkeys = _edge_keys [_i : _i + _INSERT_BULK_SIZE] + async with chat_limiter: + _ebds, _ = await asyncio.wait_for( + thread_pool_exec(embd_mdl.encode, _btexts), + timeout=_timeout, + ) + await thread_pool_exec(_write_embed_cache_batch, embd_mdl.llm_name, _bkeys, _ebds) + logging.debug( + "set_graph edge pre-warm: wrote batch %d/%d (%d items)", + _i // _INSERT_BULK_SIZE + 1, + (len(_uncached_edge_items) + _INSERT_BULK_SIZE - 1) // _INSERT_BULK_SIZE, + len(_btexts), + ) + if callback: + callback(msg=f"Batch-embedded {len(_uncached_edge_items)} edge descriptions " + f"({(len(_uncached_edge_items) + _INSERT_BULK_SIZE - 1) // _INSERT_BULK_SIZE} " + f"batches of {_INSERT_BULK_SIZE}).") + # ── end batch pre-warm ────────────────────────────────────────────────────── + tasks = [] for ii, (from_node, to_node) in enumerate(change.added_updated_edges): edge_attrs = graph.get_edge_data(from_node, to_node) @@ -663,6 +795,7 @@ async def set_graph(tenant_id: str, kb_id: str, embd_mdl, graph: nx.Graph, chang def is_continuous_subsequence(subseq, seq): + """Return True if *subseq* appears as a contiguous sub-path within tuple *seq*.""" def find_all_indexes(tup, value): indexes = [] start = 0 @@ -684,6 +817,7 @@ def is_continuous_subsequence(subseq, seq): def merge_tuples(list1, list2): + """Extend each path tuple in *list1* by matching continuations found in *list2*.""" result = [] for tup in list1: last_element = tup[-1] @@ -740,6 +874,7 @@ def n_neighbor(graph: nx.Graph, node, n_hop: int = 2): async def get_entity_type2samples(idxnms, kb_ids: list): + """Return a mapping of entity type → sample entity names fetched from the document store.""" es_res = await settings.retriever.search({"knowledge_graph_kwd": "ty2ents", "kb_id": kb_ids, "size": 10000, "fields": ["content_with_weight"]},idxnms,kb_ids) res = defaultdict(list) @@ -758,6 +893,7 @@ async def get_entity_type2samples(idxnms, kb_ids: list): def flat_uniq_list(arr, key): + """Flatten and deduplicate the values at *key* across a list of dicts.""" res = [] for a in arr: a = a[key] @@ -769,6 +905,7 @@ def flat_uniq_list(arr, key): async def rebuild_graph(tenant_id, kb_id, exclude_rebuild=None): + """Reconstruct the full knowledge-graph for *kb_id* from its stored subgraph chunks.""" graph = nx.Graph() flds = ["knowledge_graph_kwd", "content_with_weight", "source_id"] bs = 256