beekeeper-ai 0.6.6__py3-none-any.whl → 1.0.1__py3-none-any.whl

beekeeper/retrievers/__init__.py

@@ -1,5 +0,0 @@
-from beekeeper.retrievers.watson_discovery import WatsonDiscoveryRetriever
-
-__all__ = [
-    "WatsonDiscoveryRetriever",
-]

beekeeper/retrievers/watson_discovery.py

@@ -1,121 +0,0 @@
-from logging import getLogger
-from typing import List
-
-from beekeeper.core.document import Document, DocumentWithScore
-
-logger = getLogger(__name__)
-
-
-class WatsonDiscoveryRetriever:
-    """Provides functionality to interact with IBM Watson Discovery for querying documents.
-
-    See https://cloud.ibm.com/docs/discovery-data?topic=discovery-data-getting-started for more info.
-
-    Args:
-        url (str): Watson Discovery instance url.
-        api_key (str): Watson Discovery API key.
-        project_id (str): Watson Discovery project_id.
-        version (str, optional): Watson Discovery API version. Defaults to ``2023-03-31``.
-        disable_passages (bool, optional): Return the full document instead of passages (only enable this if all documents are short). Defaults to ``False``.
-
-    **Example**
-
-    .. code-block:: python
-
-        from beekeeper.retrievers import WatsonDiscoveryRetriever
-
-        doc_retriever = WatsonDiscoveryRetriever(url="your_url",
-                                                api_key="your_api_key",
-                                                project_id="your_project_id")
-    """
-
-    def __init__(self,
-                 url: str,
-                 api_key: str,
-                 project_id: str,
-                 version: str = "2023-03-31",
-                 disable_passages: bool = False
-                 ) -> None:
-        try:
-            from ibm_cloud_sdk_core.authenticators import IAMAuthenticator
-            from ibm_watson import DiscoveryV2
-
-        except ImportError:
-            raise ImportError("ibm-watson package not found, please install it with `pip install ibm-watson`")
-
-        self.disable_passages = disable_passages
-        self.project_id = project_id
-
-        try:
-            authenticator = IAMAuthenticator(api_key)
-            self._client = DiscoveryV2(authenticator=authenticator,
-                                       version=version)
-
-            self._client.set_service_url(url)
-        except Exception as e:
-            logger.error(f"Error connecting to IBM Watson Discovery: {e}")
-            raise
-
-    def query(self, query: str, filter: str = None, top_k: int = 4) -> List[DocumentWithScore]:
-        """Search your data in the Discovery API and return a list of documents.
-
-        Args:
-            query (str): Query text.
-            filter (str, optional): Searches for documents that match the filter. Use Discovery Query Language syntax.
-            top_k (int, optional): Number of top results to return. Defaults to ``4``.
-
-        **Example**
-
-        .. code-block:: python
-
-            docs = doc_retriever.query("What's Beekeeper?")
-        """
-        from ibm_watson.discovery_v2 import QueryLargePassages
-        return_fields = ["extracted_metadata.filename", "extracted_metadata.file_type"]
-
-        if not self.disable_passages:
-            return_fields.append("passages")
-        else:
-            return_fields.append("text")
-
-        discovery_results = self._client.query(
-            project_id=self.project_id,
-            natural_language_query=query,
-            count=top_k,
-            return_=return_fields,
-            filter=filter,
-            passages=QueryLargePassages(enabled=not self.disable_passages,
-                                        per_document=False,
-                                        count=top_k,
-                                        find_answers=False,
-                                        characters=600)
-        ).get_result()
-
-        docs_and_scores = []
-
-        if not self.disable_passages and len(discovery_results["passages"]) > 0:
-            # If not `disable_passages`, always use discovery passages (recommended)
-            for passage in discovery_results["passages"]:
-                document_id_target = passage["document_id"]
-                document = [doc for doc in discovery_results["results"] if doc["document_id"] == document_id_target]
-
-                docs_and_scores.append(DocumentWithScore(
-                    document=Document(
-                        text=passage["passage_text"],
-                        metadata={"collection_id": passage["collection_id"]} | document[0]["extracted_metadata"]),
-                    score=passage["passage_score"] / 100))
-
-        elif discovery_results["matching_results"] > 0:
-            # If `disable_passages`, use document text (not recommended,
-            # make sure that all documents are short to not exceed the model context window)
-            logger.warning("Not recommended to disable passages. Make sure that all documents are short to not "
-                            "exceed the model context window.")
-            for document in discovery_results["results"]:
-                docs_and_scores.append(DocumentWithScore(
-                    document=Document(
-                        text=" ".join(document["text"]),
-                        metadata={"collection_id": document["result_metadata"]["collection_id"]} | document[
-                            "extracted_metadata"]),
-                    score=document["result_metadata"]["confidence"]))
-
-        return docs_and_scores

