update LLILLM

memory_scope/models/base_generation_model.py

@@ -1,8 +1,15 @@
-from llama_index.llms.dashscope import DashScope as DashScopeLLM
+from typing import List, Dict
 
+from llama_index.llms.dashscope import DashScope as DashScopeLLM
+from llama_index.core.base.llms.types import ChatMessage
+from llama_index.core.base.llms.types import (
+    ChatResponse,
+    CompletionResponse,
+)
 from memory_scope.models import MODEL_REGISTRY
 from memory_scope.models.base_model import BaseModel
 from memory_scope.models.response import ModelResponse, ModelResponseGen
+from memory_scope.utils.timer import Timer
 
 
 class BaseGenerationModel(BaseModel):
@@ -22,3 +29,69 @@ class BaseGenerationModel(BaseModel):
 
     async def _async_call(self, **kwargs) -> ModelResponse:
         pass
+
+
+class LLILLM(BaseGenerateModel):
+
+    def before_call(self, **kwargs) -> None:
+        prompt: str = kwargs.pop("prompt", "")
+        messages: List[Dict[str, str]] = kwargs.pop("messages", [])
+
+        if prompt:
+            input_text = prompt
+            input_type = 'prompt'
+            llama_input = input_text
+        elif messages:
+            input_text = messages
+            input_type = 'messages'
+            llama_input = [ChatMessage(
+                role=x['role'], content=x['content']
+            ) for x in input_text]
+        else:
+            raise RuntimeError("prompt and messages is both empty!")
+
+        self.data = {
+            input_type: llama_input,
+        }
+
+    def after_call(
+            self, call_result: ChatResponse | CompletionResponse) -> ModelResponse | ModelResponseGen:
+        if isinstance(call_result, CompletionResponse):
+            content = call_result.text
+        elif isinstance(call_result, ChatResponse):
+            content = call_result.message.content
+        else:
+            raise NotImplementedError
+        
+        return ModelResponse(text=content,
+                             model_type="LLM")
+    
+    def _call(self, stream: bool = False, **kwargs) -> ModelResponse | ModelResponseGen:
+        if model_name is None:
+            model_name = self.model_name
+
+        self.before_call(model_name=model_name, **kwargs)
+
+        with Timer(self.__class__.__name__, log_time=False) as t:
+            assert "prompt" in self.data or "messages" in self.data
+
+            try:
+                if 'prompt' in self.data:
+                    if stream:
+                        response = self.llm.stream_complete(**self.data)
+                    else:
+                        response = self.llm.complete(**self.data)
+                else:
+                    if stream:
+                        response = self.llm.stream_chat(**self.data)
+                    else:
+                        response = self.llm.chat(**self.data)    
+                results = self.after_call(response)
+                results.details = response
+            except Exception as e:
+                results = ModelResponse(model_type="LLM",
+                                        status=False,
+                                        details=e)
+
+            return results
+

memory_scope/models/response.py

@@ -8,13 +8,14 @@ class ModelResponse(BaseModel):
 
     embedding_results: Dict[int, List[float]] | List[float] = Field([], description="")
 
-    rank_scores: Dict[int, float] = Field({}, description="")
+    rank_scores: Dict[int, float] = Field({}, description="The rank scores of each documents.")
 
-    model_type: str = Field("", description="")
+    model_type: str = Field("", description="One of LLM, EMB, RANK.")
 
-    status: bool = Field(True, description="")
+    status: bool = Field(True, description="Indicates whether the model call was successful.")
 
-    details: str = Field("", description="")
+    details: str = Field("", description=("The details information for model call, \
+                         usually for storage of raw response or failure messages."))
 
 
 ModelResponseGen = Generator[ModelResponse, None, None]

old/elastic_search_client.py

@@ -358,6 +358,7 @@ class LLIElasticSearch(object):
         self.es_index_name = es_index_name
         self.content_key = content_key
         self.embedding_client: LLIEmbedding = embedding_client
+        # using local es for debug convenient
         self.es_client = ElasticsearchStore(index_name="my_index",
                                             es_url="http://localhost:9200",
                                             retrieval_strategy=AsyncDenseVectorStrategy(hybrid=True))