Adds LCEL native support

Author: vmesel

docs/plugins.md

@@ -51,4 +51,8 @@ PLUGINS=plugins.whats_audio_synth.main,
 
 ### Using Langchain's LCEL structure
 
-If you want to use Langchain's LCEL structure, you can use the `DialogLcelLLM` setting the variable to the 'dialog.llm.lcel_default.DialogLcelLLM'.
+If you want to use Langchain's LCEL structure, you can use the default LCEL implementation available in the file `dialog/llm/agents/lcel.py`. To use it, you need to add the following environment variable to your .env file:
+
+```bash
+LLM_CLASS=dialog.llm.agents.lcel.runnable
+```

src/dialog/llm/__init__.py

@@ -3,12 +3,17 @@
 
 from typing import Type
 
-from .default import DialogLLM
 from dialog_lib.agents.abstract import AbstractLLM
 from dialog.settings import Settings
 
+from langchain.schema.runnable import RunnablePassthrough
+from langchain_core.runnables.base import RunnableSequence
 
-def get_llm_class() -> Type[AbstractLLM]:
+from .agents.default import *
+from .agents.lcel import *
+
+
+def get_llm_class():
     if Settings().LLM_CLASS is None:
         return DialogLLM
 
@@ -20,5 +25,34 @@ def get_llm_class() -> Type[AbstractLLM]:
         llm_class = getattr(llm_module, class_name)
     except Exception as e:
         logging.info(f"Failed to load LLM class {Settings().LLM_CLASS}. Using default LLM. Exception: {e}")
+        raise
+
+    if isinstance(llm_class, AbstractLLM):
+        return llm_class, "AbstractLLM"
+
+    elif 'langchain_core.runnables' in str(type(llm_class)):
+        return llm_class, "LCELRunnable"
+
+    logging.info(f"Type for LLM class is: {type(llm_class)}")
+
+    return DialogLLM, "AbstractLLM"
+
+
+def process_user_message(message, chat_id=None):
+    LLM, llm_type = get_llm_class()
+    if llm_type == "AbstractLLM":
+        llm_instance = LLM(config=Settings().PROJECT_CONFIG, session_id=chat_id)
+        ai_message = llm_instance.process(message.message)
+
+    elif llm_type == "LCELRunnable":
+        ai_message = LLM.invoke(
+            {"input": message.message},
+            {"configurable": {
+                "session_id": chat_id,
+                "model": Settings().PROJECT_CONFIG.get("model_name", "gpt-3.5-turbo"),
+                **Settings().PROJECT_CONFIG,
+            }}
+        )
+        ai_message = {"text": ai_message.content}
 
-    return llm_class or DialogLLM
+    return ai_message

src/dialog/llm/agents/__init__.py

@@ -0,0 +1,56 @@
+
+from dialog.db import get_session
+from dialog.settings import Settings
+from dialog_lib.memory import generate_memory_instance
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
+from langchain_openai.chat_models import ChatOpenAI
+from langchain_core.runnables import ConfigurableField
+from langchain_core.runnables.history import RunnableWithMessageHistory
+from langchain.output_parsers.json import SimpleJsonOutputParser
+
+# Here we define the model that we will be using as a base for our agent
+# as well as the model_name and temperature (some of these parameterers
+# are exclusive to OpenAI instances).
+chat_model = ChatOpenAI(
+    model_name="gpt-3.5-turbo",
+    temperature=0,
+    openai_api_key=Settings().OPENAI_API_KEY,
+).configurable_fields(
+    model_name=ConfigurableField(
+        id="model_name",
+        name="GPT Model",
+        description="The GPT model to use"
+    ),
+    temperature=ConfigurableField(
+        id="temperature",
+        name="Temperature",
+        description="The temperature to use"
+    )
+)
+
+prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            Settings().PROJECT_CONFIG.get("prompt").get("prompt", "What can I help you with today?")
+        ),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{input}"),
+    ]
+)
+
+def get_memory_instance(session_id):
+    return generate_memory_instance(
+        session_id=session_id,
+        dbsession=next(get_session()),
+        database_url=Settings().DATABASE_URL
+    )
+
+chain = prompt | chat_model
+
+runnable = RunnableWithMessageHistory(
+    chain,
+    get_memory_instance,
+    input_messages_key='input',
+    history_messages_key="chat_history"
+)

src/dialog/llm/default.py renamed to src/dialog/llm/agents/default.py

@@ -3,7 +3,7 @@
 import logging
 from pydantic import parse_obj_as
 
-from dialog.llm import get_llm_class
+from dialog.llm import process_user_message
 from dialog_lib.db.memory import get_messages
 from dialog_lib.db.models import Chat as ChatEntity
 from dialog.schemas import ChatModel, SessionModel, SessionsModel
@@ -41,9 +41,7 @@ async def post_message(chat_id: str, message: ChatModel, session: Session = Depe
             detail="Chat ID not found",
         )
     start_time = datetime.datetime.now()
-    LLM = get_llm_class()
-    llm_instance = LLM(config=Settings().PROJECT_CONFIG, session_id=chat_id)
-    ai_message = llm_instance.process(message.message)
+    ai_message = process_user_message(message, chat_id)
     duration = datetime.datetime.now() - start_time
     logging.info(f"Request processing time for chat_id {chat_id}: {duration}")
     return {"message": ai_message["text"]}
@@ -56,9 +54,7 @@ async def ask_question_to_llm(message: ChatModel, session: Session = Depends(get
     using memory.
     """
     start_time = datetime.datetime.now()
-    LLM = get_llm_class()
-    llm_instance = LLM(config=Settings().PROJECT_CONFIG)
-    ai_message = llm_instance.process(message.message)
+    ai_message = process_user_message(message, chat_id=None)
     duration = datetime.datetime.now() - start_time
     logging.info(f"Request processing time: {duration}")
     return {"message": ai_message["text"]}

src/dialog/llm/agents/lcel.py

@@ -43,6 +43,6 @@ def chat_session(dbsession):
 
 @pytest.fixture
 def llm_mock(mocker):
-    llm_mock = mocker.patch('dialog.routers.dialog.get_llm_class')
+    llm_mock = mocker.patch('dialog.routers.dialog.process_user_message')
     llm_mock.process.return_value = {"text": "Hello"}
     return llm_mock

src/dialog/llm/lcel_default.py

@@ -3,7 +3,8 @@
 
 from dialog.llm import get_llm_class
 from dialog_lib.agents.abstract import AbstractLLM
-from dialog.llm.default import DialogLLM
+from dialog.llm.agents.default import DialogLLM
+from dialog.llm import DialogLLM
 
 
 def test_abstract_llm_for_invalid_config():
@@ -38,10 +39,11 @@ def test_get_llm_class_get_default_class():
 
 def test_get_llm_class_get_custom_class():
     os.environ["LLM_CLASS"] = "dialog_lib.agents.abstract.AbstractLLM"
-    llm_class = get_llm_class()
-    assert llm_class == AbstractLLM
+    llm_class, llm_type = get_llm_class()
+    assert llm_class == DialogLLM
+    assert llm_type == "AbstractLLM"
 
 def test_get_llm_class_with_invalid_class():
-    os.environ["LLM_CLASS"] = "dialogl.llm.invalid_llm.InvalidLLM"
-    llm_class = get_llm_class()
-    assert llm_class == DialogLLM
+    os.environ["LLM_CLASS"] = "dialog.llm.invalid_llm.InvalidLLM"
+    with pytest.raises(ModuleNotFoundError):
+        llm_class, llm_type = get_llm_class()