Consolidates correctness and groundeness evaluation requests

Author: mafaldasv-dev

src/elastic_evals/evaluators/__init__.py

@@ -6,6 +6,7 @@
 
 from .correctness import (
     create_correctness_analysis_evaluator,
+    create_correctness_evaluators,
     create_quantitative_correctness_evaluators,
 )
 from .criteria import (
@@ -15,6 +16,7 @@
 )
 from .groundedness import (
     create_groundedness_analysis_evaluator,
+    create_groundedness_evaluators,
     create_quantitative_groundedness_evaluator,
 )
 from .input_tokens import create_input_tokens_evaluator
@@ -33,8 +35,10 @@
     "KibanaEvaluatorConfig",
     "KibanaSubScore",
     "create_correctness_analysis_evaluator",
+    "create_correctness_evaluators",
     "create_criteria_evaluator",
     "create_groundedness_analysis_evaluator",
+    "create_groundedness_evaluators",
     "create_input_tokens_evaluator",
     "create_latency_evaluator",
     "create_output_tokens_evaluator",

src/elastic_evals/evaluators/correctness/__init__.py

@@ -6,10 +6,12 @@
 
 from .evaluator import (
     create_correctness_analysis_evaluator,
+    create_correctness_evaluators,
     create_quantitative_correctness_evaluators,
 )
 
 __all__ = [
     "create_correctness_analysis_evaluator",
+    "create_correctness_evaluators",
     "create_quantitative_correctness_evaluators",
 ]

src/elastic_evals/evaluators/correctness/evaluator.py

@@ -47,39 +47,61 @@ def create_quantitative_correctness_evaluators(
     instrumentation_profile: InstrumentationProfile = "elastic-inference",
     log: logging.Logger | None = None,
 ) -> list[Evaluator]:
