This code is an implementation of a chatbot using VertexAI chat model and Langchain.
from chatbot.bot import Bot
from chatbot.memory import MemoryTypes
from chatbot.models import ModelTypes
from chatbot.chain import ChainManager
chatbot/__init__.py
btrcm00/chatbot-with-langchain
import json
from typing import AsyncIterator
from fastapi import FastAPI
from fastapi.encoders import jsonable_encoder
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from langchain.callbacks.tracers.log_stream import RunLogPatch

from chatbot import Bot
from chatbot import MemoryTypes, ModelTypes
from chatbot.common.objects import ChatRequest

bot = Bot(memory=MemoryTypes.CUSTOM_MEMORY, model=ModelTypes.VERTEX)

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
)


async def transform_stream_for_client(
        stream: AsyncIterator[RunLogPatch],
) -> AsyncIterator[str]:
    async for chunk in stream:
        yield f"event: data\ndata: {json.dumps(jsonable_encoder(chunk))}\n\n"
    yield "event: end\n\n"


@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
    sentence = request.message
    chat_history = request.history
    if chat_history:
        bot.chain.add_message_to_memory(
            human_message=chat_history[-1]["human"],
            ai_message=chat_history[-1]["ai"],
            conversation_id=request.conversation_id
        )

    chain_stream = bot.chain.chain_stream(
        input=sentence,
        conversation_id=request.conversation_id
    )
    return StreamingResponse(
        transform_stream_for_client(chain_stream),
        headers={"Content-Type": "text/event-stream"},
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8080)

chatbot/app.py
btrcm00/chatbot-with-langchain
import asyncio
import random
from queue import Queue
from typing import Optional, Dict

from langchain.prompts import PromptTemplate

from chatbot.memory import MemoryTypes
from chatbot.models import ModelTypes
from chatbot.common.config import Config, BaseObject
from chatbot.chain import ChainManager
from chatbot.prompt import BOT_PERSONALITY
from chatbot.common.objects import Message


class Bot(BaseObject):
    def __init__(
            self,
            config: Config = None,
            prompt_template: PromptTemplate = None,
            memory: Optional[MemoryTypes] = None,
            model: Optional[ModelTypes] = None,
            chain_kwargs: Optional[dict] = None,
            memory_kwargs: Optional[dict] = None,
            model_kwargs: Optional[dict] = None,
            bot_personality: str = BOT_PERSONALITY
    ):
        """
        Conversation chatbot
        :param config: System configuration
        :type config: Config
        :param prompt_template: Prompt template
        :type prompt_template: PromptTemplate
        :param memory: Conversation memory type
        :type memory: MemoryTypes
        :param model: Model type
        :type model: ModelTypes
        :param chain_kwargs: Keyword arguments for LangChain's chain. Default with verbose
        :type chain_kwargs: Dict
        :param memory_kwargs: Keyword arguments for Memory. Default with ConversationWindowBuffer arguments
        :type memory_kwargs: Dict
        :param model_kwargs: Keyword arguments for Model. Default with VertexAI arguments
        :type model_kwargs: Dict
        """
        super().__init__()
        self.config = config if config is not None else Config()
        self.chain = ChainManager(
            config=self.config,
            prompt_template=prompt_template,
            bot_personality=bot_personality,
            memory=memory,
            model=model,
            chain_kwargs=chain_kwargs if chain_kwargs else {"verbose": True},
            memory_kwargs=memory_kwargs if memory_kwargs else self.default_memory_kwargs,
            model_kwargs=model_kwargs if model_kwargs else self.default_model_kwargs
        )
        self.input_queue = Queue(maxsize=6)

    @property
    def default_model_kwargs(self):
        return {
            "max_output_tokens": 512,
            "temperature": 0.2,
            "top_p": 0.8,
            "top_k": 40,
            # "streaming": True
        }

    @property
    def default_memory_kwargs(self):
        return {
            "k": 2
        }

    def set_personality(
            self,
            bot_personalities: str = BOT_PERSONALITY
    ):
        self.chain.set_personality_prompt(
            personality_prompt=bot_personalities,
            user_personality=""
        )

    def reset_history(self, conversation_id: str = None):
        self.chain.reset_history(conversation_id=conversation_id)

    def predict(self, sentence: str, conversation_id: str = None):
        message = Message(message=sentence, role=self.config.human_prefix)
        return asyncio.run(self.chain(message, conversation_id=conversation_id))

    def send(self):
        conversation_id = str(random.randint(10000, 99999))
        while True:
            sentence = input("User:")
            output = self.predict(sentence, conversation_id=conversation_id)
            print("BOT:", output[0].message)


