From 49ce064a1145083e45bd9f4ed206d5584ab2b537 Mon Sep 17 00:00:00 2001 From: GGrassia Date: Fri, 28 Nov 2025 15:38:33 +0100 Subject: [PATCH] fix (embedding): fixed query endpoint --- lightrag/operate.py | 37 +++++++++++++++++++++++++++++-------- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/lightrag/operate.py b/lightrag/operate.py index 13262d82..13715e94 100644 --- a/lightrag/operate.py +++ b/lightrag/operate.py @@ -2831,13 +2831,22 @@ async def _perform_kg_search( query_embedding = None if query and (kg_chunk_pick_method == "VECTOR" or chunks_vdb): embedding_func_config = text_chunks_db.embedding_func - if embedding_func_config and embedding_func_config.func: + if embedding_func_config: try: - query_embedding = await embedding_func_config.func([query]) - query_embedding = query_embedding[ - 0 - ] # Extract first embedding from batch result - logger.debug("Pre-computed query embedding for all vector operations") + # Handle both EmbeddingFunc objects and plain callable functions + from .utils import EmbeddingFunc + if isinstance(embedding_func_config, EmbeddingFunc): + embedding_func = embedding_func_config.func + else: + # It's a plain callable function + embedding_func = embedding_func_config + + if embedding_func: + query_embedding = await embedding_func([query]) + query_embedding = query_embedding[ + 0 + ] # Extract first embedding from batch result + logger.debug("Pre-computed query embedding for all vector operations") except Exception as e: logger.warning(f"Failed to pre-compute query embedding: {e}") query_embedding = None @@ -3748,7 +3757,13 @@ async def _find_related_text_unit_from_entities( kg_chunk_pick_method = "WEIGHT" else: try: - actual_embedding_func = embedding_func_config.func + # Handle both EmbeddingFunc objects and plain callable functions + from .utils import EmbeddingFunc + if isinstance(embedding_func_config, EmbeddingFunc): + actual_embedding_func = embedding_func_config.func + else: + # It's a plain callable function + actual_embedding_func = embedding_func_config selected_chunk_ids = None if actual_embedding_func: @@ -4043,7 +4058,13 @@ async def _find_related_text_unit_from_relations( kg_chunk_pick_method = "WEIGHT" else: try: - actual_embedding_func = embedding_func_config.func + # Handle both EmbeddingFunc objects and plain callable functions + from .utils import EmbeddingFunc + if isinstance(embedding_func_config, EmbeddingFunc): + actual_embedding_func = embedding_func_config.func + else: + # It's a plain callable function + actual_embedding_func = embedding_func_config if actual_embedding_func: selected_chunk_ids = await pick_by_vector_similarity(