beekeeper/text_splitters/__init__.py

@@ -1,9 +0,0 @@
-from beekeeper.text_splitters.semantic import SemanticSplitter
-from beekeeper.text_splitters.sentence import SentenceSplitter
-from beekeeper.text_splitters.token import TokenTextSplitter
-
-__all__ = [
-    "SentenceSplitter",
-    "SemanticSplitter",
-    "TokenTextSplitter",
-]

beekeeper/text_splitters/semantic.py

@@ -1,139 +0,0 @@
-import re
-from typing import List, Literal, Tuple
-
-import numpy as np
-from pydantic.v1 import BaseModel
-
-from beekeeper.core.document import Document
-from beekeeper.core.embeddings import BaseEmbedding
-from beekeeper.core.utils.pairwise import cosine_similarity
-
-
-class SemanticSplitter(BaseModel):
-    """Python class designed to split text into chunks using semantic understanding.
-
-    Credit to Greg Kamradt's notebook: `5 Levels Of Text Splitting <https://github.com/FullStackRetrieval-com/RetrievalTutorials/blob/main/tutorials/LevelsOfTextSplitting/5_Levels_Of_Text_Splitting.ipynb>`_.
-
-    Args:
-        embed_model (BaseEmbedding):
-        buffer_size (int, optional): Size of the buffer for semantic chunking. Default is ``1``.
-        breakpoint_threshold_amount (int, optional): Threshold percentage for detecting breakpoints. Default is ``95``.
-        device (str, optional): Device to use for processing. Currently supports "cpu" and "cuda". Default is ``cpu``.
-
-    **Example**
-
-    .. code-block:: python
-
-        from beekeeper.embeddings import HuggingFaceEmbedding
-        from beekeeper.text_splitters import SemanticSplitter
-
-        embedding = HuggingFaceEmbedding()
-        text_splitter = SemanticSplitter(embed_model=embedding)
-    """
-
-    embed_model: BaseEmbedding
-    buffer_size: int = 1
-    breakpoint_threshold_amount: int = 95
-    device: Literal["cpu", "cuda"] = "cpu"
-
-    class Config:
-        arbitrary_types_allowed = True
-
-    def _combine_sentences(self, sentences: List[dict]) -> List[dict]:
-        """Combine sentences with neighbors based on buffer size."""
-        for i in range(len(sentences)):
-            combined_sentence = ""
-            
-            # Add previous sentences based on buffer size
-            for j in range(i - self.buffer_size, i):
-                if j >= 0:
-                    combined_sentence += sentences[j]["sentence"] + " "
-
-            # Add the current sentence
-            combined_sentence += sentences[i]["sentence"]
-
-            # Add next sentences based on buffer size
-            for j in range(i + 1, i + 1 + self.buffer_size):
-                if j < len(sentences):
-                    combined_sentence += " " + sentences[j]["sentence"]
-
-            sentences[i]["combined_sentence"] = combined_sentence
-        
-        return sentences
-        
-    def _calculate_cosine_distances(self, single_sentences_list: List[str]) -> Tuple[List[float], List[dict]]:
-        _sentences = [{"sentence": x, "index": i} for i, x in enumerate(single_sentences_list)]
-        
-        sentences = self._combine_sentences(_sentences)
-        embeddings = self.embed_model.get_texts_embedding(
-            [x["combined_sentence"] for x in sentences]
-        )
-        
-        for i, sentence in enumerate(sentences):
-            sentence["combined_sentence_embedding"] = embeddings[i]
-            
-        distances = []
-        for i in range(len(sentences) - 1):
-            embedding_current = sentences[i]["combined_sentence_embedding"]
-            embedding_next = sentences[i + 1]["combined_sentence_embedding"]
-            
-            similarity = cosine_similarity(embedding_current, embedding_next)
-            
-            distance = 1 - similarity
-            distances.append(distance)
-            
-            # Store distance in the dictionary
-            sentences[i]["distance_to_next"] = distance
-        
-        return distances, sentences
-    
-    def _calculate_breakpoint(self, distances: List[float]) -> List:
-        distance_threshold = np.percentile(distances, self.breakpoint_threshold_amount)
-        
-        return [i for i, x in enumerate(distances) if x > distance_threshold]
-        
-    def from_text(self, text: str) -> List[str]:
-        """Split text into chunks.
-        
-        Args:
-            text (str): Input text to split.
-        """
-        single_sentences_list = re.split(r"(?<=[.?!])\s+", text)
-        distances, sentences = self._calculate_cosine_distances(single_sentences_list)
-        
-        indices_above_thresh = self._calculate_breakpoint(distances)
-        
-        chunks = []
-        start_index = 0
-
-        for index in indices_above_thresh:
-            # Slice the sentence_dicts from the current start index to the end index
-            group = sentences[start_index: index + 1]
-            combined_text = " ".join([d["sentence"] for d in group])
-            chunks.append(combined_text)
-            
-            # Update the start index for the next group
-            start_index = index + 1
-            
-        # The last group, if any sentences remain
-        if start_index < len(sentences):
-            combined_text = " ".join([d["sentence"] for d in sentences[start_index:]])
-            chunks.append(combined_text)
-        
-        return chunks
-        
-    def from_documents(self, documents: List[Document]) -> List[Document]:
-        """Split documents into chunks.
-        
-        Args:
-            documents (List[Document]): List of `Document` objects to split.
-        """
-        chunks = []
-
-        for document in documents:
-            texts = self.from_text(document.get_content())
-
-            for text in texts:
-                chunks.append(Document(text=text, metadata=document.get_metadata()))
-
-        return chunks

