yangdx
commited on
Commit
·
fbe600f
1
Parent(s):
edbf071
Improve context only mode for Ollama api
Browse files
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 |
-
|
110 |
-
"/
|
111 |
-
"/
|
112 |
-
"/
|
113 |
-
"/
|
114 |
-
"/
|
115 |
-
"/
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
|
|
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,
|