[Feature] Add support for Groq LLMs (#1284)

embedchain/factory.py

@@ -23,6 +23,7 @@ class LlmFactory:
         "google": "embedchain.llm.google.GoogleLlm",
         "aws_bedrock": "embedchain.llm.aws_bedrock.AWSBedrockLlm",
         "mistralai": "embedchain.llm.mistralai.MistralAILlm",
+        "groq": "embedchain.llm.groq.GroqLlm",
     }
     provider_to_config_class = {
         "embedchain": "embedchain.config.llm.base.BaseLlmConfig",

embedchain/llm/groq.py

@@ -0,0 +1,43 @@
+import os
+from typing import Optional
+
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.schema import HumanMessage, SystemMessage
+
+try:
+    from langchain_groq import ChatGroq
+except ImportError:
+    raise ImportError("Groq requires extra dependencies. Install with `pip install langchain-groq`") from None
+
+
+from embedchain.config import BaseLlmConfig
+from embedchain.helpers.json_serializable import register_deserializable
+from embedchain.llm.base import BaseLlm
+
+
+@register_deserializable
+class GroqLlm(BaseLlm):
+    def __init__(self, config: Optional[BaseLlmConfig] = None):
+        super().__init__(config=config)
+
+    def get_llm_model_answer(self, prompt) -> str:
+        response = self._get_answer(prompt, self.config)
+        return response
+
+    def _get_answer(self, prompt: str, config: BaseLlmConfig) -> str:
+        messages = []
+        if config.system_prompt:
+            messages.append(SystemMessage(content=config.system_prompt))
+        messages.append(HumanMessage(content=prompt))
+        api_key = config.api_key or os.environ["GROQ_API_KEY"]
+        kwargs = {
+            "model_name": config.model or "mixtral-8x7b-32768",
+            "temperature": config.temperature,
+            "groq_api_key": api_key,
+        }
+        if config.stream:
+            callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
+            chat = ChatGroq(**kwargs, streaming=config.stream, callbacks=callbacks, api_key=api_key)
+        else:
+            chat = ChatGroq(**kwargs)
+        return chat.invoke(messages).content

embedchain/llm/openai.py

@@ -58,8 +58,7 @@ class OpenAILlm(BaseLlm):
         messages: list[BaseMessage],
     ) -> str:
         from langchain.output_parsers.openai_tools import JsonOutputToolsParser
-        from langchain_core.utils.function_calling import \
-            convert_to_openai_tool
+        from langchain_core.utils.function_calling import convert_to_openai_tool
 
         openai_tools = [convert_to_openai_tool(tools)]
         chat = chat.bind(tools=openai_tools).pipe(JsonOutputToolsParser())

embedchain/utils/misc.py

@@ -406,9 +406,11 @@ def validate_config(config_data):
                     "aws_bedrock",
                     "mistralai",
                     "vllm",
+                    "groq",
                 ),
                 Optional("config"): {
                     Optional("model"): str,
+                    Optional("model_name"): str,
                     Optional("number_documents"): int,
                     Optional("temperature"): float,
                     Optional("max_tokens"): int,