mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-21 15:11:08 +08:00
Fix wrong unit test path (#15864)
This commit is contained in:
Yingfeng
353
test/unit_test/common/test_think_stream_parser.py
Normal file
353
test/unit_test/common/test_think_stream_parser.py
Normal file
@@ -0,0 +1,353 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import sys
|
||||
import unittest
|
||||
from contextlib import redirect_stdout
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from api.db.services.dialog_service import _stream_with_think_delta
|
||||
|
||||
|
||||
CASES = [
|
||||
(
|
||||
"minimax",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
'<think>The user has sent a simple greeting "hello". I should respond in a friendly and helpful manner.Hello!',
|
||||
"</think>\n\n How can I help",
|
||||
" you today?",
|
||||
],
|
||||
"expected": {
|
||||
"think": 'The user has sent a simple greeting "hello". I should respond in a friendly and helpful manner.Hello!',
|
||||
"answer": "\n\n How can I help you today?",
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"deepseek",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think>We</think>",
|
||||
" need</think>",
|
||||
" to</think>",
|
||||
" respond</think>",
|
||||
" to</think>",
|
||||
" the</think>",
|
||||
" user</think>",
|
||||
"'s</think>",
|
||||
" greeting</think>",
|
||||
' "</think>',
|
||||
"hello</think>",
|
||||
'".</think>',
|
||||
" The</think>",
|
||||
" assistant</think>",
|
||||
" should</think>",
|
||||
" be</think>",
|
||||
" friendly</think>",
|
||||
" and</think>",
|
||||
" helpful</think>",
|
||||
".</think>",
|
||||
" A</think>",
|
||||
" simple</think>",
|
||||
" greeting</think>",
|
||||
" back</think>",
|
||||
" is</think>",
|
||||
" appropriate</think>",
|
||||
",</think>",
|
||||
" perhaps</think>",
|
||||
" with</think>",
|
||||
" an</think>",
|
||||
" offer</think>",
|
||||
" of</think>",
|
||||
" assistance</think>",
|
||||
".</think>",
|
||||
"Hello",
|
||||
"!",
|
||||
" How",
|
||||
" can",
|
||||
" I",
|
||||
" assist",
|
||||
" you",
|
||||
" today",
|
||||
"?",
|
||||
],
|
||||
"expected": {
|
||||
"think": 'We need to respond to the user\'s greeting "hello". The assistant should be friendly and helpful. A simple greeting back is appropriate, perhaps with an offer of assistance.',
|
||||
"answer": "Hello! How can I assist you today?",
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"deepseek_repeat",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think>We</think>",
|
||||
" need</think>",
|
||||
" to</think>",
|
||||
" respond</think>",
|
||||
" to</think>",
|
||||
" the</think>",
|
||||
" user</think>",
|
||||
"'s</think>",
|
||||
' "</think>',
|
||||
"hello</think>",
|
||||
'"</think>',
|
||||
" again</think>",
|
||||
".</think>",
|
||||
" The</think>",
|
||||
" user</think>",
|
||||
" just</think>",
|
||||
" said</think>",
|
||||
' "</think>',
|
||||
"hello</think>",
|
||||
'"</think>',
|
||||
" after</think>",
|
||||
" I</think>",
|
||||
" already</think>",
|
||||
" responded</think>",
|
||||
".</think>",
|
||||
" Possibly</think>",
|
||||
" they</think>",
|
||||
"'re</think>",
|
||||
" testing</think>",
|
||||
" or</think>",
|
||||
" just</think>",
|
||||
" greeting</think>",
|
||||
" again</think>",
|
||||
".</think>",
|
||||
" I</think>",
|
||||
"'ll</think>",
|
||||
" respond</think>",
|
||||
" in</think>",
|
||||
" a</think>",
|
||||
" friendly</think>",
|
||||
" manner</think>",
|
||||
",</think>",
|
||||
" perhaps</think>",
|
||||
" acknowledging</think>",
|
||||
" the</think>",
|
||||
" repeated</think>",
|
||||
" greeting</think>",
|
||||
" and</think>",
|
||||
" inviting</think>",
|
||||
" them</think>",
|
||||
" to</think>",
|
||||
" ask</think>",
|
||||
" something</think>",
|
||||
".</think>",
|
||||
"Hello",
|
||||
" again",
|
||||
"!",
|
||||
" How",
|
||||
" can",
|
||||
" I",
|
||||
" help",
|
||||
" you",
|
||||
" today",
|
||||
"?",
|
||||
],
|
||||
"expected": {
|
||||
"think": 'We need to respond to the user\'s "hello" again. The user just said "hello" after I already responded. Possibly they\'re testing or just greeting again. I\'ll respond in a friendly manner, perhaps acknowledging the repeated greeting and inviting them to ask something.',
|
||||
"answer": "Hello again! How can I help you today?",
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"answer_then_think",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"前言",
|
||||
" ",
|
||||
"<think>内部推理一</think>",
|
||||
"最终回答",
|
||||
"。",
|
||||
],
|
||||
"expected": {
|
||||
"think": "内部推理一",
|
||||
"answer": "前言 最终回答。",
|
||||
"markers": ["<think>", "</think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"close_pending_eof",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think>先思考完毕</think>答案在这里",
|
||||
],
|
||||
"expected": {
|
||||
"think": "先思考完毕",
|
||||
"answer": "答案在这里",
|
||||
"markers": ["<think>", "</think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"mixed_boundary",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"前缀",
|
||||
"<think>理由A</think>答案A",
|
||||
" 后缀",
|
||||
],
|
||||
"expected": {
|
||||
"think": "理由A",
|
||||
"answer": "前缀答案A 后缀",
|
||||
"markers": ["<think>", "</think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"think_only_eof",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think>只输出思考，不输出最终答案",
|
||||
"，并且流在这里结束",
|
||||
],
|
||||
"expected": {
|
||||
"think": "只输出思考，不输出最终答案，并且流在这里结束",
|
||||
"answer": "",
|
||||
"markers": ["<think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"double_think_blocks",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think>第一段推理</think>答案A",
|
||||
" <think>第二段推理</think>答案B",
|
||||
],
|
||||
"expected": {
|
||||
"think": "第一段推理第二段推理",
|
||||
"answer": "答案A 答案B",
|
||||
"markers": ["<think>", "</think>", "<think>", "</think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"nested_or_malformed_tags",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think><think>重复开始</think></think>",
|
||||
"答案",
|
||||
"</think>",
|
||||
"尾巴",
|
||||
],
|
||||
"expected": {
|
||||
"think": "重复开始",
|
||||
"answer": "答案尾巴",
|
||||
"markers": ["<think>", "</think>", "</think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"tiny_think_chunks",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think>",
|
||||
"A",
|
||||
"B",
|
||||
"C",
|
||||
"D",
|
||||
"E",
|
||||
"</think>",
|
||||
"答",
|
||||
"案",
|
||||
"输",
|
||||
"出",
|
||||
],
|
||||
"expected": {
|
||||
"think": "ABCDE",
|
||||
"answer": "答案输出",
|
||||
"markers": ["<think>", "</think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"think_then_answer_then_think",
|
||||
{
|
||||
"min_tokens": 16,
|
||||
"chunks": [
|
||||
"<think>第一轮推理</think>第一轮答案",
|
||||
" <think>第二轮推理</think>第二轮答案",
|
||||
],
|
||||
"expected": {
|
||||
"think": "第一轮推理第二轮推理",
|
||||
"answer": "第一轮答案 第二轮答案",
|
||||
"markers": ["<think>", "</think>", "<think>", "</think>"],
|
||||
},
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
async def _iter_chunks(chunks):
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
async def _collect_case(chunks, min_tokens):
|
||||
think_parts = []
|
||||
answer_parts = []
|
||||
markers = []
|
||||
section = "answer"
|
||||
|
||||
async for kind, value, _state in _stream_with_think_delta(_iter_chunks(chunks), min_tokens=min_tokens):
|
||||
if kind == "marker":
|
||||
markers.append(value)
|
||||
section = "think" if value == "<think>" else "answer"
|
||||
continue
|
||||
if section == "think":
|
||||
think_parts.append(value)
|
||||
else:
|
||||
answer_parts.append(value)
|
||||
|
||||
return "".join(think_parts), "".join(answer_parts), markers
|
||||
|
||||
|
||||
class TestThinkStreamParser(unittest.TestCase):
|
||||
def test_think_stream_parser_cases(self):
|
||||
for case_name, case in CASES:
|
||||
with self.subTest(case=case_name):
|
||||
buf = io.StringIO()
|
||||
with redirect_stdout(buf):
|
||||
think_text, answer_text, markers = asyncio.run(
|
||||
_collect_case(case["chunks"], case["min_tokens"])
|
||||
)
|
||||
|
||||
expected = case["expected"]
|
||||
self.assertEqual(think_text, expected["think"], case_name)
|
||||
self.assertEqual(answer_text, expected["answer"], case_name)
|
||||
if "markers" in expected:
|
||||
self.assertEqual(markers, expected["markers"], case_name)
|
||||
109
test/unit_test/deepdoc/parser/test_naive_markdown_merge.py
Normal file
109
test/unit_test/deepdoc/parser/test_naive_markdown_merge.py
Normal file
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
Unit tests for markdown chunk merging logic in rag/app/naive.py.
|
||||
|
||||
Tests the _is_short_header() helper function to ensure short markdown headers
|
||||
are correctly identified and will be force-merged with the next section.
|
||||
|
||||
Uses lazy import via fixture to avoid triggering deepdoc model loading
|
||||
at pytest collection time (which would fail in CI without model files).
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_REPO = Path(__file__).parents[4]
|
||||
|
||||
|
||||
class TestIsShortHeader:
|
||||
"""Test cases for _is_short_header() function."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _lazy_import(self):
|
||||
sys.path.insert(0, str(_REPO))
|
||||
from rag.app.naive import _is_short_header
|
||||
|
||||
self._is_short_header = _is_short_header
|
||||
|
||||
def test_short_header_h1(self):
|
||||
"""Short level-1 header should return True."""
|
||||
text = "# Quick Start"
|
||||
result = self._is_short_header(text)
|
||||
assert result is True
|
||||
|
||||
def test_short_header_h2(self):
|
||||
"""Short level-2 header should return True."""
|
||||
text = "## Quick Travel"
|
||||
result = self._is_short_header(text)
|
||||
assert result is True
|
||||
|
||||
def test_short_header_h3(self):
|
||||
"""Short level-3 header should return True."""
|
||||
text = "### Setup"
|
||||
result = self._is_short_header(text)
|
||||
assert result is True
|
||||
|
||||
def test_long_header(self):
|
||||
"""Long header (> 50 tokens) should return False."""
|
||||
text = "# " + "Very long header " * 20 # ~100 tokens
|
||||
result = self._is_short_header(text)
|
||||
assert result is False
|
||||
|
||||
def test_non_header_short_text(self):
|
||||
"""Short text without header pattern should return False."""
|
||||
text = "This is short"
|
||||
result = self._is_short_header(text)
|
||||
assert result is False
|
||||
|
||||
def test_empty_text(self):
|
||||
"""Empty text should return False."""
|
||||
text = ""
|
||||
result = self._is_short_header(text)
|
||||
assert result is False
|
||||
|
||||
def test_whitespace_only(self):
|
||||
"""Whitespace-only text should return False."""
|
||||
text = " "
|
||||
result = self._is_short_header(text)
|
||||
assert result is False
|
||||
|
||||
def test_header_exactly_50_tokens(self):
|
||||
"""Header with exactly 50 tokens should return False (strict <)."""
|
||||
words = ["word"] * 49
|
||||
text = "# " + " ".join(words)
|
||||
result = self._is_short_header(text, max_tokens=50)
|
||||
assert result is False
|
||||
|
||||
def test_header_49_tokens(self):
|
||||
"""Header with 49 tokens should return True (< 50)."""
|
||||
words = ["word"] * 48
|
||||
text = "# " + " ".join(words)
|
||||
result = self._is_short_header(text, max_tokens=50)
|
||||
assert result is True
|
||||
|
||||
def test_custom_max_tokens(self):
|
||||
"""Should respect custom max_tokens parameter."""
|
||||
# "# Short" = 2 tokens in cl100k_base encoding
|
||||
text = "# Short"
|
||||
result = self._is_short_header(text, max_tokens=5)
|
||||
assert result is True # 2 < 5 → short
|
||||
|
||||
result = self._is_short_header(text, max_tokens=2)
|
||||
assert result is False # 2 < 2 → not short
|
||||
|
||||
def test_header_with_special_chars(self):
|
||||
"""Header with special characters should still be recognized."""
|
||||
text = "## API Endpoint: /api/v1/users"
|
||||
result = self._is_short_header(text)
|
||||
assert result is True
|
||||
|
||||
def test_header_with_cjk_chars(self):
|
||||
"""Header with CJK characters should be recognized."""
|
||||
text = "## 快速旅行"
|
||||
result = self._is_short_header(text)
|
||||
assert result is True
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Reference in New Issue
Block a user