ai-parrot 0.1.0__cp311-cp311-manylinux_2_28_x86_64.whl

parrot/stores/milvus.py

@@ -0,0 +1,540 @@
+from typing import Optional, Union, Any
+import asyncio
+import uuid
+import torch
+from pymilvus import (
+    MilvusClient,
+    Collection,
+    FieldSchema,
+    CollectionSchema,
+    DataType,
+    connections,
+    db
+)
+from pymilvus.exceptions import MilvusException
+from langchain_milvus import Milvus  # pylint: disable=import-error, E0611
+from langchain.memory import VectorStoreRetrieverMemory
+from .abstract import AbstractStore
+from ..conf import (
+    MILVUS_HOST,
+    MILVUS_PROTOCOL,
+    MILVUS_PORT,
+    MILVUS_URL,
+    MILVUS_TOKEN,
+    MILVUS_USER,
+    MILVUS_PASSWORD,
+    MILVUS_SECURE,
+    MILVUS_SERVER_NAME,
+    MILVUS_CA_CERT,
+    MILVUS_SERVER_CERT,
+    MILVUS_SERVER_KEY,
+    MILVUS_USE_TLSv2
+)
+
+
+class MilvusStore(AbstractStore):
+    """MilvusStore class.
+
+    Milvus is a Vector Database multi-layered.
+
+    Args:
+        host (str): Milvus host.
+        port (int): Milvus port.
+        url (str): Milvus URL.
+    """
+
+    def __init__(self, embeddings = None, **kwargs):
+        super().__init__(embeddings, **kwargs)
+        self.use_bge: bool = kwargs.pop("use_bge", False)
+        self.fastembed: bool = kwargs.pop("use_fastembed", False)
+        self.database: str = kwargs.pop('database', '')
+        self.collection = kwargs.pop('collection_name', '')
+        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.host = kwargs.pop("host", MILVUS_HOST)
+        self.port = kwargs.pop("port", MILVUS_PORT)
+        self.protocol = kwargs.pop("protocol", MILVUS_PROTOCOL)
+        self.create_database: bool = kwargs.pop('create_database', True)
+        self.url = kwargs.pop("url", MILVUS_URL)
+        if not self.url:
+            self.url = f"{self.protocol}://{self.host}:{self.port}"
+        else:
+            # Extract host and port from URL
+            if not self.host:
+                self.host = self.url.split("://")[-1].split(":")[0]
+            if not self.port:
+                self.port = int(self.url.split(":")[-1])
+        self.token = kwargs.pop("token", MILVUS_TOKEN)
+        # user and password (if required)
+        self.user = kwargs.pop("user", MILVUS_USER)
+        self.password = kwargs.pop("password", MILVUS_PASSWORD)
+        # SSL/TLS
+        self._secure: bool = kwargs.pop('secure', MILVUS_SECURE)
+        self._server_name: str = kwargs.pop('server_name', MILVUS_SERVER_NAME)
+        self._cert: str = kwargs.pop('server_pem_path', MILVUS_SERVER_CERT)
+        self._ca_cert: str = kwargs.pop('ca_pem_path', MILVUS_CA_CERT)
+        self._cert_key: str = kwargs.pop('client_key_path', MILVUS_SERVER_KEY)
+        # Any other argument will be passed to the Milvus client
+        self.kwargs = {
+            "uri": self.url,
+            "host": self.host,
+            "port": self.port,
+            **kwargs
+        }
+        if self.token:
+            self.kwargs['token'] = self.token
+        if self.user:
+            self.kwargs['token'] = f"{self.user}:{self.password}"
+        # SSL Security:
+        if self._secure is True:
+            args = {
+                "secure": self._secure,
+                "server_name": self._server_name
+            }
+            if self._cert:
+                if MILVUS_USE_TLSv2 is True:
+                    args['client_pem_path'] = self._cert
+                    args['client_key_path'] = self._cert_key
+                else:
+                    args["server_pem_path"] = self._cert
+            if self._ca_cert:
+                args['ca_pem_path'] = self._ca_cert
+            self.kwargs = {**self.kwargs, **args}
+        # 1. Check if database exists:
+        if self.database:
+            self.kwargs['db_name'] = self.database
+            self.use_database(self.database, create=True)
+
+    def connect(self, client_id: str = None):
+        # 1. Set up a pyMilvus default connection
+        # Unique connection:
+        if not client_id:
+            client_id = "uri-connection"
+        self.conn = connections.connect(
+            alias=client_id,
+            **self.kwargs
+        )
+        self._client = MilvusClient(
+            **self.kwargs
+        )
+        self._connected = True
+        return self._client, client_id
+
+    def close(self, client_id: str = "uri-connection"):
+        connections.disconnect(alias=client_id)
+        self._connected = False
+        self._client.close()
+
+    def create_db(self, db_name: str, **kwargs) -> bool:
+        args = {
+            "uri": self.url,
+            "host": self.host,
+            "port": self.port,
+            **kwargs
+        }
+        print('ARGS >', args)
+        try:
+            conn = connections.connect(**args)
+            db.create_database(db_name)
+            self.logger.notice(
+                f"Database {db_name} created successfully."
+            )
+        except Exception as e:
+            raise ValueError(
+                f"Error creating database: {e}"
+            )
+        finally:
+            connections.disconnect(alias="uri-connection")
+
+    def use_database(self, db_name: str, create: bool = False) -> None:
+        try:
+            conn = connections.connect(**self.kwargs)
+        except MilvusException as exc:
+            if "database not found" in exc.message:
+                args = self.kwargs.copy()
+                del args['db_name']
+                self.create_db(db_name, **args)
+                # re-connect:
+                conn = connections.connect(**self.kwargs)
+        if db_name not in db.list_database():
+            if self.create_database is True or create is True:
+                try:
+                    db.create_database(db_name)
+                    self.logger.notice(
+                        f"Database {db_name} created successfully."
+                    )
+                except Exception as e:
+                    raise ValueError(
+                        f"Error creating database: {e}"
+                    )
+            else:
+                raise ValueError(
+                    f"Database {db_name} does not exist."
+                )
+        connections.disconnect(alias='default')
+
+    def setup_vector(self):
+        self.vector = Milvus(
+            self._embed_,
+            consistency_level='Bounded',
+            connection_args={**self.kwargs},
+            collection_name=self.collection,
+        )
+        return self.vector
+
+    def get_vectorstore(self):
+        return self.get_vector()
+
+    def collection_exists(self, collection_name: str) -> bool:
+        if collection_name in self._client.list_collections():
+            return True
+        return False
+
+    def check_state(self, collection_name: str) -> dict:
+        return self._client.get_load_state(collection_name=collection_name)
+
+    async def delete_collection(self, collection: str = None) -> dict:
+        self._client.drop_collection(
+            collection_name=collection
+        )
+
+    async def create_collection(
+        self,
+        collection_name: str,
+        document: Any = None,
+        dimension: int = 768,
+        index_type: str = None,
+        metric_type: str = None,
+        schema_type: str = 'default',
+        metadata_field: str = None,
+        **kwargs
+    ) -> dict:
+        """create_collection.
+
+        Create a Schema (Milvus Collection) on the Current Database.
+
+        Args:
+            collection_name (str): Collection Name.
+            document (Any): List of Documents.
+            dimension (int, optional): Vector Dimension. Defaults to 768.
+            index_type (str, optional): Default index type of Vector Field. Defaults to "HNSW".
+            metric_type (str, optional): Default Metric for Vector Index. Defaults to "L2".
+            schema_type (str, optional): Description of Model. Defaults to 'default'.
+
+        Returns:
+            dict: _description_
+        """
+        # Check if collection exists:
+        if self.collection_exists(collection_name):
+            self.logger.warning(
+                f"Collection {collection_name} already exists."
+            )
+            return None
+        idx_params = {}
+        if not index_type:
+            index_type = self._index_type
+        if index_type == 'HNSW':
+            idx_params = {
+                "M": 36,
+                "efConstruction": 1024
+            }
+        elif index_type in ('IVF_FLAT', 'SCANN', 'IVF_SQ8'):
+            idx_params = {
+                "nlist": 1024
+            }
+        elif index_type in ('IVF_PQ'):
+            idx_params = {
+                "nlist": 1024,
+                "m": 16
+            }
+        if not metric_type:
+            metric_type = self._metric_type  # default metric type
+        if schema_type == 'default':
+            # Default Collection for all loaders:
+            schema = MilvusClient.create_schema(
+                auto_id=False,
+                enable_dynamic_field=True,
+                description=collection_name
+            )
+            schema.add_field(
+                field_name="pk",
+                datatype=DataType.INT64,
+                is_primary=True,
+                auto_id=True,
+                max_length=100
+            )
+            schema.add_field(
+                field_name="index",
+                datatype=DataType.VARCHAR,
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="url",
+                datatype=DataType.VARCHAR,
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="source",
+                datatype=DataType.VARCHAR,
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="filename",
+                datatype=DataType.VARCHAR,
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="question",
+                datatype=DataType.VARCHAR,
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="answer",
+                datatype=DataType.VARCHAR,
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="source_type",
+                datatype=DataType.VARCHAR,
+                max_length=128
+            )
+            schema.add_field(
+                field_name="type",
+                datatype=DataType.VARCHAR,
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="text",
+                datatype=DataType.VARCHAR,
+                description="Text",
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="summary",
+                datatype=DataType.VARCHAR,
+                description="Summary (refine resume)",
+                max_length=65535
+            )
+            schema.add_field(
+                field_name="vector",
+                datatype=DataType.FLOAT_VECTOR,
+                dim=dimension,
+                description="vector"
+            )
+            # schema.add_field(
+            #     field_name="embedding",
+            #     datatype=DataType.FLOAT_VECTOR,
+            #     dim=dimension,
+            #     description="Binary Embeddings"
+            # )
+            schema.add_field(
+                field_name="document_meta",
+                datatype=DataType.JSON,
+                description="Custom Metadata information"
+            )
+            index_params = self._client.prepare_index_params()
+            index_params.add_index(
+                field_name="pk",
+                index_type="STL_SORT"
+            )
+            index_params.add_index(
+                field_name="text",
+                index_type="marisa-trie"
+            )
+            index_params.add_index(
+                field_name="summary",
+                index_type="marisa-trie"
+            )
+            index_params.add_index(
+                field_name="vector",
+                index_type=index_type,
+                metric_type=metric_type,
+                params=idx_params
+            )
+            self._client.create_collection(
+                collection_name=collection_name,
+                schema=schema,
+                index_params=index_params,
+                num_shards=2
+            )
+            await asyncio.sleep(2)
+            res = self._client.get_load_state(
+                collection_name=collection_name
+            )
+            return None
+        else:
+            self._client.create_collection(
+                collection_name=collection_name,
+                dimension=dimension
+            )
+            if metadata_field:
+                kwargs['metadata_field'] = metadata_field
+            # Here using drop_old=True to force recreate based on the first document
+            docstore = Milvus.from_documents(
+                [document],  # Only the first document
+                self._embed_,
+                connection_args={**self.kwargs},
+                collection_name=collection_name,
+                drop_old=True,
+                # consistency_level='Session',
+                primary_field='pk',
+                text_field='text',
+                vector_field='vector',
+                **kwargs
+            )
+            return docstore
+
+    async def load_documents(
+        self,
+        documents: list,
+        collection: str = None,
+        upsert: bool = False,
+        attribute: str = 'source_type',
+        metadata_field: str = None,
+        **kwargs
+    ):
+        if not collection:
+            collection = self.collection
+        try:
+            tensor = torch.randn(1000, 1000).cuda()
+        except Exception:
+            tensor = None
+        if upsert is True:
+            # get first document
+            doc = documents[0]
+            # getting source type:
+            doc_type = doc.metadata.get('attribute', None)
+            if attribute:
+                deleted = self._client.delete(
+                    collection_name=collection,
+                    filter=f'{attribute} == "{doc_type}"'
+                )
+                self.logger.notice(
+                    f"Deleted documents with {attribute} {attribute}: {deleted}"
+                )
+        if metadata_field:
+            # document_meta
+            kwargs['metadata_field'] = metadata_field
+        docstore = Milvus.from_documents(
+            documents,
+            self._embed_,
+            connection_args={**self.kwargs},
+            collection_name=collection,
+            consistency_level='Bounded',
+            drop_old=False,
+            primary_field='pk',
+            text_field='text',
+            vector_field='vector',
+            **kwargs
+        )
+        del tensor
+        return docstore
+
+    def upsert(self, payload: dict, collection: str = None) -> None:
+        pass
+
+    def insert(
+        self,
+        payload: Union[dict, list],
+        collection: str = None
+    ) -> dict:
+        if collection is None:
+            collection = self.collection
+        result = self._client.insert(
+            collection_name=collection,
+            data=payload
+        )
+        collection.flush()
+        return result
+
+    def get_vector(
+        self,
+        collection: str = None,
+        metric_type: str = None,
+        nprobe: int = 200,
+        metadata_field: str = None
+    ) -> Milvus:
+        if not metric_type:
+            metric_type = self._metric_type
+        if not collection:
+            collection = self.collection
+        _search = {
+            "search_params": {
+                "metric_type": metric_type,
+                "params": {"nprobe": nprobe},
+            }
+        }
+        if metadata_field:
+            # document_meta
+            _search['metadata_field'] = metadata_field
+        _embed_ = self.create_embedding(
+            model_name=self.embedding_name
+        )
+        return Milvus(
+            embedding_function=_embed_,
+            collection_name=collection,
+            consistency_level='Bounded',
+            connection_args={
+                **self.kwargs
+            },
+            primary_field='pk',
+            text_field='text',
+            vector_field='vector',
+            **_search
+        )
+
+    def similarity_search(self, query: str, collection: str = None, limit: int = 2) -> list:
+        if collection is None:
+            collection = self.collection
+        if self._embed_ is None:
+            _embed_ = self.create_embedding(
+                model_name=self.embedding_name
+            )
+        else:
+            _embed_ = self._embed_
+        vector_db = Milvus(
+            embedding_function=_embed_,
+            collection_name=collection,
+            consistency_level='Bounded',
+            connection_args={
+                **self.kwargs
+            },
+            primary_field='pk',
+            text_field='text',
+            vector_field='vector'
+        )
+        return vector_db.similarity_search(query, k=limit)
+
+    def search(
+        self,
+        payload: Union[dict, list],
+        collection: str = None,
+        limit: Optional[int] = None
+    ) -> list:
+        args = {}
+        if collection is None:
+            collection = self.collection
+        if limit is not None:
+            args = {"limit": limit}
+        if isinstance(payload, dict):
+            payload = [payload]
+        result = self._client.search(
+            collection_name=collection,
+            data=payload,
+            **args
+        )
+        return result
+
+    def memory_retriever(self, num_results: int  = 5) -> VectorStoreRetrieverMemory:
+        vectordb = Milvus.from_documents(
+            {},
+            self._embed_,
+            connection_args={**self.kwargs}
+        )
+        retriever = Milvus.as_retriever(vectordb, search_kwargs=dict(k=num_results))
+        return VectorStoreRetrieverMemory(retriever=retriever)
+
+    def save_context(self, memory: VectorStoreRetrieverMemory, context: list) -> None:
+        for val in context:
+            memory.save_context(val)

