agentic-base 0.1.0__tar.gz

agentic_base-0.1.0/PKG-INFO

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: agentic-base
+Version: 0.1.0
+Summary: 
+Author: jalal
+Author-email: jalalkhaldi3@gmail.com
+Requires-Python: >=3.11,<3.13
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: faiss-cpu (>=1.13.2,<2.0.0)
+Requires-Dist: langchain (>=1.2.10,<2.0.0)
+Requires-Dist: minio (>=7.2.20,<8.0.0)
+Requires-Dist: numpy (>=2.4.2,<3.0.0)
+Requires-Dist: openai (>=2.21.0,<3.0.0)
+Requires-Dist: polars (>=1.38.1,<2.0.0)
+Requires-Dist: pydantic-settings (>=2.13.0,<3.0.0)
+Requires-Dist: qdrant-client (>=1.16.2,<2.0.0)
+Requires-Dist: rank-bm25 (>=0.2.2,<0.3.0)
+Requires-Dist: transformers (>=5.2.0,<6.0.0)
+Description-Content-Type: text/markdown
+
+

agentic_base-0.1.0/agentic_base/connector/__init__.py

@@ -0,0 +1,71 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Generic, Union
+
+import polars as pl
+
+from agentic_base.mixins import FromConfigMixin
+from agentic_base.types import TCDatasetConnector
+
+if TYPE_CHECKING:
+    from agentic_base.dataset import Dataset, TextDataset
+
+_logger = logging.getLogger(__name__)
+
+
+class DatasetConnector(
+    FromConfigMixin[TCDatasetConnector],
+    ABC,
+    Generic[TCDatasetConnector],
+):
+    def __init__(self, config: TCDatasetConnector):
+        super().__init__(config)
+
+        _logger.info(
+            f"Initializing dataset connector | "
+            f"class={self.__class__.__name__} | "
+            f"module={self.__class__.__module__}"
+        )
+
+    @abstractmethod
+    def _load(self) -> Union[pl.DataFrame, pl.LazyFrame]:
+        """
+        Load raw data as Polars DataFrame or LazyFrame.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def to(self, ds: "Dataset[Any]") -> None:
+        raise NotImplementedError()
+
+    def load(self) -> "Dataset[Union[pl.DataFrame, pl.LazyFrame]]":
+        from agentic_base.dataset.polars import PolarsDataset
+
+        _logger.info(f"Loading dataset | connector={self.__class__.__name__}")
+
+        df = self._load()
+        self._log_polars_info(df)
+
+        return PolarsDataset.from_polars(df)
+
+    def load_text(self) -> "TextDataset[Union[pl.DataFrame, pl.LazyFrame]]":
+        from agentic_base.dataset.polars import PolarsTextDataset
+
+        _logger.info(f"Loading text dataset | connector={self.__class__.__name__}")
+
+        df = self._load()
+        self._log_polars_info(df)
+
+        return PolarsTextDataset.from_polars(df)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _log_polars_info(self, df: Union[pl.DataFrame, pl.LazyFrame]) -> None:
+        """
+        Log dataset structure without forcing materialization.
+        """
+        if isinstance(df, pl.DataFrame):
+            schema = df.schema
+            _logger.info(f"Loaded DataFrame | " f"columns={len(schema)} | " f"schema={list(schema.items())}")

agentic_base-0.1.0/agentic_base/connector/minio.py

@@ -0,0 +1,46 @@
+import io
+from typing import TYPE_CHECKING, Any, Union
+
+import polars as pl
+from minio import Minio
+
+from agentic_base.connector import DatasetConnector
+from agentic_base.connector.settings import MinioDatasetConnectorSettings
+
+if TYPE_CHECKING:
+    from agentic_base.dataset import Dataset
+
+
+class MinioDatasetConnector(DatasetConnector[MinioDatasetConnectorSettings]):
+
+    def __init__(self, config: MinioDatasetConnectorSettings):
+        super().__init__(config)
+        self.client = Minio(
+            self.config.endpoint.replace("http://", "").replace("https://", ""),
+            access_key=self.config.access_key.get_secret_value(),
+            secret_key=self.config.secret_key.get_secret_value(),
+            secure=self.config.endpoint.startswith("https://"),
+        )
+
+    def _load(self) -> Union[pl.DataFrame, pl.LazyFrame]:
+        response = self.client.get_object(self.config.bucket, self.config.key)
+        try:
+            buffer = io.BytesIO(response.read())
+        finally:
+            response.close()
+            response.release_conn()
+        df = pl.read_parquet(buffer)
+        return df
+
+    def to(self, ds: "Dataset[Any]") -> None:
+        df = ds.polars
+        buffer = io.BytesIO()
+        df.write_parquet(buffer)
+        buffer.seek(0)
+        self.client.put_object(
+            bucket_name=self.config.bucket,
+            object_name=self.config.key,
+            data=buffer,
+            length=buffer.getbuffer().nbytes,
+            content_type="application/octet-stream",
+        )

agentic_base-0.1.0/agentic_base/connector/parquet.py

@@ -0,0 +1,21 @@
+from typing import TYPE_CHECKING, Any, Union
+
+import polars as pl
+
+from agentic_base.connector import DatasetConnector
+from agentic_base.connector.settings import ParquetDatasetConnectorSettings
+
+if TYPE_CHECKING:
+    from agentic_base.dataset import Dataset
+
+
+class ParquetDatasetConnector(DatasetConnector[ParquetDatasetConnectorSettings]):
+
+    def __init__(self, config: ParquetDatasetConnectorSettings):
+        super().__init__(config)
+
+    def _load(self) -> Union[pl.DataFrame, pl.LazyFrame]:
+        return pl.scan_parquet(self.config.path) if self.config.lazy else pl.read_parquet(self.config.path)
+
+    def to(self, ds: "Dataset[Any]") -> None:
+        ds.polars.write_parquet(self.config.path)

agentic_base-0.1.0/agentic_base/connector/settings.py

@@ -0,0 +1,22 @@
+from pydantic import SecretStr
+from pydantic_settings import SettingsConfigDict
+
+from agentic_base.settings import FromConfigMixinSettings
+
+
+class DatasetConnectorSettings(FromConfigMixinSettings):
+    pass
+
+
+class ParquetDatasetConnectorSettings(DatasetConnectorSettings):
+    path: str
+    lazy: bool
+
+
+class MinioDatasetConnectorSettings(DatasetConnectorSettings):
+    endpoint: str
+    bucket: str
+    key: str
+    access_key: SecretStr
+    secret_key: SecretStr
+    model_config = SettingsConfigDict(env_prefix='MINIO_', extra="ignore")

agentic_base-0.1.0/agentic_base/constants.py

@@ -0,0 +1 @@
+CONFIG_PATH = "/config/config.yaml"

agentic_base-0.1.0/agentic_base/dataset/__init__.py

@@ -0,0 +1,155 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    Iterator,
+    Literal,
+    Optional,
+    Self,
+    Tuple,
+    Union,
+    overload,
+)
+
+import polars as pl
+from polars._typing import ColumnNameOrSelector, IntoExpr, IntoExprColumn
+
+from agentic_base.dataset.mixins import TextDatasetMixin
+from agentic_base.types import TDataset
+from agentic_base.utils import _get_minio_connector, _get_parquet_connector
+
+if TYPE_CHECKING:
+    pass
+
+_logger = logging.getLogger(__name__)
+
+
+class Dataset(ABC, Generic[TDataset]):
+    def __init__(self, service: TDataset):
+        super().__init__()
+        self.service = service
+        self._polars: Optional[pl.DataFrame] = None
+        self._lazy_polars: Optional[pl.LazyFrame] = None
+
+    @classmethod
+    @abstractmethod
+    def from_polars(cls, df: Union[pl.DataFrame, pl.LazyFrame]) -> Self:
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_polars(self) -> pl.DataFrame:
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_lazy_polars(self) -> pl.LazyFrame:
+        raise NotImplementedError
+
+    @classmethod
+    def from_parquet(
+        cls,
+        path: str,
+        *,
+        lazy: bool = True,
+    ) -> "Dataset[Union[pl.DataFrame, pl.LazyFrame]]":
+        return _get_parquet_connector(path, lazy=lazy).load()
+
+    def to_parquet(self, path: str, *, lazy: bool = True) -> None:
+        return _get_parquet_connector(path, lazy=lazy).to(self)
+
+    @classmethod
+    def from_minio(
+        cls,
+        bucket: str,
+        key: str,
+        endpoint: str,
+        access_key: str,
+        secret_key: str,
+    ) -> "Dataset[Union[pl.DataFrame, pl.LazyFrame]]":
+        return _get_minio_connector(bucket, key, endpoint, access_key, secret_key).load()
+
+    def to_minio(self, bucket: str, key: str, endpoint: str, access_key: str, secret_key: str) -> None:
+        return _get_minio_connector(bucket, key, endpoint, access_key, secret_key).to(self)
+
+    @property
+    def polars(self) -> pl.DataFrame:
+        if self._polars is None:
+            self._polars = self.to_polars()
+        return self._polars
+
+    @property
+    def lazy_polars(self) -> pl.LazyFrame:
+        if self._lazy_polars is None:
+            self._lazy_polars = self.to_lazy_polars()
+        return self._lazy_polars
+
+    @property
+    def polars_shape(self) -> Tuple[int, int]:
+        return self.polars.shape
+
+    def with_columns(
+        self,
+        *exprs: IntoExpr | Iterable[IntoExpr],
+        **named_exprs: IntoExpr,
+    ) -> Self:
+        df = self.polars.with_columns(*exprs, **named_exprs)
+        return self.__class__.from_polars(df)
+
+    def filter(
+        self,
+        *predicates: (IntoExprColumn | Iterable[IntoExprColumn] | bool | list[bool]),
+        **constraints: Any,
+    ) -> Self:
+        return self.__class__.from_polars(self.polars.filter(*predicates, **constraints))
+
+    def drop(
+        self,
+        *columns: ColumnNameOrSelector | Iterable[ColumnNameOrSelector],
+        strict: bool = True,
+    ) -> Self:
+        return self.__class__.from_polars(self.polars.drop(*columns, strict=strict))
+
+    def __len__(self) -> int:
+        return self.polars_shape[0]
+
+    @overload
+    def iter_rows(self, *, named: Literal[False] = ..., buffer_size: int = ...) -> Iterator[Tuple[Any, ...]]: ...
+
+    @overload
+    def iter_rows(self, *, named: Literal[True], buffer_size: int = ...) -> Iterator[Dict[str, Any]]: ...
+
+    def iter_rows(
+        self, *, named: Literal[False, True] = False, buffer_size: int = 512
+    ) -> Iterator[Tuple[Any, ...]] | Iterator[Dict[str, Any]]:
+        df_stream = self.lazy_polars.collect(streaming=True)  # type: ignore[call-overload]
+        return df_stream.iter_rows(named=named, buffer_size=buffer_size)  # type: ignore[no-any-return]
+
+
+class TextDataset(Dataset[TDataset], TextDatasetMixin[TDataset], Generic[TDataset]):
+
+    def __init__(self, service: TDataset):
+        super().__init__(service)
+        self._validate_schema()
+
+    @classmethod
+    def from_parquet(
+        cls,
+        path: str,
+        *,
+        lazy: bool = True,
+    ) -> "TextDataset[Union[pl.DataFrame, pl.LazyFrame]]":
+        return _get_parquet_connector(path, lazy=lazy).load_text()
+
+    @classmethod
+    def from_minio(
+        cls,
+        bucket: str,
+        key: str,
+        endpoint: str,
+        access_key: str,
+        secret_key: str,
+    ) -> "TextDataset[Union[pl.DataFrame, pl.LazyFrame]]":
+        return _get_minio_connector(bucket, key, endpoint, access_key, secret_key).load_text()

agentic_base-0.1.0/agentic_base/dataset/mixins.py

@@ -0,0 +1,114 @@
+import json
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Protocol,
+    Tuple,
+    Union,
+    runtime_checkable,
+)
+
+import polars as pl
+from langchain_core.documents import Document
+
+from agentic_base.types import TDataset
+
+if TYPE_CHECKING:
+
+    from agentic_base.dataset import TextDataset
+
+
+@runtime_checkable
+class SupportsPolarsDataset(Protocol[TDataset]):
+    @property
+    def polars(self) -> pl.DataFrame: ...
+
+    def to_polars(self) -> pl.DataFrame: ...
+
+    @classmethod
+    def from_polars(cls, df: Union[pl.DataFrame, pl.LazyFrame]) -> "TextDataset[TDataset]":
+        raise NotImplementedError
+
+
+class TextDatasetMixin(SupportsPolarsDataset[TDataset], Generic[TDataset]):
+    REQUIRED_COLUMNS = {"page_content", "metadata"}
+
+    def _validate_schema(self) -> None:
+        df = self.polars
+        missing = self.REQUIRED_COLUMNS - set(df.columns)
+        if missing:
+            raise ValueError(
+                f"{self.__class__.__name__} requires columns {self.REQUIRED_COLUMNS}, " f"but missing {missing}"
+            )
+
+    def iter_documents(self) -> Iterator[Document]:
+        df = self.to_polars()
+        for row in df.iter_rows(named=True):
+            yield Document(
+                page_content=row["page_content"],
+                metadata=row["metadata"],
+            )
+
+    @classmethod
+    def from_records(
+        cls,
+        records: List[Tuple[str, Dict[str, Any]]],
+    ) -> "TextDataset[TDataset]":
+        if not records:
+            raise ValueError("records must be non-empty")
+        df = pl.DataFrame(
+            {
+                "page_content": [text for text, _ in records],
+                "metadata": [meta for _, meta in records],
+            }
+        )
+        return cls.from_polars(df)
+
+    def dump_documents(
+        self,
+        out_dir: str,
+        prefix: str = "doc",
+        ext: str = ".md",
+        encoding: str = "utf-8",
+    ) -> None:
+        out_dir_path = Path(out_dir)
+        out_dir_path.mkdir(parents=True, exist_ok=True)
+
+        for i, row in enumerate(self.polars.iter_rows(named=True)):
+            path = out_dir_path / f"{prefix}_{i:05d}{ext}"
+            page_content = row.get("page_content", "")
+            metadata = row.get("metadata", {})
+            try:
+                metadata_str = json.dumps(metadata, indent=2, ensure_ascii=False)
+            except TypeError:
+                metadata_str = str(metadata)
+            content = "###PAGE_CONTENT###\n" f"{page_content}\n\n" "###METADATA###\n" f"{metadata_str}\n"
+            path.write_text(content, encoding=encoding)
+
+    def to_langchain_documents(
+        self,
+    ) -> List[Document]:
+        docs: list[Document] = []
+        for row in self.polars.iter_rows(named=True):
+            docs.append(
+                Document(
+                    page_content=row["page_content"],
+                    metadata=row["metadata"],
+                )
+            )
+        return docs
+
+    @classmethod
+    def from_documents(cls, docs: List[Document]) -> "TextDataset[TDataset]":
+        df = pl.DataFrame(
+            {
+                "page_content": [d.page_content for d in docs],
+                "metadata": [d.metadata or {} for d in docs],
+            }
+        )
+        return cls.from_polars(df)

agentic_base-0.1.0/agentic_base/dataset/polars.py

@@ -0,0 +1,43 @@
+from typing import Self, Union
+
+import polars as pl
+
+from agentic_base.dataset import Dataset, TextDataset
+
+
+class PolarsDataset(Dataset[Union[pl.DataFrame, pl.LazyFrame]]):
+    def __init__(self, service: Union[pl.DataFrame, pl.LazyFrame]):
+        super().__init__(service)
+
+    @classmethod
+    def from_polars(cls, df: Union[pl.DataFrame, pl.LazyFrame]) -> Self:
+        return cls(df)
+
+    def to_polars(self) -> pl.DataFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service.collect()
+        return self.service
+
+    def to_lazy_polars(self) -> pl.LazyFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service
+        return self.service.lazy()
+
+
+class PolarsTextDataset(TextDataset[Union[pl.DataFrame, pl.LazyFrame]]):
+    def __init__(self, service: Union[pl.DataFrame, pl.LazyFrame]):
+        super().__init__(service)
+
+    @classmethod
+    def from_polars(cls, df: Union[pl.DataFrame, pl.LazyFrame]) -> Self:
+        return cls(df)
+
+    def to_polars(self) -> pl.DataFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service.collect()
+        return self.service
+
+    def to_lazy_polars(self) -> pl.LazyFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service
+        return self.service.lazy()

agentic_base-0.1.0/agentic_base/dataset/preprocess/__init__.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from typing import Any, Generic
+
+from agentic_base.dataset.polars import TextDataset
+from agentic_base.mixins import FromConfigMixin
+from agentic_base.types import TCTextPreprocessor, TCTokenCounter
+
+
+class TextPreprocessor(FromConfigMixin[TCTextPreprocessor], ABC, Generic[TCTextPreprocessor]):
+
+    def __init__(self, config: TCTextPreprocessor) -> None:
+        self.config = config
+
+    @abstractmethod
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        raise NotImplementedError
+
+
+class TokenCounter(FromConfigMixin[TCTokenCounter], ABC, Generic[TCTokenCounter]):
+    def __init__(self, config: TCTokenCounter) -> None:
+        self.config = config
+
+    @abstractmethod
+    def count(self, text: str) -> int:
+        raise NotImplementedError()
+
+    def count_batch(self, texts: list[str]) -> list[int]:
+        return [self.count(t) for t in texts]

agentic_base-0.1.0/agentic_base/dataset/preprocess/preprocess.py

@@ -0,0 +1,98 @@
+from typing import Any, List
+
+import polars as pl
+
+from agentic_base.dataset import TextDataset
+from agentic_base.dataset.preprocess import TextPreprocessor, TokenCounter
+from agentic_base.dataset.settings import (
+    MaxTokenFilterSettings,
+    MinTokenFilterSettings,
+    PreprocessPipelineSettings,
+    QuantileTokenFilterSettings,
+    SigmaBandTokenFilterSettings,
+)
+from agentic_base.utils import load_class
+
+
+class MinTokenFilter(TextPreprocessor[MinTokenFilterSettings[Any]]):
+
+    def __init__(self, config: MinTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        return ds.__class__.from_polars(
+            ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+            .filter(pl.col("n_tokens") >= self.config.min_tokens)
+            .drop("n_tokens")
+        )
+
+
+class MaxTokenFilter(TextPreprocessor[MaxTokenFilterSettings[Any]]):
+    def __init__(self, config: MaxTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        return ds.__class__.from_polars(
+            ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+            .filter(pl.col("n_tokens") <= self.config.max_tokens)
+            .drop("n_tokens")
+        )
+
+
+class QuantileTokenFilter(TextPreprocessor[QuantileTokenFilterSettings[Any]]):
+    def __init__(self, config: QuantileTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        df = ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+        cutoff = df.select(pl.col("n_tokens").quantile(self.config.q)).item()
+        return ds.__class__.from_polars(df.filter(pl.col("n_tokens") <= cutoff).drop("n_tokens"))
+
+
+class SigmaBandTokenFilter(TextPreprocessor[SigmaBandTokenFilterSettings[Any]]):
+    def __init__(self, config: SigmaBandTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        df = ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+        mu, sigma = df.select(
+            pl.col("n_tokens").mean(),
+            pl.col("n_tokens").std(),
+        ).row(0)
+        return ds.__class__.from_polars(
+            df.filter(pl.col("n_tokens").is_between(mu - self.config.z * sigma, mu + self.config.z * sigma)).drop(
+                "n_tokens"
+            )
+        )
+
+
+class PreprocessPipeline(TextPreprocessor[PreprocessPipelineSettings]):
+
+    def __init__(self, config: PreprocessPipelineSettings):
+        super().__init__(config)
+        self.steps = self._load_steps()
+
+    def _load_steps(self) -> List[TextPreprocessor[Any]]:
+        steps = []
+        for step_config in self.config.steps:
+            step: TextPreprocessor[Any] = load_class(step_config.module_path)(step_config)
+            steps.append(step)
+        return steps
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        current = ds
+        for step in self.steps:
+            current = step.apply(current)
+        return current

agentic_base-0.1.0/agentic_base/dataset/preprocess/token_counter.py

@@ -0,0 +1,37 @@
+from transformers import AutoTokenizer
+
+from agentic_base.dataset.preprocess import TokenCounter
+from agentic_base.dataset.settings import (
+    HeuristicTokenCounterSettings,
+    HuggingFaceTokenCounterSettings,
+)
+
+
+class HeuristicTokenCounter(TokenCounter[HeuristicTokenCounterSettings]):
+
+    def __init__(self, config: HeuristicTokenCounterSettings) -> None:
+        super().__init__(config)
+
+    def count(self, text: str) -> int:
+        return max(1, int(len(text) / self.config.chars_per_token))
+
+
+class HuggingFaceTokenCounter(TokenCounter[HuggingFaceTokenCounterSettings]):
+    def __init__(self, config: HuggingFaceTokenCounterSettings) -> None:
+        super().__init__(config)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config.name,
+            revision=self.config.revision,
+            trust_remote_code=True,
+        )
+
+    def count(self, text: str) -> int:
+        return len(
+            self.tokenizer(
+                text,
+                add_special_tokens=self.config.add_special_tokens,
+                truncation=False,
+                return_attention_mask=False,
+                return_token_type_ids=False,
+            )["input_ids"]
+        )