-    return kibana_evaluators(
+    return create_correctness_evaluators(
+        client=client,
+        connector_id=connector_id,
+        instrumentation_profile=instrumentation_profile,
+        log=log,
+    )[1:]
+
+
+def create_correctness_evaluators(
+    *,
+    client: KibanaEvaluatorsClient,
+    connector_id: str,
+    instrumentation_profile: InstrumentationProfile = "elastic-inference",
+    log: logging.Logger | None = None,
+) -> list[Evaluator]:
+    quantitative = kibana_evaluators(
         [_config(connector_id)],
         client=client,
         instrumentation_profile=instrumentation_profile,
         log=log,
     )
+    return [_analysis_evaluator(quantitative[0]), *quantitative]
 
 
 def create_correctness_analysis_evaluator(
     *,
-    inference_client: KibanaInferenceClient,
+    inference_client: KibanaInferenceClient | None = None,
+    client: KibanaEvaluatorsClient | None = None,
+    connector_id: str | None = None,
     log: logging.Logger,
     instrumentation_profile: InstrumentationProfile = "elastic-inference",
 ) -> Evaluator:
-    client = KibanaEvaluatorsClient(
-        kibana_url=inference_client.kibana_url,
-        api_key=inference_client.api_key,
-        timeout=inference_client.timeout,
-    )
-    factuality = kibana_evaluators(
-        [
-            KibanaEvaluatorConfig(
-                name="correctness",
-                kind="LLM",
-                connector_id=inference_client.connector_id,
-                sub_scores=(KibanaSubScore(key="factuality", evaluator_name="correctness"),),
-            )
-        ],
+    if client is None:
+        if inference_client is None:
+            raise ValueError("client and connector_id are required")
+        client = KibanaEvaluatorsClient(
+            kibana_url=inference_client.kibana_url,
+            api_key=inference_client.api_key,
+            timeout=inference_client.timeout,
+        )
+        connector_id = inference_client.connector_id
+    elif inference_client is not None:
+        raise ValueError("Pass either client or inference_client, not both")
+    elif not connector_id:
+        raise ValueError("connector_id is required")
+
+    return create_correctness_evaluators(
         client=client,
+        connector_id=connector_id,
         instrumentation_profile=instrumentation_profile,
         log=log,
     )[0]
 
+
+def _analysis_evaluator(factuality: Evaluator) -> Evaluator:
     async def evaluate(params: EvaluatorParams) -> EvaluationResult:
         result = await factuality.evaluate(params)
         if result.label in {"error", "unavailable"}:
@@ -94,7 +116,11 @@ async def evaluate(params: EvaluatorParams) -> EvaluationResult:
             metadata=result.metadata,
         )
 
-    return SimpleEvaluator(name=QUALITATIVE_EVALUATOR_NAME, kind="LLM", evaluate=evaluate)
+    return SimpleEvaluator(
+        name=QUALITATIVE_EVALUATOR_NAME,
+        kind="LLM",
+        evaluate=evaluate,
+    )
 
 
 def _analysis_explanation(summary: Any, fallback: str | None) -> str | None:

src/elastic_evals/evaluators/groundedness/__init__.py

@@ -6,10 +6,12 @@
 
 from .evaluator import (
     create_groundedness_analysis_evaluator,
+    create_groundedness_evaluators,
     create_quantitative_groundedness_evaluator,
 )
 
 __all__ = [
     "create_groundedness_analysis_evaluator",
+    "create_groundedness_evaluators",
     "create_quantitative_groundedness_evaluator",
 ]

src/elastic_evals/evaluators/groundedness/evaluator.py

@@ -29,7 +29,22 @@ def create_quantitative_groundedness_evaluator(
     instrumentation_profile: InstrumentationProfile = "elastic-inference",
     log: logging.Logger | None = None,
 ) -> Evaluator:
-    return kibana_evaluators(
+    return create_groundedness_evaluators(
+        client=client,
+        connector_id=connector_id,
+        instrumentation_profile=instrumentation_profile,
+        log=log,
+    )[1]
+
+
+def create_groundedness_evaluators(
+    *,
+    client: KibanaEvaluatorsClient,
+    connector_id: str,
+    instrumentation_profile: InstrumentationProfile = "elastic-inference",
+    log: logging.Logger | None = None,
+) -> list[Evaluator]:
+    quantitative = kibana_evaluators(
         [
             KibanaEvaluatorConfig(
                 name="groundedness",
@@ -47,26 +62,40 @@ def create_quantitative_groundedness_evaluator(
         instrumentation_profile=instrumentation_profile,
         log=log,
     )[0]
+    return [_analysis_evaluator(quantitative), quantitative]
 
 
 def create_groundedness_analysis_evaluator(
     *,
-    inference_client: KibanaInferenceClient,
+    inference_client: KibanaInferenceClient | None = None,
+    client: KibanaEvaluatorsClient | None = None,
+    connector_id: str | None = None,
     log: logging.Logger,
     instrumentation_profile: InstrumentationProfile = "elastic-inference",
 ) -> Evaluator:
-    client = KibanaEvaluatorsClient(
-        kibana_url=inference_client.kibana_url,
-        api_key=inference_client.api_key,
-        timeout=inference_client.timeout,
-    )
-    groundedness = create_quantitative_groundedness_evaluator(
+    if client is None:
+        if inference_client is None:
+            raise ValueError("client and connector_id are required")
+        client = KibanaEvaluatorsClient(
+            kibana_url=inference_client.kibana_url,
+            api_key=inference_client.api_key,
+            timeout=inference_client.timeout,
+        )
+        connector_id = inference_client.connector_id
+    elif inference_client is not None:
+        raise ValueError("Pass either client or inference_client, not both")
+    elif not connector_id:
+        raise ValueError("connector_id is required")
+
+    return create_groundedness_evaluators(
         client=client,
-        connector_id=inference_client.connector_id,
+        connector_id=connector_id,
         instrumentation_profile=instrumentation_profile,
         log=log,
-    )
+    )[0]
+
 
+def _analysis_evaluator(groundedness: Evaluator) -> Evaluator:
     async def evaluate(params: EvaluatorParams) -> EvaluationResult:
         result = await groundedness.evaluate(params)
         if result.label in {"error", "unavailable"}:

src/elastic_evals/evaluators/kibana.py

@@ -8,6 +8,7 @@
 
 import asyncio
 import logging
+import weakref
 from dataclasses import dataclass
 from typing import Any, Literal, Sequence
 
@@ -49,6 +50,12 @@ class _ScoreSelector:
     evaluator_name: str
 
 
+@dataclass(frozen=True)
+class _CachedEvaluation:
+    request_key: str
+    task: asyncio.Task[EvaluateResponse]
+
+
 class _EvaluationBatch:
     def __init__(
         self,
@@ -60,21 +67,35 @@ def __init__(
         self._client = client
         self._configs = tuple(configs)
         self._instrumentation_profile = instrumentation_profile
-        self._evaluations: dict[str, asyncio.Task[EvaluateResponse]] = {}
+        self._evaluations: weakref.WeakKeyDictionary[object, _CachedEvaluation] = weakref.WeakKeyDictionary()
         self._lock = asyncio.Lock()
 
     async def evaluate(self, params: EvaluatorParams) -> EvaluateResponse:
         trace_id = _trace_id(params)
         if not trace_id:
             raise ValueError("A trace ID is required for Kibana evaluators")
 
+        request = self._build_request(params, trace_id)
+        request_key = request.model_dump_json(exclude_none=True)
+        scope = params._evaluation_scope
+
         async with self._lock:
-            evaluation = self._evaluations.get(trace_id)
-            if evaluation is None:
-                evaluation = asyncio.create_task(self._client.evaluate(self._build_request(params, trace_id)))
-                self._evaluations[trace_id] = evaluation
+            cached = self._evaluations.get(scope)
+            if cached is None or cached.request_key != request_key:
+                cached = _CachedEvaluation(
+                    request_key=request_key,
+                    task=asyncio.create_task(self._client.evaluate(request)),
+                )
+                self._evaluations[scope] = cached
 
-        return await asyncio.shield(evaluation)
+        try:
+            return await asyncio.shield(cached.task)
+        except Exception:
+            async with self._lock:
+                current = self._evaluations.get(scope)
+                if current is cached:
+                    del self._evaluations[scope]
+            raise
 
     def _build_request(self, params: EvaluatorParams, trace_id: str) -> EvaluateRequest:
         reference_data: dict[str, Any] | None
@@ -189,7 +210,7 @@ def kibana_evaluators(
     instrumentation_profile: InstrumentationProfile = "elastic-inference",
     log: logging.Logger | None = None,
 ) -> list[Evaluator]:
-    """Create Python evaluators backed by one Kibana request per trace."""
+    """Create Python evaluators backed by one Kibana request per evaluation."""
     _validate_configs(configs)
     immutable_configs = tuple(configs)
     selectors = _build_selectors(immutable_configs)

src/elastic_evals/executor/client.py

@@ -163,15 +163,16 @@ async def task_runner() -> TaskOutput:
 
                 log_evaluation_start(example_index, repetition, len(evaluators))
 
+                params = EvaluatorParams(
+                    input=example.input,
+                    output=task_output,
+                    expected=example.output,
+                    metadata=example.metadata,
+                    trace_id=task_trace_id,
+                )
+
                 for evaluator in evaluators:
                     log_evaluator_start(evaluator.name, example_index, repetition)
-                    params = EvaluatorParams(
-                        input=example.input,
-                        output=task_output,
-                        expected=example.output,
-                        metadata=example.metadata,
-                        trace_id=task_trace_id,
-                    )
 
                     async def evaluator_runner() -> Any:
                         return await evaluator.evaluate(params)

src/elastic_evals/types.py

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Generic, Literal, Protocol, TypeAlias, TypeVar
 
 from pydantic import BaseModel
@@ -20,6 +20,10 @@
 TaskOutput: TypeAlias = Any
 
 
+class _EvaluationScope:
+    pass
+
+
 class Example(BaseModel, Generic[TInput, TExpected, TMetadata]):
     input: TInput
     output: TExpected | None = None
@@ -64,6 +68,7 @@ class EvaluatorParams(Generic[TInput, TExpected, TMetadata, TTaskOutput]):
     expected: TExpected | None
     metadata: TMetadata
     trace_id: str | None = None
+    _evaluation_scope: object = field(default_factory=_EvaluationScope, compare=False, repr=False)
 
 
 class Evaluator(Protocol, Generic[TInput, TExpected, TMetadata, TTaskOutput]):