ai-parrot 0.3.4__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/models.py

@@ -0,0 +1,372 @@
+from typing import Union, Optional
+import uuid
+import time
+from datetime import datetime
+from pathlib import Path, PurePath
+from enum import Enum
+from langchain_core.agents import AgentAction
+
+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)
+    response: str = Field(required=False)
+    intermediate_steps: list = Field(default_factory=list)
+    chat_history: list = Field(repr=True, default_factory=list)
+    source_documents: list = Field(required=False, default_factory=list)
+
+    def __post_init__(self) -> None:
+        if self.intermediate_steps:
+            steps = []
+            for item, result in self.intermediate_steps:
+                if isinstance(item, AgentAction):
+                    # convert into dictionary:
+                    steps.append(
+                        {
+                            "tool": item.tool,
+                            "tool_input": item.tool_input,
+                            "result": result,
+                            "log": str(item.log)
+                        }
+                    )
+            if steps:
+                self.intermediate_steps = steps
+
+
+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)
+    # Bot/Agent type
+    bot_type: str = Field(default='chatbot', required=False)
+    # 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/__init__.py

@@ -0,0 +1,48 @@
+from typing import Any
+from .abstract import AbstractStore
+from ..exceptions import ConfigError  # pylint: disable=E0611
+try:
+    from .qdrant import QdrantStore
+    QDRANT_ENABLED = True
+except (ModuleNotFoundError, ImportError):
+    QDRANT_ENABLED = False
+
+try:
+    from .milvus import MilvusStore
+    MILVUS_ENABLED = True
+except (ModuleNotFoundError, ImportError):
+    MILVUS_ENABLED = False
+
+
+def get_vectordb(vector_db: str, embeddings: Any, **kwargs) -> AbstractStore:
+    if vector_db in ('QdrantStore', 'qdrant'):
+        if QDRANT_ENABLED is True:
+            ## TODO: support pluggable vector store
+            return QdrantStore(  # pylint: disable=E0110
+                embeddings=embeddings,
+                **kwargs
+            )
+        else:
+            raise ConfigError(
+                (
+                    "Qdrant is enabled but not installed, "
+                    "Hint: Please install with pip install -e .[qdrant]"
+                )
+            )
+    if vector_db in ('milvus', 'MilvusStore'):
+        if MILVUS_ENABLED is True:
+            return MilvusStore(
+                embeddings=embeddings,
+                **kwargs
+            )
+        else:
+            raise ConfigError(
+                (
+                    "Milvus is enabled but not installed, "
+                    "Hint: Please install with pip install -e .[milvus]"
+                )
+            )
+    else:
+        raise ValueError(
+            f"Vector Database {vector_db} not supported"
+        )

parrot/stores/abstract.py

@@ -0,0 +1,171 @@
+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 = kwargs.pop('embedding_name', 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
+        self._client_id = 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)