API reference¶
Generated from the apogee-ai-data source with mkdocstrings. Every symbol below is exported from apogee_ai_data, so it is part of the supported public surface.
Application · DTOs¶
BenchDTO
dataclass
¶
ChunkDTO
dataclass
¶
ChunkDTO(source: str, loader: str = 'text', chunker: str = 'fixed', size: int = 512, overlap: int = 64)
IndexDTO
dataclass
¶
IndexDTO(source: str, loader: str = 'text', chunker: str = 'fixed', size: int = 512, overlap: int = 64, embedder: str = 'hashing', dim: int = 128, store_path: str | None = None)
LoadDTO
dataclass
¶
RetrieveDTO
dataclass
¶
RetrieveDTO(source: str, query: str, k: int = 5, retriever: str = 'dense', loader: str = 'text', chunker: str = 'fixed', size: int = 512, overlap: int = 64, embedder: str = 'hashing', dim: int = 128, filters: dict[str, str] = dict())
filters
class-attribute
instance-attribute
¶
Application · Use cases¶
BenchRetrievalUseCase
¶
Sintetiza um corpus + queries e mede recall@k para dense vs bm25 vs hybrid.
execute
async
¶
execute(dto: BenchDTO) -> dict[str, float]
Source code in apogee_ai_data/application/use_cases/bench_retrieval_use_case.py
async def execute(self, dto: BenchDTO) -> dict[str, float]:
rng = random.Random(42)
documents: list[Document] = []
gold: list[tuple[int, str]] = []
for i in range(dto.documents):
topic = _TOPICS[i % len(_TOPICS)]
noise = rng.choice(_TOPICS)
text = (
f"Documento {i}: {topic}. Detalhes adicionais sobre {noise}. "
+ ("lorem ipsum " * 30)
)
doc_id = f"doc-{i:03d}"
documents.append(Document(id=doc_id, text=text))
gold.append((i, doc_id))
chunker = FixedSizeChunker(size=dto.chunk_size, overlap=32)
chunks = chunker.split(documents)
chunk_to_doc = {c.id: c.document_id for c in chunks}
embedder = HashingTextEmbedder(dim=128)
store = InMemoryVectorStore(dim=128)
embeddings = await embedder.embed_many([c.text for c in chunks])
await store.upsert_many(
(c.id, e.vector, {"document_id": c.document_id})
for c, e in zip(chunks, embeddings)
)
bm25 = BM25Retriever()
bm25.index(chunks)
dense = DenseRetriever(store, embedder)
hybrid = HybridRetriever(dense, bm25)
recall = {"dense": 0, "bm25": 0, "hybrid": 0}
latency = {"dense": 0.0, "bm25": 0.0, "hybrid": 0.0}
for q_idx in range(dto.queries):
doc_idx, expected_id = gold[q_idx % len(gold)]
topic = _TOPICS[doc_idx % len(_TOPICS)]
words = topic.split()
rng.shuffle(words)
q = RetrievalQuery(text=" ".join(words[: max(2, len(words) // 2)]), k=5)
for name, retriever in (
("dense", dense),
("bm25", bm25),
("hybrid", hybrid),
):
start = time.perf_counter()
result = await retriever.retrieve(q)
latency[name] += (time.perf_counter() - start) * 1000.0
hit_docs = {chunk_to_doc.get(h.chunk_id, "") for h in result.hits}
if expected_id in hit_docs:
recall[name] += 1
n = max(dto.queries, 1)
return {
"queries": float(dto.queries),
"documents": float(dto.documents),
"recall_dense": recall["dense"] / n,
"recall_bm25": recall["bm25"] / n,
"recall_hybrid": recall["hybrid"] / n,
"avg_latency_ms_dense": latency["dense"] / n,
"avg_latency_ms_bm25": latency["bm25"] / n,
"avg_latency_ms_hybrid": latency["hybrid"] / n,
}
ChunkDocumentsUseCase
¶
execute
async
¶
Source code in apogee_ai_data/application/use_cases/chunk_documents_use_case.py
IndexDocumentsUseCase
¶
execute
async
¶
execute(dto: IndexDTO)
Source code in apogee_ai_data/application/use_cases/index_documents_use_case.py
async def execute(self, dto: IndexDTO):
loader = build_loader(dto.loader)
chunker = build_chunker(dto.chunker, size=dto.size, overlap=dto.overlap)
embedder = build_embedder(dto.embedder, dim=dto.dim)
documents = list(await loader.load(dto.source))
chunks = chunker.split(documents)
if not chunks:
return InMemoryVectorStore(dim=embedder.dim)
store: object
if dto.store_path:
store = JsonVectorStore(dto.store_path, dim=embedder.dim)
else:
store = InMemoryVectorStore(dim=embedder.dim)
embeddings = await embedder.embed_many([c.text for c in chunks])
await store.upsert_many( # type: ignore[attr-defined]
(
chunk.id,
emb.vector,
{"text": chunk.text, **chunk.metadata},
)
for chunk, emb in zip(chunks, embeddings)
)
return store
LoadDocumentsUseCase
¶
RetrieveUseCase
¶
Loads → chunks → embeds → indexes → retrieves in one shot.
Convenient for CLI / one-off retrieval; not optimal for repeated queries (where you'd index once and reuse the store).
execute
async
¶
execute(dto: RetrieveDTO) -> RetrievalResult
Source code in apogee_ai_data/application/use_cases/retrieve_use_case.py
async def execute(self, dto: RetrieveDTO) -> RetrievalResult:
loader = build_loader(dto.loader)
chunker = build_chunker(dto.chunker, size=dto.size, overlap=dto.overlap)
embedder = build_embedder(dto.embedder, dim=dto.dim)
documents = list(await loader.load(dto.source))
chunks = chunker.split(documents)
if not chunks:
return RetrievalResult(query=dto.query, hits=(), retriever=dto.retriever)
query = RetrievalQuery(text=dto.query, k=dto.k, filters=dto.filters)
if dto.retriever == "bm25":
sparse = BM25Retriever()
sparse.index(chunks)
return await sparse.retrieve(query)
store = InMemoryVectorStore(dim=embedder.dim)
embeddings = await embedder.embed_many([c.text for c in chunks])
await store.upsert_many(
(chunk.id, emb.vector, {"text": chunk.text, **chunk.metadata})
for chunk, emb in zip(chunks, embeddings)
)
if dto.retriever == "dense":
return await DenseRetriever(store, embedder).retrieve(query)
if dto.retriever == "hybrid":
sparse = BM25Retriever()
sparse.index(chunks)
dense = DenseRetriever(store, embedder)
return await HybridRetriever(dense, sparse).retrieve(query)
raise ValueError(f"unknown retriever: {dto.retriever}")
Domain¶
Chunk
dataclass
¶
ChunkPolicy
dataclass
¶
ChunkPolicy(kind: ChunkerKind = FIXED, size: int = 512, overlap: int = 64, separators: tuple[str, ...] = ('\n\n', '\n', '. ', ' '))
separators
class-attribute
instance-attribute
¶
Dataset
dataclass
¶
Dataset(name: str, documents: tuple[Document, ...] = (), metadata: dict[str, str] = dict())
Document
dataclass
¶
Document(id: str, text: str, kind: DocumentKind = TEXT, source: str | None = None, metadata: dict[str, str] = dict())
metadata
class-attribute
instance-attribute
¶
RetrievalHit
dataclass
¶
RetrievalQuery
dataclass
¶
RetrievalResult
dataclass
¶
RetrievalResult(query: str, hits: tuple[RetrievalHit, ...] = (), retriever: str = 'dense', latency_ms: float = 0.0, metadata: dict[str, str] = dict())
metadata
class-attribute
instance-attribute
¶
Domain · Enums¶
ChunkerKind
¶
DocumentKind
¶
Bases: str, Enum
EmbedderKind
¶
RetrieverKind
¶
VectorStoreKind
¶
Bases: str, Enum
Domain · Exceptions¶
DataError
¶
Bases: Exception
Base for apogee-ai-data errors.
LoaderError
¶
Domain · Protocols (ports)¶
IChunker
¶
ILoader
¶
IRetriever
¶
Bases: Protocol
retrieve
async
¶
retrieve(query: RetrievalQuery) -> RetrievalResult
IVectorStore
¶
Bases: Protocol
upsert
async
¶
upsert_many
async
¶
query
async
¶
query(vector: tuple[float, ...] | list[float], k: int = 5, filters: dict[str, str] | None = None) -> list[RetrievalHit]
delete
async
¶
count
async
¶
Infrastructure¶
BM25Retriever
¶
Self-contained BM25 over an in-memory corpus of chunks.
Source code in apogee_ai_data/infrastructure/retrievers/bm25_retriever.py
index
¶
index(chunks: Iterable[Chunk]) -> None
Source code in apogee_ai_data/infrastructure/retrievers/bm25_retriever.py
def index(self, chunks: Iterable[Chunk]) -> None:
self._chunks = list(chunks)
self._tokens = [_tokenize(c.text) for c in self._chunks]
self._df = Counter()
for tokens in self._tokens:
for term in set(tokens):
self._df[term] += 1
self._n = len(self._chunks)
self._avgdl = (
sum(len(t) for t in self._tokens) / self._n if self._n else 0.0
)
retrieve
async
¶
retrieve(query: RetrievalQuery) -> RetrievalResult
Source code in apogee_ai_data/infrastructure/retrievers/bm25_retriever.py
async def retrieve(self, query: RetrievalQuery) -> RetrievalResult:
if not self._n:
return RetrievalResult(query=query.text, hits=(), retriever=self.name)
start = time.perf_counter()
q_terms = _tokenize(query.text)
scored: list[RetrievalHit] = []
for chunk, tokens in zip(self._chunks, self._tokens):
if query.filters and not all(
chunk.metadata.get(key) == val
for key, val in query.filters.items()
):
continue
tf = Counter(tokens)
dl = len(tokens) or 1
score = 0.0
for term in q_terms:
if term not in tf:
continue
idf = self._idf(term)
f = tf[term]
num = f * (self._k1 + 1)
den = f + self._k1 * (
1 - self._b + self._b * dl / (self._avgdl or 1)
)
score += idf * num / den
if score > 0:
scored.append(
RetrievalHit(
chunk_id=chunk.id,
score=score,
metadata={"text": chunk.text, **chunk.metadata},
)
)
scored.sort(key=lambda h: h.score, reverse=True)
elapsed = (time.perf_counter() - start) * 1000.0
return RetrievalResult(
query=query.text,
hits=tuple(scored[: query.k]),
retriever=self.name,
latency_ms=elapsed,
)
CodeLoader
¶
Walks a directory and yields one Document per source file.
Source code in apogee_ai_data/infrastructure/loaders/code_loader.py
load
async
¶
load(source: str) -> Iterable[Document]
Source code in apogee_ai_data/infrastructure/loaders/code_loader.py
async def load(self, source: str) -> Iterable[Document]:
root = Path(source)
if not root.exists():
raise LoaderError(source, "path does not exist")
if root.is_file():
paths = [root]
else:
paths = sorted(p for p in root.rglob("*") if p.is_file())
documents: list[Document] = []
for path in paths:
if path.suffix.lower() not in self._exts:
continue
try:
text = path.read_text(encoding="utf-8")
except (OSError, UnicodeDecodeError):
continue
documents.append(
Document(
id=str(path.relative_to(root) if root.is_dir() else path.name),
text=text,
kind=DocumentKind.CODE,
source=str(path),
metadata={"language": path.suffix.lstrip(".")},
)
)
return documents
DenseRetriever
¶
Source code in apogee_ai_data/infrastructure/retrievers/dense_retriever.py
retrieve
async
¶
retrieve(query: RetrievalQuery) -> RetrievalResult
Source code in apogee_ai_data/infrastructure/retrievers/dense_retriever.py
async def retrieve(self, query: RetrievalQuery) -> RetrievalResult:
start = time.perf_counter()
embedding = await self._embedder.embed(query.text)
hits = await self._store.query(
embedding.vector, k=query.k, filters=query.filters or None
)
elapsed = (time.perf_counter() - start) * 1000.0
return RetrievalResult(
query=query.text,
hits=tuple(hits),
retriever=self.name,
latency_ms=elapsed,
)
FaissVectorStore
¶
Lazy FAISS adapter — install via extras=faiss.
Uses IndexFlatIP (assumes vectors are L2-normalised → cosine).
Source code in apogee_ai_data/infrastructure/vector_stores/faiss_vector_store.py
upsert
async
¶
Source code in apogee_ai_data/infrastructure/vector_stores/faiss_vector_store.py
async def upsert(self, chunk_id: str, vector, metadata=None) -> None:
import numpy as np # type: ignore
self._ensure_index()
vec = np.asarray([list(vector)], dtype="float32")
if chunk_id in self._ids:
await self.delete(chunk_id)
self._index.add(vec) # type: ignore[union-attr]
self._ids.append(chunk_id)
self._metadata.append(dict(metadata or {}))
upsert_many
async
¶
query
async
¶
query(vector, k: int = 5, filters=None) -> list[RetrievalHit]
Source code in apogee_ai_data/infrastructure/vector_stores/faiss_vector_store.py
async def query(self, vector, k: int = 5, filters=None) -> list[RetrievalHit]:
import numpy as np # type: ignore
self._ensure_index()
if not self._ids:
return []
vec = np.asarray([list(vector)], dtype="float32")
scores, idxs = self._index.search(vec, k=min(k, len(self._ids))) # type: ignore[union-attr]
out: list[RetrievalHit] = []
for score, idx in zip(scores[0], idxs[0]):
if idx < 0:
continue
md = self._metadata[idx]
if filters and not all(md.get(key) == val for key, val in filters.items()):
continue
out.append(
RetrievalHit(
chunk_id=self._ids[idx],
score=float(score),
metadata=dict(md),
)
)
return out
delete
async
¶
Source code in apogee_ai_data/infrastructure/vector_stores/faiss_vector_store.py
count
async
¶
FixedSizeChunker
¶
Chunks by character count with optional overlap.
Source code in apogee_ai_data/infrastructure/chunkers/fixed_size_chunker.py
split
¶
Source code in apogee_ai_data/infrastructure/chunkers/fixed_size_chunker.py
def split(self, documents: Iterable[Document]) -> list[Chunk]:
out: list[Chunk] = []
for doc in documents:
text = doc.text
stride = self._size - self._overlap
if not text:
continue
for index, start in enumerate(range(0, len(text), stride)):
piece = text[start : start + self._size]
if not piece:
break
out.append(
Chunk(
id=f"{doc.id}::{index}",
document_id=doc.id,
text=piece,
index=index,
metadata=dict(doc.metadata),
)
)
if start + self._size >= len(text):
break
return out
HashingTextEmbedder
¶
HtmlLoader
¶
Strips HTML to plain text. Uses bs4 if available, else regex fallback.
load
async
¶
load(source: str) -> Iterable[Document]
Source code in apogee_ai_data/infrastructure/loaders/html_loader.py
async def load(self, source: str) -> Iterable[Document]:
path = Path(source)
if not path.is_file():
raise LoaderError(source, "file not found")
raw = path.read_text(encoding="utf-8", errors="ignore")
try:
from bs4 import BeautifulSoup # type: ignore
text = BeautifulSoup(raw, "html.parser").get_text(separator=" ")
except ImportError:
text = _TAG_RE.sub(" ", raw)
text = _WS_RE.sub(" ", text).strip()
if not text:
raise LoaderError(source, "no text extracted")
return [
Document(
id=path.name,
text=text,
kind=DocumentKind.HTML,
source=str(path),
)
]
HybridRetriever
¶
Combines dense + BM25 results via Reciprocal Rank Fusion.
Source code in apogee_ai_data/infrastructure/retrievers/hybrid_retriever.py
retrieve
async
¶
retrieve(query: RetrievalQuery) -> RetrievalResult
Source code in apogee_ai_data/infrastructure/retrievers/hybrid_retriever.py
async def retrieve(self, query: RetrievalQuery) -> RetrievalResult:
start = time.perf_counter()
dense = await self._dense.retrieve(query)
sparse = await self._sparse.retrieve(query)
scores: dict[str, float] = {}
metadata: dict[str, dict[str, str]] = {}
for rank, hit in enumerate(dense.hits, 1):
scores[hit.chunk_id] = scores.get(hit.chunk_id, 0.0) + 1.0 / (
self._k + rank
)
metadata.setdefault(hit.chunk_id, dict(hit.metadata))
for rank, hit in enumerate(sparse.hits, 1):
scores[hit.chunk_id] = scores.get(hit.chunk_id, 0.0) + 1.0 / (
self._k + rank
)
metadata.setdefault(hit.chunk_id, dict(hit.metadata))
merged = sorted(
(
RetrievalHit(
chunk_id=cid,
score=score,
metadata=metadata.get(cid, {}),
)
for cid, score in scores.items()
),
key=lambda h: h.score,
reverse=True,
)
elapsed = (time.perf_counter() - start) * 1000.0
return RetrievalResult(
query=query.text,
hits=tuple(merged[: query.k]),
retriever=self.name,
latency_ms=elapsed,
metadata={"dense": str(len(dense.hits)), "sparse": str(len(sparse.hits))},
)
InMemoryVectorStore
¶
Source code in apogee_ai_data/infrastructure/vector_stores/in_memory_vector_store.py
upsert
async
¶
Source code in apogee_ai_data/infrastructure/vector_stores/in_memory_vector_store.py
async def upsert(
self,
chunk_id: str,
vector,
metadata: dict[str, str] | None = None,
) -> None:
vec = tuple(float(v) for v in vector)
if len(vec) != self.dim:
raise VectorStoreError(
f"expected dim={self.dim}, got {len(vec)}"
)
self._entries[chunk_id] = _Entry(
chunk_id=chunk_id, vector=vec, metadata=dict(metadata or {})
)
upsert_many
async
¶
query
async
¶
query(vector, k: int = 5, filters: dict[str, str] | None = None) -> list[RetrievalHit]
Source code in apogee_ai_data/infrastructure/vector_stores/in_memory_vector_store.py
async def query(
self,
vector,
k: int = 5,
filters: dict[str, str] | None = None,
) -> list[RetrievalHit]:
if k <= 0:
raise VectorStoreError("k must be positive")
vec = tuple(float(v) for v in vector)
if len(vec) != self.dim:
raise VectorStoreError(f"expected dim={self.dim}, got {len(vec)}")
scored: list[RetrievalHit] = []
for entry in self._entries.values():
if filters and not all(
entry.metadata.get(key) == val for key, val in filters.items()
):
continue
score = cosine(vec, entry.vector)
scored.append(
RetrievalHit(
chunk_id=entry.chunk_id,
score=score,
metadata=dict(entry.metadata),
)
)
scored.sort(key=lambda h: h.score, reverse=True)
return scored[:k]
delete
async
¶
count
async
¶
JsonVectorStore
¶
JsonlLoader
¶
Source code in apogee_ai_data/infrastructure/loaders/jsonl_loader.py
load
async
¶
load(source: str) -> Iterable[Document]
Source code in apogee_ai_data/infrastructure/loaders/jsonl_loader.py
async def load(self, source: str) -> Iterable[Document]:
path = Path(source)
if not path.is_file():
raise LoaderError(source, "file not found")
documents: list[Document] = []
with path.open("r", encoding="utf-8") as fh:
for lineno, line in enumerate(fh, 1):
line = line.strip()
if not line:
continue
try:
payload = json.loads(line)
except json.JSONDecodeError as exc:
raise LoaderError(
source, f"line {lineno}: {exc.msg}"
) from exc
if self._text_field not in payload:
raise LoaderError(
source, f"line {lineno} missing field {self._text_field!r}"
)
doc_id = str(payload.get(self._id_field, f"{path.name}:{lineno}"))
documents.append(
Document(
id=doc_id,
text=str(payload[self._text_field]),
kind=DocumentKind.JSONL,
source=str(path),
metadata={
k: str(v)
for k, v in payload.items()
if k not in (self._text_field, self._id_field)
},
)
)
return documents
MarkdownLoader
¶
PdfLoader
¶
Lazy adapter for pypdf — install via extras=pdf.
load
async
¶
load(source: str) -> Iterable[Document]
Source code in apogee_ai_data/infrastructure/loaders/pdf_loader.py
async def load(self, source: str) -> Iterable[Document]:
try:
from pypdf import PdfReader # type: ignore
except ImportError as exc: # pragma: no cover
raise LoaderError(
source, "install with `pip install apogee-ai-data[pdf]`"
) from exc
path = Path(source)
if not path.is_file():
raise LoaderError(source, "file not found")
reader = PdfReader(str(path))
text_parts = [page.extract_text() or "" for page in reader.pages]
text = "\n\n".join(t for t in text_parts if t.strip())
if not text:
raise LoaderError(source, "no text extracted")
return [
Document(
id=path.name,
text=text,
kind=DocumentKind.PDF,
source=str(path),
metadata={"pages": str(len(reader.pages))},
)
]
RecursiveChunker
¶
RecursiveChunker(size: int = 512, overlap: int = 64, separators: tuple[str, ...] = ('\n\n', '\n', '. ', ' '))
Splits along separators (paragraph→line→sentence→word) up to size.
Source code in apogee_ai_data/infrastructure/chunkers/recursive_chunker.py
def __init__(
self,
size: int = 512,
overlap: int = 64,
separators: tuple[str, ...] = ("\n\n", "\n", ". ", " "),
) -> None:
if size <= 0:
raise ChunkerError("size must be positive")
if overlap < 0 or overlap >= size:
raise ChunkerError("overlap must be in [0, size)")
if not separators:
raise ChunkerError("at least one separator is required")
self._size = size
self._overlap = overlap
self._separators = separators
split
¶
Source code in apogee_ai_data/infrastructure/chunkers/recursive_chunker.py
def split(self, documents: Iterable[Document]) -> list[Chunk]:
out: list[Chunk] = []
for doc in documents:
pieces = self._split_recursive(doc.text)
for index, piece in enumerate(pieces):
if not piece:
continue
out.append(
Chunk(
id=f"{doc.id}::{index}",
document_id=doc.id,
text=piece,
index=index,
metadata=dict(doc.metadata),
)
)
return out
SentenceTransformersEmbedder
¶
Lazy adapter for sentence-transformers — install via extras.
Source code in apogee_ai_data/infrastructure/embedders/sentence_transformers_embedder.py
embed
async
¶
embed(text: str) -> Embedding
Source code in apogee_ai_data/infrastructure/embedders/sentence_transformers_embedder.py
embed_many
async
¶
embed_many(texts: Iterable[str]) -> list[Embedding]
Source code in apogee_ai_data/infrastructure/embedders/sentence_transformers_embedder.py
async def embed_many(self, texts: Iterable[str]) -> list[Embedding]:
self._ensure_model()
items = list(texts)
vecs = self._model.encode(items, normalize_embeddings=True) # type: ignore[union-attr]
return [
Embedding(vector=tuple(float(v) for v in row), model=self._model_name)
for row in vecs
]
TextLoader
¶
Source code in apogee_ai_data/infrastructure/loaders/text_loader.py
load
async
¶
load(source: str) -> Iterable[Document]
Source code in apogee_ai_data/infrastructure/loaders/text_loader.py
async def load(self, source: str) -> Iterable[Document]:
path = Path(source)
if not path.is_file():
raise LoaderError(source, "file not found")
try:
text = path.read_text(encoding=self._encoding)
except OSError as exc:
raise LoaderError(source, str(exc)) from exc
return [
Document(
id=path.name,
text=text,
kind=DocumentKind.TEXT,
source=str(path),
)
]
TokenChunker
¶
Approximates tokens via whitespace splits — no tiktoken dependency.
Source code in apogee_ai_data/infrastructure/chunkers/token_chunker.py
split
¶
Source code in apogee_ai_data/infrastructure/chunkers/token_chunker.py
def split(self, documents: Iterable[Document]) -> list[Chunk]:
out: list[Chunk] = []
stride = self._tokens - self._overlap
for doc in documents:
words = doc.text.split()
if not words:
continue
for index, start in enumerate(range(0, len(words), stride)):
window = words[start : start + self._tokens]
if not window:
break
out.append(
Chunk(
id=f"{doc.id}::{index}",
document_id=doc.id,
text=" ".join(window),
index=index,
metadata=dict(doc.metadata),
)
)
if start + self._tokens >= len(words):
break
return out