Parent-child chunks: мелкие чанки для поиска, крупные — для контекста

from __future__ import annotations

import uuid
from dataclasses import dataclass, field
from typing import Protocol


@dataclass
class Document:
    page_content: str
    metadata: dict = field(default_factory=dict)


# ─────────────────────────────────────────────────────────────
# Протокол сплиттера — любой класс с методом split_documents
# ─────────────────────────────────────────────────────────────

class Splitter(Protocol):
    def split_documents(self, docs: list[Document]) -> list[Document]: ...


# ─────────────────────────────────────────────────────────────
# ParentChildSplitter
# ─────────────────────────────────────────────────────────────

class ParentChildSplitter:
    """
    Создаёт двухуровневую структуру чанков.

    Args:
        parent_splitter: сплиттер для создания родительских чанков.
            Определяет размер единицы контекста для LLM.
        child_splitter: сплиттер для нарезки детей из каждого родителя.
            Определяет размер единицы поиска в векторной БД.

    Returns (split_documents):
        child_docs: список дочерних чанков с parent_id в metadata.
            → индексировать в векторную БД.
        parent_store: dict {parent_id: Document}.
            → хранить в docstore для lookup при retrieval.
    """

    def __init__(
        self,
        parent_splitter: Splitter,
        child_splitter: Splitter,
    ) -> None:
        self.parent_splitter = parent_splitter
        self.child_splitter = child_splitter

    def split_documents(
        self,
        documents: list[Document],
    ) -> tuple[list[Document], dict[str, Document]]:
        """
        Возвращает (child_docs, parent_store).
        child_docs  — загружать в vectorstore.
        parent_store — загружать в docstore.
        """
        child_docs: list[Document] = []
        parent_store: dict[str, Document] = {}

        for doc in documents:
            # Шаг 1: нарезаем документ на родительские чанки
            parents = self.parent_splitter.split_documents([doc])

            for parent in parents:
                # Шаг 2: присваиваем родителю уникальный ID
                parent_id = str(uuid.uuid4())
                parent.metadata["doc_id"] = parent_id
                parent_store[parent_id] = parent

                # Шаг 3: нарезаем родителя на дочерние чанки
                children = self.child_splitter.split_documents([parent])

                for child in children:
                    # Шаг 4: каждый ребёнок ссылается на родителя
                    child.metadata["parent_id"] = parent_id
                    # Сохраняем source из исходного документа
                    if "source" in doc.metadata:
                        child.metadata["source"] = doc.metadata["source"]
                    child_docs.append(child)

        return child_docs, parent_store

class InMemoryDocStore:
    """Простое in-memory хранилище документов."""

    def __init__(self) -> None:
        self._store: dict[str, Document] = {}

    def mset(self, pairs: list[tuple[str, Document]]) -> None:
        """Записать несколько документов сразу."""
        for key, doc in pairs:
            self._store[key] = doc

    def mget(self, keys: list[str]) -> list[Document | None]:
        """Получить документы по ключам. None если ключ не найден."""
        return [self._store.get(k) for k in keys]

    def __len__(self) -> int:
        return len(self._store)

    def __contains__(self, key: str) -> bool:
        return key in self._store


class RedisDocStore:
    """
    Redis-based docstore для production.
    Сериализует Document в JSON; родители доступны между перезапусками.
    """

    def __init__(self, redis_url: str = "redis://localhost:6379", ttl: int | None = None) -> None:
        import redis, json
        self._r = redis.from_url(redis_url)
        self._ttl = ttl
        self._json = json

    def mset(self, pairs: list[tuple[str, Document]]) -> None:
        pipe = self._r.pipeline()
        for key, doc in pairs:
            payload = self._json.dumps({
                "page_content": doc.page_content,
                "metadata": doc.metadata,
            })
            if self._ttl:
                pipe.setex(key, self._ttl, payload)
            else:
                pipe.set(key, payload)
        pipe.execute()

    def mget(self, keys: list[str]) -> list[Document | None]:
        raw_list = self._r.mget(keys)
        result = []
        for raw in raw_list:
            if raw is None:
                result.append(None)
            else:
                data = self._json.loads(raw)
                result.append(Document(
                    page_content=data["page_content"],
                    metadata=data["metadata"],
                ))
        return result

