[Feature] Add support for vllm as llm source (#1149)

embedchain/app.py

@@ -9,14 +9,9 @@ from typing import Any, Dict, Optional
 import requests
 import yaml
 
-from embedchain.cache import (
-    Config,
-    ExactMatchEvaluation,
-    SearchDistanceEvaluation,
-    cache,
-    gptcache_data_manager,
-    gptcache_pre_function,
-)
+from embedchain.cache import (Config, ExactMatchEvaluation,
+                              SearchDistanceEvaluation, cache,
+                              gptcache_data_manager, gptcache_pre_function)
 from embedchain.client import Client
 from embedchain.config import AppConfig, CacheConfig, ChunkerConfig
 from embedchain.constants import SQLITE_PATH

embedchain/config/llm/base.py

@@ -73,7 +73,7 @@ class BaseLlmConfig(BaseConfig):
         callbacks: Optional[List] = None,
         api_key: Optional[str] = None,
         endpoint: Optional[str] = None,
-        model_kwargs: Optional[Dict[str, Any]] = {},
+        model_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """
         Initializes a configuration class instance for the LLM.
@@ -115,6 +115,8 @@ class BaseLlmConfig(BaseConfig):
         :type model_kwargs: Optional[Dict[str, Any]], optional
         :param callbacks: Langchain callback functions to use, defaults to None
         :type callbacks: Optional[List], optional
+        :param query_type: The type of query to use, defaults to None
+        :type query_type: Optional[str], optional
         :raises ValueError: If the template is not valid as template should
         contain $context and $query (and optionally $history)
         :raises ValueError: Stream is not boolean
@@ -142,6 +144,7 @@ class BaseLlmConfig(BaseConfig):
         self.api_key = api_key
         self.endpoint = endpoint
         self.model_kwargs = model_kwargs
+
         if type(prompt) is str:
             prompt = Template(prompt)
 

embedchain/embedchain.py

@@ -7,7 +7,12 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 from dotenv import load_dotenv
 from langchain.docstore.document import Document
 
-from embedchain.cache import adapt, get_gptcache_session, gptcache_data_convert, gptcache_update_cache_callback
+from embedchain.cache import (
+    adapt,
+    get_gptcache_session,
+    gptcache_data_convert,
+    gptcache_update_cache_callback,
+)
 from embedchain.chunkers.base_chunker import BaseChunker
 from embedchain.config import AddConfig, BaseLlmConfig, ChunkerConfig
 from embedchain.config.base_app_config import BaseAppConfig

embedchain/llm/base.py

@@ -4,7 +4,9 @@ from typing import Any, Dict, Generator, List, Optional
 from langchain.schema import BaseMessage as LCBaseMessage
 
 from embedchain.config import BaseLlmConfig
-from embedchain.config.llm.base import DEFAULT_PROMPT, DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE, DOCS_SITE_PROMPT_TEMPLATE
+from embedchain.config.llm.base import (DEFAULT_PROMPT,
+                                        DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE,
+                                        DOCS_SITE_PROMPT_TEMPLATE)
 from embedchain.helpers.json_serializable import JSONSerializable
 from embedchain.memory.base import ChatHistory
 from embedchain.memory.message import ChatMessage

embedchain/llm/vllm.py

@@ -0,0 +1,40 @@
+from typing import Iterable, Optional, Union
+
+from langchain.callbacks.manager import CallbackManager
+from langchain.callbacks.stdout import StdOutCallbackHandler
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain_community.llms import VLLM as BaseVLLM
+
+from embedchain.config import BaseLlmConfig
+from embedchain.helpers.json_serializable import register_deserializable
+from embedchain.llm.base import BaseLlm
+
+
+@register_deserializable
+class VLLM(BaseLlm):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config=config)
+        if self.config.model is None:
+            self.config.model = "mosaicml/mpt-7b"
+
+    def get_llm_model_answer(self, prompt):
+        return self._get_answer(prompt=prompt, config=self.config)
+
+    @staticmethod
+    def _get_answer(prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]:
+        callback_manager = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()]
+
+        # Prepare the arguments for BaseVLLM
+        llm_args = {
+            "model": config.model,
+            "temperature": config.temperature,
+            "top_p": config.top_p,
+            "callback_manager": CallbackManager(callback_manager),
+        }
+
+        # Add model_kwargs if they are not None
+        if config.model_kwargs is not None:
+            llm_args.update(config.model_kwargs)
+
+        llm = BaseVLLM(**llm_args)
+        return llm(prompt)

embedchain/vectordb/zilliz.py

@@ -6,7 +6,15 @@ from embedchain.helpers.json_serializable import register_deserializable
 from embedchain.vectordb.base import BaseVectorDB
 
 try:
-    from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, MilvusClient, connections, utility
+    from pymilvus import (
+        Collection,
+        CollectionSchema,
+        DataType,
+        FieldSchema,
+        MilvusClient,
+        connections,
+        utility,
+    )
 except ImportError:
     raise ImportError(
         "Zilliz requires extra dependencies. Install with `pip install --upgrade embedchain[milvus]`"