beekeeper-ai 0.6.1__py3-none-any.whl

beekeeper/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.6.1"

beekeeper/core/document/__init__.py

@@ -0,0 +1,6 @@
+from beekeeper.core.document.schema import Document, DocumentWithScore
+
+__all__ = [
+    "Document",
+    "DocumentWithScore"
+]

beekeeper/core/document/schema.py

@@ -0,0 +1,97 @@
+import uuid
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from pydantic.v1 import BaseModel, Field, validator
+
+if TYPE_CHECKING:
+    from langchain_core.documents import Document as LangChainDocument
+
+
+class BaseDocument(ABC, BaseModel):
+    """Generic abstract interface for retrievable documents."""
+
+    doc_id: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Unique ID of the document.")
+    metadata: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="A flat dictionary of metadata fields.")
+
+    @validator("metadata", pre=True)
+    def _validate_metadata(cls, v) -> Dict:
+        if v is None:
+            return {}
+        return v
+
+    @abstractmethod
+    def get_content(self) -> str:
+        """Get document content."""
+
+    @abstractmethod
+    def get_metadata(self) -> str:
+        """Get metadata."""
+
+
+class Document(BaseDocument):
+    """Generic interface for data document."""
+
+    text: str = Field(default="", description="Text content of the document.")
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "Document"
+
+    def get_content(self) -> str:
+        """Get the text content."""
+        return self.text
+
+    def get_metadata(self) -> dict:
+        """Get metadata."""
+        return self.metadata
+
+    @classmethod
+    def from_langchain_format(cls, doc: "LangChainDocument") -> "Document":
+        """
+        Convert a document from LangChain  format.
+
+        Args:
+            doc (LangChainDocument): Document in LangChain format.
+        """
+        return cls(text=doc.page_content, metadata=doc.metadata)
+
+
+@dataclass
+class DocumentWithScore:
+    document: BaseDocument
+    score: Optional[float] = None
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "DocumentWithScore"
+
+    def get_score(self) -> float:
+        """Get score."""
+        if self.score is None:
+            return 0.0
+        else:
+            return self.score
+
+    # #### pass through methods to BaseDocument ####
+    @property
+    def doc_id(self) -> str:
+        return self.document.doc_id
+
+    @property
+    def text(self) -> str:
+        if isinstance(self.document, Document):
+            return self.document.text
+        else:
+            raise ValueError("Must be a Document to get text")
+
+    def get_content(self) -> str:
+        return self.document.get_content()
+
+    def get_metadata(self) -> str:
+        return self.document.get_metadata()

beekeeper/core/document_loaders/__init__.py

@@ -0,0 +1,5 @@
+from beekeeper.core.document_loaders.base import BaseLoader
+
+__all__ = [
+    "BaseLoader",
+]

beekeeper/core/document_loaders/base.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from pydantic.v1 import BaseModel
+
+from beekeeper.core.document import Document
+
+
+class BaseLoader(ABC, BaseModel):
+    """An interface for document loader."""
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "BaseLoader"
+
+    @abstractmethod
+    def load_data(self, extra_info: Optional[dict] = None) -> List[Document]:
+        """Loads data."""
+
+    def load(self) -> List[Document]:
+        return self.load_data()
+
+    def lazy_load(self) -> List[Document]:
+        return self.load_data()

beekeeper/core/embeddings/__init__.py

@@ -0,0 +1,6 @@
+from beekeeper.core.embeddings.base import BaseEmbedding, Embedding
+
+__all__ = [
+    "BaseEmbedding",
+    "Embedding",
+]

