fischer-agentkit/src/agentkit/memory/contextual_retrieval.py

"""ContextualChunker - 上下文增强分块

在嵌入前为每个文档块添加 LLM 生成的上下文前缀，
解决分块后上下文丢失问题（Anthropic Contextual Retrieval）。
"""

from __future__ import annotations

import hashlib
import logging
from dataclasses import dataclass
from typing import Any

from agentkit.memory.embedder import EmbeddingCache

logger = logging.getLogger(__name__)


@dataclass
class ContextualChunk:
    """带上下文前缀的文档块"""

    original_content: str
    context_prefix: str
    enhanced_content: str
    chunk_index: int
    metadata: dict[str, Any]

    @property
    def content(self) -> str:
        """获取增强后的完整内容"""
        return self.enhanced_content


CONTEXT_PROMPT_TEMPLATE = """\
Given the full document below and a specific chunk from it, write a brief context that helps someone understand what this chunk is about in the broader document. Output ONLY the context, no explanations.

<document>
{document}
</document>

<chunk>
{chunk}
</chunk>

Context:"""


class ContextualChunker:
    """上下文增强分块器

    为每个文档块生成 LLM 上下文前缀，增强检索质量。

    工作流程：
    1. 接收文档和分块列表
    2. 对每个块，调用 LLM 生成简洁上下文语句
    3. 将上下文前缀添加到原始内容前
    4. 缓存结果避免重复计算

    成本优化：
    - 文档级 Prompt Caching（同一文档的多个块共享文档前缀）
    - EmbeddingCache 缓存上下文生成结果
    - 批处理（batch_size）
    """

    def __init__(
        self,
        llm_gateway: Any = None,
        cache: EmbeddingCache | None = None,
        batch_size: int = 8,
        max_context_length: int = 200,
        prompt_template: str = CONTEXT_PROMPT_TEMPLATE,
    ):
        """
        Args:
            llm_gateway: LLM Gateway 实例，用于生成上下文
            cache: 嵌入缓存，用于缓存上下文生成结果
            batch_size: 批处理大小
            max_context_length: 上下文最大字符长度
            prompt_template: 上下文生成 prompt 模板
        """
        self._llm_gateway = llm_gateway
        self._cache = cache
        self._batch_size = batch_size
        self._max_context_length = max_context_length
        self._prompt_template = prompt_template
        self._context_cache: dict[str, str] = {}

    async def enhance_chunks(
        self,
        document: str,
        chunks: list[str],
        metadata: dict[str, Any] | None = None,
    ) -> list[ContextualChunk]:
        """为文档块添加上下文前缀

        Args:
            document: 完整文档内容
            chunks: 文档分块列表
            metadata: 附加元数据

        Returns:
            增强后的 ContextualChunk 列表
        """
        if not chunks:
            return []

        if not self._llm_gateway:
            # No LLM available — return chunks without context
            logger.info("No LLM gateway configured, skipping contextual enhancement")
            return [
                ContextualChunk(
                    original_content=chunk,
                    context_prefix="",
                    enhanced_content=chunk,
                    chunk_index=i,
                    metadata=metadata or {},
                )
                for i, chunk in enumerate(chunks)
            ]

        result: list[ContextualChunk] = []

        # Process in batches
        for batch_start in range(0, len(chunks), self._batch_size):
            batch = chunks[batch_start : batch_start + self._batch_size]
            batch_results = await self._process_batch(document, batch, batch_start, metadata)
            result.extend(batch_results)

        return result

    async def _process_batch(
        self,
        document: str,
        chunks: list[str],
        start_index: int,
        metadata: dict[str, Any] | None,
    ) -> list[ContextualChunk]:
        """处理一批文档块"""
        results: list[ContextualChunk] = []

        for i, chunk in enumerate(chunks):
            chunk_index = start_index + i
            chunk_meta = dict(metadata or {})
            chunk_meta["chunk_index"] = chunk_index

            # Check cache
            cache_key = self._make_cache_key(document, chunk)
            if cache_key in self._context_cache:
                context = self._context_cache[cache_key]
            else:
                context = await self._generate_context(document, chunk)
                self._context_cache[cache_key] = context

            # Truncate context if too long
            if len(context) > self._max_context_length:
                context = context[: self._max_context_length]

            # Build enhanced content
            if context:
                enhanced = f"{context}\n{chunk}"
            else:
                enhanced = chunk

            chunk_meta["context_prefix"] = context
            chunk_meta["has_context"] = bool(context)

            results.append(
                ContextualChunk(
                    original_content=chunk,
                    context_prefix=context,
                    enhanced_content=enhanced,
                    chunk_index=chunk_index,
                    metadata=chunk_meta,
                )
            )

        return results

    async def _generate_context(self, document: str, chunk: str) -> str:
        """使用 LLM 为单个块生成上下文"""
        # Truncate document for prompt efficiency
        doc_preview = document[:3000] if len(document) > 3000 else document
        chunk_preview = chunk[:1000] if len(chunk) > 1000 else chunk

        prompt = self._prompt_template.format(
            document=doc_preview,
            chunk=chunk_preview,
        )

        try:
            response = await self._llm_gateway.chat(
                messages=[{"role": "user", "content": prompt}],
                model="default",
            )
            context = response.content.strip()
            return context
        except Exception as e:
            logger.warning(f"Context generation failed for chunk: {e}")
            return ""

    @staticmethod
    def _make_cache_key(document: str, chunk: str) -> str:
        """生成缓存键"""
        content = f"{document[:500]}:{chunk[:500]}"
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def clear_cache(self) -> None:
        """清除上下文缓存"""
        self._context_cache.clear()