beekeeper/text_splitters/sentence.py

@@ -1,107 +0,0 @@
-from typing import List
-
-from beekeeper.core.document import Document
-from beekeeper.core.text_splitters.utils import (
-    merge_splits,
-    split_by_char,
-    split_by_fns,
-    split_by_regex,
-    split_by_sentence_tokenizer,
-    split_by_sep,
-    tokenizer,
-)
-
-
-class SentenceSplitter:
-    """Designed to split input text into smaller chunks,
-    particularly useful for processing large documents or texts, tries to keep sentences and paragraphs together.
-
-    Args:
-        chunk_size (int, optional): Size of each chunk. Default is ``512``.
-        chunk_overlap (int, optional): Amount of overlap between chunks. Default is ``256``.
-        separator (str, optional): Separators used for splitting into words. Default is ``" "``
-
-    **Example**
-
-    .. code-block:: python
-
-        from beekeeper.text_splitters import SentenceSplitter
-
-        text_splitter = SentenceSplitter()
-    """
-
-    def __init__(self,
-                 chunk_size: int = 512,
-                 chunk_overlap: int = 256,
-                 separator=" "
-                 ) -> None:
-
-        if chunk_overlap > chunk_size:
-            raise ValueError(
-                f"Got a larger `chunk_overlap` ({chunk_overlap}) than `chunk_size` "
-                f"({chunk_size}). `chunk_overlap` should be smaller."
-            )
-
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-
-        self._split_fns = [
-            split_by_sep("\n\n\n"),
-            split_by_sentence_tokenizer()
-        ]
-        self._sub_split_fns = [
-            split_by_regex("[^,.;？！]+[,.;？！]?"),
-            split_by_sep(separator),
-            split_by_char()
-        ]
-
-    def from_text(self, text: str) -> List[str]:
-        """Split text into chunks.
-        
-        Args:
-            text (str): Input text to split.
-
-        **Example**
-
-        .. code-block:: python
-
-            chunks = text_splitter.from_text("Beekeeper is a data framework to load any data in one line of code and connect with AI applications.")
-        """
-        splits = self._split(text)
-
-        return merge_splits(splits, self.chunk_size, self.chunk_overlap)
-
-    def from_documents(self, documents: List[Document]) -> List[Document]:
-        """Split documents into chunks.
-
-        Args:
-            documents (List[Document]): List of `Document` objects to split.
-        """
-        chunks = []
-
-        for document in documents:
-            texts = self.from_text(document.get_content())
-
-            for text in texts:
-                chunks.append(Document(text=text, metadata=document.get_metadata()))
-
-        return chunks
-
-    def _split(self, text: str) -> List[dict]:
-
-        text_len = len(tokenizer(text))
-        if text_len <= self.chunk_size:
-            return [{"text": text, "is_sentence": True, "token_size": text_len}]
-
-        text_splits = []
-        text_splits_by_fns, is_sentence = split_by_fns(text, self._split_fns, self._sub_split_fns)
-
-        for text_split_by_fns in text_splits_by_fns:
-            split_len = len(tokenizer(text_split_by_fns))
-            if split_len <= self.chunk_size:
-                text_splits.append({"text": text_split_by_fns, "is_sentence": is_sentence, "token_size": split_len})
-            else:
-                recursive_text_splits = self._split(text_split_by_fns)
-                text_splits.extend(recursive_text_splits)
-
-        return text_splits

