Feat/llm responses

chatsky/conditions/llm.py

@@ -0,0 +1,20 @@
+from chatsky.llm.methods import BaseMethod
+from chatsky.core import BaseCondition, Context, Pipeline
+
+
+class LLMCondition(BaseCondition):
+    model_name: str
+    prompt: str
+    method: BaseMethod
+    pipeline: Pipeline
+
+    async def call(self, ctx: Context) -> bool:
+        """
+        Basic function for using LLM in condition cases.
+
+        :param model_name: Key of the model from the `Pipeline.models` dictionary.
+        :param prompt: Prompt for the model to use on users input.
+        :param method: Method that takes models output and returns boolean.
+        """
+        model = self.pipeline.models[self.model_name]
+        return await model.condition(self.prompt, self.method)

chatsky/core/pipeline.py

@@ -78,6 +78,10 @@ class Pipeline(BaseModel, extra="forbid", arbitrary_types_allowed=True):
     """
     Slots configuration.
     """
+    models: Dict = Field(default_factory=dict)
+    """
+    LLM models.
+    """
     messenger_interface: MessengerInterface = Field(default_factory=CLIMessengerInterface)
     """
     A `MessengerInterface` instance for this pipeline.
@@ -116,6 +120,7 @@ def __init__(
         *,
         default_priority: float = None,
         slots: GroupSlot = None,
+        models: dict = None,
         messenger_interface: MessengerInterface = None,
         context_storage: Union[DBContextStorage, dict] = None,
         pre_services: ServiceGroupInitTypes = None,
@@ -133,6 +138,7 @@ def __init__(
             "fallback_label": fallback_label,
             "default_priority": default_priority,
             "slots": slots,
+            "models": models,
             "messenger_interface": messenger_interface,
             "context_storage": context_storage,
             "pre_services": pre_services,

chatsky/llm/__init__.py

@@ -0,0 +1,16 @@
+try:
+    from langchain_openai import ChatOpenAI
+    from langchain_anthropic import ChatAnthropic
+    from langchain_google_vertexai import ChatVertexAI
+    from langchain_cohere import ChatCohere
+    from langchain_mistralai import ChatMistralAI
+    from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+    from langchain_core.output_parsers import StrOutputParser
+    from langchain_core.language_models.chat_models import BaseChatModel
+except ImportError:
+    raise ImportError("Langchain is not available. Please install it with `pip install chatsky[llm]`.")
+
+from chatsky.llm.filters import BaseFilter, FromTheModel, IsImportant
+from chatsky.llm.methods import BaseMethod, LogProb, LLMResult
+from chatsky.llm.llm_api import LLM_API
+from chatsky.llm.utils import message_to_langchain, __attachment_to_content

chatsky/llm/filters.py

@@ -0,0 +1,56 @@
+"""
+Filters.
+---------
+This module contains a collection of basic functions for history filtering to avoid cluttering LLMs context window.
+"""
+
+from chatsky.core.message import Message
+from chatsky.core.context import Context
+from pydantic import BaseModel
+import abc
+
+
+class BaseFilter(BaseModel, abc.ABC):
+    """
+    Base class for all message history filters.
+    """
+
+    @abc.abstractmethod
+    def __call__(self, ctx: Context, request: Message, response: Message, model_name: str) -> bool:
+        """
+        :param Context ctx: Context object.
+        :param Message request: Request message.
+        :param Message response: Response message.
+        :param str model_name: Name of the model in the Pipeline.models.
+        """
+        raise NotImplementedError
+
+
+class IsImportant(BaseFilter):
+    """
+    Filter that checks if the "important" field in a Message.misc is True.
+    """
+
+    def __call__(
+        self, ctx: Context = None, request: Message = None, response: Message = None, model_name: str = None
+    ) -> bool:
+        if request and request.misc["important"]:
+            return True
+        if response and response.misc["important"]:
+            return True
+        return False
+
+
+class FromTheModel(BaseFilter):
+    """
+    Filter that checks if the message was sent by the model.
+    """
+
+    def __call__(
+        self, ctx: Context = None, request: Message = None, response: Message = None, model_name: str = None
+    ) -> bool:
+        if request is not None and request.annotation["__generated_by_model__"] == model_name:
+            return True
+        elif response is not None and response.annotation["__generated_by_model__"] == model_name:
+            return True
+        return False

chatsky/llm/llm_api.py

@@ -0,0 +1,95 @@
+"""
+LLM responses.
+---------
+Wrapper around langchain.
+"""
+
+try:
+    from langchain_openai import ChatOpenAI
+    from langchain_anthropic import ChatAnthropic
+    from langchain_google_vertexai import ChatVertexAI
+    from langchain_cohere import ChatCohere
+    from langchain_mistralai import ChatMistralAI
+    from langchain_core.output_parsers import StrOutputParser
+    from langchain_core.language_models.chat_models import BaseChatModel
+
+    langchain_available = True
+except ImportError:
+    langchain_available = False
+
+
+from chatsky.core.message import Message
+from chatsky.core.context import Context
+from chatsky.core.pipeline import Pipeline
+from chatsky.llm.methods import BaseMethod
+
+from typing import Union, Type, Optional
+from pydantic import BaseModel
+
+from chatsky.llm.utils import message_to_langchain
+
+
+class LLM_API:
+    """
+    This class acts as a wrapper for all LLMs from langchain
+    and handles message exchange between remote model and chatsky classes.
+    """
+
+    def __init__(
+        self,
+        model: BaseChatModel,
+        system_prompt: Optional[str] = "",
+    ) -> None:
+        """
+        :param model: Model object.
+        :param system_prompt: System prompt for the model.
+        """
+        self.__check_imports()
+        self.model: BaseChatModel = model
+        self.name = ""
+        self.parser = StrOutputParser()
+        self.system_prompt = system_prompt
+
+    def __check_imports(self):
+        if not langchain_available:
+            raise ImportError("Langchain is not available. Please install it with `pip install chatsky[llm]`.")
+
+    async def __get_llm_response(self, history: list = [""], message_schema: BaseModel = None):
+        if message_schema is None:
+            result = await self.parser.ainvoke(await self.model.ainvoke(history))
+        else:
+            structured_model = self.model.with_structured_output(message_schema)
+            result = Message.model_validate(await structured_model.ainvoke(history))
+        return result
+
+    async def respond(
+        self, history: list = [""], message_schema: Union[None, Type[Message], Type[BaseModel]] = None
+    ) -> Message:
+
+        result = await self.__get_llm_response(history, message_schema)
+
+        if message_schema is None:
+            result = Message(text=result)
+        elif issubclass(message_schema, Message):
+            # Case if the message_schema desribes Message structure
+            result = Message.model_validate(result)
+        elif issubclass(message_schema, BaseModel):
+            # Case if the message_schema desribes Message.text structure
+            result = Message(text=str(result))
+
+        if result.annotations:
+            result.annotations["__generated_by_model__"] = self.name
+        else:
+            result.annotations = {"__generated_by_model__": self.name}
+        return result
+
+    async def condition(self, prompt: str, method: BaseMethod, return_schema=None):
+        async def process_input(ctx: Context, _: Pipeline) -> bool:
+            condition_history = [
+                await message_to_langchain(Message(prompt), pipeline=_, source="system"),
+                await message_to_langchain(ctx.last_request, pipeline=_, source="human"),
+            ]
+            result = method(ctx, await self.model.agenerate([condition_history], logprobs=True, top_logprobs=10))
+            return result
+
+        return process_input

chatsky/llm/methods.py

@@ -0,0 +1,70 @@
+"""
+LLM methods.
+---------
+In this file stored unified functions for some basic condition cases
+including regex search, semantic distance (cosine) etc.
+"""
+
+from chatsky.core.context import Context
+from pydantic import BaseModel
+from langchain_core.outputs.llm_result import LLMResult
+import abc
+
+
+class BaseMethod(BaseModel, abc.ABC):
+    """
+    Base class to evaluate models response as condition.
+    """
+
+    @abc.abstractmethod
+    async def __call__(self, ctx: Context, model_result: LLMResult) -> bool:
+        raise NotImplementedError
+
+    async def model_result_to_text(self, model_result: LLMResult) -> str:
+        """
+        Converts raw model generation to a string.
+        """
+        return model_result.generations[0][0].text
+
+
+class Contains(BaseMethod):
+    """
+    Simple method to check if a string contains a pattern.
+
+    :param str pattern: pattern to check
+
+    :return: True if pattern is contained in model result
+    :rtype: bool
+    """
+
+    pattern: str
+
+    async def __call__(self, ctx: Context, model_result: LLMResult) -> bool:
+        text = await self.model_result_to_text(model_result)
+        return bool(self.pattern.lower() in text.lower())
+
+
+class LogProb(BaseMethod):
+    """
+    Method to check whether a target token's log probability is higher then a threshold.
+
+    :param str target_token: token to check (e.g. `"TRUE"`)
+    :param float threshold: threshold to bypass. by default `-0.5`
+
+    :return: True if logprob is higher then threshold
+    :rtype: bool
+    """
+
+    target_token: str
+    threshold: float = -0.5
+
+    async def __call__(self, ctx: Context, model_result: LLMResult) -> bool:
+        try:
+            result = model_result.generations[0][0].generation_info["logprobs"]["content"][0]["top_logprobs"]
+        except ValueError:
+            raise ValueError("LogProb method can only be applied to OpenAI models.")
+        for tok in result:
+            if tok["token"] == self.target_token and tok["logprob"] > self.threshold:
+                return True
+
+        return False

chatsky/llm/utils.py

@@ -0,0 +1,57 @@
+import base64
+from chatsky.core.message import Image, Message
+from chatsky.core.pipeline import Pipeline
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+
+
+async def message_to_langchain(message: Message, pipeline: Pipeline, source: str = "human", max_size: int = 1000):
+    """
+    Creates a langchain message from a ~chatsky.script.core.message.Message object.
+
+    :param Message message: ~chatsky.script.core.message.Message object.
+    :param Pipeline pipeline: ~chatsky.pipeline.Pipeline object.
+    :param str source: Source of a message [`human`, `ai`, `system`]. Defaults to "human".
+    :param int max_size: Maximum size of the message in symbols. 
+    If exceed the limit will raise ValueError. Is not affected by system prompt size.
+
+    :return: Langchain message object.
+    :rtype: HumanMessage|AIMessage|SystemMessage
+    """
+    if len(message.text) > max_size:
+        raise ValueError("Message is too long.")
+
+    if message.text is None:
+        message.text = ""
+    content = [{"type": "text", "text": message.text}]
+
+    if message.attachments:
+        for image in message.attachments:
+            if isinstance(image, Image):
+                content.append(
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": await __attachment_to_content(image, pipeline.messenger_interface)},
+                    }
+                )
+
+    if source == "human":
+        return HumanMessage(content=content)
+    elif source == "ai":
+        return AIMessage(content=content)
+    elif source == "system":
+        return SystemMessage(content=content)
+    else:
+        raise ValueError("Invalid source name. Only `human`, `ai` and `system` are supported.")
+
+
+async def __attachment_to_content(attachment: Image, iface) -> str:
+    """
+    Helper function to convert image to base64 string.
+    """
+    image_bytes = await attachment.get_bytes(iface)
+    image_b64 = base64.b64encode(image_bytes).decode("utf-8")
+    extension = str(attachment.source).split(".")[-1]
+    if image_b64 == "" or extension is None:
+        raise ValueError("Data image is not accessible.")
+    image_b64 = f"data:image/{extension};base64,{image_b64}"
+    return image_b64