llama-index

1
from abc import ABC, abstractmethod
2
from typing import Any, Dict, List, Optional
3

4
from llama_index.legacy.bridge.pydantic import BaseModel, Field
5

6

7
class RetrievalMetricResult(BaseModel):
8
    """Metric result.
9

10
    Attributes:
11
        score (float): Score for the metric
12
        metadata (Dict[str, Any]): Metadata for the metric result
13

14
    """
15

16
    score: float = Field(..., description="Score for the metric")
17
    metadata: Dict[str, Any] = Field(
18
        default_factory=dict, description="Metadata for the metric result"
19
    )
20

21
    def __str__(self) -> str:
22
        """String representation."""
23
        return f"Score: {self.score}\nMetadata: {self.metadata}"
24

25
    def __float__(self) -> float:
26
        """Float representation."""
27
        return self.score
28

29

30
class BaseRetrievalMetric(BaseModel, ABC):
31
    """Base class for retrieval metrics."""
32

33
    metric_name: str
34

35
    @abstractmethod
36
    def compute(
37
        self,
38
        query: Optional[str] = None,
39
        expected_ids: Optional[List[str]] = None,
40
        retrieved_ids: Optional[List[str]] = None,
41
        expected_texts: Optional[List[str]] = None,
42
        retrieved_texts: Optional[List[str]] = None,
43
        **kwargs: Any,
44
    ) -> RetrievalMetricResult:
45
        """Compute metric.
46

47
        Args:
48
            query (Optional[str]): Query string
49
            expected_ids (Optional[List[str]]): Expected ids
50
            retrieved_ids (Optional[List[str]]): Retrieved ids
51
            **kwargs: Additional keyword arguments
52

53
        """
54

55
    class Config:
56
        arbitrary_types_allowed = True
57