beekeeper/text_splitters/token.py

@@ -1,101 +0,0 @@
-from typing import List
-
-from beekeeper.core.document import Document
-from beekeeper.core.text_splitters.utils import (
-    merge_splits,
-    split_by_char,
-    split_by_fns,
-    split_by_sep,
-    tokenizer,
-)
-
-
-class TokenTextSplitter:
-    r"""This is the simplest splitting method. Designed to split input text into smaller chunks looking at word tokens.
-
-    Args:
-        chunk_size (int, optional): Size of each chunk. Default is ``512``.
-        chunk_overlap (int, optional): Amount of overlap between chunks. Default is ``256``.
-        separator (str, optional): Separators used for splitting into words. Default is ``\\n\\n``.
-
-    **Example**
-
-    .. code-block:: python
-
-        from beekeeper.text_splitters import TokenTextSplitter
-
-        text_splitter = TokenTextSplitter()
-    """
-
-    def __init__(self,
-                 chunk_size: int = 512,
-                 chunk_overlap: int = 256,
-                 separator="\n\n") -> None:
-
-        if chunk_overlap > chunk_size:
-            raise ValueError(
-                f"Got a larger `chunk_overlap` ({chunk_overlap}) than `chunk_size` "
-                f"({chunk_size}). `chunk_overlap` should be smaller."
-            )
-
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-
-        self._split_fns = [
-            split_by_sep(separator)
-        ]
-
-        self._sub_split_fns = [
-            split_by_char()
-        ]
-
-    def from_text(self, text: str) -> List[str]:
-        """Split text into chunks.
-
-        Args:
-            text (str): Input text to split.
-
-        **Example**
-
-        .. code-block:: python
-
-            chunks = text_splitter.from_text("Beekeeper is a data framework to load any data in one line of code and connect with AI applications.")
-        """
-        splits = self._split(text)
-
-        return merge_splits(splits, self.chunk_size, self.chunk_overlap)
-
-    def from_documents(self, documents: List[Document]) -> List[Document]:
-        """Split documents into chunks.
-
-        Args:
-            documents (List[Document]): List of `Document` objects to split.
-        """
-        chunks = []
-
-        for document in documents:
-            texts = self.from_text(document.get_content())
-
-            for text in texts:
-                chunks.append(Document(text=text, metadata=document.get_metadata()))
-
-        return chunks
-
-    def _split(self, text: str) -> List[dict]:
-
-        text_len = len(tokenizer(text))
-        if text_len <= self.chunk_size:
-            return [{"text": text, "is_sentence": True, "token_size": text_len}]
-
-        text_splits = []
-        text_splits_by_fns, is_sentence = split_by_fns(text, self._split_fns, self._sub_split_fns)
-
-        for text_split_by_fns in text_splits_by_fns:
-            split_len = len(tokenizer(text_split_by_fns))
-            if split_len <= self.chunk_size:
-                text_splits.append({"text": text_split_by_fns, "is_sentence": False, "token_size": split_len})
-            else:
-                recursive_text_splits = self._split(text_split_by_fns)
-                text_splits.extend(recursive_text_splits)
-
-        return text_splits

