diff --git a/memory_scope/models/base_generation_model.py b/memory_scope/models/base_generation_model.py index de38466f..0639f29a 100644 --- a/memory_scope/models/base_generation_model.py +++ b/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 + diff --git a/memory_scope/models/response.py b/memory_scope/models/response.py index 11890183..5c435042 100644 --- a/memory_scope/models/response.py +++ b/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] diff --git a/old/elastic_search_client.py b/old/elastic_search_client.py index eb8490bd..1907bb65 100644 --- a/old/elastic_search_client.py +++ b/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))