Feat/llm responses

chatsky/conditions/__init__.py

@@ -11,3 +11,4 @@
 )
 from chatsky.conditions.slots import SlotsExtracted
 from chatsky.conditions.service import ServiceFinished
+from chatsky.conditions.llm import LLMCondition

chatsky/conditions/llm.py

@@ -0,0 +1,32 @@
+"""
+LLM Conditions
+--------------
+This module provides LLM-based conditions.
+"""
+
+from chatsky.llm.methods import BaseMethod
+from chatsky.core import BaseCondition, Context
+
+
+class LLMCondition(BaseCondition):
+    """
+    LLM-based condition.
+    Uses prompt to produce result from model and evaluates the result using given method.
+    """
+
+    model_name: str
+    """
+    Key of the model in the :py:attr:`~chatsky.core.pipeline.Pipeline.models` dictionary.
+    """
+    prompt: str
+    """
+    Condition prompt.
+    """
+    method: BaseMethod
+    """
+    Method that takes model's output and returns boolean.
+    """
+
+    async def call(self, ctx: Context) -> bool:
+        model = ctx.pipeline.models[self.model_name]
+        return await model.condition(ctx, self.prompt, self.method)

chatsky/core/pipeline.py

@@ -29,6 +29,7 @@
 from chatsky.core.service.actor import Actor
 from chatsky.core.node_label import AbsoluteNodeLabel, AbsoluteNodeLabelInitTypes
 from chatsky.core.script_parsing import JSONImporter, Path
+from chatsky.llm.llm_api import LLM_API
 
 logger = logging.getLogger(__name__)
 
@@ -78,6 +79,10 @@ class Pipeline(BaseModel, extra="forbid", arbitrary_types_allowed=True):
     """
     Slots configuration.
     """
