yangdx
commited on
Commit
·
36d19a0
1
Parent(s):
a726a93
Add conversation history support to chat API
Browse files- Added HISTORY_TURNS env variable
- Updated chat request data creation
- Modified server to handle history
- Added history to test cases
- .env.example +1 -0
- lightrag/api/lightrag_server.py +22 -4
- test_lightrag_ollama_chat.py +50 -7
.env.example
CHANGED
@@ -43,6 +43,7 @@ MAX_ASYNC=4
|
|
43 |
MAX_TOKENS=32768
|
44 |
EMBEDDING_DIM=1024
|
45 |
MAX_EMBED_TOKENS=8192
|
|
|
46 |
|
47 |
# Security (empty for no key)
|
48 |
LIGHTRAG_API_KEY=your-secure-api-key-here
|
|
|
43 |
MAX_TOKENS=32768
|
44 |
EMBEDDING_DIM=1024
|
45 |
MAX_EMBED_TOKENS=8192
|
46 |
+
#HISTORY_TURNS=3
|
47 |
|
48 |
# Security (empty for no key)
|
49 |
LIGHTRAG_API_KEY=your-secure-api-key-here
|
lightrag/api/lightrag_server.py
CHANGED
@@ -470,6 +470,13 @@ def parse_args() -> argparse.Namespace:
|
|
470 |
help="Enable automatic scanning when the program starts",
|
471 |
)
|
472 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
473 |
args = parser.parse_args()
|
474 |
|
475 |
return args
|
@@ -1576,8 +1583,9 @@ def create_app(args):
|
|
1576 |
if not messages:
|
1577 |
raise HTTPException(status_code=400, detail="No messages provided")
|
1578 |
|
1579 |
-
# Get the last message as query
|
1580 |
query = messages[-1].content
|
|
|
1581 |
|
1582 |
# Check for query prefix
|
1583 |
cleaned_query, mode = parse_query_mode(query)
|
@@ -1585,9 +1593,19 @@ def create_app(args):
|
|
1585 |
start_time = time.time_ns()
|
1586 |
prompt_tokens = estimate_tokens(cleaned_query)
|
1587 |
|
1588 |
-
|
1589 |
-
|
1590 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1591 |
|
1592 |
if request.stream:
|
1593 |
from fastapi.responses import StreamingResponse
|
|
|
470 |
help="Enable automatic scanning when the program starts",
|
471 |
)
|
472 |
|
473 |
+
parser.add_argument(
|
474 |
+
"--history-turns",
|
475 |
+
type=int,
|
476 |
+
default=get_env_value("HISTORY_TURNS", None, int),
|
477 |
+
help="Number of conversation history turns to include (default: from env or None)",
|
478 |
+
)
|
479 |
+
|
480 |
args = parser.parse_args()
|
481 |
|
482 |
return args
|
|
|
1583 |
if not messages:
|
1584 |
raise HTTPException(status_code=400, detail="No messages provided")
|
1585 |
|
1586 |
+
# Get the last message as query and previous messages as history
|
1587 |
query = messages[-1].content
|
1588 |
+
conversation_history = messages[:-1] # 所有之前的消息作为历史记录
|
1589 |
|
1590 |
# Check for query prefix
|
1591 |
cleaned_query, mode = parse_query_mode(query)
|
|
|
1593 |
start_time = time.time_ns()
|
1594 |
prompt_tokens = estimate_tokens(cleaned_query)
|
1595 |
|
1596 |
+
# 构建 query_param
|
1597 |
+
param_dict = {
|
1598 |
+
"mode": mode,
|
1599 |
+
"stream": request.stream,
|
1600 |
+
"only_need_context": False,
|
1601 |
+
"conversation_history": conversation_history,
|
1602 |
+
}
|
1603 |
+
|
1604 |
+
# 如果设置了 history_turns,添加到参数中
|
1605 |
+
if args.history_turns is not None:
|
1606 |
+
param_dict["history_turns"] = args.history_turns
|
1607 |
+
|
1608 |
+
query_param = QueryParam(**param_dict)
|
1609 |
|
1610 |
if request.stream:
|
1611 |
from fastapi.responses import StreamingResponse
|
test_lightrag_ollama_chat.py
CHANGED
@@ -189,19 +189,32 @@ def get_base_url(endpoint: str = "chat") -> str:
|
|
189 |
|
190 |
|
191 |
def create_chat_request_data(
|
192 |
-
content: str,
|
|
|
|
|
|
|
|
|
193 |
) -> Dict[str, Any]:
|
194 |
"""Create chat request data
|
195 |
Args:
|
196 |
content: User message content
|
197 |
stream: Whether to use streaming response
|
198 |
model: Model name
|
|
|
|
|
199 |
Returns:
|
200 |
Dictionary containing complete chat request data
|
201 |
"""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
202 |
return {
|
203 |
"model": model or CONFIG["server"]["model"],
|
204 |
-
"messages":
|
205 |
"stream": stream,
|
206 |
}
|
207 |
|
@@ -259,11 +272,25 @@ def run_test(func: Callable, name: str) -> None:
|
|
259 |
def test_non_stream_chat() -> None:
|
260 |
"""Test non-streaming call to /api/chat endpoint"""
|
261 |
url = get_base_url()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
262 |
data = create_chat_request_data(
|
263 |
-
CONFIG["test_cases"]["basic"]["query"],
|
|
|
|
|
|
|
264 |
)
|
265 |
-
|
266 |
-
# Send request
|
267 |
response = make_request(url, data)
|
268 |
|
269 |
# Print response
|
@@ -297,9 +324,25 @@ def test_stream_chat() -> None:
|
|
297 |
The last message will contain performance statistics, with done set to true.
|
298 |
"""
|
299 |
url = get_base_url()
|
300 |
-
data = create_chat_request_data(CONFIG["test_cases"]["basic"]["query"], stream=True)
|
301 |
|
302 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
303 |
response = make_request(url, data, stream=True)
|
304 |
|
305 |
if OutputControl.is_verbose():
|
|
|
189 |
|
190 |
|
191 |
def create_chat_request_data(
|
192 |
+
content: str,
|
193 |
+
stream: bool = False,
|
194 |
+
model: str = None,
|
195 |
+
conversation_history: List[Dict[str, str]] = None,
|
196 |
+
history_turns: int = None,
|
197 |
) -> Dict[str, Any]:
|
198 |
"""Create chat request data
|
199 |
Args:
|
200 |
content: User message content
|
201 |
stream: Whether to use streaming response
|
202 |
model: Model name
|
203 |
+
conversation_history: List of previous conversation messages
|
204 |
+
history_turns: Number of history turns to include
|
205 |
Returns:
|
206 |
Dictionary containing complete chat request data
|
207 |
"""
|
208 |
+
messages = conversation_history or []
|
209 |
+
if history_turns is not None and conversation_history:
|
210 |
+
messages = messages[
|
211 |
+
-2 * history_turns :
|
212 |
+
] # Each turn has 2 messages (user + assistant)
|
213 |
+
messages.append({"role": "user", "content": content})
|
214 |
+
|
215 |
return {
|
216 |
"model": model or CONFIG["server"]["model"],
|
217 |
+
"messages": messages,
|
218 |
"stream": stream,
|
219 |
}
|
220 |
|
|
|
272 |
def test_non_stream_chat() -> None:
|
273 |
"""Test non-streaming call to /api/chat endpoint"""
|
274 |
url = get_base_url()
|
275 |
+
|
276 |
+
# Example conversation history
|
277 |
+
conversation_history = [
|
278 |
+
{"role": "user", "content": "你好"},
|
279 |
+
{"role": "assistant", "content": "你好!我是一个AI助手,很高兴为你服务。"},
|
280 |
+
{"role": "user", "content": "西游记里有几个主要人物?"},
|
281 |
+
{
|
282 |
+
"role": "assistant",
|
283 |
+
"content": "西游记的主要人物有唐僧、孙悟空、猪八戒、沙和尚这四位主角。",
|
284 |
+
},
|
285 |
+
]
|
286 |
+
|
287 |
+
# Send request with conversation history and history turns
|
288 |
data = create_chat_request_data(
|
289 |
+
CONFIG["test_cases"]["basic"]["query"],
|
290 |
+
stream=False,
|
291 |
+
conversation_history=conversation_history,
|
292 |
+
history_turns=2, # Only include last 2 turns
|
293 |
)
|
|
|
|
|
294 |
response = make_request(url, data)
|
295 |
|
296 |
# Print response
|
|
|
324 |
The last message will contain performance statistics, with done set to true.
|
325 |
"""
|
326 |
url = get_base_url()
|
|
|
327 |
|
328 |
+
# Example conversation history
|
329 |
+
conversation_history = [
|
330 |
+
{"role": "user", "content": "你好"},
|
331 |
+
{"role": "assistant", "content": "你好!我是一个AI助手,很高兴为你服务。"},
|
332 |
+
{"role": "user", "content": "西游记里有几个主要人物?"},
|
333 |
+
{
|
334 |
+
"role": "assistant",
|
335 |
+
"content": "西游记的主要人物有唐僧、孙悟空、猪八戒、沙和尚这四位主角。",
|
336 |
+
},
|
337 |
+
]
|
338 |
+
|
339 |
+
# Send request with conversation history and history turns
|
340 |
+
data = create_chat_request_data(
|
341 |
+
CONFIG["test_cases"]["basic"]["query"],
|
342 |
+
stream=True,
|
343 |
+
conversation_history=conversation_history,
|
344 |
+
history_turns=2, # Only include last 2 turns
|
345 |
+
)
|
346 |
response = make_request(url, data, stream=True)
|
347 |
|
348 |
if OutputControl.is_verbose():
|