add llamaindex support

memory_scope/db/elastic_search_client.py

@@ -1,11 +1,14 @@
 from elasticsearch import Elasticsearch
 from elasticsearch.helpers import bulk
 
-from common.dash_embedding_client import DashEmbeddingClient
-from common.logger import Logger
+from memory_scope.models.dash_embedding_client import DashEmbeddingClient, LLIEmbedding
 from constants.common_constants import ES_ENV_URL_DICT
 from enumeration.env_type import EnvType
-
+from memory_scope.utils.logger import Logger
+from llama_index.core import VectorStoreIndex, StorageContext, ServiceContext
+from llama_index.vector_stores.elasticsearch import ElasticsearchStore
+from llama_index.core.schema import TextNode
+from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy
 
 class ElasticSearchClient(object):
     def __init__(self,
@@ -339,3 +342,73 @@ class ElasticSearchClient(object):
             self.print_hits(hits)
 
         return hits
+    
+ 
+class LLIElasticSearch(object):
+    def __init__(self,
+                 es_index_name: str,
+                 embedding_client: LLIEmbedding | None = None,
+                 retrieve_topk: int = 3,
+                 content_key: str = "text",
+                 ):
+        self.es_index_name = es_index_name
+        self.content_key = content_key
+        self.embedding_client: LLIEmbedding = embedding_client
+        self.es_client = ElasticsearchStore(index_name="my_index",
+                                            es_url="http://localhost:9200",
+                                            retrieval_strategy=AsyncDenseVectorStrategy(hybrid=True))
+        
+        self.service_context = ServiceContext.from_defaults(embed_model=self.embedding_client, llm=None)
+        self.storage_context = StorageContext.from_defaults(vector_store=self.es_client)
+        self.index = VectorStoreIndex(storage_context=self.storage_context, 
+                                      service_context=self.service_context)
+
+        self.retriever = self.index.as_retriever(similarity_top_k=retrieve_topk)
+        self.logger = Logger.get_logger()
+
+    def log_index_info(self, ):
+        pass
+
+    def print_hits(self, hits: list):
+        for hit in hits:
+            print_kwargs = {
+                "_id": hit['_id'],
+                "_score": hit['_score'],
+            }
+            for k, v in hit['_source'].items():
+                # 不打印vector
+                if k == self.vector_key:
+                    v = len(v)
+                print_kwargs[k] = v
+            self.logger.info(" ".join([f"{k}={v}" for k, v in print_kwargs.items()]))
+
+    def similar_search(self,
+                       text: str, 
+                       size: int, ):
+        
+        ret_nodes = self.retriever.retrieve(text)
+        return ret_nodes
+    
+    def insert_batch(self, doc_list:list[str]):
+        node_list = []
+        for doc in doc_list:
+            assert "_id" in doc and "_source" in doc
+            content = doc["_source"]["text"]
+            doc["_source"].pop("text")
+            meta = doc["_source"]
+            node = TextNode(text=content, metadata=meta)
+            node.node_id(doc['_id'])
+            node_list.append(node)
+        self.index.insert_nodes(node_list)
+
+
+    def insert(self, _id: str, body: dict):
+        assert body and self.content_key in body, f"body={body} is illegal!"        
+        content = body[self.content_key]
+        body.pop(self.content_key)
+        meta = body
+        node = TextNode(text=content, metadata=meta)
+        self.index.insert_nodes([node])
+
+
+  

memory_scope/models/__init__.py

@@ -0,0 +1,18 @@
+from utils.registry import Registry
+
+from llama_index.embeddings.dashscope import (
+    DashScopeEmbedding,
+)
+# from llama_index.postprocessor.dashscope_rerank import DashScopeRerank
+
+from llama_index.llms.dashscope import DashScope # type: ignore
+
+
+EMB = Registry('embedding')
+EMB.register_module(DashScopeEmbedding)
+
+# RERANKER = Registry('reranker')
+# RERANKER.register_module(DashScopeRerank)
+
+LLM = Registry('llm')
+LLM.register_module(DashScope)

memory_scope/models/dash_client.py

@@ -89,3 +89,36 @@ class DashClient(object):
                 time.sleep(self.retry_sleep_time)
 
         return None
+
+
+class LLIClient(object):
+
+    def __init__(self,
+                 model_name: str,
+                 timeout: int = None,
+                 max_retry_count: int = 2,
+                 retry_sleep_time: float = 1.0,
+                 **kwargs):
+
+        self.model_name: str = model_name
+        self.timeout: int = timeout
+        self.max_retry_count: int = max_retry_count
+        self.retry_sleep_time: float = retry_sleep_time
+        self.kwargs: dict = kwargs
+
+        self.data = {}
+        self.logger = Logger.get_logger()
+    
+
+    def before_call(self, **kwargs):
+        pass
+
+    def after_call(self, **kwargs):
+        pass
+
+    def call_once(self, **kwargs):
+        pass
+
+    def call(self, **kwargs):
+        pass
+    

memory_scope/models/dash_embedding_client.py

@@ -1,8 +1,16 @@
 from typing import List, Dict
 
 import dashscope
+import time 
+
+from models import EMB
+from models.dash_client import DashClient, LLIClient
+
+from typing import List, Dict
+from utils.registry import build_from_cfg
+from utils.timer import Timer
+
 
-from common.dash_client import DashClient
 from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT
 from enumeration.dash_api_enum import DashApiEnum
 
@@ -41,3 +49,55 @@ class DashEmbeddingClient(DashClient):
         if len(embedding_results) == 1:
             embedding_results = list(embedding_results.values())[0]
         return embedding_results
+
+class LLIEmbedding(LLIClient):
+
+    def __init__(self, method, model_name, **kwargs):
+        super(LLIEmbedding, self).__init__(model_name, **kwargs)
+        self.config = {
+                       "method": method,
+                       "model_name": model_name,
+                       **kwargs}
+        self.embedder = build_from_cfg(self.config, EMB)
+
+    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, emb: Dict[int, List[float]], **kwargs) -> Dict[int, List[float]] | List[float]:
+        embedding_results = {}
+        for idx, e in enumerate(emb):
+            embedding_results[idx] = e
+
+        if len(embedding_results) == 1:
+            embedding_results = list(embedding_results.values())[0]
+        return embedding_results
+
+
+    def call_once(self, model_name: str = None, retry_cnt: int = 0, **kwargs):
+        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:
+            self.logger.debug(f"data={self.data} timeout={self.timeout}")
+            try:
+                results = self.embedder.get_text_embedding_batch(**self.data)  
+                results = self.after_call(results)
+                return results, True
+            except Exception as e:
+                self.logger.debug(f"Get Error in Embedding: {e}")
+                return None, False
+        
+
+    def call(self, model_name: str = None, **kwargs):
+        for i in range(self.max_retry_count):
+            result, flag = self.call_once(model_name=model_name, retry_cnt=i, **kwargs)
+            if flag:
+                return result
+            else:
+                time.sleep(self.retry_sleep_time)
+        return None

