gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve context only mode for Ollama api

Browse Source
This commit is contained in:
yangdx
2025-04-08 18:41:00 +08:00
parent 3c2d0462bb
commit fa8f6d615a
1 changed files with 23 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
lightrag/api/routers/ollama_api.py
View File

@@ -101,27 +101,37 @@ def estimate_tokens(text: str) -> int:
return len(tokens) return len(tokens)
def parse_query_mode(query: str) -> tuple[str, SearchMode]: def parse_query_mode(query: str) -> tuple[str, SearchMode, bool]:
"""Parse query prefix to determine search mode """Parse query prefix to determine search mode
Returns tuple of (cleaned_query, search_mode) Returns tuple of (cleaned_query, search_mode, only_need_context)
""" """
# 定义前缀映射，包含模式和是否只需要上下文
mode_map = { mode_map = {
"/local ": SearchMode.local, # 原有的前缀
"/global ": SearchMode.global_, # global_ is used because 'global' is a Python keyword "/local ": (SearchMode.local, False),
"/naive ": SearchMode.naive, "/global ": (SearchMode.global_, False), # global_ is used because 'global' is a Python keyword
"/hybrid ": SearchMode.hybrid, "/naive ": (SearchMode.naive, False),
"/mix ": SearchMode.mix, "/hybrid ": (SearchMode.hybrid, False),
"/bypass ": SearchMode.bypass, "/mix ": (SearchMode.mix, False),
"/context": SearchMode.context, "/bypass ": (SearchMode.bypass, False),
"/context": (SearchMode.hybrid, True), # context模式使用hybrid模式，并设置only_need_context为True
# 新增的前缀
"/localcontext": (SearchMode.local, True),
"/globalcontext": (SearchMode.global_, True),
"/hybridcontext": (SearchMode.hybrid, True),
"/naivecontext": (SearchMode.naive, True),
"/mixcontext": (SearchMode.mix, True),
} }
for prefix, mode in mode_map.items(): for prefix, (mode, only_need_context) in mode_map.items():
if query.startswith(prefix): if query.startswith(prefix):
# After removing prefix an leading spaces # After removing prefix an leading spaces
cleaned_query = query[len(prefix) :].lstrip() cleaned_query = query[len(prefix) :].lstrip()
return cleaned_query, mode return cleaned_query, mode, only_need_context
return query, SearchMode.hybrid # 默认使用hybrid模式，不需要上下文
return query, SearchMode.hybrid, False
class OllamaAPI: class OllamaAPI:
@@ -351,17 +361,11 @@ class OllamaAPI:
] ]
# Check for query prefix # Check for query prefix
cleaned_query, mode = parse_query_mode(query) cleaned_query, mode, only_need_context = parse_query_mode(query)
start_time = time.time_ns() start_time = time.time_ns()
prompt_tokens = estimate_tokens(cleaned_query) prompt_tokens = estimate_tokens(cleaned_query)
if mode == SearchMode.context:
mode = SearchMode.hybrid
only_need_context = True
else:
only_need_context = False
param_dict = { param_dict = {
"mode": mode, "mode": mode,
"stream": request.stream, "stream": request.stream,