Spaces:
Sleeping
Sleeping
File size: 2,761 Bytes
c96b98a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | from typing import List, Optional
from phi.document.chunking.strategy import ChunkingStrategy
from phi.document.base import Document
from phi.model.openai import OpenAIChat
from phi.model.base import Model
from phi.model.message import Message
class AgenticChunking(ChunkingStrategy):
"""Chunking strategy that uses an LLM to determine natural breakpoints in the text"""
def __init__(self, model: Optional[Model] = None, max_chunk_size: int = 5000):
self.model = model or OpenAIChat()
self.max_chunk_size = max_chunk_size
def chunk(self, document: Document) -> List[Document]:
"""Split text into chunks using LLM to determine natural breakpoints based on context"""
if len(document.content) <= self.max_chunk_size:
return [document]
chunks: List[Document] = []
remaining_text = self.clean_text(document.content)
chunk_meta_data = document.meta_data
chunk_number = 1
while remaining_text:
# Ask model to find a good breakpoint within max_chunk_size
prompt = f"""Analyze this text and determine a natural breakpoint within the first {self.max_chunk_size} characters.
Consider semantic completeness, paragraph boundaries, and topic transitions.
Return only the character position number of where to break the text:
{remaining_text[: self.max_chunk_size]}"""
try:
response = self.model.response([Message(role="user", content=prompt)])
if response and response.content:
break_point = min(int(response.content.strip()), self.max_chunk_size)
else:
break_point = self.max_chunk_size
except Exception:
# Fallback to max size if model fails
break_point = self.max_chunk_size
# Extract chunk and update remaining text
chunk = remaining_text[:break_point].strip()
meta_data = chunk_meta_data.copy()
meta_data["chunk"] = chunk_number
chunk_id = None
if document.id:
chunk_id = f"{document.id}_{chunk_number}"
elif document.name:
chunk_id = f"{document.name}_{chunk_number}"
meta_data["chunk_size"] = len(chunk)
chunks.append(
Document(
id=chunk_id,
name=document.name,
meta_data=meta_data,
content=chunk,
)
)
chunk_number += 1
remaining_text = remaining_text[break_point:].strip()
if not remaining_text:
break
return chunks
|