Fix linting
This commit is contained in:
yangdx
parent
b08b8a6a6a
commit
b848ca49e6
4 changed files with 44 additions and 64 deletions
|
|
@ -224,8 +224,12 @@ def create_query_routes(rag, api_key: Optional[str] = None, top_k: int = 60):
|
|||
if request.include_references:
|
||||
try:
|
||||
# Use aquery_data to get reference list independently
|
||||
data_param = request.to_query_params(False) # Non-streaming for data
|
||||
data_result = await rag.aquery_data(request.query, param=data_param)
|
||||
data_param = request.to_query_params(
|
||||
False
|
||||
) # Non-streaming for data
|
||||
data_result = await rag.aquery_data(
|
||||
request.query, param=data_param
|
||||
)
|
||||
if isinstance(data_result, dict) and "data" in data_result:
|
||||
reference_list = data_result["data"].get("references", [])
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -828,6 +828,7 @@ class DeletionResult:
|
|||
|
||||
# Unified Query Result Data Structures for Reference List Support
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
"""
|
||||
|
|
@ -839,6 +840,7 @@ class QueryResult:
|
|||
raw_data: Complete structured data including references and metadata
|
||||
is_streaming: Whether this is a streaming result
|
||||
"""
|
||||
|
||||
content: Optional[str] = None
|
||||
response_iterator: Optional[AsyncIterator[str]] = None
|
||||
raw_data: Optional[Dict[str, Any]] = None
|
||||
|
|
@ -879,6 +881,7 @@ class QueryContextResult:
|
|||
context: LLM context string
|
||||
raw_data: Complete structured data including reference_list
|
||||
"""
|
||||
|
||||
context: str
|
||||
raw_data: Dict[str, Any]
|
||||
|
||||
|
|
|
|||
|
|
@ -2118,7 +2118,7 @@ class LightRAG:
|
|||
query_result = QueryResult(
|
||||
content=response if not param.stream else None,
|
||||
response_iterator=response if param.stream else None,
|
||||
is_streaming=param.stream
|
||||
is_streaming=param.stream,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown mode {param.mode}")
|
||||
|
|
@ -2301,10 +2301,7 @@ class LightRAG:
|
|||
[], # no references
|
||||
"bypass",
|
||||
)
|
||||
query_result = QueryResult(
|
||||
content="",
|
||||
raw_data=empty_raw_data
|
||||
)
|
||||
query_result = QueryResult(content="", raw_data=empty_raw_data)
|
||||
else:
|
||||
raise ValueError(f"Unknown mode {data_param.mode}")
|
||||
|
||||
|
|
|
|||
|
|
@ -2386,8 +2386,7 @@ async def kg_query(
|
|||
# Return different content based on query parameters
|
||||
if query_param.only_need_context and not query_param.only_need_prompt:
|
||||
return QueryResult(
|
||||
content=context_result.context,
|
||||
raw_data=context_result.raw_data
|
||||
content=context_result.context, raw_data=context_result.raw_data
|
||||
)
|
||||
|
||||
# Build system prompt
|
||||
|
|
@ -2405,10 +2404,7 @@ async def kg_query(
|
|||
|
||||
if query_param.only_need_prompt:
|
||||
prompt_content = "\n\n".join([sys_prompt, "---User Query---", user_query])
|
||||
return QueryResult(
|
||||
content=prompt_content,
|
||||
raw_data=context_result.raw_data
|
||||
)
|
||||
return QueryResult(content=prompt_content, raw_data=context_result.raw_data)
|
||||
|
||||
# Call LLM
|
||||
tokenizer: Tokenizer = global_config["tokenizer"]
|
||||
|
|
@ -2466,16 +2462,13 @@ async def kg_query(
|
|||
),
|
||||
)
|
||||
|
||||
return QueryResult(
|
||||
content=response,
|
||||
raw_data=context_result.raw_data
|
||||
)
|
||||
return QueryResult(content=response, raw_data=context_result.raw_data)
|
||||
else:
|
||||
# Streaming response (AsyncIterator)
|
||||
return QueryResult(
|
||||
response_iterator=response,
|
||||
raw_data=context_result.raw_data,
|
||||
is_streaming=True
|
||||
is_streaming=True,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3478,10 +3471,7 @@ async def _build_query_context(
|
|||
f"[_build_query_context] Raw data entities: {len(raw_data.get('data', {}).get('entities', []))}, relationships: {len(raw_data.get('data', {}).get('relationships', []))}, chunks: {len(raw_data.get('data', {}).get('chunks', []))}"
|
||||
)
|
||||
|
||||
return QueryContextResult(
|
||||
context=context,
|
||||
raw_data=raw_data
|
||||
)
|
||||
return QueryContextResult(context=context, raw_data=raw_data)
|
||||
|
||||
|
||||
async def _get_node_data(
|
||||
|
|
@ -4157,10 +4147,7 @@ async def naive_query(
|
|||
"naive",
|
||||
)
|
||||
empty_raw_data["message"] = "No relevant document chunks found."
|
||||
return QueryResult(
|
||||
content=PROMPTS["fail_response"],
|
||||
raw_data=empty_raw_data
|
||||
)
|
||||
return QueryResult(content=PROMPTS["fail_response"], raw_data=empty_raw_data)
|
||||
|
||||
# Calculate dynamic token limit for chunks
|
||||
max_total_tokens = getattr(
|
||||
|
|
@ -4275,10 +4262,7 @@ async def naive_query(
|
|||
"""
|
||||
|
||||
if query_param.only_need_context and not query_param.only_need_prompt:
|
||||
return QueryResult(
|
||||
content=context_content,
|
||||
raw_data=raw_data
|
||||
)
|
||||
return QueryResult(content=context_content, raw_data=raw_data)
|
||||
|
||||
user_query = (
|
||||
"\n\n".join([query, query_param.user_prompt])
|
||||
|
|
@ -4294,10 +4278,7 @@ async def naive_query(
|
|||
|
||||
if query_param.only_need_prompt:
|
||||
prompt_content = "\n\n".join([sys_prompt, "---User Query---", user_query])
|
||||
return QueryResult(
|
||||
content=prompt_content,
|
||||
raw_data=raw_data
|
||||
)
|
||||
return QueryResult(content=prompt_content, raw_data=raw_data)
|
||||
|
||||
len_of_prompts = len(tokenizer.encode(query + sys_prompt))
|
||||
logger.debug(
|
||||
|
|
@ -4354,14 +4335,9 @@ async def naive_query(
|
|||
),
|
||||
)
|
||||
|
||||
return QueryResult(
|
||||
content=response,
|
||||
raw_data=raw_data
|
||||
)
|
||||
return QueryResult(content=response, raw_data=raw_data)
|
||||
else:
|
||||
# Streaming response (AsyncIterator)
|
||||
return QueryResult(
|
||||
response_iterator=response,
|
||||
raw_data=raw_data,
|
||||
is_streaming=True
|
||||
response_iterator=response, raw_data=raw_data, is_streaming=True
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue