File size: 6,050 Bytes
8e874f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | """
Main query handler for the QAFD-RAG system.
This module provides the main entry point for knowledge graph queries,
handling keyword extraction, context building, and response generation.
"""
import re
import json
from ..base import (
BaseGraphStorage,
BaseKVStorage,
BaseVectorStorage,
TextChunkSchema,
QueryParam,
)
from ..prompts import PROMPTS
from ..utils import (
logger,
compute_args_hash,
handle_cache,
save_to_cache,
CacheData,
)
from .context import build_query_context
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:
"""
Main query function for the QAFD-RAG system.
This function orchestrates the entire query pipeline:
1. Check cache for existing response
2. Extract keywords from query using LLM
3. Build context from knowledge graph
4. Generate response using LLM
Parameters:
-----------
query : str
The user's query string
knowledge_graph_inst : BaseGraphStorage
Knowledge graph storage instance
entities_vdb : BaseVectorStorage
Entity vector database
relationships_vdb : BaseVectorStorage
Relationships vector database
text_chunks_db : BaseKVStorage[TextChunkSchema]
Text chunks database
query_param : QueryParam
Query parameters including mode, top_k, etc.
global_config : dict
Global configuration including LLM functions
hashing_kv : BaseKVStorage, optional
Cache storage for responses
Returns:
--------
str
The generated response or context (based on query_param settings)
"""
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 ["local", "global", "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.debug("kw_prompt result: %s", 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:
logger.warning("JSON parsing error: %s - %s", 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"]
# Validate keywords based on mode
if query_param.mode == "local" and ll_keywords == []:
logger.warning("low_level_keywords is empty for local mode")
return PROMPTS["fail_response"]
elif query_param.mode == "global" and hl_keywords == []:
logger.warning("high_level_keywords is empty for global mode")
return PROMPTS["fail_response"]
elif query_param.mode == "hybrid" and ll_keywords == [] and hl_keywords == []:
logger.warning("Both low_level_keywords and high_level_keywords are empty for hybrid mode")
return PROMPTS["fail_response"]
# Convert lists to strings
if ll_keywords:
ll_keywords = ", ".join(ll_keywords)
else:
ll_keywords = ""
if hl_keywords:
hl_keywords = ", ".join(hl_keywords)
else:
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,
global_config,
)
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
|