diff --git a/lightrag/api/routers/ollama_api.py b/lightrag/api/routers/ollama_api.py index 86534da4..dee0e8b2 100644 --- a/lightrag/api/routers/ollama_api.py +++ b/lightrag/api/routers/ollama_api.py @@ -101,27 +101,37 @@ def estimate_tokens(text: str) -> int: 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 - Returns tuple of (cleaned_query, search_mode) + Returns tuple of (cleaned_query, search_mode, only_need_context) """ + # 定义前缀映射,包含模式和是否只需要上下文 mode_map = { - "/local ": SearchMode.local, - "/global ": SearchMode.global_, # global_ is used because 'global' is a Python keyword - "/naive ": SearchMode.naive, - "/hybrid ": SearchMode.hybrid, - "/mix ": SearchMode.mix, - "/bypass ": SearchMode.bypass, - "/context": SearchMode.context, + # 原有的前缀 + "/local ": (SearchMode.local, False), + "/global ": (SearchMode.global_, False), # global_ is used because 'global' is a Python keyword + "/naive ": (SearchMode.naive, False), + "/hybrid ": (SearchMode.hybrid, False), + "/mix ": (SearchMode.mix, False), + "/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): # After removing prefix an leading spaces 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: @@ -351,17 +361,11 @@ class OllamaAPI: ] # 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() prompt_tokens = estimate_tokens(cleaned_query) - if mode == SearchMode.context: - mode = SearchMode.hybrid - only_need_context = True - else: - only_need_context = False - param_dict = { "mode": mode, "stream": request.stream,