if __name__ == "__main__":
    config = Config()
    bot = Bot(config=config)
    bot.send()

chatbot/bot.py
btrcm00/chatbot-with-langchain
from typing import Optional, Union
from operator import itemgetter
from langchain.prompts import PromptTemplate
from langchain import hub
from langchain.callbacks.tracers.langchain import wait_for_all_tracers
from langchain.schema.output_parser import StrOutputParser
from langchain.schema.runnable import RunnableLambda

from chatbot.common.config import BaseObject, Config
from chatbot.prompt import *
from chatbot.common.objects import Message, MessageTurn
from chatbot.utils import ChatbotCache, BotAnonymizer
from chatbot.memory import MemoryTypes, MEM_TO_CLASS
from chatbot.models import ModelTypes, MODEL_TO_CLASS


class ChainManager(BaseObject):
    def __init__(
            self,
            config: Config = None,
            memory: Optional[MemoryTypes] = None,
            model: Optional[ModelTypes] = None,
            prompt_template: PromptTemplate = None,
            chain_kwargs: Optional[dict] = None,
            memory_kwargs: Optional[dict] = None,
            model_kwargs: Optional[dict] = None,
            bot_personality: Optional[str] = BOT_PERSONALITY
    ):
        super().__init__()
        self.config = config if config is not None else Config()

        self._memory = self.get_memory(memory_type=memory, parameters=memory_kwargs)
        self._base_model = self.get_model(model_type=model, parameters=model_kwargs)
        if prompt_template:
            self._prompt = prompt_template
        else:
            partial_variables = {
                "bot_personality": bot_personality or BOT_PERSONALITY,
                "user_personality": "",
            }
            self._init_prompt_template(partial_variables=partial_variables)
        chain_kwargs = chain_kwargs or {}
        self._cache = ChatbotCache.create(config=self.config)
        self.anonymizer = BotAnonymizer(config=self.config)
        self._init_chain(**chain_kwargs)

    def get_memory(
            self,
            memory_type: Optional[MemoryTypes] = None,
            parameters: Optional[dict] = None
    ):
        if memory_type is None:
            memory_type = MemoryTypes.BASE_MEMORY
        if memory_type is not None:
            if memory_type not in MEM_TO_CLASS:
                raise ValueError(
                    f"Got unknown memory type: {memory_type}. "
                    f"Valid types are: {MEM_TO_CLASS.keys()}."
                )
            memory_class = MEM_TO_CLASS[memory_type]
        else:
            raise ValueError(
                "Somehow both `memory` is None, "
                "this should never happen."
            )
        return memory_class(config=self.config, **parameters)

    def get_model(
            self,
            model_type: Optional[ModelTypes] = None,
            model_name: str = None,
            parameters: Optional[dict] = None
    ):
        if model_type is None:
            model_type = ModelTypes.VERTEX

        if model_type is not None:
            if model_type not in MODEL_TO_CLASS:
                raise ValueError(
                    f"Got unknown model type: {model_type}. "
                    f"Valid types are: {MODEL_TO_CLASS.keys()}."
                )
            model_class = MODEL_TO_CLASS[model_type]
        else:
            raise ValueError(
                "Somehow both `memory` is None, "
                "this should never happen."
            )

        if model_type in [ModelTypes.VERTEX, ModelTypes.OPENAI]:
            if not model_name:
                model_name = self.config.base_model_name
            return model_class(model_name=model_name, **parameters)
        return model_class(**parameters, return_messages=True)

    @property
    def memory(self):
        return self._memory

    def reset_history(self, conversation_id: str = None):
        self.memory.clear(conversation_id=conversation_id)

    def _init_chain(self, **kwargs):
        anonymizer_runnable = self.anonymizer.get_runnable_anonymizer()
        self.chain = ({
                          "input": itemgetter("input"),
                          "history": itemgetter("conversation_id") | RunnableLambda(self.memory.load_history)
                      }
                      | anonymizer_runnable
                      | self._prompt
                      | self._base_model
                      | StrOutputParser()
                      | self.anonymizer.anonymizer.deanonymize).with_config(run_name="GenerateResponse")

        if not isinstance(self._memory, MEM_TO_CLASS[MemoryTypes.CUSTOM_MEMORY]):
            self.chain = self.chain | self.memory.memory

    def _init_prompt_template(self, partial_variables=None):
        _prompt: PromptTemplate = hub.pull("minhi/personality-chatbot-prompt")
        self._prompt = _prompt.partial(**partial_variables)

    def set_personality_prompt(self, bot_personality: str = BOT_PERSONALITY, user_personality: str = ""):
        partial_variables = {
            "bot_personality": bot_personality,
            "user_personality": user_personality
        }
        self._init_prompt_template(partial_variables)
        self._init_chain()

    async def _predict(self, message: Message, conversation_id: str):
        try:
            output = self.chain.invoke({"input": message.message, "conversation_id": conversation_id})
            output = Message(message=output, role=self.config.ai_prefix)
            return output
        finally:
            # Wait for invoke chain finish before push to Langsmith
            wait_for_all_tracers()

    def chain_stream(self, input: str, conversation_id: str):
        # history = self.memory.load_history(conversation_id=conversation_id)
        return self.chain.astream_log(
            {"input": input, "conversation_id": conversation_id},
            include_names=["FindDocs"]
        )

    def add_message_to_memory(
            self,
            human_message: Union[Message, str],
            ai_message: Union[Message, str],
            conversation_id: str
    ):
        if isinstance(human_message, str):
            human_message = Message(message=human_message, role=self.config.human_prefix)
        if isinstance(ai_message, str):
            ai_message = Message(message=ai_message, role=self.config.ai_prefix)

        turn = MessageTurn(
            human_message=human_message,
            ai_message=ai_message,
            conversation_id=conversation_id
        )
        self.memory.add_message(turn)

    async def __call__(self, message: Message, conversation_id: str):
        # history = self.memory.load_history(conversation_id=conversation_id)
        output: Message = await self._predict(message=message, conversation_id=conversation_id)
        self.add_message_to_memory(human_message=message, ai_message=output, conversation_id=conversation_id)
        return output

chatbot/chain.py
btrcm00/chatbot-with-langchain
SERP_API_TOKEN = "SERP_API_TOKEN"
MONGO_USERNAME = "MONGO_USERNAME"
MONGO_PASSWORD = "MONGO_PASSWORD"
MONGO_DATABASE = "MONGO_DATABASE"
MONGO_CLUSTER = "MONGO_CLUSTER"
MONGO_COLLECTION = "MONGO_COLLECTION"
MONGO_CONNECTION_STRING = "MONGO_CONNECTION_STRING"
CREDENTIALS_FILE = "CREDENTIALS_FILE"
AI_PREFIX = "AI_PREFIX"
HUMAN_PREFIX = "HUMAN_PREFIX"
MEMORY_KEY = "MEMORY_KEY"


chatbot/common/common_keys.py
btrcm00/chatbot-with-langchain
import logging
import os
import json
import vertexai
import urllib.parse

from chatbot.common.constants import CHAT_MODEL_NAME
from chatbot.common.common_keys import *


class Singleton(type):
    _instances = {}

    def __call__(cls, *args, **kwargs):
        if cls not in cls._instances:
            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
        return cls._instances[cls]


class BaseSingleton(metaclass=Singleton):
    pass


class BaseObject(BaseSingleton):
    def __init__(self, **kwargs):
        self.logger = logging.getLogger(self.class_name())

    @classmethod
    def class_name(cls):
        return cls.__name__


class Config(BaseObject):
    def __init__(
            self,
            credentials: str = None,
            cache_type: str = None,
            base_model_name: str = None,
            serp_api_token: str = None,
            model_max_input_size: int = 5,
            waiting_time: float = 2,
            memory_connection_string: str = None,
            memory_database_name: str = None,
            memory_collection_name: str = None,
            session_id: str = None,
            mongo_username: str = None,
            mongo_password: str = None,
            mongo_cluster: str = None,
            memory_window_size: int = 5
    ):
        super().__init__()
        self.credentials = credentials if credentials is not None else os.getenv(CREDENTIALS_FILE,
                                                                                 "./secure/vertex.json")
        self.init_env()
        self.serp_api_token = serp_api_token if serp_api_token is not None else ""
        self.cache_type = cache_type if cache_type is not None else "in_memory"
        self.base_model_name = base_model_name if base_model_name is not None else CHAT_MODEL_NAME
        self.model_max_input_size = model_max_input_size
        self.waiting_time = waiting_time

        self.mongo_username = mongo_username if mongo_username is not None else os.getenv(MONGO_USERNAME)
        self.mongo_password = mongo_password if mongo_password is not None else os.getenv(MONGO_PASSWORD)
        if self.mongo_password:
            self.mongo_password = urllib.parse.quote_plus(self.mongo_password)
        if self.mongo_username:
            self.mongo_username = urllib.parse.quote_plus(self.mongo_username)
        self.mongo_cluster = mongo_cluster if mongo_cluster is not None else os.getenv(MONGO_CLUSTER)
        self.memory_database_name = memory_database_name if memory_database_name is not None \
            else os.getenv(MONGO_DATABASE, "langchain_bot")
        self.memory_collection_name = memory_collection_name if memory_collection_name is not None \
            else os.getenv(MONGO_COLLECTION, "chatbot")
        self.memory_connection_string = memory_connection_string if memory_connection_string is not None \
            else os.getenv(MONGO_CONNECTION_STRING,
                           f"mongodb+srv://{self.mongo_username}:{self.mongo_password}@{self.mongo_cluster}.xnkswcg.mongodb.net")
        self.session_id = session_id if session_id is not None else "chatbot"
        self.memory_window_size = memory_window_size if memory_window_size is not None else 5
        self.ai_prefix = os.getenv(AI_PREFIX, "AI")
        self.human_prefix = os.getenv(HUMAN_PREFIX, "Human")
        self.memory_key = os.getenv(MEMORY_KEY, "history")

    def init_env(self):
        credential_data = json.load(open(self.credentials, "r"))
        project = credential_data["project_id"]
        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.credentials
        os.environ["GOOGLE_CLOUD_PROJECT"] = project
        vertexai.init(
            project=project,
            location="us-central1"
        )

chatbot/common/config.py
btrcm00/chatbot-with-langchain
TEXT_MODEL_NAME = "text-bison@001"
CHAT_MODEL_NAME = "chat-bison@001"

ANONYMIZED_FIELDS = ['US_BANK_NUMBER', 'US_DRIVER_LICENSE', 'US_ITIN', 'US_PASSPORT', 'US_SSN']
NLP_CONFIG = {
    "nlp_engine_name": "spacy",
    "models": [
        {"lang_code": "vi", "model_name": "vi_core_news_lg"},
        {"lang_code": "en", "model_name": "en_core_web_md"}
    ],
}

chatbot/common/constants.py
btrcm00/chatbot-with-langchain
from pydantic import BaseModel, Field
from typing import Optional, List, Dict


class Message(BaseModel):
    message: str = Field(description="User message")
    role: str = Field(description="Message role in conversation")


class MessageTurn(BaseModel):
    human_message: Message = Field(description="Message of human")
    ai_message: Message = Field(description="Message of AI")
    conversation_id: str = Field(description="The id of user in this turn")


class ChatRequest(BaseModel):
    message: str
    history: Optional[List[Dict[str, str]]]
    conversation_id: Optional[str]


def messages_from_dict(message: dict) -> str:
    human_message = message["human_message"]
    ai_message = message["ai_message"]

    human_message = Message(message=human_message["message"], role=human_message["role"])
    ai_message = Message(message=ai_message["message"], role=ai_message["role"])
    return f"{human_message.role}: {human_message.message}\n{ai_message.role}: {ai_message.message}"

chatbot/common/objects.py
btrcm00/chatbot-with-langchain
from chatbot.memory.base_memory import BaseChatbotMemory
from chatbot.memory.mongo_memory import MongoChatbotMemory
from chatbot.memory.custom_memory import CustomMongoChatbotMemory
from chatbot.memory.memory_types import MemoryTypes, MEM_TO_CLASS
chatbot/memory/__init__.py
btrcm00/chatbot-with-langchain
from typing import Optional

import langchain.schema.messages
from langchain.memory import ConversationBufferWindowMemory, ChatMessageHistory

from chatbot.common.config import BaseObject, Config
from chatbot.common.objects import MessageTurn


class BaseChatbotMemory(BaseObject):
    __slots__ = ["_base_memory", "_memory"]

    def __init__(
            self,
            config: Config = None,
            chat_history_class=ChatMessageHistory,
            memory_class=ConversationBufferWindowMemory,
            chat_history_kwargs: Optional[dict] = None,
            **kwargs
    ):
        """
        Base chatbot memory
        :param config: Config object
        :param chat_history_class: LangChain's chat history class
        :param memory_class: LangChain's memory class
        :param kwargs: Memory class kwargs
        """
        super().__init__()
        self.config = config if config is not None else Config()
        self._params = kwargs
        self.chat_history_kwargs = chat_history_kwargs or {}
        self._base_memory_class = chat_history_class
        self._memory = memory_class(**self.params)
        self._user_memory = dict()

    @property
    def params(self):
        if self._params:
            return self._params
        else:
            return {
                "ai_prefix": self.config.ai_prefix,
                "human_prefix": self.config.human_prefix,
                "memory_key": self.config.memory_key,
                "k": self.config.memory_window_size
            }

    @property
    def memory(self):
        return self._memory

    @property
    def user_memory(self):
        return self._user_memory

    def clear(self, conversation_id: str):
        if conversation_id in self.user_memory:
            memory = self.user_memory.pop(conversation_id)
            memory.clear()

    def load_history(self, conversation_id: str) -> str:
        if conversation_id not in self._user_memory:
            memory = self._base_memory_class(**self.chat_history_kwargs)
            self.memory.chat_memory = memory
            self.user_memory[conversation_id] = memory
            return ""

        self.memory.chat_memory = self.user_memory.get(conversation_id)
        return self._memory.load_memory_variables({})["history"]

    def add_message(self, message_turn: MessageTurn):
        pass

chatbot/memory/base_memory.py
btrcm00/chatbot-with-langchain
import json
from typing import List
from pymongo import MongoClient, errors

from chatbot.common.config import Config, BaseObject
from chatbot.common.objects import MessageTurn, messages_from_dict


class BaseCustomMongoChatbotMemory(BaseObject):
    def __init__(
            self,
            config: Config = None,
            connection_string: str = None,
            session_id: str = None,
            database_name: str = None,
            collection_name: str = None,
            k: int = 5,
            **kwargs
    ):
        super(BaseCustomMongoChatbotMemory, self).__init__()
        self.config = config if config is not None else Config()

        self.connection_string = connection_string
        self.session_id = session_id
        self.database_name = database_name
        self.collection_name = collection_name

        try:
            self.client: MongoClient = MongoClient(connection_string)
        except errors.ConnectionFailure as error:
            self.logger.error(error)

        self.db = self.client[database_name]
        self.collection = self.db[collection_name]
        self.collection.create_index("SessionId")
        self.k = k

    def add_message(self, message_turn: MessageTurn):
        conversation_id = message_turn.conversation_id
        try:
            self.logger.info(f"Save 1 message turn of conversation <{conversation_id}>")
            self.collection.insert_one(
                {
                    "ConversationId": conversation_id,
                    "SessionId": self.session_id,
                    "History": json.dumps(message_turn.dict(), ensure_ascii=False),
                }
            )
        except errors.WriteError as err:
            self.logger.error(err)

    def clear_history(self, conversation_id: str = None):
        try:
            self.logger.info(f"Deleting the history of conversation <{conversation_id}> at session <{self.session_id}>")
            if conversation_id is None:
                self.logger.warning(f"You are deleting all collection with session: {self.session_id}")
                self.collection.delete_many({"SessionId": self.session_id})
            else:
                self.collection.delete_many({"SessionId": self.session_id, "ConversationId": conversation_id})
        except errors.WriteError as err:
            self.logger.error(err)

    def load_history(self, conversation_id: str) -> str:
        """Retrieve the messages from MongoDB"""
        from pymongo import errors
        cursor = None
        try:
            cursor = self.collection.find({"SessionId": self.session_id, "ConversationId": conversation_id})
        except errors.OperationFailure as error:
            self.logger.error(error)

        if cursor:
            items = [json.loads(document["History"]) for document in cursor][-self.k:]
        else:
            items = []

        messages: List[str] = [messages_from_dict(item) for item in items]
        return "\n".join(messages)


class CustomMongoChatbotMemory(BaseObject):
    def __init__(self, config: Config = None, **kwargs):
        super(CustomMongoChatbotMemory, self).__init__()
        self.memory = BaseCustomMongoChatbotMemory(
            connection_string=config.memory_connection_string,
            session_id=config.session_id,
            database_name=config.memory_database_name,
            collection_name=config.memory_collection_name,
            **kwargs
        )

    def clear(self, conversation_id: str = None):
        self.memory.clear_history(conversation_id=conversation_id)

    def load_history(self, conversation_id: str):
        return self.memory.load_history(conversation_id)

    def add_message(self, message_turn: MessageTurn):
        self.memory.add_message(message_turn)

chatbot/memory/custom_memory.py
btrcm00/chatbot-with-langchain
from enum import Enum
from chatbot.memory import MongoChatbotMemory, BaseChatbotMemory, CustomMongoChatbotMemory


class MemoryTypes(str, Enum):
    """Enumerator with the Memory types."""
    BASE_MEMORY = "base-memory"
    MONGO_MEMORY = "mongodb-memory"
    CUSTOM_MEMORY = "custom-memory"


MEM_TO_CLASS = {
    "mongodb-memory": MongoChatbotMemory,
    "base-memory": BaseChatbotMemory,
    "custom-memory": CustomMongoChatbotMemory
}

chatbot/memory/memory_types.py
btrcm00/chatbot-with-langchain
import logging
from langchain.memory import MongoDBChatMessageHistory

from chatbot.common.config import Config
from chatbot.memory.base_memory import BaseChatbotMemory

logger = logging.getLogger(__name__)


class MongoChatbotMemory(BaseChatbotMemory):
    def __init__(self, config: Config = None, **kwargs):
        config = config if config is not None else Config()
        super(MongoChatbotMemory, self).__init__(
            config=config,
            chat_history_class=MongoDBChatMessageHistory,
            chat_history_kwargs={
                "connection_string": config.memory_connection_string,
                "session_id": config.session_id,
                "database_name": config.memory_database_name,
                "collection_name": config.memory_collection_name
            }
        )

chatbot/memory/mongo_memory.py
btrcm00/chatbot-with-langchain
from chatbot.models.model_types import ModelTypes, MODEL_TO_CLASS

chatbot/models/__init__.py
btrcm00/chatbot-with-langchain
from enum import Enum
from langchain.chat_models import ChatVertexAI, ChatOpenAI
from langchain.llms import LlamaCpp


class ModelTypes(str, Enum):
    OPENAI = "OPENAI"
    VERTEX = "VERTEX"
    LLAMA_CPP = "LLAMA-CPP"


MODEL_TO_CLASS = {
    "OPENAI": ChatOpenAI,
    "VERTEX": ChatVertexAI,
    "LLAMA-CPP": LlamaCpp
}

chatbot/models/model_types.py
btrcm00/chatbot-with-langchain
BOT_PERSONALITY = """
Gender: Male
Name: MinHi
Age: 23
Personality trait: Extrovert
Conversational style: cheerful, don't like to talk a lot
Conversation domain: tán gẫu, tâm sự
Other info: developed by Cong Minh and Tran Hieu"""

LLAMA_PROMPT = """
[INST] <<SYS>>
You are a helpful Vietnamese assistant with given personalities. 
Help the user answer any questions or chat with them about anything they want. 
<</SYS>>
Below are descriptions of your personality traits, talk to users according to those traits:
{personality}
Current conversation:
{history}
Human: {input}
AI: [/INST]
"""

chatbot/prompt.py
btrcm00/chatbot-with-langchain
from typing import Optional
from langchain.tools import BaseTool
from langchain.callbacks.manager import CallbackManagerForToolRun
from langchain.utilities import SerpAPIWrapper


class CustomSearchTool(BaseTool):
    name = "Custom search"
    description = "Useful for when you need to answer questions about current or newest events, which you are not trained before"

    def __init__(self):
        super().__init__()
        self.params = {
            "engine": "google",
            "gl": "us",
            "hl": "vi",
        }
        self._search = SerpAPIWrapper(params=self.params)

    def _run(
            self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """Use the tool."""
        return self._search.run(query)

chatbot/tools/serp.py
btrcm00/chatbot-with-langchain
from .chain_cache import ChatbotCache
from .anonymizer import BotAnonymizer
chatbot/utils/__init__.py
btrcm00/chatbot-with-langchain
import langdetect
from langchain_experimental.data_anonymizer import PresidioReversibleAnonymizer, PresidioAnonymizer
from langchain.schema import runnable

from chatbot.common.config import BaseObject, Config
from chatbot.common.constants import ANONYMIZED_FIELDS, NLP_CONFIG


class BotAnonymizer(BaseObject):
    def __init__(self, config: Config = None):
        super(BotAnonymizer, self).__init__()
        self.config = config if config is not None else Config()
        self._anonymizer = PresidioReversibleAnonymizer(languages_config=NLP_CONFIG)

    @property
    def anonymizer(self):
        return self._anonymizer

    @property
    def supported_lang(self):
        return ["vi", "en"]

    def _detect_lang(self, input_dict: dict) -> dict:
        language = langdetect.detect(input_dict["input"])
        if language not in self.supported_lang:
            self.logger.warning(
                f"Detected language is not supported in this Chatbot, it only support {self.supported_lang}, but detected {language}")
            language = None
        return {"language": language, **input_dict}

    def anonymize_func(self, input_dict: dict):
        if input_dict["language"] is None:
            return {
                "input": input_dict["input"],
                "history": input_dict["history"]
            }
        return {
            "input": self.anonymizer.anonymize(input_dict["input"], input_dict["language"]),
            "history": self.anonymizer.anonymize(input_dict["history"], input_dict["language"])
        }

    def get_runnable_anonymizer(self):
        return runnable.RunnableLambda(self._detect_lang) | runnable.RunnableLambda(self.anonymize_func)

chatbot/utils/anonymizer.py
btrcm00/chatbot-with-langchain
import langchain
from langchain.cache import InMemoryCache, RedisCache

from chatbot.common.config import BaseObject, Config

CACHE_TYPE = {
    "in_memory": InMemoryCache,
    "redis": RedisCache
}


class ChatbotCache(BaseObject):
    @classmethod
    def create(cls, config: Config = None):
        config = config if config is not None else Config()
        langchain.llm_cache = CACHE_TYPE[config.cache_type]()

chatbot/utils/chain_cache.py
btrcm00/chatbot-with-langchain
import os

from langchain.prompts import PromptTemplate
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

from gradio_ui.base_ui import BaseGradioUI
from chatbot import MemoryTypes
from chatbot.models import ModelTypes
from chatbot.common.config import Config
from chatbot.bot import Bot
from chatbot.prompt import LLAMA_PROMPT, BOT_PERSONALITY


if __name__ == "__main__":
    GGML_MODEL_PATH = os.environ["GGML_MODEL_PATH"]
    config = Config()
    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])

    partial_variables = {"personality": BOT_PERSONALITY}
    prompt_template = PromptTemplate(
        template=LLAMA_PROMPT,
        input_variables=["history", "input"],
        partial_variables=partial_variables
    )

    bot = Bot(
        config,
        prompt_template,
        model=ModelTypes.LLAMA_CPP,
        memory=MemoryTypes.CUSTOM_MEMORY,
        model_kwargs={
            "model_path": GGML_MODEL_PATH,
            "n_ctx": 512,
            "temperature": 0.75,
            "max_tokens": 512,
            "top_p": 0.95,
            "callback_manager": callback_manager,
            "verbose": True
        }
    )
    demo = BaseGradioUI(bot=bot)
    demo.start_demo()