class ParentChildRetriever:
    """
    Ищет по дочерним чанкам, возвращает родительские.

    Args:
        vectorstore: хранилище с дочерними чанками (поддерживает asimilarity_search).
        docstore: хранилище с родительскими чанками (mget по parent_id).
        k: сколько уникальных родителей вернуть.
        search_k_multiplier: ищем k * multiplier дочерних чанков,
            чтобы после дедупликации набрать k уникальных родителей.
    """

    def __init__(
        self,
        vectorstore,
        docstore: InMemoryDocStore | RedisDocStore,
        k: int = 4,
        search_k_multiplier: int = 3,
    ) -> None:
        self.vectorstore = vectorstore
        self.docstore = docstore
        self.k = k
        self._search_k = k * search_k_multiplier

    async def aretrieve(self, query: str) -> list[Document]:
        """Async retrieval: query → дочерние → родительские."""
        # 1. Ищем дочерние чанки
        child_results = await self.vectorstore.asimilarity_search(
            query, k=self._search_k
        )

        # 2. Собираем уникальные parent_id (порядок = релевантность)
        seen: dict[str, None] = {}
        for doc in child_results:
            pid = doc.metadata.get("parent_id")
            if pid and pid not in seen:
                seen[pid] = None
            if len(seen) >= self.k:
                break

        parent_ids = list(seen.keys())

        # 3. Загружаем родителей из docstore
        parents = self.docstore.mget(parent_ids)
        return [p for p in parents if p is not None]

    def retrieve(self, query: str) -> list[Document]:
        """Sync wrapper для использования вне async-контекста."""
        import asyncio
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                import concurrent.futures
                with concurrent.futures.ThreadPoolExecutor() as pool:
                    future = pool.submit(asyncio.run, self.aretrieve(query))
                    return future.result()
            return loop.run_until_complete(self.aretrieve(query))
        except RuntimeError:
            return asyncio.run(self.aretrieve(query))

"""
Полный пример: индексируем документацию FastAPI через parent-child chunking,
затем задаём вопрос и получаем ответ с богатым контекстом.

Зависимости: chromadb, openai
"""

import asyncio
from pathlib import Path

import chromadb
from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
from openai import AsyncOpenAI


# ── Собственные сплиттеры (из предыдущих уроков) ──

class RecursiveCharacterSplitter:
    DEFAULT_SEPS = ["\n\n", "\n", ". ", " ", ""]

    def __init__(self, chunk_size=2000, chunk_overlap=100, separators=None):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.seps = separators or self.DEFAULT_SEPS

    def split_documents(self, docs):
        result = []
        for doc in docs:
            for chunk in self._split(doc.page_content, self.seps):
                result.append(Document(page_content=chunk, metadata=dict(doc.metadata)))
        return result

    def _split(self, text, seps):
        if len(text) <= self.chunk_size:
            return [text] if text.strip() else []
        sep = next((s for s in seps if s in text), seps[-1])
        parts = text.split(sep)
        chunks, buf = [], ""
        for part in parts:
            candidate = buf + (sep if buf else "") + part
            if len(candidate) <= self.chunk_size:
                buf = candidate
            else:
                if buf:
                    chunks.append(buf)
                    buf = buf[max(0, len(buf) - self.chunk_overlap):] + (sep if buf else "") + part
                else:
                    sub_seps = seps[seps.index(sep)+1:] if sep in seps[1:] else [""]
                    chunks.extend(self._split(part, sub_seps))
        if buf.strip():
            chunks.append(buf)
        return chunks


# ── Инициализация компонентов ──