beekeeper/core/embeddings/base.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from typing import List, Literal
+
+import numpy as np
+
+from beekeeper.core.utils.pairwise import cosine_similarity
+
+Embedding = List[float]
+
+
+class BaseEmbedding(ABC):
+    """An interface for embedding models."""
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "BaseEmbedding"
+
+    @abstractmethod
+    def get_query_embedding(self, query: str) -> Embedding:
+        """Get query embedding."""
+
+    @abstractmethod
+    def get_texts_embedding(self, texts: List[str]) -> List[Embedding]:
+        """Get text embeddings."""
+
+    @abstractmethod
+    def get_documents_embedding(self, documents: List[str]) -> List[Embedding]:
+        """Get documents embeddings."""
+
+    def embed_documents(self, texts: List[str]) -> List[Embedding]:
+        return self.get_texts_embedding(texts=texts)
+
+    @staticmethod
+    def similarity(embedding1: Embedding, embedding2: Embedding,
+                   mode: Literal["cosine", "dot_product", "euclidean"] = "cosine"):
+        """Get embedding similarity."""
+        if mode == "euclidean":
+            return -float(np.linalg.norm(np.array(embedding1) - np.array(embedding2)))
+
+        elif mode == "dot_product":
+            return np.dot(embedding1, embedding2)
+
+        else:
+            return cosine_similarity(embedding1, embedding2)

beekeeper/core/text_splitters/utils.py

@@ -0,0 +1,142 @@
+from typing import Callable, List, Tuple
+
+
+def tokenizer(text: str) -> List:
+    try:
+        import tiktoken
+    except ImportError:
+        raise ImportError("tiktoken package not found, please install it with `pip install tiktoken`")
+
+    enc = tiktoken.get_encoding("cl100k_base")
+    return enc.encode(text)
+
+
+def split_by_sep(sep) -> Callable[[str], List[str]]:
+    """Split text by separator."""
+    return lambda text: text.split(sep)
+
+
+def split_by_regex(regex: str) -> Callable[[str], List[str]]:
+    """Split text by regex."""
+    import re
+
+    return lambda text: re.findall(regex, text)
+
+
+def split_by_char() -> Callable[[str], List[str]]:
+    """Split text by character."""
+    return lambda text: list(text)
+
+
+def split_by_sentence_tokenizer() -> Callable[[str], List[str]]:
+    try:
+        import nltk
+    except ImportError:
+        raise ImportError("nltk package not found, please install it with `pip install nltk`")
+
+    sentence_tokenizer = nltk.tokenize.PunktSentenceTokenizer()
+    return lambda text: _split_by_sentence_tokenizer(text, sentence_tokenizer)
+
+
+def _split_by_sentence_tokenizer(text: str, sentence_tokenizer) -> List[str]:
+    """Get the spans and then return the sentences.
+
+    Using the start index of each span
+    Instead of using end, use the start of the next span
+    """
+    spans = list(sentence_tokenizer.span_tokenize(text))
+    sentences = []
+    for i, span in enumerate(spans):
+        start = span[0]
+        if i < len(spans) - 1:
+            end = spans[i + 1][0]
+        else:
+            end = len(text)
+        sentences.append(text[start:end])
+    return sentences
+
+
+def split_by_fns(text: str,
+                 split_fns: List[Callable],
+                 sub_split_fns: List[Callable] = None) -> Tuple[List[str], bool]:
+    """Split text by defined list of split functions."""
+    if not split_fns:
+        raise ValueError("Must provide a `split_fns` parameter")
+
+    for split_fn in split_fns:
+        splits = split_fn(text)
+        if len(splits) > 1:
+            return splits, True
+
+    if sub_split_fns: # noqa: RET503
+        for split_fn in sub_split_fns: # noqa: RET503
+            splits = split_fn(text)
+            if len(splits) > 1:
+                return splits, False
+
+
+def merge_splits(splits: List[dict],
+                 chunk_size: int,
+                 chunk_overlap: int) -> List[str]:
+    """Merge splits into chunks."""
+    chunks: List[str] = []
+    cur_chunk: List[Tuple[str, int]] = []
+    cur_chunk_len = 0
+    last_chunk: List[Tuple[str, int]] = []
+    new_chunk = True
+
+    def close_chunk() -> None:
+        nonlocal chunks, cur_chunk, last_chunk, cur_chunk_len, new_chunk
+
+        chunks.append("".join([text for text, length in cur_chunk]))
+        last_chunk = cur_chunk
+        cur_chunk = []
+        cur_chunk_len = 0
+        new_chunk = True
+
+        # add overlap to the next chunk using previous chunk
+        if len(last_chunk) > 0:
+            last_index = len(last_chunk) - 1
+            while (
+                    last_index >= 0
+                    and cur_chunk_len + last_chunk[last_index][1] <= chunk_overlap):
+                text, length = last_chunk[last_index]
+                cur_chunk_len += length
+                cur_chunk.insert(0, (text, length))
+                last_index -= 1
+
+    def postprocess_chunks(_chunks: List[str]) -> List[str]:
+        """Post-process chunks."""
+        post_chunks = []
+        for _chunk in _chunks:
+            stripped_chunk = _chunk.strip()
+            if stripped_chunk == "":
+                continue
+            post_chunks.append(stripped_chunk)
+        return post_chunks
+
+    while len(splits) > 0:
+        cur_split = splits[0]
+
+        if cur_split["token_size"] > chunk_size:
+            raise ValueError("Got a split size that exceeded chunk size")
+
+        if cur_chunk_len + cur_split["token_size"] > chunk_size and not new_chunk:
+            close_chunk()
+        else:
+            if (cur_split["is_sentence"]
+                    or cur_chunk_len + cur_split["token_size"] <= chunk_size
+                    or new_chunk):  # If `new_chunk`, always add at least one split
+
+                cur_chunk_len += cur_split["token_size"]
+                cur_chunk.append((cur_split["text"], cur_split["token_size"]))
+                splits.pop(0)
+                new_chunk = False
+            else:
+                close_chunk()
+
+    if not new_chunk:
+        chunk = "".join([text for text, length in cur_chunk])
+        chunks.append(chunk)
+
+    return postprocess_chunks(chunks)

