cognee/evals/qa_context_provider_utils.py
alekszievr 3494521cae
Support 4 different rag options in eval (#439)
* QA eval dataset as argument, with hotpot and 2wikimultihop as options. Json schema validation for datasets.

* Load dataset file by filename, outsource utilities

* restructure metric selection

* Add comprehensiveness, diversity and empowerment metrics

* add promptfoo as an option

* refactor RAG solution in eval;2C

* LLM as a judge metrics implemented in a uniform way

* Use requests.get instead of wget

* clean up promptfoo config template

* minor fixes

* get promptfoo path instead of hardcoding

* minor fixes

* Add LLM as a judge prompts

* Support 4 different rag options in eval

* Minor refactor and logger usage
2025-01-15 15:34:13 +01:00

59 lines
2 KiB
Python

import cognee
from cognee.api.v1.search import SearchType
from cognee.infrastructure.databases.vector import get_vector_engine
from cognee.modules.retrieval.brute_force_triplet_search import brute_force_triplet_search
from cognee.tasks.completion.graph_query_completion import retrieved_edges_to_string
async def get_raw_context(instance: dict) -> str:
return instance["context"]
async def cognify_instance(instance: dict):
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
for title, sentences in instance["context"]:
await cognee.add("\n".join(sentences), dataset_name="QA")
await cognee.cognify("QA")
async def get_context_with_cognee(instance: dict) -> str:
await cognify_instance(instance)
insights = await cognee.search(SearchType.INSIGHTS, query_text=instance["question"])
summaries = await cognee.search(SearchType.SUMMARIES, query_text=instance["question"])
search_results = insights + summaries
search_results_str = "\n".join([context_item["text"] for context_item in search_results])
return search_results_str
async def get_context_with_simple_rag(instance: dict) -> str:
await cognify_instance(instance)
vector_engine = get_vector_engine()
found_chunks = await vector_engine.search("document_chunk_text", instance["question"], limit=5)
search_results_str = "\n".join([context_item.payload["text"] for context_item in found_chunks])
return search_results_str
async def get_context_with_brute_force_triplet_search(instance: dict) -> str:
await cognify_instance(instance)
found_triplets = await brute_force_triplet_search(instance["question"], top_k=5)
search_results_str = retrieved_edges_to_string(found_triplets)
return search_results_str
qa_context_providers = {
"no_rag": get_raw_context,
"cognee": get_context_with_cognee,
"simple_rag": get_context_with_simple_rag,
"brute_force": get_context_with_brute_force_triplet_search,
}