+    models: Dict[str, LLM_API] = Field(default_factory=dict)
+    """
+    LLM models.
+    """
     messenger_interface: MessengerInterface = Field(default_factory=CLIMessengerInterface)
     """
     A `MessengerInterface` instance for this pipeline.
@@ -116,6 +121,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 +139,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,3 @@
+from chatsky.llm.filters import BaseHistoryFilter, FromModel, IsImportant
+from chatsky.llm.methods import BaseMethod, LogProb, Contains
+from chatsky.llm.llm_api import LLM_API

chatsky/llm/_langchain_imports.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+try:
+    from langchain_core.output_parsers import StrOutputParser
+    from langchain_core.language_models.chat_models import BaseChatModel
+    from langchain_core.messages.base import BaseMessage
+    from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+    from langchain_core.outputs.llm_result import LLMResult
+
+    langchain_available = True
+except ImportError:
+    StrOutputParser = Any
+    BaseChatModel = Any
+    BaseMessage = Any
+    HumanMessage = Any
+    SystemMessage = Any
+    AIMessage = Any
+    LLMResult = Any
+
+    langchain_available = False
+
+
+def check_langchain_available():
+    if not langchain_available:
+        raise ImportError("Langchain is not available. Please install it with `pip install chatsky[llm]`.")

chatsky/llm/filters.py

@@ -0,0 +1,93 @@
+"""
+Filters
+---------
+This module contains a collection of basic functions for history filtering to avoid cluttering LLMs context window.
+"""
+
+import abc
+from enum import Enum
+from pydantic import BaseModel
+
+from logging import Logger
+
+from chatsky.core.message import Message
+from chatsky.core.context import Context
+
+logger = Logger(name=__name__)
+
+
+class BaseHistoryFilter(BaseModel, abc.ABC):
+    """
+    Base class for all message history filters.
+    """
+
+    class Return(Enum):
+        Request = 1
+        Response = 2
+        Turn = 3
+
+    def call(self, ctx: Context, request: Message, response: Message, model_name: str):
+        return self.Return.Turn
+
+    def __call__(self, ctx: Context, request: Message, response: Message, model_name: str) -> bool:
+        """
+        :param ctx: Context object.
+        :param request: Request message.
+        :param response: Response message.
+        :param model_name: Name of the model in the Pipeline.models.
+        """
+        try:
+            result = self.call(ctx, request, response, model_name)
+        except Exception as exc:
+            logger.warning(exc)
+            return []
+        if result == self.Return.Turn:
+            return [request, response]
+        if result == self.Return.Response:
+            return [response]
+        if result == self.Return.Request:
+            return [request]
+        return []
+
+
+class MessageFilter(BaseHistoryFilter):
+    @abc.abstractmethod
+    def call(self, ctx, message, model_name):
+        raise NotImplemented
+
+    def __call__(self, ctx, request, response, model_name):
+        return self.call(ctx, request, model_name) + self.call(ctx, response, model_name)
+
+
+class IsImportant(BaseHistoryFilter):
+    """
+    Filter that checks if the "important" field in a Message.misc is True.
+    """
+
+    def call(self, ctx: Context, request: Message, response: Message, model_name: str) -> bool:
+        if request is not None and request.misc is not None and request.misc.get("important", None):
+            return self.Return.Request
+        if response is not None and response.misc is not None and response.misc.get("important", None):
+            return self.Return.Response
+        return False
+
+
+class FromModel(BaseHistoryFilter):
+    """
+    Filter that checks if the message was sent by the model.
+    """
+
+    def call(self, ctx: Context, request: Message, response: Message, model_name: str) -> bool:
+        if (
+            request is not None
+            and request.annotations is not None
+            and request.annotations.get("__generated_by_model__") == model_name
+        ):
+            return self.Return.Request
+        elif (
+            response is not None
+            and response.annotations is not None
+            and response.annotations.get("__generated_by_model__") == model_name
+        ):
+            return self.Return.Response
+        return False

chatsky/llm/llm_api.py

@@ -0,0 +1,68 @@
+"""
+LLM responses.
+--------------
+Wrapper around langchain.
+"""
+
+from typing import Union, Type, Optional
+from pydantic import BaseModel
+import logging
+from chatsky.core.message import Message
+from chatsky.core.context import Context
+from chatsky.llm.methods import BaseMethod
+from chatsky.llm.utils import message_to_langchain
+from chatsky.llm._langchain_imports import StrOutputParser, BaseChatModel, BaseMessage, check_langchain_available
+
+
+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.
+        """
+        check_langchain_available()
+        self.model: BaseChatModel = model
+        self.parser = StrOutputParser()
+        self.system_prompt = system_prompt
+
+    async def respond(
+        self,
+        history: list[BaseMessage],
+        message_schema: Union[None, Type[Message], Type[BaseModel]] = None,
+    ) -> Message:
+
+        if message_schema is None:
+            result = await self.parser.ainvoke(await self.model.ainvoke(history))
+            return Message(text=result)
+        elif issubclass(message_schema, Message):
+            # Case if the message_schema describes Message structure
+            structured_model = self.model.with_structured_output(message_schema, method="json_mode")
+            model_result = await structured_model.ainvoke(history)
+            logging.debug(f"Generated response: {model_result}")
+            return Message.model_validate(model_result)
+        elif issubclass(message_schema, BaseModel):
+            # Case if the message_schema describes Message.text structure
+            structured_model = self.model.with_structured_output(message_schema)
+            model_result = await structured_model.ainvoke(history)
+            return Message(text=message_schema.model_validate(model_result).model_dump_json())
+        else:
+            raise ValueError
+
+    async def condition(
+        self, ctx: Context, prompt: str, method: BaseMethod, return_schema: Optional[BaseModel] = None
+    ) -> bool:
+        condition_history = [
+            await message_to_langchain(Message(prompt), ctx=ctx, source="system"),
+            await message_to_langchain(ctx.last_request, ctx=ctx, source="human"),
+        ]
+        result = await method(ctx, await self.model.agenerate([condition_history], logprobs=True, top_logprobs=10))
+        return result

chatsky/llm/methods.py

@@ -0,0 +1,79 @@
+"""
+LLM methods
+-----------
+This module provides basic methods to support LLM conditions.
+These methods return bool values based on LLM result.
+"""
+
+import abc
+
+from pydantic import BaseModel
+
+from chatsky.core.context import Context
+from chatsky.llm._langchain_imports import LLMResult
+
+
+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:
+        """
+        Determine if result of an LLM invocation satisfies the condition of this method.
+
+        :param ctx: Current dialog context.
+        :param model_result: Result of langchain model's invoke.
+
+        """
+        raise NotImplementedError
+
+    async def model_result_to_text(self, model_result: LLMResult) -> str:
+        """
+        Extract text from raw model result.
+        """
+        return model_result.generations[0][0].text
+
+
+class Contains(BaseMethod):
+    """
+    Simple method to check if a string contains a pattern.
+
+    :param pattern: pattern that will be searched in model_result.
+    """
+
+    pattern: str
+
+    async def __call__(self, ctx: Context, model_result: LLMResult) -> bool:
+        """
+        :return: True if pattern is contained in model_result.
+        """
+        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`
+    """
+
+    target_token: str
+    threshold: float = -0.5
+
+    async def __call__(self, ctx: Context, model_result: LLMResult) -> bool:
+        """
+        :return: True if logprob of the token is higher then threshold.
+        """
+        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