beekeeper/vector_stores/__init__.py

@@ -1,7 +0,0 @@
-from beekeeper.vector_stores.chroma import ChromaVectorStore
-from beekeeper.vector_stores.elasticsearch import ElasticsearchVectorStore
-
-__all__ = [
-    "ChromaVectorStore",
-    "ElasticsearchVectorStore",
-]

beekeeper/vector_stores/chroma.py

@@ -1,115 +0,0 @@
-import uuid
-from logging import getLogger
-from typing import List, Literal
-
-from beekeeper.core.document import Document, DocumentWithScore
-from beekeeper.core.embeddings import BaseEmbedding
-
-logger = getLogger(__name__)
-
-
-class ChromaVectorStore:
-    """Chroma is the AI-native open-source vector database. Embeddings are stored within a ChromaDB collection.
-
-    Args:
-        embed_model (BaseEmbedding):
-        collection_name (str, optional): Name of the ChromaDB collection.
-        distance_strategy (str, optional): Distance strategy for similarity search. Currently supports "cosine", "ip" and "l2". Defaults to ``cosine``.
-
-    **Example**
-
-    .. code-block:: python
-
-        from beekeeper.embeddings import HuggingFaceEmbedding
-        from beekeeper.vector_stores import ChromaVectorStore
-
-        embedding = HuggingFaceEmbedding()
-        vector_db = ChromaVectorStore(embed_model=embedding)
-    """
-
-    def __init__(self, embed_model: BaseEmbedding,
-                 collection_name: str = None,
-                 distance_strategy: Literal["cosine", "ip", "l2"] = "cosine") -> None:
-        try:
-            import chromadb
-            import chromadb.config
-
-        except ImportError:
-            raise ImportError("chromadb package not found, please install it with `pip install chromadb`")
-
-        self._embed_model = embed_model
-        self._client_settings = chromadb.config.Settings()
-        self._client = chromadb.Client(self._client_settings)
-
-        if collection_name is None:
-            collection_name = "auto-generated-" + str(uuid.uuid4())[:8]
-            logger.info(f"collection_name: {collection_name}")
-
-        self._collection = self._client.get_or_create_collection(
-            name=collection_name,
-            embedding_function=None,
-            metadata={"hnsw:space": distance_strategy}
-        )
-
-    def add_documents(self, documents: List[Document]) -> List:
-        """Add documents to the ChromaDB collection.
-
-        Args:
-            documents (List[Document]): List of `Document` objects to add to the collection.
-        """
-        embeddings = []
-        metadatas = []
-        ids = []
-        chroma_documents = []
-
-        for doc in documents:
-            embeddings.append(self._embed_model.get_query_embedding(doc.get_content()))
-            metadatas.append(doc.get_metadata() if doc.get_metadata() else None)
-            ids.append(doc.doc_id if doc.doc_id else str(uuid.uuid4()))
-            chroma_documents.append(doc.get_content())
-
-        self._collection.add(embeddings=embeddings,
-                             ids=ids,
-                             metadatas=metadatas,
-                             documents=chroma_documents)
-
-        return ids
-
-    def query(self, query: str, top_k: int = 4) -> List[DocumentWithScore]:
-        """Performs a similarity search for top-k most similar documents.
-
-        Args:
-            query (str): Query text.
-            top_k (int, optional): Number of top results to return. Defaults to ``4``.
-        """
-        query_embedding = self._embed_model.get_query_embedding(query)
-
-        results = self._collection.query(
-            query_embeddings=query_embedding,
-            n_results=top_k
-        )
-
-        return [
-            DocumentWithScore(document=Document(
-                doc_id=result[0],
-                text=result[1],
-                metadata=result[2]
-            ), score=result[3])
-            for result in zip(
-                results["ids"][0],
-                results["documents"][0],
-                results["metadatas"][0],
-                results["distances"][0],
-            )
-        ]
-
-    def delete_documents(self, ids: List[str] = None) -> None:
-        """Delete documents from the ChromaDB collection.
-
-        Args:
-            ids (List[str]): List of `Document` IDs to delete. Defaults to ``None``.
-        """
-        if not ids:
-            raise ValueError("No ids provided to delete.")
-
-        self._collection.delete(ids=ids)