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 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
- query_param = QueryParam(
1589
- mode=mode, stream=request.stream, only_need_context=False
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, stream: bool = False, model: str = None
 
 
 
 
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": [{"role": "user", "content": content}],
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"], stream=False
 
 
 
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
- # Send request and get streaming response
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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():