wehub-resource-sync
29 lines
794 B
Python
29 lines
794 B
Python
import os
|
|
|
|
import pytest
|
|
|
|
from ragas.testset import TestsetGenerator
|
|
|
|
|
|
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
|
|
def test_testset_generation_e2e():
|
|
# generate kg
|
|
from langchain_community.document_loaders import DirectoryLoader
|
|
|
|
loader = DirectoryLoader("./docs", glob="**/*.md")
|
|
docs = loader.load()
|
|
|
|
# choose llm
|
|
from ragas.embeddings import embedding_factory
|
|
from ragas.llms import llm_factory
|
|
|
|
generator_llm = llm_factory("gpt-4o")
|
|
generator_embeddings = embedding_factory()
|
|
|
|
generator = TestsetGenerator(
|
|
llm=generator_llm,
|
|
embedding_model=generator_embeddings, # type: ignore
|
|
)
|
|
dataset = generator.generate_with_langchain_docs(docs, testset_size=3)
|
|
assert dataset is not None
|