wehub-resource-sync
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
from backend.app.main import app
|
|
from fastapi.testclient import TestClient
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture
|
|
def client():
|
|
return TestClient(app)
|
|
|
|
|
|
# Define parameter combinations for query_rag
|
|
query_rag_params = [
|
|
{
|
|
"llm_name": "gemini-2.0-flash",
|
|
"temperature": 0.2,
|
|
"similarity_top_k": 2,
|
|
"retrieval_strategy": "auto_merging",
|
|
"use_hyde": True,
|
|
"use_refine": True,
|
|
"use_node_rerank": True,
|
|
"use_react": False,
|
|
"qa_followup": True,
|
|
"hybrid_retrieval": True,
|
|
"query": "What were Google's Q1 Earnings?",
|
|
"evaluate_response": True,
|
|
"eval_model_name": "gemini-2.0-flash",
|
|
"embedding_model_name": "text-embedding-005",
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("payload", query_rag_params)
|
|
def test_query_rag(client, payload):
|
|
response = client.post("/query_rag", json=payload)
|
|
assert response.status_code == 200
|
|
|
|
|
|
eval_batch_params = [
|
|
{
|
|
"llm_name": "gemini-2.0-flash",
|
|
"temperature": 0.2,
|
|
"similarity_top_k": 5,
|
|
"retrieval_strategy": "parent",
|
|
"use_hyde": True,
|
|
"use_refine": True,
|
|
"use_node_rerank": False,
|
|
"use_react": False,
|
|
"eval_model_name": "gemini-2.0-flash",
|
|
"embedding_model_name": "text-embedding-005",
|
|
"input_eval_dataset_bucket_uri": "rag-llm-bucket/test_ground_truth.csv",
|
|
"bq_eval_results_table_id": "eval_results.eval_results_table",
|
|
"ragas_metrics": ["faithfulness", "answer_relevancy"],
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("payload", eval_batch_params)
|
|
def test_eval_batch(client, payload):
|
|
response = client.post("/eval_batch", json=payload)
|
|
assert response.status_code == 200
|