embedding_fn = OpenAIEmbeddingFunction(
    api_key="sk-...",
    model_name="text-embedding-3-small",
)

chroma_client = chromadb.Client()
child_collection = chroma_client.create_collection(
    name="children",
    embedding_function=embedding_fn,
)

docstore = InMemoryDocStore()

parent_splitter = RecursiveCharacterSplitter(chunk_size=2000, chunk_overlap=200)
child_splitter  = RecursiveCharacterSplitter(chunk_size=400,  chunk_overlap=50)

splitter = ParentChildSplitter(
    parent_splitter=parent_splitter,
    child_splitter=child_splitter,
)


# ── Адаптер Chroma для retriever'а ──

class ChromaVectorStore:
    def __init__(self, collection):
        self.col = collection

    def add_documents(self, docs: list[Document]) -> None:
        self.col.add(
            ids=[str(i) for i in range(len(docs))],
            documents=[d.page_content for d in docs],
            metadatas=[d.metadata for d in docs],
        )

    async def asimilarity_search(self, query: str, k: int = 10) -> list[Document]:
        results = self.col.query(query_texts=[query], n_results=k)
        docs = []
        for content, meta in zip(results["documents"][0], results["metadatas"][0]):
            docs.append(Document(page_content=content, metadata=meta))
        return docs


# ── Индексирование ──

def index_documents(file_paths: list[str]) -> None:
    raw_docs = []
    for path in file_paths:
        content = Path(path).read_text(encoding="utf-8")
        raw_docs.append(Document(page_content=content, metadata={"source": path}))

    child_docs, parent_dict = splitter.split_documents(raw_docs)

    # Сохраняем родителей в docstore
    docstore.mset(list(parent_dict.items()))

    # Индексируем детей в Chroma
    vectorstore = ChromaVectorStore(child_collection)
    vectorstore.add_documents(child_docs)

    print(f"Проиндексировано:")
    print(f"  Родительских чанков: {len(parent_dict)}")
    print(f"  Дочерних чанков:     {len(child_docs)}")
    print(f"  Среднее детей/родитель: {len(child_docs)/len(parent_dict):.1f}")


# ── Retrieval + генерация ──

async def answer_question(question: str) -> str:
    vectorstore = ChromaVectorStore(child_collection)
    retriever = ParentChildRetriever(
        vectorstore=vectorstore,
        docstore=docstore,
        k=3,
    )

    # Получаем крупные родительские чанки
    parent_docs = await retriever.aretrieve(question)

    context = "\n\n---\n\n".join(d.page_content for d in parent_docs)

    client = AsyncOpenAI()
    response = await client.chat.completions.create(
        model="claude-sonnet-4-6",
        messages=[
            {"role": "system", "content": "Ты помощник по документации. Отвечай только на основе предоставленного контекста."},
            {"role": "user", "content": f"Контекст:\n{context}\n\nВопрос: {question}"},
        ],
    )
    return response.choices[0].message.content


# ── Запуск ──

if __name__ == "__main__":
    # Индексируем
    index_documents(list(Path("docs/").rglob("*.md")))

    # Задаём вопрос
    answer = asyncio.run(answer_question("Как обновить истёкший JWT токен?"))
    print(answer)

from document_aware_chunking import MarkdownHeaderSplitter

# Родитель — H2-секция (весь раздел целиком)
parent_splitter = MarkdownHeaderSplitter(
    headers_to_split_on=[("##", "h2")],
    max_chunk_size=5000,      # допускаем большие родители
    inject_context=True,      # breadcrumb в начало
    strip_headers=False,      # заголовок остаётся в тексте родителя
)

# Ребёнок — H3-подсекция или абзацы (если нет H3)
child_splitter = MarkdownHeaderSplitter(
    headers_to_split_on=[("###", "h3")],
    max_chunk_size=600,
    inject_context=True,
)

splitter = ParentChildSplitter(
    parent_splitter=parent_splitter,
    child_splitter=child_splitter,
)

