From 7432450e46f58c13897d864cefd2cb5222a904ac Mon Sep 17 00:00:00 2001
From: Exploding-Soda <easonherobrine@outlook.com>
Date: Tue, 8 Apr 2025 02:16:40 +0800
Subject: [PATCH] Add /context query prefix to Ollama model simulation for
 LightRAG Server

---
 lightrag/api/routers/ollama_api.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/lightrag/api/routers/ollama_api.py b/lightrag/api/routers/ollama_api.py
index 6d9ea8e6..86534da4 100644
--- a/lightrag/api/routers/ollama_api.py
+++ b/lightrag/api/routers/ollama_api.py
@@ -23,6 +23,7 @@ class SearchMode(str, Enum):
     hybrid = "hybrid"
     mix = "mix"
     bypass = "bypass"
+    context = "context"
 
 
 class OllamaMessage(BaseModel):
@@ -111,6 +112,7 @@ def parse_query_mode(query: str) -> tuple[str, SearchMode]:
         "/hybrid ": SearchMode.hybrid,
         "/mix ": SearchMode.mix,
         "/bypass ": SearchMode.bypass,
+        "/context": SearchMode.context,
     }
 
     for prefix, mode in mode_map.items():
@@ -354,10 +356,16 @@ class OllamaAPI:
                 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,
-                    "only_need_context": False,
+                    "only_need_context": only_need_context,
                     "conversation_history": conversation_history,
                     "top_k": self.top_k,
                 }