beekeeper/core/utils/pairwise.py

@@ -0,0 +1,20 @@
+from typing import List
+
+import numpy as np
+
+Matrix = List[float]
+
+def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:
+    """Row-wise cosine similarity between two equal-width matrices."""
+    X = np.array(X)
+    Y = np.array(Y)
+    
+    if X.shape[0] != Y.shape[0]:
+        raise ValueError(
+            f"Number of rows in X and Y must be the same. X has shape {X.shape} "
+            f"and Y has shape {Y.shape}."
+        )
+    product = np.dot(X, Y)
+    norm = np.linalg.norm(X) * np.linalg.norm(Y)
+    
+    return product / norm

beekeeper/document_loaders/__init__.py

@@ -0,0 +1,17 @@
+from beekeeper.document_loaders.directory import DirectoryLoader
+from beekeeper.document_loaders.docx import DocxLoader
+from beekeeper.document_loaders.html import HTMLLoader
+from beekeeper.document_loaders.json import JSONLoader
+from beekeeper.document_loaders.pdf import PDFLoader
+from beekeeper.document_loaders.s3 import S3Loader
+from beekeeper.document_loaders.watson_discovery import WatsonDiscoveryLoader
+
+__all__ = [
+    "DirectoryLoader",
+    "DocxLoader",
+    "HTMLLoader",
+    "JSONLoader",
+    "PDFLoader",
+    "S3Loader",
+    "WatsonDiscoveryLoader",
+]

beekeeper/document_loaders/directory.py