memory_scope/models/dash_generate_client.py

@@ -2,10 +2,21 @@ from typing import List, Dict
 
 import dashscope
 
-from common.dash_client import DashClient
+from models.dash_client import DashClient, LLIClient
 from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT
 from enumeration.dash_api_enum import DashApiEnum
 
+import time
+from typing import List, Dict
+from utils.timer import Timer
+from models import LLM
+from utils.registry import build_from_cfg
+from llama_index.core.base.llms.types import ChatMessage
+from llama_index.core.base.llms.types import (
+    ChatResponse,
+    CompletionResponse,
+)
+
 
 class DashGenerateClient(DashClient):
     """
@@ -43,3 +54,76 @@ class DashGenerateClient(DashClient):
             return output["choices"][0]["message"]["content"]
         else:
             raise NotImplementedError
+
+
+class LLILLM(LLIClient):
+
+    def __init__(self, method, model_name: str, **kwargs):
+        super(LLILLM, self).__init__(model_name, **kwargs)
+        self.config = {
+                       "method": method,
+                       "model_name": model_name,
+                       **kwargs}
+        self.llm = build_from_cfg(self.config, LLM)
+
+
+    def before_call(self, model_name: str = None, **kwargs):
+        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, response_obj, **kwargs):
+        self.logger.debug(f"response_obj={response_obj}")
+        if isinstance(response_obj, CompletionResponse):
+            return response_obj.text
+        elif isinstance(response_obj, ChatResponse):
+            return response_obj.message.content
+        else:
+            raise NotImplementedError
+        
+    
+    def call_once(self, model_name: str = None, retry_cnt: int = 0, **kwargs):
+        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:
+            self.logger.debug(f"data={self.data} timeout={self.timeout}")
+            if True:
+          # try:
+                if 'prompt' in self.data:
+                    results = self.llm.complete(**self.data)
+                else:
+                    results = self.llm.chat(**self.data)    
+                results = self.after_call(results)
+                return results, True
+          #  except:
+          #      return None, False
+        
+
+    def call(self, model_name: str = None, **kwargs):
+        for i in range(self.max_retry_count):
+            result, flag = self.call_once(model_name=model_name, retry_cnt=i, **kwargs)
+            print("dashscope llm results:",result)
+            if flag:
+                return result
+            else:
+                time.sleep(self.retry_sleep_time)
+        return None

memory_scope/models/dash_rerank_client.py

@@ -2,11 +2,21 @@ from typing import List
 
 import dashscope
 
-from common.dash_client import DashClient
+from models.dash_client import DashClient, LLIClient
 from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT
 from enumeration.dash_api_enum import DashApiEnum
 
 
+import time
+from typing import List
+from models import RERANKER
+from utils.timer import Timer
+
+from utils.registry import build_from_cfg
+from llama_index.core.data_structs import Node
+from llama_index.core.schema import NodeWithScore # type: ignore
+
+
 class DashReRankClient(DashClient):
     """
     url: https://help.aliyun.com/document_detail/2780059.html
