dataknobs-bots 0.2.4__py3-none-any.whl

dataknobs_bots/knowledge/retrieval/__init__.py

@@ -0,0 +1,23 @@
+"""Retrieval utilities for RAG knowledge bases.
+
+This module provides post-retrieval processing for improving RAG quality,
+including chunk merging, context formatting, and result optimization.
+"""
+
+from dataknobs_bots.knowledge.retrieval.formatter import (
+    ContextFormatter,
+    FormatterConfig,
+)
+from dataknobs_bots.knowledge.retrieval.merger import (
+    ChunkMerger,
+    MergedChunk,
+    MergerConfig,
+)
+
+__all__ = [
+    "ChunkMerger",
+    "MergedChunk",
+    "MergerConfig",
+    "ContextFormatter",
+    "FormatterConfig",
+]

dataknobs_bots/knowledge/retrieval/formatter.py

@@ -0,0 +1,249 @@
+"""Context formatting utilities for RAG retrieval.
+
+This module provides formatting for retrieved chunks to optimize
+LLM context window usage and improve comprehension.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from dataknobs_bots.knowledge.retrieval.merger import MergedChunk
+
+
+@dataclass
+class FormatterConfig:
+    """Configuration for context formatting.
+
+    Attributes:
+        small_chunk_threshold: Max chars for "small" chunks (full heading path)
+        medium_chunk_threshold: Max chars for "medium" chunks (last 2 headings)
+        include_scores: Whether to include similarity scores
+        include_source: Whether to include source file information
+        group_by_source: Whether to group chunks by source file
+    """
+
+    small_chunk_threshold: int = 200
+    medium_chunk_threshold: int = 800
+    include_scores: bool = False
+    include_source: bool = True
+    group_by_source: bool = False
+
+
+class ContextFormatter:
+    """Formats retrieved chunks for LLM context with dynamic heading inclusion.
+
+    This formatter applies intelligent heading inclusion based on content
+    size to optimize token usage while maintaining context clarity:
+    - Small chunks: Full heading path (need context)
+    - Medium chunks: Last 2 heading levels
+    - Large chunks: No headings (content is self-contained)
+
+    Example:
+        ```python
+        formatter = ContextFormatter(FormatterConfig(
+            small_chunk_threshold=200,
+            include_scores=True
+        ))
+
+        # Format standard results
+        context = formatter.format(results)
+
+        # Format merged chunks
+        context = formatter.format_merged(merged_chunks)
+        ```
+    """
+
+    def __init__(self, config: FormatterConfig | None = None):
+        """Initialize the context formatter.
+
+        Args:
+            config: Formatter configuration, uses defaults if not provided
+        """
+        self.config = config or FormatterConfig()
+
+    def format(self, results: list[dict[str, Any]]) -> str:
+        """Format search results for LLM context.
+
+        Args:
+            results: Search results from RAGKnowledgeBase.query()
+
+        Returns:
+            Formatted context string
+        """
+        if not results:
+            return ""
+
+        if self.config.group_by_source:
+            return self._format_grouped_by_source(results)
+
+        formatted_chunks = []
+        for i, result in enumerate(results, 1):
+            formatted = self._format_result(result, i)
+            formatted_chunks.append(formatted)
+
+        return "\n\n---\n\n".join(formatted_chunks)
+
+    def format_merged(self, merged_chunks: list[MergedChunk]) -> str:
+        """Format merged chunks for LLM context.
+
+        Args:
+            merged_chunks: Merged chunks from ChunkMerger
+
+        Returns:
+            Formatted context string
+        """
+        if not merged_chunks:
+            return ""
+
+        # Convert to result format and use standard formatting
+        results = []
+        for chunk in merged_chunks:
+            results.append({
+                "text": chunk.text,
+                "source": chunk.source,
+                "heading_path": chunk.heading_display,
+                "similarity": chunk.avg_similarity,
+                "metadata": {
+                    "headings": chunk.heading_path,
+                    "content_length": chunk.content_length,
+                },
+            })
+
+        return self.format(results)
+
+    def _format_result(self, result: dict[str, Any], index: int) -> str:
+        """Format a single result with dynamic heading inclusion.
+
+        Args:
+            result: Search result dictionary
+            index: Result index for numbering
+
+        Returns:
+            Formatted chunk string
+        """
+        text = result.get("text", "")
+        source = result.get("source", "")
+        similarity = result.get("similarity", 0.0)
+        metadata = result.get("metadata", {})
+
+        # Get heading information
+        headings = metadata.get("headings", [])
+        if not headings:
+            heading_path = result.get("heading_path", "")
+            if isinstance(heading_path, str) and heading_path:
+                headings = heading_path.split(" > ")
+
+        # Determine content length for heading decision
+        content_length = metadata.get("content_length", len(text))
+
+        # Get headings to display based on content size
+        display_headings = self._get_display_headings(headings, content_length)
+
+        # Build formatted chunk
+        lines = []
+
+        # Add index and heading
+        if display_headings:
+            heading_str = " > ".join(display_headings)
+            if self.config.include_scores:
+                lines.append(f"[{index}] [{similarity:.2f}] {heading_str}")
+            else:
+                lines.append(f"[{index}] {heading_str}")
+        else:
+            if self.config.include_scores:
+                lines.append(f"[{index}] [{similarity:.2f}]")
+            else:
+                lines.append(f"[{index}]")
+
+        # Add content
+        lines.append(text.strip())
+
+        # Add source
+        if self.config.include_source and source:
+            lines.append(f"(Source: {source})")
+
+        return "\n".join(lines)
+
+    def _get_display_headings(
+        self,
+        headings: list[str],
+        content_length: int,
+    ) -> list[str]:
+        """Get headings to display based on content length.
+
+        Implements dynamic heading inclusion:
+        - Small chunks: Full heading path
+        - Medium chunks: Last 2 heading levels
+        - Large chunks: No headings
+
+        Args:
+            headings: Full heading path
+            content_length: Length of content in characters
+
+        Returns:
+            List of headings to display
+        """
+        if not headings:
+            return []
+
+        if content_length < self.config.small_chunk_threshold:
+            # Small chunks: include full heading path
+            return headings
+        elif content_length < self.config.medium_chunk_threshold:
+            # Medium chunks: include last 2 heading levels
+            return headings[-2:] if len(headings) > 2 else headings
+        else:
+            # Large chunks: omit headings (content is self-contained)
+            return []
+
+    def _format_grouped_by_source(self, results: list[dict[str, Any]]) -> str:
+        """Format results grouped by source file.
+
+        Args:
+            results: Search results
+
+        Returns:
+            Formatted context string with source grouping
+        """
+        from collections import defaultdict
+
+        # Group by source
+        groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        for result in results:
+            source = result.get("source", "unknown")
+            groups[source].append(result)
+
+        # Format each group
+        formatted_groups = []
+        chunk_index = 1
+
+        for source, source_results in groups.items():
+            group_lines = [f"## Source: {source}"]
+
+            for result in source_results:
+                formatted = self._format_result(result, chunk_index)
+                # Remove source line since we're grouping
+                lines = formatted.split("\n")
+                lines = [line for line in lines if not line.startswith("(Source:")]
+                group_lines.append("\n".join(lines))
+                chunk_index += 1
+
+            formatted_groups.append("\n\n".join(group_lines))
+
+        return "\n\n---\n\n".join(formatted_groups)
+
+    def wrap_for_prompt(self, context: str, tag: str = "knowledge_base") -> str:
+        """Wrap formatted context in XML tags for prompt injection.
+
+        Args:
+            context: Formatted context string
+            tag: Tag name to wrap with
+
+        Returns:
+            Context wrapped in XML tags
+        """
+        if not context:
+            return ""
+        return f"<{tag}>\n{context}\n</{tag}>"

dataknobs_bots/knowledge/retrieval/merger.py

@@ -0,0 +1,279 @@
+"""Chunk merging utilities for RAG retrieval optimization.
+
+This module provides functionality to merge adjacent chunks that share
+the same heading path, improving context coherence for LLM consumption.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class MergerConfig:
+    """Configuration for chunk merging.
+
+    Attributes:
+        max_merged_size: Maximum size of merged chunk content in characters
+        preserve_order: Whether to preserve positional ordering within groups
+    """
+
+    max_merged_size: int = 2000
+    preserve_order: bool = True
+
+
+@dataclass
+class MergedChunk:
+    """A merged chunk combining multiple related chunks.
+
+    Attributes:
+        text: Combined text content
+        source: Source file path
+        heading_path: Shared heading path
+        heading_display: Formatted heading display string
+        chunks: Original chunks that were merged
+        avg_similarity: Average similarity score of merged chunks
+        content_length: Total content length
+    """
+
+    text: str
+    source: str
+    heading_path: list[str]
+    heading_display: str
+    chunks: list[dict[str, Any]]
+    avg_similarity: float
+    content_length: int
+
+
+class ChunkMerger:
+    """Merges adjacent chunks sharing the same heading path.
+
+    This merger groups search results by their heading path and source,
+    then combines them into coherent context units while respecting
+    size limits.
+
+    Example:
+        ```python
+        merger = ChunkMerger(MergerConfig(max_merged_size=2000))
+        results = await kb.query("How do I configure auth?", k=10)
+        merged = merger.merge(results)
+
+        for chunk in merged:
+            print(f"[{chunk.avg_similarity:.2f}] {chunk.heading_display}")
+            print(chunk.text)
+        ```
+    """
+
+    def __init__(self, config: MergerConfig | None = None):
+        """Initialize the chunk merger.
+
+        Args:
+            config: Merger configuration, uses defaults if not provided
+        """
+        self.config = config or MergerConfig()
+
+    def merge(self, results: list[dict[str, Any]]) -> list[MergedChunk]:
+        """Merge search results by shared heading path.
+
+        Groups chunks by (source, heading_path) and merges those that
+        share identical heading paths. Chunks are ordered by their
+        position within the document.
+
+        Args:
+            results: Search results from RAGKnowledgeBase.query()
+                Each result should have:
+                - text: Chunk content
+                - source: Source file
+                - heading_path: Heading hierarchy string or list
+                - similarity: Similarity score
+                - metadata: Full chunk metadata
+
+        Returns:
+            List of MergedChunk objects, sorted by average similarity
+        """
+        if not results:
+            return []
+
+        # Group chunks by (source, heading_path)
+        groups: dict[tuple[str, tuple[str, ...]], list[dict[str, Any]]] = defaultdict(list)
+
+        for result in results:
+            source = result.get("source", "")
+            heading_path = self._normalize_heading_path(result)
+            key = (source, tuple(heading_path))
+            groups[key].append(result)
+
+        # Merge each group
+        merged_chunks = []
+        for (source, heading_path_tuple), chunks in groups.items():
+            heading_path = list(heading_path_tuple)
+
+            # Sort by position if available
+            if self.config.preserve_order:
+                chunks = self._sort_by_position(chunks)
+
+            # Merge chunks respecting size limit
+            merged = self._merge_chunk_group(chunks, source, heading_path)
+            merged_chunks.extend(merged)
+
+        # Sort by average similarity (descending)
+        merged_chunks.sort(key=lambda c: c.avg_similarity, reverse=True)
+
+        return merged_chunks
+
+    def _normalize_heading_path(self, result: dict[str, Any]) -> list[str]:
+        """Extract and normalize heading path from result.
+
+        Args:
+            result: Search result dictionary
+
+        Returns:
+            List of heading strings
+        """
+        # Try to get from metadata first (may have list format)
+        metadata = result.get("metadata", {})
+        headings = metadata.get("headings", [])
+        if headings:
+            return headings
+
+        # Fall back to heading_path string
+        heading_path = result.get("heading_path", "")
+        if isinstance(heading_path, list):
+            return heading_path
+        elif heading_path:
+            return heading_path.split(" > ")
+
+        return []
+
+    def _sort_by_position(self, chunks: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Sort chunks by their position in the document.
+
+        Args:
+            chunks: List of chunk results
+
+        Returns:
+            Sorted list
+        """
+        def get_position(chunk: dict[str, Any]) -> int:
+            metadata = chunk.get("metadata", {})
+            # Try chunk_index first, then line_number
+            return metadata.get("chunk_index", metadata.get("line_number", 0))
+
+        return sorted(chunks, key=get_position)
+
+    def _merge_chunk_group(
+        self,
+        chunks: list[dict[str, Any]],
+        source: str,
+        heading_path: list[str],
+    ) -> list[MergedChunk]:
+        """Merge a group of chunks with the same heading path.
+
+        Combines chunks until max_merged_size is reached, then starts
+        a new merged chunk. Overflow chunks are returned as separate
+        merged chunks.
+
+        Args:
+            chunks: Chunks to merge
+            source: Source file path
+            heading_path: Shared heading path
+
+        Returns:
+            List of merged chunks
+        """
+        if not chunks:
+            return []
+
+        merged_results = []
+        current_chunks: list[dict[str, Any]] = []
+        current_size = 0
+
+        for chunk in chunks:
+            chunk_text = chunk.get("text", "")
+            chunk_size = len(chunk_text)
+
+            # Check if adding this chunk would exceed the limit
+            if current_size + chunk_size > self.config.max_merged_size and current_chunks:
+                # Save current merge and start new one
+                merged_results.append(
+                    self._create_merged_chunk(current_chunks, source, heading_path)
+                )
+                current_chunks = []
+                current_size = 0
+
+            current_chunks.append(chunk)
+            current_size += chunk_size
+
+        # Don't forget the last group
+        if current_chunks:
+            merged_results.append(
+                self._create_merged_chunk(current_chunks, source, heading_path)
+            )
+
+        return merged_results
+
+    def _create_merged_chunk(
+        self,
+        chunks: list[dict[str, Any]],
+        source: str,
+        heading_path: list[str],
+    ) -> MergedChunk:
+        """Create a MergedChunk from a list of chunks.
+
+        Args:
+            chunks: Chunks to combine
+            source: Source file path
+            heading_path: Shared heading path
+
+        Returns:
+            MergedChunk object
+        """
+        # Combine text with double newline separator
+        texts = [chunk.get("text", "") for chunk in chunks]
+        combined_text = "\n\n".join(text.strip() for text in texts if text.strip())
+
+        # Calculate average similarity
+        similarities = [chunk.get("similarity", 0.0) for chunk in chunks]
+        avg_similarity = sum(similarities) / len(similarities) if similarities else 0.0
+
+        # Build heading display
+        heading_display = " > ".join(heading_path) if heading_path else ""
+
+        return MergedChunk(
+            text=combined_text,
+            source=source,
+            heading_path=heading_path,
+            heading_display=heading_display,
+            chunks=chunks,
+            avg_similarity=avg_similarity,
+            content_length=len(combined_text),
+        )
+
+    def to_result_list(self, merged_chunks: list[MergedChunk]) -> list[dict[str, Any]]:
+        """Convert merged chunks back to result list format.
+
+        Useful for compatibility with existing code that expects
+        the standard result format.
+
+        Args:
+            merged_chunks: List of merged chunks
+
+        Returns:
+            List of result dictionaries
+        """
+        results = []
+        for merged in merged_chunks:
+            results.append({
+                "text": merged.text,
+                "source": merged.source,
+                "heading_path": merged.heading_display,
+                "similarity": merged.avg_similarity,
+                "metadata": {
+                    "headings": merged.heading_path,
+                    "content_length": merged.content_length,
+                    "merged_count": len(merged.chunks),
+                },
+            })
+        return results

dataknobs_bots/memory/__init__.py

@@ -0,0 +1,56 @@
+"""Memory implementations for DynaBot."""
+
+from typing import Any
+
+from .base import Memory
+from .buffer import BufferMemory
+from .vector import VectorMemory
+
+__all__ = ["Memory", "BufferMemory", "VectorMemory", "create_memory_from_config"]
+
+
+async def create_memory_from_config(config: dict[str, Any]) -> Memory:
+    """Create memory instance from configuration.
+
+    Args:
+        config: Memory configuration with 'type' field and type-specific params
+
+    Returns:
+        Configured Memory instance
+
+    Raises:
+        ValueError: If memory type is not recognized
+
+    Example:
+        ```python
+        # Buffer memory
+        config = {
+            "type": "buffer",
+            "max_messages": 10
+        }
+        memory = await create_memory_from_config(config)
+
+        # Vector memory
+        config = {
+            "type": "vector",
+            "backend": "faiss",
+            "dimension": 1536,
+            "embedding_provider": "openai",
+            "embedding_model": "text-embedding-3-small"
+        }
+        memory = await create_memory_from_config(config)
+        ```
+    """
+    memory_type = config.get("type", "buffer").lower()
+
+    if memory_type == "buffer":
+        return BufferMemory(max_messages=config.get("max_messages", 10))
+
+    elif memory_type == "vector":
+        return await VectorMemory.from_config(config)
+
+    else:
+        raise ValueError(
+            f"Unknown memory type: {memory_type}. "
+            f"Available types: buffer, vector"
+        )

dataknobs_bots/memory/base.py

@@ -0,0 +1,38 @@
+"""Base memory interface for bot memory implementations."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Memory(ABC):
+    """Abstract base class for memory implementations."""
+
+    @abstractmethod
+    async def add_message(
+        self, content: str, role: str, metadata: dict[str, Any] | None = None
+    ) -> None:
+        """Add message to memory.
+
+        Args:
+            content: Message content
+            role: Message role (user, assistant, system, etc.)
+            metadata: Optional metadata for the message
+        """
+        pass
+
+    @abstractmethod
+    async def get_context(self, current_message: str) -> list[dict[str, Any]]:
+        """Get relevant context for current message.
+
+        Args:
+            current_message: The current message to get context for
+
+        Returns:
+            List of relevant message dictionaries
+        """
+        pass
+
+    @abstractmethod
+    async def clear(self) -> None:
+        """Clear all memory."""
+        pass

dataknobs_bots/memory/buffer.py

@@ -0,0 +1,58 @@
+"""Buffer memory implementation for simple FIFO message storage."""
+
+from collections import deque
+from typing import Any
+
+from .base import Memory
+
+
+class BufferMemory(Memory):
+    """Simple buffer memory keeping last N messages.
+
+    This implementation uses a fixed-size buffer that keeps the most recent
+    messages in memory. When the buffer is full, the oldest messages are
+    automatically removed.
+
+    Attributes:
+        max_messages: Maximum number of messages to keep in buffer
+        messages: Deque containing the messages
+    """
+
+    def __init__(self, max_messages: int = 10):
+        """Initialize buffer memory.
+
+        Args:
+            max_messages: Maximum number of messages to keep
+        """
+        self.max_messages = max_messages
+        self.messages: deque[dict[str, Any]] = deque(maxlen=max_messages)
+
+    async def add_message(
+        self, content: str, role: str, metadata: dict[str, Any] | None = None
+    ) -> None:
+        """Add message to buffer.
+
+        Args:
+            content: Message content
+            role: Message role
+            metadata: Optional metadata
+        """
+        self.messages.append({"content": content, "role": role, "metadata": metadata or {}})
+
+    async def get_context(self, current_message: str) -> list[dict[str, Any]]:
+        """Get all messages in buffer.
+
+        The current_message parameter is not used in buffer memory since
+        we simply return all buffered messages in order.
+
+        Args:
+            current_message: Not used in buffer memory
+
+        Returns:
+            List of all buffered messages
+        """
+        return list(self.messages)
+
+    async def clear(self) -> None:
+        """Clear all messages from buffer."""
+        self.messages.clear()