@@ -0,0 +1,65 @@
+import glob
+import os
+from pathlib import Path
+from typing import List, Optional, Type
+
+from beekeeper.core.document import Document
+from beekeeper.core.document_loaders import BaseLoader
+
+
+def _loading_default_supported_loaders():
+    from beekeeper.document_loaders import DocxLoader, HTMLLoader, PDFLoader
+
+    return {
+    ".docx": DocxLoader,
+    ".html": HTMLLoader,
+    ".pdf": PDFLoader,
+}
+
+
+class DirectoryLoader(BaseLoader):
+    """Simple directory loader.
+
+    Args:
+        required_exts: (List[str], optional): List of file extensions to only load files with those extensions.
+        recursive (str, optional): Whether to recursively search for files. Defaults to ``False``.
+    """
+    
+    required_exts: List[str] = [".pdf", ".docx", ".html"]
+    recursive: Optional[bool] = False
+    file_loader: Optional[dict[str, Type[BaseLoader]]] = None
+
+    def load_data(self, input_dir: str, extra_info: Optional[dict] = None) -> List[Document]:
+        """Loads data from the specified directory.
+        
+        Args:
+            input_dir (str): Directory path from which to load the documents.
+        """
+        if not os.path.isdir(input_dir):
+            raise ValueError(f"`{input_dir}` is not a valid directory.")
+        
+        if self.file_loader is None:
+            self.file_loader = _loading_default_supported_loaders()
+        
+        input_dir = Path(input_dir)
+        documents = []
+        
+        pattern_prefix = "**/*" if self.recursive else ""
+
+        for extension in self.required_exts:
+            files = glob.glob(os.path.join(input_dir, pattern_prefix + extension), recursive=self.recursive)
+            
+            for file_dir in files:
+                loader_cls = self.file_loader.get(extension)
+                if loader_cls:
+                    try:
+                        #TO-DO add `file_loader_kwargs`
+                        doc = loader_cls().load_data(file_dir)
+                        documents.extend(doc)
+                    except Exception as e:
+                        raise f"Error reading {file_dir}: {e}"
+                else:
+                    #TO-DO add `unstructured file` support
+                    raise f"Unsupported file type: {extension}"
+
+        return documents

beekeeper/document_loaders/docx.py

@@ -0,0 +1,31 @@
+import os
+from pathlib import Path
+from typing import List, Optional
+
+from beekeeper.core.document import Document
+from beekeeper.core.document_loaders import BaseLoader
+
+
+class DocxLoader(BaseLoader):
+    """Microsoft Word (Docx) loader."""
+
+    def load_data(self, input_file: str, extra_info: Optional[dict] = None) -> List[Document]:
+        """Loads data from the specified directory.
+        
+        Args:
+            input_file (str): File path to load.
+        """
+        try:
+            import docx2txt  # noqa: F401
+        except ImportError:
+            raise ImportError("docx2txt package not found, please install it with `pip install docx2txt`")
+        
+        if not os.path.isfile(input_file):
+            raise ValueError(f"File `{input_file}` does not exist")
+        
+        input_file = str(Path(input_file).resolve())
+
+        text = docx2txt.process(input_file)
+        metadata = {"source": input_file}
+
+        return [Document(text=text, metadata=metadata)]

beekeeper/document_loaders/html.py

@@ -0,0 +1,77 @@
+import os
+from pathlib import Path
+from typing import List, Optional
+
+from beekeeper.core.document import Document
+from beekeeper.core.document_loaders import BaseLoader
+
+
+class HTMLLoader(BaseLoader):
+    """Load HTML file and extract text from a specific tag.
+
+    Args:
+        tag (str): HTML tag to extract. Defaults to ``section``.
+    """
+    
+    tag: str = "section"
+
+    def load_data(self, input_file: str, extra_info: Optional[dict] = None) -> List[Document]:
+        """Loads data from the specified directory.
+        
+        Args:
+            input_file (str): File path to load.
+        """
+        try:
+            from bs4 import BeautifulSoup  # noqa: F401
+        except ImportError:
+            raise ImportError("beautifulsoup4 package not found, please install it with `pip install beautifulsoup4`")
+        
+        if not os.path.isfile(input_file):
+            raise ValueError(f"File `{input_file}` does not exist")
+        
+        input_file = str(Path(input_file).resolve())
+
+        with open(input_file, encoding="utf-8") as html_file:
+            soup = BeautifulSoup(html_file, "html.parser")
+
+        tags = soup.find_all(self.tag)
+        documents = []
+
+        for tag in tags:
+            tag_text = self._extract_text_from_tag(tag)
+
+            metadata = {
+                "tag": self.tag,
+                "source": input_file,
+            }
+
+            doc = Document(
+                text=tag_text,
+                metadata=metadata,
+            )
+
+            documents.append(doc)
+
+        return documents
+
+    def _extract_text_from_tag(self, tag) -> str:
+        """Extract the text from an HTML tag, ignoring other nested tags."""
+        try:
+            from bs4 import NavigableString  # noqa: F401
+        except ImportError:
+            raise ImportError("beautifulsoup4 package not found, please install it with `pip install beautifulsoup4`")
+
+        texts = []
+
+        for elem in tag.children:
+            # Check if the element is a text node, not a tag
+            if isinstance(elem, NavigableString):
+                if elem.strip():
+                    texts.append(elem.strip())
+            # Ignore any tag that matches the main tag being processed (to avoid recursion)
+            elif elem.name == self.tag:
+                continue
+            else:
+                texts.append(elem.get_text().strip())
+
+        return "\n".join(texts)

