add rag test

This commit is contained in:
jyong 2024-03-06 14:54:06 +08:00
parent 703aefbd17
commit 52e6f458be

View File

@ -3,6 +3,8 @@ import datetime
import uuid
from typing import Optional
import pytest
from core.rag.cleaner.clean_processor import CleanProcessor
from core.rag.datasource.keyword.keyword_factory import Keyword
from core.rag.datasource.retrieval_service import RetrievalService
@ -16,9 +18,9 @@ from models.dataset import Dataset
from models.model import UploadFile
class ParagraphIndexProcessor(BaseIndexProcessor):
def extract(self) -> list[Document]:
@pytest.mark.parametrize('setup_unstructured_mock', [['partition_md', 'chunk_by_title']], indirect=True)
def extract() -> list[Document]:
file_detail = UploadFile(
tenant_id='test',
storage_type='local',
@ -40,8 +42,8 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
)
text_docs = ExtractProcessor.extract(extract_setting=extract_setting,
is_automatic=False)
is_automatic=True)
assert isinstance(text_docs, list)
return text_docs
def transform(self, documents: list[Document], **kwargs) -> list[Document]: