ai-parrot 0.1.0__cp311-cp311-manylinux_2_28_x86_64.whl

parrot/manager.py

@@ -0,0 +1,152 @@
+"""
+Chatbot Manager.
+
+Tool for instanciate, managing and interacting with Chatbot through APIs.
+"""
+from typing import Any, Dict, Type
+from importlib import import_module
+from aiohttp import web
+from navconfig.logging import logging
+from .chatbots import (
+    AbstractChatbot,
+    Chatbot,
+    HRAgent,
+    AskTROC,
+    OddieBot,
+    Cody,
+    BoseBot
+)
+from .handlers.chat import ChatHandler # , BotHandler
+from .handlers import ChatbotHandler
+from .models import ChatbotModel
+
+
+logging.getLogger(name='selenium.webdriver').setLevel(logging.WARNING)
+
+
+class ChatbotManager:
+    """ChatbotManager.
+
+    Manage chatbots and interact with them through via aiohttp App.
+
+    """
+    app: web.Application = None
+    chatbots: Dict[str, AbstractChatbot] = {}
+
+    def __init__(self) -> None:
+        self.app = None
+        self.chatbots = {}
+        self.logger = logging.getLogger(name='Parrot.Manager')
+
+    def get_chatbot_class(self, class_name: str) -> Type[AbstractChatbot]:
+        """
+        Dynamically import a chatbot class based on the class name from the relative module '.chatbots'.
+        Args:
+        class_name (str): The name of the chatbot class to be imported.
+        Returns:
+        Type[AbstractChatbot]: A chatbot class derived from AbstractChatbot.
+        """
+        module = import_module('.chatbots', __package__)
+        try:
+            return getattr(module, class_name)
+        except AttributeError:
+            raise ImportError(f"No class named '{class_name}' found in the module 'chatbots'.")
+
+
+    async def load_bots(self, app: web.Application) -> None:
+        """Load all chatbots from DB."""
+        self.logger.info("Loading chatbots from DB...")
+        db = app['database']
+        async with await db.acquire() as conn:
+            ChatbotModel.Meta.connection = conn
+            bots = await ChatbotModel.all()
+            for bot in bots:
+                self.logger.notice(
+                    f"Loading chatbot '{bot.name}'..."
+                )
+                cls_name = bot.custom_class
+                if cls_name is None:
+                    class_name = Chatbot
+                else:
+                    class_name = self.get_chatbot_class(cls_name)
+                chatbot = class_name(
+                    chatbot_id=bot.chatbot_id,
+                    name=bot.name
+                )
+                try:
+                    await chatbot.configure()
+                    self.add_chatbot(chatbot)
+                except Exception as e:
+                    self.logger.error(
+                        f"Failed to configure chatbot '{chatbot.name}': {e}"
+                    )
+        self.logger.info(
+            ":: Chatbots loaded successfully."
+        )
+
+    def create_chatbot(self, class_name: Any = None, name: str = None, **kwargs) -> AbstractChatbot:
+        """Create a chatbot and add it to the manager."""
+        if class_name is None:
+            class_name = Chatbot
+        chatbot = class_name(**kwargs)
+        chatbot.name = name
+        self.add_chatbot(chatbot)
+        if 'llm' in kwargs:
+            llm = kwargs['llm']
+            llm_name = llm.pop('name')
+            model = llm.pop('model')
+            llm = chatbot.load_llm(
+                llm_name, model=model, **llm
+            )
+            chatbot.llm = llm
+        return chatbot
+
+    def add_chatbot(self, chatbot: AbstractChatbot) -> None:
+        """Add a chatbot to the manager."""
+        self.chatbots[chatbot.name] = chatbot
+
+    def get_chatbot(self, name: str) -> AbstractChatbot:
+        """Get a chatbot by name."""
+        return self.chatbots.get(name)
+
+    def remove_chatbot(self, name: str) -> None:
+        """Remove a chatbot by name."""
+        del self.chatbots[name]
+
+    def get_chatbots(self) -> Dict[str, AbstractChatbot]:
+        """Get all chatbots."""
+        return self.chatbots
+
+    def get_app(self) -> web.Application:
+        """Get the app."""
+        if self.app is None:
+            raise RuntimeError("App is not set.")
+        return self.app
+
+    def setup(self, app: web.Application) -> web.Application:
+        if isinstance(app, web.Application):
+            self.app = app  # register the app into the Extension
+        else:
+            self.app = app.get_app()  # Nav Application
+        # register signals for startup and shutdown
+        self.app.on_startup.append(self.on_startup)
+        self.app.on_shutdown.append(self.on_shutdown)
+        # Add Manager to main Application:
+        self.app['chatbot_manager'] = self
+        ## Configure Routes
+        router = self.app.router
+        router.add_view('/api/v1/chat', ChatHandler)
+        router.add_view('/api/v1/chat/{chatbot_name}', ChatHandler)
+        ChatbotHandler.configure(self.app, '/api/v1/bots')
+        # router.add_view('/api/v1/bots', ChatbotHandler)
+        # router.add_view('/api/v1/bots/{chatbot_name}', ChatbotHandler)
+        return self.app
+
+    async def on_startup(self, app: web.Application) -> None:
+        """On startup."""
+        # configure all pre-configured chatbots:
+        await self.load_bots(app)
+
+    async def on_shutdown(self, app: web.Application) -> None:
+        """On shutdown."""
+        pass

parrot/models.py

@@ -0,0 +1,347 @@
+from typing import Union, Optional
+import uuid
+import time
+from datetime import datetime
+from pathlib import Path, PurePath
+from enum import Enum
+from datamodel import BaseModel, Field
+from datamodel.types import Text  # pylint: disable=no-name-in-module
+from asyncdb.models import Model
+
+def created_at(*args, **kwargs) -> int:
+    return int(time.time()) * 1000
+
+
+class AgentResponse(BaseModel):
+    """AgentResponse.
+    dict_keys(
+        ['input', 'chat_history', 'output', 'intermediate_steps']
+    )
+
+    Response from Chatbots.
+    """
+    question: str = Field(required=False)
+    input: str = Field(required=False)
+    output: str = Field(required=False)
+    intermediate_steps: list = Field(default_factory=list)
+
+class ChatResponse(BaseModel):
+    """ChatResponse.
+    dict_keys(
+        ['question', 'chat_history', 'answer', 'source_documents', 'generated_question']
+    )
+
+    Response from Chatbots.
+    """
+    query: str = Field(required=False)
+    result: str = Field(required=False)
+    question: str = Field(required=False)
+    generated_question: str = Field(required=False)
+    answer: str = Field(required=False)
+    response: str = Field(required=False)
+    chat_history: list = Field(repr=True, default_factory=list)
+    source_documents: list = Field(required=False, default_factory=list)
+    documents: dict = Field(required=False, default_factory=dict)
+    sid: uuid.UUID = Field(primary_key=True, required=False, default=uuid.uuid4)
+    at: int = Field(default=created_at)
+
+    def __post_init__(self) -> None:
+        if self.result and not self.answer:
+            self.answer = self.result
+        if self.question and not self.generated_question:
+            self.generated_question = self.question
+        return super().__post_init__()
+
+
+# Chatbot Model:
+class ChatbotModel(Model):
+    """Chatbot.
+    CREATE TABLE IF NOT EXISTS navigator.chatbots (
+        chatbot_id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+        name VARCHAR NOT NULL DEFAULT 'Nav',
+        description VARCHAR,
+        config_file VARCHAR DEFAULT 'config.toml',
+        company_information JSONB DEFAULT '{}'::JSONB,
+        contact_information VARCHAR,
+        contact_form VARCHAR,
+        contact_email VARCHAR,
+        company_website VARCHAR,
+        avatar TEXT,
+        enabled BOOLEAN NOT NULL DEFAULT TRUE,
+        timezone VARCHAR DEFAULT 'UTC',
+        attributes JSONB DEFAULT '{}'::JSONB,
+        role VARCHAR DEFAULT 'a Human Resources Assistant',
+        goal VARCHAR NOT NULL DEFAULT 'Bring useful information to Users.',
+        backstory VARCHAR NOT NULL DEFAULT 'I was created by a team of developers to assist with users tasks.',
+        rationale VARCHAR NOT NULL DEFAULT 'Remember to maintain a professional tone. Please provide accurate and relevant information.',
+        language VARCHAR DEFAULT 'en',
+        template_prompt VARCHAR,
+        pre_instructions JSONB DEFAULT '[]'::JSONB,
+        llm VARCHAR DEFAULT 'VertexLLM',
+        model_name VARCHAR DEFAULT 'gemini-pro',
+        model_config JSONB DEFAULT '{}'::JSONB,
+        embedding_name VARCHAR DEFAULT 'thenlper/gte-base',
+        tokenizer VARCHAR DEFAULT 'thenlper/gte-base',
+        summarize_model VARCHAR DEFAULT 'facebook/bart-large-cnn',
+        classification_model VARCHAR DEFAULT 'facebook/bart-large-cnn',
+        database JSONB DEFAULT '{"vector_database": "MilvusStore", "database": "TROC", "collection_name": "troc_information"}'::JSONB,
+        created_at TIMESTAMPTZ DEFAULT NOW(),
+        created_by INTEGER,
+        updated_at TIMESTAMPTZ DEFAULT NOW()
+    );
+    """
+    chatbot_id: uuid.UUID = Field(primary_key=True, required=False, default_factory=uuid.uuid4)
+    name: str = Field(default='Nav', required=True)
+    description: str = Field(default='Nav Chatbot', required=False)
+    config_file: str = Field(default='config.toml', required=False)
+    custom_class: str = Field(required=False)
+    company_information: dict = Field(default_factory=dict, required=False)
+    avatar: Text
+    enabled: bool = Field(required=True, default=True)
+    timezone: str = Field(required=False, max=75, default="UTC", repr=False)
+    attributes: Optional[dict] = Field(required=False, default_factory=dict)
+    # Chatbot Configuration
+    role: str = Field(
+        default="a Human Resources Assistant",
+        required=False
+    )
+    goal: str = Field(
+        default="Bring useful information to Users.",
+        required=True
+    )
+    backstory: str = Field(
+        default="I was created by a team of developers to assist with users tasks.",
+        required=True
+    )
+    rationale: str = Field(
+        default=(
+            "Remember to maintain a professional tone."
+            " Please provide accurate and relevant information."
+        ),
+        required=True
+    )
+    language: str = Field(default='en', required=False)
+    template_prompt: Union[str, PurePath] = Field(
+        default=None,
+        required=False
+    )
+    pre_instructions: list = Field(
+        default_factory=list,
+        required=False
+    )
+    # Model Configuration:
+    llm: str = Field(default='VertexLLM', required=False)
+    llm_config: dict = Field(default_factory=dict, required=False)
+    embedding_name: str = Field(default="thenlper/gte-base", required=False)
+    tokenizer: str = Field(default='thenlper/gte-base', required=False)
+    summarize_model: str = Field(default="facebook/bart-large-cnn", required=False)
+    classification_model: str = Field(default="facebook/bart-large-cnn", required=False)
+    # Database Configuration
+    database: dict = Field(default='TROC', required=False, default_factory=dict)
+    # When created
+    created_at: datetime = Field(required=False, default=datetime.now())
+    created_by: int = Field(required=False)
+    updated_at: datetime = Field(required=False, default=datetime.now())
+
+    def __post_init__(self) -> None:
+        super(ChatbotModel, self).__post_init__()
+        if isinstance(self.config_file, str):
+            self.config_file = Path(self.config_file).resolve()
+        if isinstance(self.config_file, PurePath):
+            self.config_file = str(self.config_file)
+
+    class Meta:
+        """Meta Chatbot."""
+        driver = 'pg'
+        name = "chatbots"
+        schema = "navigator"
+        strict = True
+        frozen = False
+
+
+class ChatbotUsage(Model):
+    """ChatbotUsage.
+
+    Saving information about Chatbot Usage.
+
+    -- ScyllaDB CREATE TABLE Syntax --
+    CREATE TABLE IF NOT EXISTS navigator.chatbots_usage (
+        chatbot_id TEXT,
+        user_id SMALLINT,
+        sid TEXT,
+        source_path TEXT,
+        platform TEXT,
+        origin inet,
+        user_agent TEXT,
+        question TEXT,
+        response TEXT,
+        used_at BIGINT,
+        at TEXT,
+        PRIMARY KEY ((chatbot_id, sid, at), used_at)
+    ) WITH CLUSTERING ORDER BY (used_at DESC)
+    AND default_time_to_live = 10368000;
+
+    """
+    chatbot_id: uuid.UUID = Field(primary_key=True, required=False)
+    user_id: int = Field(primary_key=True, required=False)
+    sid: uuid.UUID = Field(primary_key=True, required=False, default=uuid.uuid4)
+    source_path: str = Field(required=False, default='web')
+    platform: str = Field(required=False, default='web')
+    origin: str = Field(required=False)
+    user_agent: str = Field(required=False)
+    question: str = Field(required=False)
+    response: str = Field(required=False)
+    used_at: int = Field(required=False, default=created_at)
+    event_timestamp: datetime = Field(required=False, default=datetime.now)
+    _at: str = Field(primary_key=True, required=False)
+
+    class Meta:
+        """Meta Chatbot."""
+        driver = 'bigquery'
+        name = "chatbots_usage"
+        schema = "navigator"
+        ttl = 10368000  # 120 days in seconds
+        strict = True
+        frozen = False
+
+    def __post_init__(self) -> None:
+        if not self._at:
+            # Generate a unique session id
+            self._at = f'{self.sid}:{self.used_at}'
+        super(ChatbotUsage, self).__post_init__()
+
+
+class FeedbackType(Enum):
+    """FeedbackType."""
+    # Good Feedback
+    GOOD_COMPLETE = "Completeness"
+    GOOD_CORRECT = "Correct"
+    GOOD_FOLLOW = "Follow the instructions"
+    GOOD_UNDERSTAND = "Understandable"
+    GOOD_USEFUL = "very useful"
+    GOOD_OTHER = "Please Explain"
+    # Bad Feedback
+    BAD_DONTLIKE = "Don't like the style"
+    BAD_INCORRECT = "Incorrect"
+    BAD_NOTFOLLOW = "Didn't follow the instructions"
+    BAD_LAZY = "Being lazy"
+    BAD_NOTUSEFUL = "Not useful"
+    BAD_UNSAFE = "Unsafe or problematic"
+    BAD_OTHER = "Other"
+
+    @classmethod
+    def list_feedback(cls, feedback_category):
+        """Return a list of feedback types based on the given category (Good or Bad)."""
+        prefix = feedback_category.upper() + "_"
+        return [feedback for feedback in cls if feedback.name.startswith(prefix)]
+
+class ChatbotFeedback(Model):
+    """ChatbotFeedback.
+
+    Saving information about Chatbot Feedback.
+
+    -- ScyllaDB CREATE TABLE Syntax --
+    CREATE TABLE IF NOT EXISTS navigator.chatbots_feedback (
+        chatbot_id UUID,
+        user_id INT,
+        sid UUID,
+        at TEXT,
+        rating TINYINT,
+        like BOOLEAN,
+        dislike BOOLEAN,
+        feedback_type TEXT,
+        feedback TEXT,
+        created_at BIGINT,
+        PRIMARY KEY ((chatbot_id, user_id, sid), created_at)
+    ) WITH CLUSTERING ORDER BY (created_at DESC)
+    AND default_time_to_live = 7776000;
+
+    """
+    chatbot_id: uuid.UUID = Field(primary_key=True, required=False)
+    user_id: int = Field(required=False)
+    sid: uuid.UUID = Field(primary_key=True, required=False)
+    _at: str = Field(primary_key=True, required=False)
+    # feedback information:
+    rating: int = Field(required=False, default=0)
+    _like: bool = Field(required=False, default=False)
+    _dislike: bool = Field(required=False, default=False)
+    feedback_type: FeedbackType = Field(required=False)
+    feedback: str = Field(required=False)
+    created_at: int = Field(required=False, default=created_at)
+    expiration_timestamp: datetime = Field(required=False, default=datetime.now)
+
+    class Meta:
+        """Meta Chatbot."""
+        driver = 'bigquery'
+        name = "chatbots_feedback"
+        schema = "navigator"
+        ttl = 7776000  # 3 months in seconds
+        strict = True
+        frozen = False
+
+    def __post_init__(self) -> None:
+        if not self._at:
+            # Generate a unique session id
+            if not self.created_at:
+                self.created_at = created_at()
+            self._at = f'{self.sid}:{self.created_at}'
+        super(ChatbotFeedback, self).__post_init__()
+
+
+## Prompt Library:
+
+class PromptCategory(Enum):
+    """
+    Prompt Category.
+
+    Categorization of Prompts, as "tech", "tech-or-explain", "idea", "explain", "action", "command", "other".
+    """
+    TECH = "tech"
+    TECH_OR_EXPLAIN = "tech-or-explain"
+    IDEA = "idea"
+    EXPLAIN = "explain"
+    ACTION = "action"
+    COMMAND = "command"
+    OTHER = "other"
+
+class PromptLibrary(Model):
+    """PromptLibrary.
+
+    Saving information about Prompt Library.
+
+    -- PostgreSQL CREATE TABLE Syntax --
+    CREATE TABLE IF NOT EXISTS navigator.prompt_library (
+            prompt_id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+            chatbot_id UUID,
+            title varchar,
+            query varchar,
+            description TEXT,
+            prompt_category varchar,
+            prompt_tags varchar[],
+            created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
+            created_by INTEGER,
+            updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
+    );
+    """
+    prompt_id: uuid.UUID = Field(primary_key=True, required=False, default_factory=uuid.uuid4)
+    chatbot_id: uuid.UUID = Field(required=True)
+    title: str = Field(required=True)
+    query: str = Field(required=True)
+    description: str = Field(required=False)
+    prompt_category: str = Field(required=False, default=PromptCategory.OTHER)
+    prompt_tags: list = Field(required=False, default_factory=list)
+    created_at: datetime = Field(required=False, default=datetime.now)
+    created_by: int = Field(required=False)
+    updated_at: datetime = Field(required=False, default=datetime.now)
+
+    class Meta:
+        """Meta Prompt Library."""
+        driver = 'pg'
+        name = "prompt_library"
+        schema = "navigator"
+        strict = True
+        frozen = False
+
+
+    def __post_init__(self) -> None:
+        super(PromptLibrary, self).__post_init__()