examples/run_llama_cpp.py
btrcm00/chatbot-with-langchain
from gradio_ui.base_ui import BaseGradioUI
from chatbot import MemoryTypes, ModelTypes

if __name__ == "__main__":
    demo = BaseGradioUI(bot_memory=MemoryTypes.CUSTOM_MEMORY,bot_model=ModelTypes.VERTEX)
    demo.start_demo()

gradio_app.py
btrcm00/chatbot-with-langchain
from typing import Any, Optional
import random

import gradio as gr

from chatbot.bot import Bot
from chatbot import MemoryTypes, ModelTypes


class BaseGradioUI:
    def __init__(
            self,
            bot: Bot = None,
            bot_memory: Optional[MemoryTypes] = None,
            bot_model: Optional[ModelTypes] = None
    ):
        self.bot = bot if bot is not None else Bot(memory=bot_memory, model=bot_model)
        self._conversation_id = None

    @staticmethod
    def create_conversation_id():
        return str(random.randint(100000000, 999999999))

    def user_state(self, message: str, chat_history: Any, conversation_id):
        """Initiate user state and chat history

        Args:
            message (str): user message
            chat_history (Any): chat history
        """
        if not conversation_id:
            conversation_id = self.create_conversation_id()
        return "", chat_history + [[message, None]], conversation_id

    def respond(self, conversation_id, chat_history):
        message = chat_history[-1][0]
        result = self.bot.predict(sentence=message, conversation_id=conversation_id)
        chat_history[-1][-1] = result.message
        return chat_history

    def start_demo(self, port=8000, debug=False, share=True):
        with gr.Blocks() as demo:
            conversation_id_state = gr.State("")
            gr.Markdown("""<h1><center> LLM Assistant </center></h1>""")
            chatbot = gr.Chatbot(label="Assistant").style(height=700)

            with gr.Row():
                message = gr.Textbox(show_label=False,
                                     placeholder="Enter your prompt and press enter",
                                     visible=True)

            btn_refresh = gr.ClearButton(components=[message, chatbot],
                                         value="Refresh the conversation history")

            def clear_user_state():
                return {conversation_id_state: ""}

            message.submit(
                self.user_state,
                inputs=[message, chatbot, conversation_id_state],
                outputs=[message, chatbot, conversation_id_state],
                queue=False
            ).then(
                self.respond,
                inputs=[conversation_id_state, chatbot],
                outputs=[chatbot]
            )
            btn_refresh.click(clear_user_state, outputs=conversation_id_state)
        demo.queue()
        demo.launch(debug=debug, server_port=port, share=share)

gradio_ui/base_ui.py
btrcm00/chatbot-with-langchain