fix: import train config and add tests (#1776)

The cache feature causes import of config to fail because of the
pydantic schema not found error. Added fix and a test to prevent this in
future.

---------

Co-authored-by: Jithin James <[email protected]>
Co-authored-by: Ayush Thakur <[email protected]>

Author: 3 people

src/ragas/cache.py

@@ -5,7 +5,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, GetCoreSchemaHandler
+from pydantic_core import CoreSchema, core_schema
 
 
 class CacheInterface(ABC):
@@ -21,6 +22,17 @@ def set(self, key: str, value) -> None:
     def has_key(self, key: str) -> bool:
         pass
 
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """
+        Define how Pydantic generates a schema for BaseRagasEmbeddings.
+        """
+        return core_schema.no_info_after_validator_function(
+            cls, core_schema.is_instance_schema(cls)  # The validator function
+        )
+
 
 class DiskCacheBackend(CacheInterface):
     def __init__(self, cache_dir: str = ".cache"):

tests/conftest.py

@@ -2,9 +2,11 @@
 
 import typing as t
 
+import numpy as np
 import pytest
 from langchain_core.outputs import Generation, LLMResult
 
+from ragas.embeddings.base import BaseRagasEmbeddings
 from ragas.llms.base import BaseRagasLLM
 
 if t.TYPE_CHECKING:
@@ -46,6 +48,26 @@ async def agenerate_text(  # type: ignore
         return LLMResult(generations=[[Generation(text=prompt.to_string())]])
 
 
+class EchoEmbedding(BaseRagasEmbeddings):
+
+    async def aembed_documents(self, texts: t.List[str]) -> t.List[t.List[float]]:
+        return [np.random.rand(768).tolist() for _ in texts]
+
+    async def aembed_query(self, text: str) -> t.List[float]:
+        return [np.random.rand(768).tolist()]
+
+    def embed_documents(self, texts: t.List[str]) -> t.List[t.List[float]]:
+        return [np.random.rand(768).tolist() for _ in texts]
+
+    def embed_query(self, text: str) -> t.List[float]:
+        return [np.random.rand(768).tolist()]
+
+
 @pytest.fixture
 def fake_llm():
     return EchoLLM()
+
+
+@pytest.fixture
+def fake_embedding():
+    return EchoEmbedding()

tests/unit/test_optimizer_config.py

@@ -0,0 +1,8 @@
+def test_load_config(fake_llm, fake_embedding):
+
+    from ragas.config import DemonstrationConfig, InstructionConfig
+
+    inst_config = InstructionConfig(llm=fake_llm)
+    demo_config = DemonstrationConfig(embedding=fake_embedding)
+    assert inst_config.llm == fake_llm
+    assert demo_config.embedding == fake_embedding