beekeeper/document_loaders/json.py

@@ -0,0 +1,53 @@
+import json
+import os
+from pathlib import Path
+from typing import List, Optional
+
+from beekeeper.core.document import Document
+from beekeeper.core.document_loaders import BaseLoader
+
+
+class JSONLoader(BaseLoader):
+    """JSON loader.
+
+    Args:
+        jq_schema (str, optional): jq schema to use to extract the data from the JSON.
+    """
+
+    jq_schema: Optional[str] = None
+
+    def load_data(self, input_file: str, extra_info: Optional[dict] = None) -> List[Document]:
+        """Loads data from the specified directory.
+        
+        Args:
+            input_file (str): File path to load.
+        """
+        try:
+            import jq  # noqa: F401
+        except ImportError:
+            raise ImportError("jq package not found, please install it with `pip install jq`")
+        
+        if not os.path.isfile(input_file):
+            raise ValueError(f"File `{input_file}` does not exist")
+        
+        documents = []
+        jq_compiler = jq.compile(self.jq_schema)
+        json_file = Path(input_file).resolve().read_text(encoding="utf-8")
+        json_data = jq_compiler.input(json.loads(json_file))
+        
+        
+        for content in json_data:
+
+            if isinstance(content, str):
+                content = content
+            elif isinstance(content, dict):
+                content = json.dumps(content) if content else ""
+            else:
+                content = str(content) if content is not None else ""
+            
+            if content.strip() != "":
+                documents.append(Document(
+                    text=content,
+                    metadata={"source": str(Path(input_file).resolve())}))
+
+        return documents

beekeeper/document_loaders/pdf.py

@@ -0,0 +1,38 @@
+import logging
+import os
+from pathlib import Path
+from typing import List, Optional
+
+from beekeeper.core.document import Document
+from beekeeper.core.document_loaders import BaseLoader
+
+logging.getLogger("pypdf").setLevel(logging.ERROR)
+
+class PDFLoader(BaseLoader):
+    """PDF loader using PyPDF."""
+
+    def load_data(self, input_file: str, extra_info: Optional[dict] = None) -> List[Document]:
+        """Loads data from the specified directory.
+        
+        Args:
+            input_file (str): File path to load.
+        """
+        try:
+            import pypdf  # noqa: F401
+
+        except ImportError:
+            raise ImportError("pypdf package not found, please install it with `pip install pypdf`")
+        
+        if not os.path.isfile(input_file):
+            raise ValueError(f"File `{input_file}` does not exist")
+        
+        input_file = str(Path(input_file).resolve())
+        pdf_loader = pypdf.PdfReader(input_file)
+
+        return [
+            Document(
+                text=page.extract_text().strip(),
+                metadata={"source": input_file, "page": page_number}
+            )
+            for page_number, page in enumerate(pdf_loader.pages)
+        ]