# Если у H2-секции нет H3-подсекций, MarkdownHeaderSplitter вернёт
# весь текст как один чанк — он станет одновременно и родителем, и ребёнком.
# В этом случае ParentChildSplitter создаёт одну пару parent=child,
# что нормально: retrieval всё равно работает корректно.

# Пример: документация с глубокой иерархией
docs = [Document(page_content=Path("fastapi-docs.md").read_text())]
child_docs, parent_store = splitter.split_documents(docs)

print(f"Родителей (H2-секции): {len(parent_store)}")
print(f"Детей (H3-подсекции):  {len(child_docs)}")

# Выведем первый родитель и его детей
first_parent_id = next(iter(parent_store))
first_parent = parent_store[first_parent_id]
print(f"\nРодитель: {first_parent.metadata.get('h2', 'н/д')}")
print(f"Размер:   {len(first_parent.page_content)} символов")

its_children = [c for c in child_docs if c.metadata.get("parent_id") == first_parent_id]
print(f"Дочерних: {len(its_children)}")
for c in its_children:
    print(f"  — [{c.metadata.get('h3', 'абзац')}] {len(c.page_content)} символов")

import re
from dataclasses import dataclass, field


class SentenceWindowSplitter:
    """
    Создаёт пары: ребёнок = одно предложение,
    родитель = окно window_size предложений вокруг него.

    Каждое предложение становится и ребёнком, и центром своего родителя.
    Это означает, что родители перекрываются — нормально для docstore.
    """

    _ABBREV = re.compile(r"\b(?:т\.е|и\.е|т\.к|ст|п|гл|рис|табл|др|пр|проф|д-р)\.")

    def __init__(self, window_size: int = 3) -> None:
        self.window_size = window_size  # предложений с каждой стороны

    def split_documents(
        self,
        docs: list[Document],
    ) -> tuple[list[Document], dict[str, Document]]:
        child_docs: list[Document] = []
        parent_store: dict[str, Document] = {}

        for doc in docs:
            sentences = self._split_sentences(doc.page_content)

            for i, sentence in enumerate(sentences):
                # Окно: берём window_size предложений с каждой стороны
                start = max(0, i - self.window_size)
                end   = min(len(sentences), i + self.window_size + 1)
                window_text = " ".join(sentences[start:end])

                parent_id = str(uuid.uuid4())
                parent = Document(
                    page_content=window_text,
                    metadata={
                        "doc_id": parent_id,
                        "source": doc.metadata.get("source", ""),
                        "center_sentence": i,
                        "window_start": start,
                        "window_end": end,
                    },
                )
                parent_store[parent_id] = parent

                child = Document(
                    page_content=sentence,
                    metadata={
                        "parent_id": parent_id,
                        "source": doc.metadata.get("source", ""),
                        "sentence_idx": i,
                    },
                )
                child_docs.append(child)

        return child_docs, parent_store

    def _split_sentences(self, text: str) -> list[str]:
        """Разбивает текст на предложения, защищая аббревиатуры."""
        protected = self._ABBREV.sub(lambda m: m.group().replace(".", "‼"), text)
        parts = re.split(r"(?<=[.!?])\s+", protected)
        return [p.replace("‼", ".").strip() for p in parts if p.strip()]


# Пример использования
window_splitter = SentenceWindowSplitter(window_size=2)
docs = [Document(page_content=Path("contract.txt").read_text(), metadata={"source": "contract.txt"})]

child_docs, parent_store = window_splitter.split_documents(docs)

print(f"Предложений (дети):    {len(child_docs)}")
print(f"Окон контекста (родители): {len(parent_store)}")

# Пример пары:
child = child_docs[5]
parent = parent_store[child.metadata["parent_id"]]
print(f"\nДочерний чанк (ищем по нему):")
print(f"  «{child.page_content}»")
print(f"\nРодительское окно (отдаём LLM):")
print(f"  «{parent.page_content}»")