from pathlib import Path from typing import ClassVar from docling.document_converter import DocumentConverter from docling_core.types.doc.document import DoclingDocument from haiku.rag.utils import text_to_docling_document class FileReader: # Extensions supported by docling docling_extensions: ClassVar[list[str]] = [ ".adoc", ".asc", ".asciidoc", ".bmp", ".csv", ".docx", ".html", ".xhtml", ".jpeg", ".jpg", ".md", ".pdf", ".png", ".pptx", ".tiff", ".xlsx", ".xml", ".webp", ] # Plain text extensions that we'll read directly text_extensions: ClassVar[list[str]] = [ ".astro", ".c", ".cpp", ".css", ".go", ".h", ".hpp", ".java", ".js", ".json", ".kt", ".mdx", ".mjs", ".php", ".py", ".rb", ".rs", ".svelte", ".swift", ".ts", ".tsx", ".txt", ".vue", ".yaml", ".yml", ] # Code file extensions with their markdown language identifiers for syntax highlighting code_markdown_identifier: ClassVar[dict[str, str]] = { ".astro": "astro", ".c": "c", ".cpp": "cpp", ".css": "css", ".go": "go", ".h": "c", ".hpp": "cpp", ".java": "java", ".js": "javascript", ".json": "json", ".kt": "kotlin", ".mjs": "javascript", ".php": "php", ".py": "python", ".rb": "ruby", ".rs": "rust", ".svelte": "svelte", ".swift": "swift", ".ts": "typescript", ".tsx": "tsx", ".vue": "vue", ".yaml": "yaml", ".yml": "yaml", } extensions: ClassVar[list[str]] = docling_extensions + text_extensions @staticmethod def parse_file(path: Path) -> DoclingDocument: try: file_extension = path.suffix.lower() if file_extension in FileReader.docling_extensions: # Use docling for complex document formats converter = DocumentConverter() result = converter.convert(path) return result.document elif file_extension in FileReader.text_extensions: # Read plain text files directly content = path.read_text(encoding="utf-8") # Wrap code files (but not plain txt) in markdown code blocks for better presentation if file_extension in FileReader.code_markdown_identifier: language = FileReader.code_markdown_identifier[file_extension] content = f"```{language}\n{content}\n```" # Convert text to DoclingDocument by wrapping as markdown return text_to_docling_document(content, name=f"{path.stem}.md") else: # Fallback: try to read as text and convert to DoclingDocument content = path.read_text(encoding="utf-8") return text_to_docling_document(content, name=f"{path.stem}.md") except Exception: raise ValueError(f"Failed to parse file: {path}")