parrot/stores/qdrant.py

@@ -0,0 +1,153 @@
+from typing import Any
+from qdrant_client import QdrantClient  # pylint: disable=import-error
+from langchain_community.vectorstores import (  # pylint: disable=import-error, E0611
+    Qdrant
+)
+from .abstract import AbstractStore
+from ..conf import (
+    QDRANT_PROTOCOL,
+    QDRANT_HOST,
+    QDRANT_PORT,
+    QDRANT_USE_HTTPS,
+    QDRANT_CONN_TYPE,
+    QDRANT_URL
+)
+
+
+class QdrantStore(AbstractStore):
+    """QdrantStore class.
+
+
+    Args:
+        host (str): Qdrant host.
+        port (int): Qdrant port.
+        index_name (str): Qdrant index name.
+    """
+
+    def _create_qdrant_client(self, host, port, url, https, verify, qdrant_args):
+        """
+        Creates a Qdrant client based on the provided configuration.
+
+        Args:
+            host: Host of the Qdrant server (if using "server" connection).
+            port: Port of the Qdrant server (if using "server" connection).
+            url: URL of the Qdrant cloud service (if using "cloud" connection).
+            https: Whether to use HTTPS for the connection.
+            verify: Whether to verify the SSL certificate.
+            qdrant_args: Additional arguments for the Qdrant client.
+
+        Returns:
+            A QdrantClient object.
+        """
+        if url is not None:
+            return QdrantClient(
+            url=url,
+            port=None,
+            verify=verify,
+            **qdrant_args
+            )
+        else:
+            return QdrantClient(
+            host,
+            port=port,
+            https=https,
+            verify=verify,
+            **qdrant_args
+            )
+
+    def __init__(self, embeddings = None, **kwargs):
+        super().__init__(embeddings, **kwargs)
+        self.host = kwargs.get("host", QDRANT_HOST)
+        self.port = kwargs.get("port", QDRANT_PORT)
+        qdrant_args = kwargs.get("qdrant_args", {})
+        connection_type = kwargs.get("connection_type", QDRANT_CONN_TYPE)
+        url = kwargs.get("url", QDRANT_URL)
+        if connection_type == "server":
+            self.client = self._create_qdrant_client(
+                self.host, self.port, url, QDRANT_USE_HTTPS, False, qdrant_args
+            )
+        elif connection_type == "cloud":
+            if url is None:
+                raise ValueError(
+                    "A URL is required for 'cloud' connection"
+                )
+            self.client = self._create_qdrant_client(
+                None, None, url, False, False, qdrant_args
+            )
+        else:
+            raise ValueError(
+                f"Invalid connection type: {connection_type}"
+            )
+        if url is not None:
+            self.url = url
+        else:
+            self.url = f"{QDRANT_PROTOCOL}://{self.host}"
+            if self.port:
+                self.url += f":{self.port}"
+
+    def get_vectorstore(self):
+        if self._embed_ is None:
+            _embed_ = self.create_embedding(
+                model_name=self.embedding_name
+            )
+        else:
+            _embed_ = self._embed_
+        self.vector = Qdrant(
+            client=self.client,
+            collection_name=self.collection,
+            embeddings=_embed_,
+        )
+        return self.vector
+
+    async def load_documents(
+        self,
+        documents: list,
+        collection: str = None
+    ):
+        if collection is None:
+            collection = self.collection
+
+        docstore = Qdrant.from_documents(
+            documents,
+            self._embed_,
+            url=self.url,
+            # location=":memory:",  # Local mode with in-memory storage only
+            collection_name=collection,
+            force_recreate=False,
+        )
+        return docstore
+
+    def upsert(self, payload: dict, collection: str = None) -> None:
+        if collection is None:
+            collection = self.collection
+        self.client.upsert(
+            collection_name=collection,
+            points=self._embed_,
+            payload=payload
+        )
+
+    def search(self, payload: dict, collection: str = None) -> dict:
+        pass
+
+    async def delete_collection(self, collection: str = None) -> dict:
+        self.client.delete_collection(
+            collection_name=collection
+        )
+
+    async def create_collection(
+        self,
+        collection_name: str,
+        document: Any,
+        dimension: int = 768,
+        **kwargs
+    ) -> dict:
+        # Here using drop_old=True to force recreate based on the first document
+        docstore = Qdrant.from_documents(
+            [document],
+            self._embed_,
+            url=self.url,
+            # location=":memory:",  # Local mode with in-memory storage only
+            collection_name=collection_name,
+            force_recreate=True,
+        )
+        return docstore

parrot/tools/__init__.py

@@ -0,0 +1,16 @@
+import os
+from langchain_community.tools.yahoo_finance_news import YahooFinanceNewsTool
+from langchain_community.tools import YouTubeSearchTool
+from langchain_community.agent_toolkits import O365Toolkit
+from navconfig import config
+from .wikipedia import WikipediaTool, WikidataTool
+from .asknews import AskNewsTool
+from .duck import DuckDuckGoSearchTool, DuckDuckGoRelevantSearch
+from .weather import OpenWeather, OpenWeatherMapTool
+from .google import GoogleLocationFinder, GoogleSiteSearchTool, GoogleSearchTool
+from .zipcode import ZipcodeAPIToolkit
+from .bing import BingSearchTool
+from .stack import StackExchangeTool
+
+
+os.environ["USER_AGENT"] = "Parrot.AI/1.0"