add unittest for llm embedding

memory_scope/models/base_embedding_model.py

@@ -1,3 +1,4 @@
+from typing import List, Dict
 from llama_index.embeddings.dashscope import DashScopeEmbedding
 
 from memory_scope.models import MODEL_REGISTRY
@@ -22,3 +23,31 @@ class BaseEmbeddingModel(BaseModel):
 
     async def _async_call(self, **kwargs) -> ModelResponse:
         pass
+
+
+class LLIEmbedding(BaseEmbeddingModel):
+
+    def before_call(self, **kwargs):
+        text: str | List[str] = kwargs.pop("text", "")
+        if isinstance(text, str):
+            text = [text]
+        self.data = dict(texts=text)
+
+    def after_call(self, model_response: ModelResponse, **kwargs) -> ModelResponse:
+        embeddings = model_response.raw
+        model_response.embedding_results = embeddings
+        return model_response
+
+    def _call(self, **kwargs) -> ModelResponse:
+        results = ModelResponse()
+        try:
+            response = self.model.get_text_embedding_batch(**self.data)  
+            results.raw = response
+            results.status = True
+        except Exception as e:
+            results.details = e
+            results.status = False
+        return results
+    
+
+   

memory_scope/models/base_generation_model.py

@@ -31,7 +31,7 @@ class BaseGenerationModel(BaseModel):
         pass
 
 
-class LLILLM(BaseGenerateModel):
+class LLILLM(BaseGenerationModel):
 
     def before_call(self, **kwargs) -> None:
         prompt: str = kwargs.pop("prompt", "")

memory_scope/models/base_rank_model.py

@@ -1,3 +1,7 @@
+from typing import List, Dict
+from llama_index.core.data_structs import Node
+from llama_index.core.schema import NodeWithScore # type: ignore
+
 from memory_scope.models import MODEL_REGISTRY
 from memory_scope.models.base_model import BaseModel
 from memory_scope.models.response import ModelResponse, ModelResponseGen
@@ -11,12 +15,52 @@ class BaseRankModel(BaseModel):
     def before_call(self, **kwargs) -> None:
         pass
 
-    def after_call(self, model_response: ModelResponse | ModelResponseGen,
-                   **kwargs) -> ModelResponse | ModelResponseGen:
+    def after_call(self, **kwargs) -> ModelResponse:
         pass
 
-    def _call(self, stream: bool = False, **kwargs) -> ModelResponse | ModelResponseGen:
+    def _call(self, stream: bool = False, **kwargs) -> ModelResponse:
         pass
 
     async def _async_call(self, **kwargs) -> ModelResponse:
         pass
+
+
+class LLIReRank(BaseRankModel):
+   
+    def before_call(self, **kwargs) -> None:
+        assert "query" in kwargs or "documents" in kwargs
+        query: str = kwargs.pop("query", "")
+        documents: List[str] = kwargs.pop("documents", [])
+        
+        assert query and documents, f"query or documents is empty! query={query}, documents={len(documents)}"
+        if top_n is None:
+            top_n = len(documents)
+
+        # using -1.0 as dummy scores
+        nodes = [NodeWithScore(node=Node(text=text), score=-1.0) for text in documents] 
+
+        self.data = {
+            "nodes": nodes,
+            "query_str": query,
+        }
+        
+    def after_call(self, nodes: List[NodeWithScore]) -> ModelResponse: 
+        ranks = list()
+        for node in nodes:
+            ranks.append(dict(relevance_score=node.score,
+                                document=node.node.text))
+        results = ModelResponse(rank_scores=ranks)
+        return results
+
+    def _call(self, **kwargs) -> ModelResponse:
+        results = ModelResponse()
+        try:
+            self.before_call(**kwargs)
+            response = self.model.postprocess_nodes(**self.data)    
+            response = self.after_call(response)
+            results.rank_scores = response
+            results.status = True
+        except Exception as e:
+            results.details = e
+            results.status = False
+        return results

memory_scope/models/response.py

@@ -1,4 +1,4 @@
-from typing import Generator, List, Dict
+from typing import Generator, List, Dict, Any
 
 from pydantic import BaseModel, Field
 
@@ -6,10 +6,11 @@ from pydantic import BaseModel, Field
 class ModelResponse(BaseModel):
     text: str = Field("", description="")
 
-    embedding_results: Dict[int, List[float]] | List[float] = Field([], description="")
-
-    rank_scores: Dict[int, float] = Field({}, description="The rank scores of each documents.")
+    embedding_results:  List[List[float]] = Field([], description="")
 
+    #rank_scores: Dict[int, float] = Field({}, description="The rank scores of each documents.")
+    rank_scores: List[Dict[str, Any]] = Field({}, description="The rank scores of each documents.")
+    # [{"document": "xxx", "score": 0.5}, {{"document": "yyy", "score": 0.3}}]
     model_type: str = Field("", description="One of LLM, EMB, RANK.")
 
     status: bool = Field(True, description="Indicates whether the model call was successful.")
@@ -17,5 +18,5 @@ class ModelResponse(BaseModel):
     details: str = Field("", description=("The details information for model call, \
                          usually for storage of raw response or failure messages."))
 
-
+    raw: Any = Field("", description=("raw response from model call"))
 ModelResponseGen = Generator[ModelResponse, None, None]

memory_scope/utils/registry.py

@@ -22,6 +22,6 @@ class Registry(object):
         module_name_dict = {m.__name__: m for m in modules}
         self.module_dict.update(module_name_dict)
 
-    def __getitem__(self, module_name: str):
+    def get(self, module_name: str):
         assert module_name in self.module_dict, f'{module_name} not found in {self.name}'
         return self.module_dict[module_name]

tests/models/test_models_llm_embedding.py

@@ -0,0 +1,24 @@
+import unittest
+from memory_scope.models.base_embedding_model import LLIEmbedding
+
+class TestLLIEmbedding(unittest.TestCase):
+    """Tests for LLIEmbedding"""
+
+    def setUp(self):
+        config = {
+            "method_type": "DashScopeEmbedding",
+            "model_name": "text-embedding-v2",
+            "clazz": "models.base_embedding_model"
+        }
+        self.emb = LLIEmbedding(**config)
+    
+    def test_single_embedding(self):
+        text = "您吃了吗？"
+        embs = self.emb.call(text=text)
+
+    def test_batch_embedding(self):
+        text = ["您吃了吗？", 
+                "吃了吗您？"]
+        embs = self.emb.call(text=text)
+
+