yangdx commited on
Commit
fbe600f
·
1 Parent(s): edbf071

Improve context only mode for Ollama api

Browse files
Files changed (1) hide show
  1. lightrag/api/routers/ollama_api.py +23 -19
lightrag/api/routers/ollama_api.py CHANGED
@@ -101,27 +101,37 @@ def estimate_tokens(text: str) -> int:
101
  return len(tokens)
102
 
103
 
104
- def parse_query_mode(query: str) -> tuple[str, SearchMode]:
105
  """Parse query prefix to determine search mode
106
- Returns tuple of (cleaned_query, search_mode)
107
  """
 
108
  mode_map = {
109
- "/local ": SearchMode.local,
110
- "/global ": SearchMode.global_, # global_ is used because 'global' is a Python keyword
111
- "/naive ": SearchMode.naive,
112
- "/hybrid ": SearchMode.hybrid,
113
- "/mix ": SearchMode.mix,
114
- "/bypass ": SearchMode.bypass,
115
- "/context": SearchMode.context,
 
 
 
 
 
 
 
 
116
  }
117
 
118
- for prefix, mode in mode_map.items():
119
  if query.startswith(prefix):
120
  # After removing prefix an leading spaces
121
  cleaned_query = query[len(prefix) :].lstrip()
122
- return cleaned_query, mode
123
 
124
- return query, SearchMode.hybrid
 
125
 
126
 
127
  class OllamaAPI:
@@ -351,17 +361,11 @@ class OllamaAPI:
351
  ]
352
 
353
  # Check for query prefix
354
- cleaned_query, mode = parse_query_mode(query)
355
 
356
  start_time = time.time_ns()
357
  prompt_tokens = estimate_tokens(cleaned_query)
358
 
359
- if mode == SearchMode.context:
360
- mode = SearchMode.hybrid
361
- only_need_context = True
362
- else:
363
- only_need_context = False
364
-
365
  param_dict = {
366
  "mode": mode,
367
  "stream": request.stream,
 
101
  return len(tokens)
102
 
103
 
104
+ def parse_query_mode(query: str) -> tuple[str, SearchMode, bool]:
105
  """Parse query prefix to determine search mode
106
+ Returns tuple of (cleaned_query, search_mode, only_need_context)
107
  """
108
+ # 定义前缀映射,包含模式和是否只需要上下文
109
  mode_map = {
110
+ # 原有的前缀
111
+ "/local ": (SearchMode.local, False),
112
+ "/global ": (SearchMode.global_, False), # global_ is used because 'global' is a Python keyword
113
+ "/naive ": (SearchMode.naive, False),
114
+ "/hybrid ": (SearchMode.hybrid, False),
115
+ "/mix ": (SearchMode.mix, False),
116
+ "/bypass ": (SearchMode.bypass, False),
117
+ "/context": (SearchMode.hybrid, True), # context模式使用hybrid模式,并设置only_need_context为True
118
+
119
+ # 新增的前缀
120
+ "/localcontext": (SearchMode.local, True),
121
+ "/globalcontext": (SearchMode.global_, True),
122
+ "/hybridcontext": (SearchMode.hybrid, True),
123
+ "/naivecontext": (SearchMode.naive, True),
124
+ "/mixcontext": (SearchMode.mix, True),
125
  }
126
 
127
+ for prefix, (mode, only_need_context) in mode_map.items():
128
  if query.startswith(prefix):
129
  # After removing prefix an leading spaces
130
  cleaned_query = query[len(prefix) :].lstrip()
131
+ return cleaned_query, mode, only_need_context
132
 
133
+ # 默认使用hybrid模式,不需要上下文
134
+ return query, SearchMode.hybrid, False
135
 
136
 
137
  class OllamaAPI:
 
361
  ]
362
 
363
  # Check for query prefix
364
+ cleaned_query, mode, only_need_context = parse_query_mode(query)
365
 
366
  start_time = time.time_ns()
367
  prompt_tokens = estimate_tokens(cleaned_query)
368
 
 
 
 
 
 
 
369
  param_dict = {
370
  "mode": mode,
371
  "stream": request.stream,