parrot/stores/abstract.py

@@ -0,0 +1,170 @@
+from abc import ABC, abstractmethod
+from typing import Union, Any
+from collections.abc import Callable
+import torch
+from langchain_huggingface import (
+    HuggingFaceEmbeddings
+)
+from langchain_community.embeddings import (
+    HuggingFaceBgeEmbeddings
+)
+from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
+from navconfig.logging import logging
+from ..conf import (
+    EMBEDDING_DEVICE,
+    EMBEDDING_DEFAULT_MODEL,
+    CUDA_DEFAULT_DEVICE,
+    MAX_BATCH_SIZE
+)
+
+
+class AbstractStore(ABC):
+    """AbstractStore class.
+
+    Args:
+        embeddings (str): Embeddings.
+    """
+
+    def __init__(self, embeddings: Union[str, Callable] = None, **kwargs):
+        self.client: Callable = None
+        self.vector: Callable = None
+        self._embed_: Callable = None
+        self._connected: bool = False
+        self.use_bge: bool = kwargs.pop("use_bge", False)
+        self.fastembed: bool = kwargs.pop("use_fastembed", False)
+        self.embedding_name: str = EMBEDDING_DEFAULT_MODEL
+        self.dimension: int = kwargs.pop("dimension", 768)
+        self._metric_type: str = kwargs.pop("metric_type", 'COSINE')
+        self._index_type: str = kwargs.pop("index_type", 'IVF_FLAT')
+        self.database: str = kwargs.pop('database', '')
+        self.collection = kwargs.pop("collection_name", "my_collection")
+        self.index_name = kwargs.pop("index_name", "my_index")
+        if embeddings is not None:
+            if isinstance(embeddings, str):
+                self.embedding_name = embeddings
+            else:
+                self._embed_ = embeddings
+        self.logger = logging.getLogger(f"Store.{__name__}")
+        # client
+        self._client = None
+
+    @property
+    def connected(self) -> bool:
+        return self._connected
+
+    async def __aenter__(self):
+        try:
+            self.tensor = torch.randn(1000, 1000).cuda()
+        except RuntimeError:
+            self.tensor = None
+        if self._embed_ is None:
+            self._embed_ = self.create_embedding(
+                model_name=self.embedding_name
+            )
+        self._client, self.client_id = self.connect()
+        return self
+
+    @abstractmethod
+    def connect(self):
+        pass
+
+    def __enter__(self):
+        if self._embed_ is None:
+            self._embed_ = self.create_embedding(
+                model_name=self.embedding_name
+            )
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        # closing Embedding
+        self._embed_ = None
+        del self.tensor
+        try:
+            torch.cuda.empty_cache()
+        except RuntimeError:
+            pass
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # closing Embedding
+        self._embed_ = None
+        try:
+            torch.cuda.empty_cache()
+        except RuntimeError:
+            pass
+
+    @abstractmethod
+    def get_vector(self):
+        pass
+
+    @abstractmethod
+    async def load_documents(
+        self,
+        documents: list,
+        collection: str = None
+    ):
+        pass
+
+    @abstractmethod
+    def upsert(self, payload: dict, collection_name: str = None) -> None:
+        pass
+
+    @abstractmethod
+    def search(self, payload: dict, collection_name: str = None) -> dict:
+        pass
+
+    @abstractmethod
+    async def delete_collection(self, collection_name: str = None) -> dict:
+        pass
+
+    @abstractmethod
+    async def create_collection(self, collection_name: str, document: Any) -> dict:
+        pass
+
+    def create_embedding(
+        self,
+        model_name: str = None
+    ):
+        encode_kwargs: str = {
+            'normalize_embeddings': True,
+            "batch_size": MAX_BATCH_SIZE
+        }
+        if torch.backends.mps.is_available():
+            # Use CUDA Multi-Processing Service if available
+            device = torch.device("mps")
+        elif torch.cuda.is_available():
+            # Use CUDA GPU if available
+            device = torch.device(
+                f'cuda:{CUDA_DEFAULT_DEVICE}'
+            )
+        elif EMBEDDING_DEVICE == 'cuda':
+            device = torch.device(
+                f'cuda:{CUDA_DEFAULT_DEVICE}'
+            )
+        else:
+            device = torch.device(EMBEDDING_DEVICE)
+        model_kwargs: str = {'device': device}
+        if model_name is None:
+            model_name = EMBEDDING_DEFAULT_MODEL
+        if self.use_bge is True:
+            return HuggingFaceBgeEmbeddings(
+                model_name=model_name,
+                model_kwargs=model_kwargs,
+                encode_kwargs=encode_kwargs
+            )
+        if self.fastembed is True:
+            return FastEmbedEmbeddings(
+                model_name=model_name,
+                max_length=1024,
+                threads=4
+            )
+        return HuggingFaceEmbeddings(
+            model_name=model_name,
+            model_kwargs=model_kwargs,
+            encode_kwargs=encode_kwargs
+        )
+
+    def get_default_embedding(
+        self,
+        model_name: str = EMBEDDING_DEFAULT_MODEL
+    ):
+        return self.create_embedding(model_name=model_name)