Initial commit: RAGGER - Sistema RAG per associazione di volontariato

chunker.py

@@ -0,0 +1,109 @@
+from dataclasses import dataclass
+from typing import Optional
+from config import settings
+from logger import get_logger
+
+log = get_logger("chunker")
+
+
+@dataclass
+class Chunk:
+    text: str
+    index: int
+    source: str
+    page: Optional[int] = None
+    sheet: Optional[str] = None
+    metadata: Optional[dict] = None
+
+
+class TextChunker:
+    def __init__(
+        self,
+        chunk_size: int = None,
+        overlap: int = None
+    ):
+        self.chunk_size = chunk_size or settings.CHUNK_SIZE
+        self.overlap = overlap or settings.CHUNK_OVERLAP
+        log.debug(f"TextChunker: size={self.chunk_size}, overlap={self.overlap}")
+    
+    def chunk_text(
+        self,
+        text: str,
+        source: str,
+        page: Optional[int] = None,
+        sheet: Optional[str] = None,
+        metadata: Optional[dict] = None
+    ) -> list[Chunk]:
+        words = text.split()
+        log.debug(f"Chunking text: {len(words)} words, source={source}, page={page}, sheet={sheet}")
+        
+        if len(words) <= self.chunk_size:
+            log.debug(f"Text fits in single chunk")
+            return [Chunk(
+                text=text,
+                index=0,
+                source=source,
+                page=page,
+                sheet=sheet,
+                metadata=metadata
+            )]
+        
+        chunks = []
+        start = 0
+        idx = 0
+        
+        while start < len(words):
+            end = start + self.chunk_size
+            chunk_words = words[start:end]
+            chunk_text = " ".join(chunk_words)
+            
+            chunks.append(Chunk(
+                text=chunk_text,
+                index=idx,
+                source=source,
+                page=page,
+                sheet=sheet,
+                metadata=metadata
+            ))
+            
+            log.debug(f"Chunk {idx}: {len(chunk_words)} words, pos {start}-{end}")
+            start = end - self.overlap
+            idx += 1
+        
+        log.debug(f"Created {len(chunks)} chunks")
+        return chunks
+    
+    def chunk_document(self, doc: dict) -> list[Chunk]:
+        log.info(f"Chunking document: {doc.get('source')}")
+        all_chunks = []
+        
+        if "pages" in doc:
+            log.debug(f"Document has {len(doc['pages'])} pages/sheets")
+            for page_data in doc["pages"]:
+                page_chunks = self.chunk_text(
+                    text=page_data["text"],
+                    source=doc["source"],
+                    page=page_data.get("page"),
+                    sheet=page_data.get("sheet"),
+                    metadata={"type": doc["type"]}
+                )
+                all_chunks.extend(page_chunks)
+                if page_data.get("page"):
+                    log.debug(f"Page {page_data['page']}: {len(page_chunks)} chunks")
+                elif page_data.get("sheet"):
+                    log.debug(f"Sheet '{page_data['sheet']}': {len(page_chunks)} chunks")
+        else:
+            log.debug(f"Single page document")
+            all_chunks = self.chunk_text(
+                text=doc["full_text"],
+                source=doc["source"],
+                metadata={"type": doc["type"]}
+            )
+        
+        log.info(f"Document chunked: {len(all_chunks)} total chunks")
+        return all_chunks
+
+
+def chunk_document(doc: dict) -> list[Chunk]:
+    chunker = TextChunker()
+    return chunker.chunk_document(doc)