leaudit-platform-frontend/app/components/dify-dataset-manager/retrieve-test.tsx

import { useState } from 'react';
import {
    Input,
    Button,
    Card,
    Select,
    Slider,
    Table,
    Tag,
    Empty,
    Spin,
    message,
} from 'antd';
import { FileSearchOutlined } from '@ant-design/icons';
import type { ColumnsType } from 'antd/es/table';
import type { RetrieveRecord } from '~/api/dify-dataset/type';
import { retrieveDataset } from '~/api/dify-dataset/api/segmentApi';

interface RetrieveTestProps {
    datasetId: string;
}

/**
 * 召回测试组件
 * 用于测试知识库的检索效果
 */
export default function RetrieveTest({ datasetId }: RetrieveTestProps) {
    const [searchQuery, setSearchQuery] = useState('');
    const [retrieveResults, setRetrieveResults] = useState<RetrieveRecord[]>([]);
    const [retrieving, setRetrieving] = useState(false);
    const [searchMethod, setSearchMethod] = useState<string>('hybrid_search');
    const [topK, setTopK] = useState<number>(5);

    /**
     * 执行检索
     */
    const handleRetrieve = async () => {
        if (!searchQuery.trim()) {
            message.warning('请输入检索关键词');
            return;
        }

        if (!datasetId) {
            message.warning('知识库ID不存在');
            return;
        }

        setRetrieving(true);
        try {
            const response = await retrieveDataset(datasetId, searchQuery, {
                search_method: searchMethod as any,
                top_k: topK,
            });
            setRetrieveResults(response.records || []);
            if (response.records?.length === 0) {
                message.info('未找到匹配的结果');
            }
        } catch (err: any) {
            console.error('检索失败:', err);
            message.error(err.message || '检索失败');
        } finally {
            setRetrieving(false);
        }
    };

    // 检索结果列定义
    const columns: ColumnsType<RetrieveRecord> = [
        {
            title: '相关度',
            dataIndex: 'score',
            key: 'score',
            width: 100,
            render: (score: number) => (
                <Tag color={score > 0.8 ? 'green' : score > 0.5 ? 'orange' : 'default'}>
                    {(score * 100).toFixed(1)}%
                </Tag>
            ),
        },
        {
            title: '内容',
            key: 'content',
            render: (_, record) => (
                <div className="retrieve-result-content">
                    <div className="content-text">
                        {record.segment.content.length > 300
                            ? record.segment.content.substring(0, 300) + '...'
                            : record.segment.content}
                    </div>
                    {record.segment.answer && (
                        <div className="answer-text">
                            <strong>答案：</strong>
                            {record.segment.answer.length > 150
                                ? record.segment.answer.substring(0, 150) + '...'
                                : record.segment.answer}
                        </div>
                    )}
                </div>
            ),
        },
        {
            title: '字数',
            key: 'word_count',
            width: 80,
            render: (_, record) => record.segment.word_count,
        },
        {
            title: '命中次数',
            key: 'hit_count',
            width: 100,
            render: (_, record) => record.segment.hit_count,
        },
    ];

    return (
        <div className="retrieve-test-page">
            {/* 页面标题 */}
            <div className="page-header">
                <h1>召回测试</h1>
                <p className="page-description">
                    输入查询内容，测试知识库的检索效果
                </p>
            </div>

            {/* 检索设置 */}
            <Card className="retrieve-settings" size="small">
                <div className="search-row">
                    <Input
                        placeholder="输入检索关键词..."
                        prefix={<FileSearchOutlined />}
                        value={searchQuery}
                        onChange={(e) => setSearchQuery(e.target.value)}
                        onPressEnter={handleRetrieve}
                        className="search-input"
                    />
                    <Button
                        type="primary"
                        onClick={handleRetrieve}
                        loading={retrieving}
                    >
                        检索
                    </Button>
                </div>

                <div className="options-row">
                    <div className="option-item">
                        <span className="option-label">检索方式：</span>
                        <Select
                            value={searchMethod}
                            onChange={setSearchMethod}
                            style={{ width: 140 }}
                            options={[
                                { value: 'keyword_search', label: '关键词搜索' },
                                { value: 'semantic_search', label: '语义搜索' },
                                { value: 'full_text_search', label: '全文搜索' },
                                { value: 'hybrid_search', label: '混合搜索' },
                            ]}
                        />
                    </div>
                    <div className="option-item">
                        <span className="option-label">返回数量 (Top K)：</span>
                        <Slider
                            value={topK}
                            onChange={setTopK}
                            min={1}
                            max={20}
                            style={{ width: 120 }}
                        />
                        <span className="option-value">{topK}</span>
                    </div>
                </div>
            </Card>

            {/* 检索结果 */}
            <div className="retrieve-results">
                {retrieving ? (
                    <div className="loading-state">
                        <Spin size="large" />
                        <div className="loading-text">检索中...</div>
                    </div>
                ) : retrieveResults.length === 0 ? (
                    <Empty
                        description="请输入关键词进行检索"
                        className="empty-state"
                    />
                ) : (
                    <>
                        <div className="results-header">
                            <span>找到 {retrieveResults.length} 条结果</span>
                        </div>
                        <Table
                            columns={columns}
                            dataSource={retrieveResults}
                            rowKey={(record) => record.segment.id}
                            pagination={false}
                            size="small"
                        />
                    </>
                )}
            </div>
        </div>
    );
}