@@ -41,3 +51,69 @@ class DashReRankClient(DashClient):
 
     def after_call(self, response_obj, **kwargs):
         return response_obj["output"]["results"]
+
+
+class LLIReRank(LLIClient):
+   
+    def __init__(self, method, model_name, **kwargs):
+        super(LLIReRank, self).__init__(model_name, **kwargs)
+
+        self.config = {
+                       "method": method,
+                       "model_name": model_name,
+                       **kwargs}
+        self.reranker = build_from_cfg(self.config, RERANKER)
+
+
+    def before_call(self, model_name: str = None, **kwargs):
+        query: str = kwargs.pop("query", "")
+        documents: List[str] = kwargs.pop("documents", [])
+        top_n: int | None = kwargs.pop("top_n", None)
+        return_documents: bool = kwargs.pop("return_documents", False)
+        
+        assert query and documents, f"query or documents is empty! query={query}, documents={len(documents)}"
+        if top_n is None:
+            top_n = len(documents)
+        
+        nodes = [NodeWithScore(Node(text=text, score=1.0)) for text in documents]
+        self.reranker = self.reranker(top_n=top_n, 
+                                   return_documents=return_documents)
+
+        self.data = {
+            "nodes": nodes,
+            "query_str": query,
+        }
+        
+
+    def after_call(self, nodes, **kwargs):
+        results = []
+        for node in nodes:
+            results.append(dict(relevance_score=node.score,
+                                document=node.node.text))
+        return results
+
+
+    def call_once(self, model_name: str = None, retry_cnt: int = 0, **kwargs):
+        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:
+            self.logger.debug(f"data={self.data} timeout={self.timeout}")
+            try:
+                results = self.reranker.postprocess_nodes(*self.data)    
+                results = self.after_call(results)
+                return results, True
+            except:
+                return None, False
+        
+
+    def call(self, model_name: str = None, **kwargs):
+        for i in range(self.max_retry_count):
+            result, flag = self.call_once(model_name=model_name, retry_cnt=i, **kwargs)
+            if flag:
+                return result
+            else:
+                time.sleep(self.retry_sleep_time)
+        return None

memory_scope/utils/registry.py

@@ -0,0 +1,54 @@
+"""
+Registry for different modules.
+Init class according to the class name and verify the input parameters.
+"""
+import inspect
+
+
+class Registry:
+    def __init__(self, name):
+        self.name = name
+        self.module_dict = dict()
+
+    def register_module(self, module, module_name=None):
+        if module_name is None:
+            module_name = module.__name__
+        if module_name in self.module_dict:
+            raise KeyError(f'{module_name} is already registered in {self.name}')
+        self.module_dict[module_name] = module
+
+    def get_module(self, module_name):
+        assert module_name in self.module_dict, f'{module_name} not found in {self.name}'
+        return self.module_dict[module_name]
+
+
+def build_from_cfg(config, registry, default_args: dict = None, skip_param_check=False):
+
+    if default_args is None:
+        default_args = {}
+
+    args = config.copy()
+    method_type = args.pop('method')
+    #params = args.get("parameters", {}) or default_args
+    params = args
+    if isinstance(method_type, str):
+        obj_cls = registry.get_module(method_type)   
+    else:
+        raise TypeError(
+            f'type must be a str or valid type, but got {type(method_type)}')
+
+    allowed_params = list(inspect.signature(obj_cls.__init__).parameters.keys())
+    print(allowed_params, params)
+    if not skip_param_check:
+        filter_params = {key: value for key, value in params.items() if key in allowed_params}
+    else:
+        filter_params = params
+    # print(
+    #     f"Registry {registry.name}, "
+    #     f"allowed parameters {allowed_params}, filter parameters {filter_params}",
+    #     flush=True
+    # )
+    print(filter_params)
+    return obj_cls(**filter_params)
+
+