diff --git a/.env.example b/.env.example index 944e7851..b652847b 100644 --- a/.env.example +++ b/.env.example @@ -43,6 +43,7 @@ MAX_ASYNC=4 MAX_TOKENS=32768 EMBEDDING_DIM=1024 MAX_EMBED_TOKENS=8192 +#HISTORY_TURNS=3 # Security (empty for no key) LIGHTRAG_API_KEY=your-secure-api-key-here diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py index 533be818..b94e0d17 100644 --- a/lightrag/api/lightrag_server.py +++ b/lightrag/api/lightrag_server.py @@ -470,6 +470,13 @@ def parse_args() -> argparse.Namespace: help="Enable automatic scanning when the program starts", ) + parser.add_argument( + "--history-turns", + type=int, + default=get_env_value("HISTORY_TURNS", None, int), + help="Number of conversation history turns to include (default: from env or None)", + ) + args = parser.parse_args() return args @@ -1576,8 +1583,9 @@ def create_app(args): if not messages: raise HTTPException(status_code=400, detail="No messages provided") - # Get the last message as query + # Get the last message as query and previous messages as history query = messages[-1].content + conversation_history = messages[:-1] # 所有之前的消息作为历史记录 # Check for query prefix cleaned_query, mode = parse_query_mode(query) @@ -1585,9 +1593,19 @@ def create_app(args): start_time = time.time_ns() prompt_tokens = estimate_tokens(cleaned_query) - query_param = QueryParam( - mode=mode, stream=request.stream, only_need_context=False - ) + # 构建 query_param + param_dict = { + "mode": mode, + "stream": request.stream, + "only_need_context": False, + "conversation_history": conversation_history, + } + + # 如果设置了 history_turns,添加到参数中 + if args.history_turns is not None: + param_dict["history_turns"] = args.history_turns + + query_param = QueryParam(**param_dict) if request.stream: from fastapi.responses import StreamingResponse diff --git a/test_lightrag_ollama_chat.py b/test_lightrag_ollama_chat.py index d1e61d39..faf9485d 100644 --- a/test_lightrag_ollama_chat.py +++ b/test_lightrag_ollama_chat.py @@ -189,19 +189,32 @@ def get_base_url(endpoint: str = "chat") -> str: def create_chat_request_data( - content: str, stream: bool = False, model: str = None + content: str, + stream: bool = False, + model: str = None, + conversation_history: List[Dict[str, str]] = None, + history_turns: int = None, ) -> Dict[str, Any]: """Create chat request data Args: content: User message content stream: Whether to use streaming response model: Model name + conversation_history: List of previous conversation messages + history_turns: Number of history turns to include Returns: Dictionary containing complete chat request data """ + messages = conversation_history or [] + if history_turns is not None and conversation_history: + messages = messages[ + -2 * history_turns : + ] # Each turn has 2 messages (user + assistant) + messages.append({"role": "user", "content": content}) + return { "model": model or CONFIG["server"]["model"], - "messages": [{"role": "user", "content": content}], + "messages": messages, "stream": stream, } @@ -259,11 +272,25 @@ def run_test(func: Callable, name: str) -> None: def test_non_stream_chat() -> None: """Test non-streaming call to /api/chat endpoint""" url = get_base_url() - data = create_chat_request_data( - CONFIG["test_cases"]["basic"]["query"], stream=False - ) - # Send request + # Example conversation history + conversation_history = [ + {"role": "user", "content": "你好"}, + {"role": "assistant", "content": "你好!我是一个AI助手,很高兴为你服务。"}, + {"role": "user", "content": "西游记里有几个主要人物?"}, + { + "role": "assistant", + "content": "西游记的主要人物有唐僧、孙悟空、猪八戒、沙和尚这四位主角。", + }, + ] + + # Send request with conversation history and history turns + data = create_chat_request_data( + CONFIG["test_cases"]["basic"]["query"], + stream=False, + conversation_history=conversation_history, + history_turns=2, # Only include last 2 turns + ) response = make_request(url, data) # Print response @@ -297,9 +324,25 @@ def test_stream_chat() -> None: The last message will contain performance statistics, with done set to true. """ url = get_base_url() - data = create_chat_request_data(CONFIG["test_cases"]["basic"]["query"], stream=True) - # Send request and get streaming response + # Example conversation history + conversation_history = [ + {"role": "user", "content": "你好"}, + {"role": "assistant", "content": "你好!我是一个AI助手,很高兴为你服务。"}, + {"role": "user", "content": "西游记里有几个主要人物?"}, + { + "role": "assistant", + "content": "西游记的主要人物有唐僧、孙悟空、猪八戒、沙和尚这四位主角。", + }, + ] + + # Send request with conversation history and history turns + data = create_chat_request_data( + CONFIG["test_cases"]["basic"]["query"], + stream=True, + conversation_history=conversation_history, + history_turns=2, # Only include last 2 turns + ) response = make_request(url, data, stream=True) if OutputControl.is_verbose():