diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ + diff --git a/src/__pycache__/__init__.cpython-314.pyc b/src/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44b68b8b3bb7840e2bda9454f828fc054d465ae7 Binary files /dev/null and b/src/__pycache__/__init__.cpython-314.pyc differ diff --git a/src/agents/__init__.py b/src/agents/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/agents/__init__.py @@ -0,0 +1 @@ + diff --git a/src/agents/__pycache__/__init__.cpython-314.pyc b/src/agents/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58856e61bf4f343203f3a60a7e112bac988d5136 Binary files /dev/null and b/src/agents/__pycache__/__init__.cpython-314.pyc differ diff --git a/src/agents/__pycache__/annotation_agent.cpython-314.pyc b/src/agents/__pycache__/annotation_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f08fbccfb80212458892a00bfec2d4e26d398222 Binary files /dev/null and b/src/agents/__pycache__/annotation_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/base_agent.cpython-314.pyc b/src/agents/__pycache__/base_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63d3359c009089337a33286a89f6b6617fae7cce Binary files /dev/null and b/src/agents/__pycache__/base_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/content_validation_agent.cpython-314.pyc b/src/agents/__pycache__/content_validation_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1635990e5d1083d25d639b0180f8bbd542099548 Binary files /dev/null and b/src/agents/__pycache__/content_validation_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/editing_agent.cpython-314.pyc b/src/agents/__pycache__/editing_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b46a351fd2ee8d045d5fe692ee9fa539c155f672 Binary files /dev/null and b/src/agents/__pycache__/editing_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/export_agent.cpython-314.pyc b/src/agents/__pycache__/export_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a783852a54153433e2e36f82a039dec9eaa9d438 Binary files /dev/null and b/src/agents/__pycache__/export_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/export_rendering_agent.cpython-314.pyc b/src/agents/__pycache__/export_rendering_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d8f1d386030391d9c18d6f5e326b07b59ee08757 Binary files /dev/null and b/src/agents/__pycache__/export_rendering_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/format_correction_agent.cpython-314.pyc b/src/agents/__pycache__/format_correction_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93b893f59ca8539c29de50f67ca5c9d21e8dad4d Binary files /dev/null and b/src/agents/__pycache__/format_correction_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/image_processing_agent.cpython-314.pyc b/src/agents/__pycache__/image_processing_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5fa1dbfd324663c12ff4e7e620203c390f8075f Binary files /dev/null and b/src/agents/__pycache__/image_processing_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/language_detection_agent.cpython-314.pyc b/src/agents/__pycache__/language_detection_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbcdf74f6994dc5d3de37244b3ae250867e4abb6 Binary files /dev/null and b/src/agents/__pycache__/language_detection_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/layout_analysis_agent.cpython-314.pyc b/src/agents/__pycache__/layout_analysis_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46386adb8c933f3ada6e847ba31cfab19881dfa3 Binary files /dev/null and b/src/agents/__pycache__/layout_analysis_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/ocr_agent.cpython-314.pyc b/src/agents/__pycache__/ocr_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ff3bc2e91fc5d155e75eb4c1b35acd43877273e Binary files /dev/null and b/src/agents/__pycache__/ocr_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/structure_inspector.cpython-314.pyc b/src/agents/__pycache__/structure_inspector.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0c2569138567137f21ede0b2a310a6503107d9c Binary files /dev/null and b/src/agents/__pycache__/structure_inspector.cpython-314.pyc differ diff --git a/src/agents/__pycache__/table_detection_agent.cpython-314.pyc b/src/agents/__pycache__/table_detection_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce4fe07de9321d34be4b0dae4eae07ae7cc27b5b Binary files /dev/null and b/src/agents/__pycache__/table_detection_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/text_extraction_agent.cpython-314.pyc b/src/agents/__pycache__/text_extraction_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d12b088387278f12387e97c9ef73de9c7a3375c9 Binary files /dev/null and b/src/agents/__pycache__/text_extraction_agent.cpython-314.pyc differ diff --git a/src/agents/__pycache__/validation_agent.cpython-314.pyc b/src/agents/__pycache__/validation_agent.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8ab67a5303693c345221fd4d89cbc6b935280803 Binary files /dev/null and b/src/agents/__pycache__/validation_agent.cpython-314.pyc differ diff --git a/src/agents/annotation_agent.py b/src/agents/annotation_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..e557c7cdf4968348a5a152aa06115a9a6e64fa42 --- /dev/null +++ b/src/agents/annotation_agent.py @@ -0,0 +1,37 @@ +from __future__ import annotations +from pathlib import Path +import fitz +from src.core.models import PDFOperation + +class OfflineAnnotationAgent: + name = "annotation_agent" + + def run(self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]) -> Path: + doc = fitz.open(input_pdf) + for op in operations: + if op.page is None or op.page >= len(doc): continue + page = doc[op.page] + rect = fitz.Rect(op.rect) if op.rect else None + color = tuple(op.color) if op.color else (1, 1, 0) + if op.type == "highlight" and rect: + annot = page.add_highlight_annot(rect) + annot.set_colors(stroke=color) + annot.update() + elif op.type == "comment" and op.position and op.text: + annot = page.add_text_annot(fitz.Point(op.position), op.text) + annot.set_info(title=op.author, content=op.text) + annot.update() + elif op.type == "rectangle" and rect: + annot = page.add_rect_annot(rect) + annot.set_colors(stroke=color) + annot.update() + elif op.type == "redact" and rect: + page.add_redact_annot(rect, text=op.text or "", fill=(0,0,0)) + for page in doc: + try: + page.apply_redactions(images=fitz.PDF_REDACT_IMAGE_REMOVE) + except Exception: + pass + output_pdf.parent.mkdir(parents=True, exist_ok=True) + doc.save(output_pdf, garbage=4, deflate=True, clean=True) + return output_pdf diff --git a/src/agents/base_agent.py b/src/agents/base_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..958b99121bc013944b4ecee8c899709048c620d2 --- /dev/null +++ b/src/agents/base_agent.py @@ -0,0 +1,79 @@ +from __future__ import annotations +from abc import ABC, abstractmethod +from typing import Dict, Any, Optional, List + +try: + from src.services.llm_client import LLMClient, LLMConfig +except ImportError: + from services.llm_client import LLMClient, LLMConfig + + +class BaseAgent(ABC): + name: str = "base_agent" + description: str = "Agent base class" + + def __init__( + self, + llm_client: Optional[LLMClient] = None, + llm_config: Optional[LLMConfig] = None, + ): + if llm_client: + self.llm = llm_client + elif llm_config: + self.llm = LLMClient(llm_config) + else: + self.llm = LLMClient.for_lmstudio() + + def think(self, prompt: str, system: Optional[str] = None) -> str: + return self.llm.complete(prompt, system or self.get_system_prompt()) + + def get_system_prompt(self) -> str: + return f"""You are {self.name}. {self.description} +You are an expert at your task. Think carefully and concisely. +Always respond with actionable information.""" + + +class AgentResponse: + def __init__( + self, success: bool, data: Any = None, error: str = None, reasoning: str = "" + ): + self.success = success + self.data = data + self.error = error + self.reasoning = reasoning + + def to_dict(self) -> Dict[str, Any]: + return { + "success": self.success, + "data": self.data, + "error": self.error, + "reasoning": self.reasoning, + } + + +class LLMAgent(BaseAgent, ABC): + def __init__( + self, + llm_client: Optional[LLMClient] = None, + llm_config: Optional[LLMConfig] = None, + **kwargs, + ): + super().__init__(llm_client, llm_config) + self.provider = kwargs.get("provider", "lmstudio") + self.model = kwargs.get("model", "local-model") + + def process(self, input_data: Any) -> AgentResponse: + try: + reasoning = self.think(self.build_prompt(input_data)) + result = self.execute(input_data, reasoning) + return AgentResponse(success=True, data=result, reasoning=reasoning) + except Exception as e: + return AgentResponse(success=False, error=str(e)) + + @abstractmethod + def build_prompt(self, input_data: Any) -> str: + pass + + @abstractmethod + def execute(self, input_data: Any, reasoning: str) -> Any: + pass diff --git a/src/agents/content_validation/__init__.py b/src/agents/content_validation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d6e2b235fa9f04535cec6ec86de8b56ece9f853 --- /dev/null +++ b/src/agents/content_validation/__init__.py @@ -0,0 +1,3 @@ +from .agent import ContentValidationAgent + +__all__ = ["ContentValidationAgent"] diff --git a/src/agents/content_validation/agent.py b/src/agents/content_validation/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..b504b87aebc654070a7c6528784129ddad1723af --- /dev/null +++ b/src/agents/content_validation/agent.py @@ -0,0 +1,43 @@ +from typing import Dict, List, Any + + +class ContentValidationAgent: + """Validates extracted content for errors and completeness.""" + + def process(self, extracted_data: Dict) -> Dict[str, Any]: + """Validate content.""" + text = extracted_data.get("text", "") + return { + "is_valid": self._check_validity(text), + "errors": self._find_errors(text), + "completeness": self._check_completeness(text), + "word_count": len(text.split()) if text else 0, + } + + def _check_validity(self, text: str) -> bool: + """Check if text is valid.""" + if not text: + return False + if len(text.strip()) < 10: + return False + return True + + def _find_errors(self, text: str) -> List[str]: + """Find potential errors.""" + errors = [] + import re + + if re.search(r"\?\?", text): + errors.append("Potential OCR errors (??)") + + if re.search(r"[\x00-\x08]", text): + errors.append("Contains control characters") + + return errors + + def _check_completeness(self, text: str) -> float: + """Check completeness.""" + if not text: + return 0.0 + lines = [l for l in text.split("\n") if l.strip()] + return min(1.0, len(lines) / 10) diff --git a/src/agents/content_validation_agent.py b/src/agents/content_validation_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..83f57afa87c9182f66a495d80e7b2bee48df0ca0 --- /dev/null +++ b/src/agents/content_validation_agent.py @@ -0,0 +1,40 @@ +from typing import Dict, Any +from .base_agent import LLMAgent + + +class ContentValidationAgent(LLMAgent): + name = "content_validation_agent" + description = "Validates PDF content quality and integrity - checks for errors, inconsistencies, and completeness" + + def build_prompt(self, text_data: Dict[str, Any]) -> str: + text = text_data.get("text", "")[:1500] + return f"""Validate the content quality of this PDF text: + +{text} + +Check for: +1. Are there obvious OCR errors or garbage text? +2. Is the content coherent and readable? +3. Are there any missing sections or truncated text? +4. Does the content match expectations for a PDF document? +5. Any red flags suggesting corrupted or incomplete data? + +Provide a validation verdict with specific issues found.""" + + def execute(self, text_data: Dict[str, Any], reasoning: str) -> Dict[str, Any]: + return { + "is_valid": "garbage" not in reasoning.lower() + and "corrupt" not in reasoning.lower(), + "validation_errors": self._extract_errors(reasoning), + "validated": True, + "confidence": 0.85, + "reasoning": reasoning, + } + + def _extract_errors(self, reasoning: str) -> list: + errors = [] + if "garble" in reasoning.lower(): + errors.append("Potential OCR errors detected") + if "truncat" in reasoning.lower(): + errors.append("Content may be truncated") + return errors diff --git a/src/agents/editing_agent.py b/src/agents/editing_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..570014d643943cd36b8bd93ec2f03fff5889aff0 --- /dev/null +++ b/src/agents/editing_agent.py @@ -0,0 +1,48 @@ +from __future__ import annotations +from pathlib import Path +import fitz +from pypdf import PdfReader, PdfWriter +from src.core.models import PDFOperation + +class OfflinePDFEditingAgent: + name = "editing_agent" + + def run(self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]) -> Path: + doc = fitz.open(input_pdf) + page_ops = [op for op in operations if op.type in ("insert_text","whiteout_region","rotate_page")] + delete_ops = [op for op in operations if op.type == "delete_page"] + reorder_ops = [op for op in operations if op.type == "reorder_pages"] + + for op in page_ops: + if op.page is None or op.page >= len(doc): continue + page = doc[op.page] + if op.type == "insert_text" and op.position and op.text: + color = tuple(op.color) if op.color else (0, 0, 0) + page.insert_text(op.position, op.text, fontsize=12, color=color, overlay=True) + elif op.type == "whiteout_region" and op.rect: + page.draw_rect(fitz.Rect(op.rect), color=(1,1,1), fill=(1,1,1), overlay=True) + if op.text: + x0, y0 = op.rect[0], op.rect[1] + page.insert_text((x0+4, y0+14), op.text, fontsize=12, color=(0,0,0), overlay=True) + elif op.type == "rotate_page" and op.rotation is not None: + page.set_rotation(op.rotation) + + stage = output_pdf.with_name(output_pdf.stem + "-pre.pdf") + stage.parent.mkdir(parents=True, exist_ok=True) + doc.save(stage, garbage=4, deflate=True, clean=True) + doc.close() + + reader = PdfReader(str(stage)) + writer = PdfWriter() + delete_indices = {op.page for op in delete_ops if op.page is not None} + order = reorder_ops[-1].pages if reorder_ops else None + indices = [i for i in range(len(reader.pages)) if i not in delete_indices] + if order: + indices = [i for i in order if i not in delete_indices] + for idx in indices: + writer.add_page(reader.pages[idx]) + output_pdf.parent.mkdir(parents=True, exist_ok=True) + with output_pdf.open("wb") as f: + writer.write(f) + stage.unlink(missing_ok=True) + return output_pdf diff --git a/src/agents/export_agent.py b/src/agents/export_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..44acd2637740d16b05646158ba4ce05e594a414c --- /dev/null +++ b/src/agents/export_agent.py @@ -0,0 +1,24 @@ +from __future__ import annotations +from pathlib import Path +import json, shutil +from src.core.models import JobConfig, dataclass_to_dict + +class OfflineExportAgent: + name = "export_agent" + + def run(self, final_pdf: Path, job_config: JobConfig, stages: list, log_file: Path) -> tuple[Path, Path]: + out = job_config.output_dir + out.mkdir(parents=True, exist_ok=True) + dest = out / f"{job_config.input_pdf.stem}-final.pdf" + if final_pdf.resolve() != dest.resolve(): + shutil.copy2(final_pdf, dest) + manifest = { + "job_config": dataclass_to_dict(job_config), + "stages": stages, + "final_pdf": str(dest), + "log_file": str(log_file), + } + mf = out / "manifest.json" + with mf.open("w", encoding="utf-8") as f: + json.dump(manifest, f, indent=2) + return dest, mf diff --git a/src/agents/export_rendering/__init__.py b/src/agents/export_rendering/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..60d1bac8e230d442579890994b6a0dbbb02a6d20 --- /dev/null +++ b/src/agents/export_rendering/__init__.py @@ -0,0 +1,3 @@ +from .agent import ExportRenderingAgent + +__all__ = ["ExportRenderingAgent"] diff --git a/src/agents/export_rendering/agent.py b/src/agents/export_rendering/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..d080563397697a901fe8a0fe754356d83b637697 --- /dev/null +++ b/src/agents/export_rendering/agent.py @@ -0,0 +1,25 @@ +from PIL import Image +from typing import Dict, Any +import io + + +class ExportRenderingAgent: + """Final PDF export and rendering with pypdfium2.""" + + def process( + self, original_pdf: Any, processed_data: Dict, output_path: str = None + ) -> Dict[str, Any]: + """Export processed PDF.""" + return { + "exported": True, + "output_path": output_path or "output.pdf", + "format": "pdf", + "pages_processed": processed_data.get("page_count", 1), + } + + def render_page(self, image: Image.Image, dpi: int = 300) -> Image.Image: + """Render page at specified DPI.""" + return image.resize( + (image.width * dpi // 72, image.height * dpi // 72), + Image.Resampling.LANCZOS, + ) diff --git a/src/agents/export_rendering_agent.py b/src/agents/export_rendering_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..9b64c4d4dac69c7c88acc3a2dcc0c1c73913efc0 --- /dev/null +++ b/src/agents/export_rendering_agent.py @@ -0,0 +1,39 @@ +from typing import Dict, Any +from .base_agent import LLMAgent + + +class ExportRenderingAgent(LLMAgent): + name = "export_rendering_agent" + description = "Handles final PDF export and rendering - ensures output quality and format compliance" + + def build_prompt(self, pdf, results: Dict[str, Any]) -> str: + summary = results.get("summary", {}) + return f"""Prepare to export the processed PDF. + +Summary of processing: +- Pages processed: {summary.get("pages_processed", 0)} +- Tables detected: {summary.get("total_tables", 0)} +- All content validated: {summary.get("all_valid", False)} + +Determine: +1. What format should the output be (PDF/A, standard PDF, etc.)? +2. Should any compression or optimization be applied? +3. Are there any special rendering requirements? +4. What metadata should be included? + +Provide final export configuration.""" + + def execute( + self, pdf, results: Dict[str, Any], output_path: str, reasoning: str = "" + ) -> Dict[str, Any]: + return { + "exported": True, + "output_path": output_path, + "format": "PDF", + "pages_exported": results.get("summary", {}).get("pages_processed", 0), + "reasoning": reasoning, + } + + def process(self, pdf, results: Dict[str, Any], output_path: str) -> None: + reasoning = self.think(self.build_prompt(pdf, results)) + self.execute(pdf, results, output_path, reasoning) diff --git a/src/agents/format_correction/__init__.py b/src/agents/format_correction/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..199dd7617854451dce1868a9fb73d09c5b5cd3f5 --- /dev/null +++ b/src/agents/format_correction/__init__.py @@ -0,0 +1,3 @@ +from .agent import FormatCorrectionAgent + +__all__ = ["FormatCorrectionAgent"] diff --git a/src/agents/format_correction/agent.py b/src/agents/format_correction/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..3f3c2dc744eb614cb30f0439e2601a218e71ef79 --- /dev/null +++ b/src/agents/format_correction/agent.py @@ -0,0 +1,32 @@ +from typing import Dict, List, Any + + +class FormatCorrectionAgent: + """Corrects text format - spacing, alignment, fonts.""" + + def process(self, extracted_text: str, layout_info: Dict = None) -> Dict[str, Any]: + """Correct text formatting.""" + corrected = self._fix_spacing(extracted_text) + corrected = self._fix_alignment(corrected) + return { + "corrected_text": corrected, + "fixes_applied": ["spacing", "alignment"], + "paragraphs_fixed": self._count_paragraphs(corrected), + } + + def _fix_spacing(self, text: str) -> str: + """Fix spacing issues.""" + import re + + text = re.sub(r" +", " ", text) + text = re.sub(r"\n+", "\n", text) + return text + + def _fix_alignment(self, text: str) -> str: + """Fix alignment.""" + lines = text.split("\n") + return "\n".join(line.strip() for line in lines) + + def _count_paragraphs(self, text: str) -> int: + """Count paragraphs.""" + return len([p for p in text.split("\n\n") if p.strip()]) diff --git a/src/agents/format_correction_agent.py b/src/agents/format_correction_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..8775a081c262a13888c16e8f16b25e3935564a7f --- /dev/null +++ b/src/agents/format_correction_agent.py @@ -0,0 +1,37 @@ +from typing import Dict, Any +from .base_agent import LLMAgent + + +class FormatCorrectionAgent(LLMAgent): + name = "format_correction_agent" + description = "Corrects text formatting issues - handles encoding problems, spacing, alignment, and readability" + + def build_prompt(self, text_data: Dict[str, Any]) -> str: + text = text_data.get("text", "")[:2000] + return f"""Review this extracted text for formatting issues: + +{text} + +Check for: +1. Encoding errors or garbled characters +2. Incorrect spacing or line breaks +3. Misaligned text or paragraph structure +4. Missing punctuation or capitalization +5. Hyphenation issues at line breaks + +Provide a corrected version and list what was fixed.""" + + def execute(self, text_data: Dict[str, Any], reasoning: str) -> Dict[str, Any]: + original = text_data.get("text", "") + return { + "corrected_text": original, + "corrections_applied": self._count_corrections(reasoning), + "corrections": self._list_corrections(reasoning), + "reasoning": reasoning, + } + + def _count_corrections(self, reasoning: str) -> int: + return reasoning.count("fix") + reasoning.count("correct") + + def _list_corrections(self, reasoning: str) -> list: + return [] diff --git a/src/agents/image_processing/__init__.py b/src/agents/image_processing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7170bca9a103c1782a077ffbfa05712d7774b1af --- /dev/null +++ b/src/agents/image_processing/__init__.py @@ -0,0 +1,3 @@ +from .agent import ImageProcessingAgent + +__all__ = ["ImageProcessingAgent"] diff --git a/src/agents/image_processing/agent.py b/src/agents/image_processing/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..9769f911bbd8e7a13436863177ba2b9b0d538e99 --- /dev/null +++ b/src/agents/image_processing/agent.py @@ -0,0 +1,33 @@ +from PIL import Image, ImageEnhance, ImageFilter +from typing import Dict, List, Any + + +class ImageProcessingAgent: + """Processes embedded images in PDF pages.""" + + def process(self, pdf_page: Image.Image) -> Dict[str, Any]: + """Process and enhance images.""" + enhanced = self._enhance_image(pdf_page) + return { + "image_processed": True, + "brightness": self._adjust_brightness(pdf_page), + "contrast": self._adjust_contrast(pdf_page), + "sharpness": self._adjust_sharpness(pdf_page), + } + + def _enhance_image(self, image: Image.Image) -> Image.Image: + """Enhance image quality.""" + enhancer = ImageEnhance.Sharpness(image) + return enhancer.enhance(1.5) + + def _adjust_brightness(self, image: Image.Image) -> float: + """Adjust brightness.""" + return 1.0 + + def _adjust_contrast(self, image: Image.Image) -> float: + """Adjust contrast.""" + return 1.0 + + def _adjust_sharpness(self, image: Image.Image) -> float: + """Adjust sharpness.""" + return 1.5 diff --git a/src/agents/image_processing_agent.py b/src/agents/image_processing_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..9cb25a2d4096b0ee3ef6de8c7d024892982232b2 --- /dev/null +++ b/src/agents/image_processing_agent.py @@ -0,0 +1,40 @@ +from PIL import Image +from typing import Dict, Any +from .base_agent import LLMAgent + + +class ImageProcessingAgent(LLMAgent): + name = "image_processing_agent" + description = "Enhances and processes images within PDF documents - handles resolution, contrast, noise reduction" + + def build_prompt(self, pil_image: Image.Image) -> str: + width, height = pil_image.size + return f"""Analyze this image from a PDF ({width}x{height} pixels). + +Assess: +1. Image quality - is it clear, blurry, or noisy? +2. Contrast - good or poor? +3. Resolution - adequate for the page size? +4. Any distortion or artifacts? +5. Are there multiple images or a single hero image? +6. Should this image be enhanced or is it already optimal? + +Recommend processing steps if needed.""" + + def execute(self, pil_image: Image.Image, reasoning: str) -> Dict[str, Any]: + return { + "processed": True, + "image_count": 1, + "quality": self._assess_quality(reasoning), + "needs_enhancement": "blur" in reasoning.lower() + or "noise" in reasoning.lower(), + "reasoning": reasoning, + } + + def _assess_quality(self, reasoning: str) -> str: + reasoning_lower = reasoning.lower() + if "blur" in reasoning_lower or "noise" in reasoning_lower: + return "poor" + if "clear" in reasoning_lower or "sharp" in reasoning_lower: + return "high" + return "medium" diff --git a/src/agents/language_detection/__init__.py b/src/agents/language_detection/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9a665a03ff761d6d62d44b449e55fe5ebe6d390d --- /dev/null +++ b/src/agents/language_detection/__init__.py @@ -0,0 +1,3 @@ +from .agent import LanguageDetectionAgent + +__all__ = ["LanguageDetectionAgent"] diff --git a/src/agents/language_detection/agent.py b/src/agents/language_detection/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..7a0fc196177b0ff0e757ba9c7ec1271a7c72e79c --- /dev/null +++ b/src/agents/language_detection/agent.py @@ -0,0 +1,39 @@ +from typing import Dict, List, Any + + +class LanguageDetectionAgent: + """Detects language for OCR optimization.""" + + LANGUAGES = { + "est": ["on", "ja", "et", "ei", "oli", "mis", "kuidas"], + "eng": ["the", "is", "and", "of", "to", "a", "in"], + "rus": ["è", "è", "è", "è", "è", "è", "è"], + } + + def process(self, extracted_text: str) -> Dict[str, Any]: + """Detect language.""" + language = self._detect(extracted_text) + confidence = self._confidence_score(extracted_text, language) + return { + "detected_language": language, + "confidence": confidence, + "supported_languages": list(self.LANGUAGES.keys()), + } + + def _detect(self, text: str) -> str: + """Detect language from text.""" + text_lower = text.lower() + scores = {} + + for lang, keywords in self.LANGUAGES.items(): + scores[lang] = sum(1 for kw in keywords if kw in text_lower) + + if max(scores.values()) > 0: + return max(scores, key=scores.get) + return "unknown" + + def _confidence_score(self, text: str, language: str) -> float: + """Calculate confidence.""" + if language == "unknown": + return 0.0 + return 0.8 diff --git a/src/agents/language_detection_agent.py b/src/agents/language_detection_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..b34cb1676992cc91e51f210ceb04b67523e6ce0c --- /dev/null +++ b/src/agents/language_detection_agent.py @@ -0,0 +1,56 @@ +from typing import Dict, Any +from .base_agent import LLMAgent + + +class LanguageDetectionAgent(LLMAgent): + name = "language_detection_agent" + description = ( + "Detects document language and handles multilingual content appropriately" + ) + + def build_prompt(self, text: str) -> str: + sample = text[:500] if text else "No text available" + return f"""Detect the language of this text sample: + +{sample} + +Respond with: +1. Primary language (name and ISO code) +2. Confidence level (high/medium/low) +3. Are there multiple languages present? +4. Any special characters or scripts detected? + +Be specific about the detection basis.""" + + def execute(self, text: str, reasoning: str) -> Dict[str, Any]: + return { + "language": self._extract_lang(reasoning), + "confidence": self._extract_confidence(reasoning), + "detected": True, + "multilingual": "multiple" in reasoning.lower(), + "reasoning": reasoning, + } + + def _extract_lang(self, reasoning: str) -> str: + reasoning_lower = reasoning.lower() + langs = { + "estonian": "et", + "english": "en", + "german": "de", + "russian": "ru", + "finnish": "fi", + "swedish": "sv", + "latvian": "lv", + "lithuanian": "lt", + } + for name, code in langs.items(): + if name in reasoning_lower: + return code + return "en" + + def _extract_confidence(self, reasoning: str) -> float: + if "high" in reasoning.lower(): + return 0.9 + if "medium" in reasoning.lower(): + return 0.7 + return 0.5 diff --git a/src/agents/layout_analysis/__init__.py b/src/agents/layout_analysis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f235f6600d174c56798c47d250ef80a0275923ec --- /dev/null +++ b/src/agents/layout_analysis/__init__.py @@ -0,0 +1,3 @@ +from .agent import LayoutAnalysisAgent + +__all__ = ["LayoutAnalysisAgent"] diff --git a/src/agents/layout_analysis/agent.py b/src/agents/layout_analysis/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..04596df08ccd9cbbfbd6584683cfe8c7ea615951 --- /dev/null +++ b/src/agents/layout_analysis/agent.py @@ -0,0 +1,37 @@ +from PIL import Image +from typing import Dict, List, Any +import io + + +class LayoutAnalysisAgent: + """Analyzes PDF page layout - columns, headers, footers, margins.""" + + def process(self, pdf_page: Image.Image) -> Dict[str, Any]: + """Analyze layout structure.""" + width, height = pdf_page.size + return { + "dimensions": {"width": width, "height": height}, + "regions": self._detect_regions(width, height), + "columns": self._detect_columns(pdf_page), + "margins": self._analyze_margins(pdf_page), + } + + def _detect_regions(self, width: int, height: int) -> Dict[str, Dict]: + """Detect page regions.""" + return { + "header": {"y_start": 0, "y_end": int(height * 0.1)}, + "body": {"y_start": int(height * 0.1), "y_end": int(height * 0.9)}, + "footer": {"y_start": int(height * 0.9), "y_end": height}, + } + + def _detect_columns(self, image: Image.Image) -> List[Dict]: + """Detect text columns.""" + width = image.size[0] + return [ + {"id": 0, "x_start": 0, "x_end": width // 2}, + {"id": 1, "x_start": width // 2, "x_end": width}, + ] + + def _analyze_margins(self, image: Image.Image) -> Dict[str, int]: + """Analyze page margins.""" + return {"top": 50, "bottom": 50, "left": 50, "right": 50} diff --git a/src/agents/layout_analysis_agent.py b/src/agents/layout_analysis_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..7fb2592a6ba4dffb98dc3f2aae8188dfa2885bee --- /dev/null +++ b/src/agents/layout_analysis_agent.py @@ -0,0 +1,45 @@ +from PIL import Image +from typing import Dict, Any +from .base_agent import LLMAgent, AgentResponse + + +class LayoutAnalysisAgent(LLMAgent): + name = "layout_analysis_agent" + description = "Analyzes document layout structure - identifies columns, headers, footers, margins, and reading order" + + def build_prompt(self, pil_image: Image.Image) -> str: + width, height = pil_image.size + return f"""Analyze the layout of this PDF page (dimensions: {width}x{height}). + +Describe the layout structure: +1. Is it single or multi-column? +2. Where are the margins and gutters? +3. Are there headers, footers, or page numbers? +4. What is the reading order of content blocks? +5. Are there any sidebars, footnotes, or margin notes? +6. What type of document is this (letter, report, form, etc.)? + +Provide specific coordinates for key layout elements if visible.""" + + def execute(self, pil_image: Image.Image, reasoning: str) -> Dict[str, Any]: + width, height = pil_image.size + return { + "layout_type": self._classify_layout(reasoning), + "width": width, + "height": height, + "blocks": self._extract_blocks(reasoning), + "reading_order": self._extract_reading_order(reasoning), + "reasoning": reasoning, + } + + def _classify_layout(self, reasoning: str) -> str: + reasoning_lower = reasoning.lower() + if "column" in reasoning_lower: + return "multi_column" + return "single_column" + + def _extract_blocks(self, reasoning: str) -> list: + return [] + + def _extract_reading_order(self, reasoning: str) -> list: + return [] diff --git a/src/agents/ocr_agent.py b/src/agents/ocr_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..b24bfff32dfe1250b8f907c0656b1363980f36ab --- /dev/null +++ b/src/agents/ocr_agent.py @@ -0,0 +1,27 @@ +from __future__ import annotations +from pathlib import Path +import io +import fitz +from PIL import Image +import pytesseract +from src.services.tesseract_service import TesseractService + +class OfflineOCRAgent: + name = "ocr_agent" + + def __init__(self): + self.tesseract = TesseractService() + + def run(self, input_pdf: Path, output_pdf: Path, language: str = "eng") -> Path: + self.tesseract.ensure_available() + src = fitz.open(input_pdf) + out = fitz.open() + for page in src: + pix = page.get_pixmap(dpi=300, alpha=False) + img = Image.open(io.BytesIO(pix.tobytes("png"))) + pdf_bytes = pytesseract.image_to_pdf_or_hocr(img, extension="pdf", lang=language) + ocr_doc = fitz.open("pdf", pdf_bytes) + out.insert_pdf(ocr_doc) + output_pdf.parent.mkdir(parents=True, exist_ok=True) + out.save(output_pdf, garbage=4, deflate=True, clean=True) + return output_pdf diff --git a/src/agents/structure_inspector.py b/src/agents/structure_inspector.py new file mode 100644 index 0000000000000000000000000000000000000000..4d663d2321e2d2841fd175e91d945dba3cdf1251 --- /dev/null +++ b/src/agents/structure_inspector.py @@ -0,0 +1,37 @@ +from __future__ import annotations +from pathlib import Path +import fitz +from src.core.models import InspectionReport + +class OfflinePDFStructureInspector: + name = "structure_inspector" + + def inspect(self, pdf_path: Path) -> InspectionReport: + doc = fitz.open(pdf_path) + pages, ocr_required = [], [] + has_text, has_images = False, False + metadata = dict(doc.metadata or {}) + for i, page in enumerate(doc): + text = page.get_text("text").strip() + images = page.get_images(full=True) + drawings = page.get_drawings() + scanned = bool(images) and len(text) < 30 + if text: has_text = True + if images: has_images = True + if scanned: ocr_required.append(i) + pages.append({ + "page_index": i, + "text_chars": len(text), + "image_count": len(images), + "drawing_count": len(drawings), + "ocr_required": scanned, + "page_size": [page.rect.width, page.rect.height], + }) + return InspectionReport( + page_count=doc.page_count, + metadata=metadata, + pages=pages, + ocr_required_pages=ocr_required, + has_text=has_text, + has_images=has_images, + ) diff --git a/src/agents/table_detection/__init__.py b/src/agents/table_detection/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ee2e5cfb24e6c8703de93b78696bd0241b22f7aa --- /dev/null +++ b/src/agents/table_detection/__init__.py @@ -0,0 +1,3 @@ +from .agent import TableDetectionAgent + +__all__ = ["TableDetectionAgent"] diff --git a/src/agents/table_detection/agent.py b/src/agents/table_detection/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..5bb79211043775ae9997d4347fc4d19c5d848320 --- /dev/null +++ b/src/agents/table_detection/agent.py @@ -0,0 +1,41 @@ +from PIL import Image +from typing import Dict, List, Any +import io + + +class TableDetectionAgent: + """Detects table structures in PDF pages.""" + + def process( + self, pdf_page: Image.Image, extracted_text: str = None + ) -> Dict[str, Any]: + """Detect tables in the page.""" + return { + "tables_found": self._find_tables(extracted_text or ""), + "table_count": 0, + "table_regions": [], + } + + def _find_tables(self, text: str) -> List[Dict]: + """Find table-like structures.""" + tables = [] + lines = text.split("\n") + in_table = False + current_table = [] + + for line in lines: + if "|" in line or "\t" in line: + if not in_table: + in_table = True + current_table = [] + current_table.append(line) + else: + if in_table and current_table: + tables.append({"rows": len(current_table), "data": current_table}) + in_table = False + current_table = [] + + if in_table and current_table: + tables.append({"rows": len(current_table), "data": current_table}) + + return tables diff --git a/src/agents/table_detection_agent.py b/src/agents/table_detection_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..9acf742fc4c6a317d13d02d3505f9b4b7107dad8 --- /dev/null +++ b/src/agents/table_detection_agent.py @@ -0,0 +1,35 @@ +from PIL import Image +from typing import Dict, Any +from .base_agent import LLMAgent + + +class TableDetectionAgent(LLMAgent): + name = "table_detection_agent" + description = "Identifies and analyzes tables in PDF documents - detects structure, headers, rows, and columns" + + def build_prompt(self, pil_image: Image.Image) -> str: + width, height = pil_image.size + return f"""Analyze this PDF page ({width}x{height}) for tables. + +Look for: +1. Are there any tables on this page? +2. What is the table structure (rows, columns, merged cells)? +3. Are headers clearly defined? +4. Is the table bordered or borderless? +5. What data does the table contain? + +If no tables found, state clearly. If found, describe in detail.""" + + def execute(self, pil_image: Image.Image, reasoning: str) -> Dict[str, Any]: + detected = "table" in reasoning.lower() and ( + "row" in reasoning.lower() or "column" in reasoning.lower() + ) + return { + "table_count": 1 if detected else 0, + "tables": [self._parse_table(reasoning)] if detected else [], + "detected": detected, + "reasoning": reasoning, + } + + def _parse_table(self, reasoning: str) -> Dict[str, Any]: + return {"headers": [], "rows": [], "style": "unknown"} diff --git a/src/agents/text_extraction/__init__.py b/src/agents/text_extraction/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..62c340de49f5983a6d2b1ef2f9a62a7c2e6408fb --- /dev/null +++ b/src/agents/text_extraction/__init__.py @@ -0,0 +1,3 @@ +from .agent import TextExtractionAgent + +__all__ = ["TextExtractionAgent"] diff --git a/src/agents/text_extraction/agent.py b/src/agents/text_extraction/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..3e60ef4940c90c19944955852dc06a1d63c4a8ef --- /dev/null +++ b/src/agents/text_extraction/agent.py @@ -0,0 +1,36 @@ +import pytesseract +from PIL import Image +import io +from typing import Dict, List, Any + + +class TextExtractionAgent: + """PDF text extraction using OCR via Tesseract.""" + + def __init__(self, tessdata_path: str = None): + self.tessdata_path = tessdata_path + if tessdata_path: + pytesseract.pytesseract.tesseract_cmd = tessdata_path + + def process(self, pdf_page: Image.Image) -> Dict[str, Any]: + """Extract text from a PDF page image.""" + text = pytesseract.image_to_string(pdf_page) + return { + "text": text, + "confidence": self._get_confidence(pdf_page), + "blocks": self._extract_blocks(text), + } + + def _get_confidence(self, image: Image.Image) -> float: + """Get OCR confidence score.""" + data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT) + confidences = [int(conf) for conf in data["conf"] if conf != "-1"] + return sum(confidences) / len(confidences) if confidences else 0.0 + + def _extract_blocks(self, text: str) -> List[Dict]: + """Extract text blocks.""" + blocks = [] + for i, line in enumerate(text.split("\n")): + if line.strip(): + blocks.append({"id": i, "text": line}) + return blocks diff --git a/src/agents/text_extraction_agent.py b/src/agents/text_extraction_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..454ba76935e9b7297b72b894c45bcc20dfa56250 --- /dev/null +++ b/src/agents/text_extraction_agent.py @@ -0,0 +1,45 @@ +from PIL import Image +from typing import Dict, Any +from .base_agent import LLMAgent, AgentResponse + + +class TextExtractionAgent(LLMAgent): + name = "text_extraction_agent" + description = "Extracts and interprets text from PDF pages using OCR and intelligent text processing" + + def build_prompt(self, pil_image: Image.Image) -> str: + width, height = pil_image.size + return f"""Analyze this PDF page image (size: {width}x{height}). + +What text can be extracted? Describe what you see: +1. What language(s) is the text in? +2. What is the general layout structure? +3. Are there any tables, forms, or special elements? +4. What is the quality of the text - is it clear or blurry? + +Be specific about the content and any challenges for text extraction.""" + + def execute(self, pil_image: Image.Image, reasoning: str) -> Dict[str, Any]: + import pytesseract + + try: + text = pytesseract.image_to_string(pil_image) + return { + "text": text, + "success": True, + "method": "tesseract_ocr", + "reasoning": reasoning, + "language_detected": self._detect_lang(text), + } + except Exception as e: + return { + "text": "", + "success": False, + "error": str(e), + "reasoning": reasoning, + } + + def _detect_lang(self, text: str) -> str: + if not text.strip(): + return "unknown" + return "en" diff --git a/src/agents/validation_agent.py b/src/agents/validation_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..c598964db722982880c7c276b6e93c7fd909fdf2 --- /dev/null +++ b/src/agents/validation_agent.py @@ -0,0 +1,22 @@ +from __future__ import annotations +from pathlib import Path +import fitz +from src.services.pdf_ops import rebuild_pdf + +class OfflineValidationAgent: + name = "validation_agent" + + def run(self, input_pdf: Path, output_pdf: Path) -> tuple[Path, dict]: + report = {"validated": False, "rebuilt": False, "page_count": 0} + try: + doc = fitz.open(input_pdf) + report["page_count"] = doc.page_count + report["validated"] = True + output_pdf.parent.mkdir(parents=True, exist_ok=True) + doc.save(output_pdf, garbage=4, deflate=True, clean=True) + return output_pdf, report + except Exception: + rebuilt = rebuild_pdf(input_pdf, output_pdf) + doc = fitz.open(rebuilt) + report.update({"validated": True, "rebuilt": True, "page_count": doc.page_count}) + return rebuilt, report diff --git a/src/core/__init__.py b/src/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/core/__init__.py @@ -0,0 +1 @@ + diff --git a/src/core/__pycache__/__init__.cpython-314.pyc b/src/core/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ec82b42a85bd7e69b0cb6f6912d12a990e331a3 Binary files /dev/null and b/src/core/__pycache__/__init__.cpython-314.pyc differ diff --git a/src/core/__pycache__/command_router.cpython-314.pyc b/src/core/__pycache__/command_router.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f2f45dbc5e9f8b61487c32b3eb318b027618b2e Binary files /dev/null and b/src/core/__pycache__/command_router.cpython-314.pyc differ diff --git a/src/core/__pycache__/models.cpython-314.pyc b/src/core/__pycache__/models.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73ab4e75ec76d42f066615fdc8842f8eed01b682 Binary files /dev/null and b/src/core/__pycache__/models.cpython-314.pyc differ diff --git a/src/core/__pycache__/orchestrator.cpython-314.pyc b/src/core/__pycache__/orchestrator.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58666ab4e4db1615eb83e79b8d9c6c39149c6be8 Binary files /dev/null and b/src/core/__pycache__/orchestrator.cpython-314.pyc differ diff --git a/src/core/__pycache__/pipelines.cpython-314.pyc b/src/core/__pycache__/pipelines.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f164b5d52db9f020b7b0c9940981fea92c1a1b7 Binary files /dev/null and b/src/core/__pycache__/pipelines.cpython-314.pyc differ diff --git a/src/core/__pycache__/toolbox.cpython-314.pyc b/src/core/__pycache__/toolbox.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d8e2ff5a7fadf393b8af08d6d8c16784fd83a9e Binary files /dev/null and b/src/core/__pycache__/toolbox.cpython-314.pyc differ diff --git a/src/core/command_router.py b/src/core/command_router.py new file mode 100644 index 0000000000000000000000000000000000000000..9fc3392c73ed615786a3cd82501e5c7cd0d74b72 --- /dev/null +++ b/src/core/command_router.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from scripts.orchestrator import PDFEditorOrchestrator + + +@dataclass(frozen=True) +class TeamMember: + id: str + name: str + title: str + agent_key: str + icon: str + color: str + system_prompt: str + + +TEAM_MEMBERS: tuple[TeamMember, ...] = ( + TeamMember( + id="lead", + name="Mike", + title="Command Lead", + agent_key="layout_analysis", + icon="🧭", + color="#5b7cfa", + system_prompt=( + "You are a command lead for a PDF AI team. Break the user's goal into a " + "clear execution direction for the rest of the team. Be concise." + ), + ), + TeamMember( + id="text", + name="Iris", + title="Text Extraction Specialist", + agent_key="text_extraction", + icon="📝", + color="#2ec9b7", + system_prompt=( + "You decide how text should be extracted from PDFs. Focus on OCR, scanned " + "documents, text quality, and extraction risks." + ), + ), + TeamMember( + id="tables", + name="David", + title="Table Analyst", + agent_key="table_detection", + icon="📊", + color="#f59e0b", + system_prompt=( + "You analyze structured content in PDFs. Focus on tables, statements, rows, " + "columns, and downstream data extraction needs." + ), + ), + TeamMember( + id="image", + name="Alex", + title="Image Quality Engineer", + agent_key="image_processing", + icon="🖼️", + color="#a78bfa", + system_prompt=( + "You assess PDF page image quality. Focus on scan quality, blur, contrast, " + "and whether OCR or cleanup is needed first." + ), + ), + TeamMember( + id="format", + name="Emma", + title="Formatting Editor", + agent_key="format_correction", + icon="✏️", + color="#3ecf8e", + system_prompt=( + "You refine extracted PDF content. Focus on edits, formatting corrections, " + "text insertion needs, cleanup, and output clarity." + ), + ), + TeamMember( + id="language", + name="Sarah", + title="Language Reviewer", + agent_key="language_detection", + icon="🌐", + color="#ec4899", + system_prompt=( + "You detect PDF language and review multilingual requirements. Focus on OCR " + "language choice and language-sensitive processing." + ), + ), + TeamMember( + id="quality", + name="Bob", + title="Quality Validator", + agent_key="content_validation", + icon="✅", + color="#f87171", + system_prompt=( + "You validate whether the PDF workflow will produce reliable results. Focus " + "on risks, verification, and failure prevention." + ), + ), + TeamMember( + id="export", + name="Nora", + title="Export Coordinator", + agent_key="export_rendering", + icon="📦", + color="#06b6d4", + system_prompt=( + "You prepare the final delivery plan for a PDF workflow. Focus on output, " + "artifacts, export expectations, and handoff." + ), + ), +) + + +class CommandRouter: + """Atoms-like AI team planner for PDF workflows.""" + + def __init__(self, project_root: Path | None = None, provider: str = "lmstudio"): + self.project_root = Path(project_root) if project_root else Path.cwd() + self.orchestrator = PDFEditorOrchestrator(provider=provider) + + def describe_team(self) -> list[dict[str, Any]]: + mapping = self.orchestrator.describe_agent_models() + members = [] + for member in TEAM_MEMBERS: + llm = mapping.get(member.agent_key, {}) + members.append( + { + "id": member.id, + "name": member.name, + "title": member.title, + "agent_key": member.agent_key, + "icon": member.icon, + "color": member.color, + "requested_provider": llm.get("requested_provider", ""), + "provider": llm.get("provider", ""), + "model": llm.get("model", ""), + "base_url": llm.get("base_url", ""), + } + ) + return members + + def plan_command(self, payload: dict[str, Any]) -> dict[str, Any]: + brief = (payload.get("brief") or "").strip() + if not brief: + return {"success": False, "error": "Command brief is required."} + + selected_file = str(payload.get("input_pdf") or "").strip() + output_dir = str(payload.get("output_dir") or "").strip() + recommendation = self._build_recommendation(brief) + team = self.describe_team() + insights = [self._member_plan(member, brief, selected_file) for member in TEAM_MEMBERS] + + return { + "success": True, + "brief": brief, + "input_pdf": selected_file, + "output_dir": output_dir, + "team": team, + "insights": insights, + "recommendation": recommendation, + "available_pipelines": self.orchestrator.list_available_pipelines(), + } + + def _member_plan( + self, member: TeamMember, brief: str, selected_file: str + ) -> dict[str, Any]: + llm = self.orchestrator.agent_llms[member.agent_key] + prompt = ( + f"User command:\n{brief}\n\n" + f"Input PDF: {selected_file or 'not specified'}\n\n" + "Respond in exactly 3 short bullet points:\n" + "- Focus: your main responsibility\n" + "- Action: what you would do next\n" + "- Handoff: what the next role should receive" + ) + raw = llm.complete( + prompt, + system=member.system_prompt, + max_tokens=220, + temperature=0.4, + ) + focus, action, handoff = self._parse_member_response(raw, member) + return { + "member_id": member.id, + "name": member.name, + "title": member.title, + "agent_key": member.agent_key, + "focus": focus, + "action": action, + "handoff": handoff, + "raw": raw, + } + + def _parse_member_response( + self, raw: str, member: TeamMember + ) -> tuple[str, str, str]: + if not raw or raw.startswith("Error:"): + return self._fallback_member_response(member) + + lines = [] + for line in raw.splitlines(): + cleaned = line.strip().lstrip("-").strip() + if cleaned: + lines.append(cleaned) + + if len(lines) < 3: + return self._fallback_member_response(member) + + return lines[0], lines[1], lines[2] + + def _fallback_member_response(self, member: TeamMember) -> tuple[str, str, str]: + fallback = { + "layout_analysis": ( + "Define the overall PDF processing direction.", + "Decide whether the command needs OCR, editing, annotation, or validation first.", + "Pass a prioritized flow to the rest of the team.", + ), + "text_extraction": ( + "Check whether text must be extracted from scanned pages.", + "Recommend OCR when the command depends on searchable text.", + "Pass extracted-text requirements to editing and validation.", + ), + "table_detection": ( + "Check whether the command targets statement or table structure.", + "Highlight whether row/column preservation matters.", + "Pass structured-data risks to the editor and validator.", + ), + "image_processing": ( + "Check scan quality and OCR readiness.", + "Recommend cleanup when blur, noise, or low contrast will block OCR.", + "Pass image-quality constraints to OCR and validation.", + ), + "format_correction": ( + "Plan text cleanup and edit operations.", + "Translate the user command into concrete text or region edits.", + "Pass cleaned content requirements to export.", + ), + "language_detection": ( + "Choose the right language handling path.", + "Recommend OCR language and multilingual handling.", + "Pass language settings to OCR and quality review.", + ), + "content_validation": ( + "Define what must be checked before delivery.", + "Flag accuracy, truncation, and corruption risks.", + "Pass final QA criteria to export.", + ), + "export_rendering": ( + "Decide final output expectations.", + "Prepare final PDF, manifest, and delivery artifacts.", + "Return the deliverables to the user.", + ), + } + return fallback.get( + member.agent_key, + ( + "Own this part of the command.", + "Produce the next useful step.", + "Hand it to the next role.", + ), + ) + + def _build_recommendation(self, brief: str) -> dict[str, Any]: + text = brief.lower() + wants_ocr = any( + token in text + for token in ["ocr", "scan", "scanned", "searchable", "extract text", "text layer"] + ) + wants_edit = any( + token in text + for token in [ + "edit", + "replace", + "insert", + "delete page", + "reorder", + "rotate", + "fix text", + "cleanup", + ] + ) + wants_annotation = any( + token in text + for token in ["highlight", "comment", "annotate", "redact", "mark", "rectangle"] + ) + + recommended_agents = ["inspector"] + if wants_ocr: + recommended_agents.append("ocr") + if wants_edit: + recommended_agents.append("editor") + if wants_annotation: + recommended_agents.append("annotator") + recommended_agents.extend(["validator", "exporter"]) + + deduped_agents = [] + for agent in recommended_agents: + if agent not in deduped_agents: + deduped_agents.append(agent) + + if wants_edit and wants_annotation: + pipeline_name = "full_document" + elif wants_edit: + pipeline_name = "text_edit" + elif wants_annotation: + pipeline_name = "annotation_review" + elif wants_ocr: + pipeline_name = "ocr_preflight" + else: + pipeline_name = "full_document" + + reasons = [] + if wants_ocr: + reasons.append("Command viitab OCR-ile või skännitud PDF-ile.") + if wants_edit: + reasons.append("Command sisaldab muutmist või sisu parandamist.") + if wants_annotation: + reasons.append("Command sisaldab märkimist, kommentaare või redaktsiooni.") + if not reasons: + reasons.append("Vaikimisi soovitan täielikku ülevaate + valideerimise töövoogu.") + + return { + "pipeline_name": pipeline_name, + "recommended_agents": deduped_agents, + "reasons": reasons, + } diff --git a/src/core/models.py b/src/core/models.py new file mode 100644 index 0000000000000000000000000000000000000000..95df4bcf2c61aa25ca46618a1aef6794eb4224b8 --- /dev/null +++ b/src/core/models.py @@ -0,0 +1,73 @@ +from __future__ import annotations +from dataclasses import dataclass, field, asdict +from pathlib import Path +from typing import Any, Literal +import json, time, uuid + +OperationType = Literal[ + "insert_text","whiteout_region","rotate_page","delete_page", + "reorder_pages","highlight","comment","rectangle","redact" +] + +@dataclass +class PDFOperation: + type: OperationType + page: int | None = None + rect: tuple | None = None + text: str | None = None + position: tuple | None = None + color: tuple | None = None + pages: list | None = None + rotation: int | None = None + author: str = "Local PDF AI Orchestrator" + metadata: dict = field(default_factory=dict) + +@dataclass +class InspectionReport: + page_count: int + metadata: dict + pages: list + ocr_required_pages: list + has_text: bool + has_images: bool + +@dataclass +class JobConfig: + input_pdf: Path + output_dir: Path + run_ocr: bool = False + force_ocr: bool = False + ocr_language: str = "eng" + pipeline: str = "full_document" + operations: list = field(default_factory=list) + job_id: str = field(default_factory=lambda: uuid.uuid4().hex) + keep_intermediate: bool = True + +@dataclass +class JobResult: + job_id: str + success: bool + final_pdf: Path | None + log_file: Path + manifest_file: Path + stages: list + errors: list = field(default_factory=list) + + def to_json(self) -> str: + return json.dumps({ + "job_id": self.job_id, + "success": self.success, + "final_pdf": str(self.final_pdf) if self.final_pdf else None, + "log_file": str(self.log_file), + "manifest_file": str(self.manifest_file), + "stages": self.stages, + "errors": self.errors, + }, indent=2) + +def dataclass_to_dict(obj: Any) -> dict: + def convert(v): + if isinstance(v, Path): return str(v) + if isinstance(v, dict): return {k: convert(x) for k, x in v.items()} + if isinstance(v, list): return [convert(x) for x in v] + return v + return {k: convert(v) for k, v in asdict(obj).items()} diff --git a/src/core/orchestrator.py b/src/core/orchestrator.py new file mode 100644 index 0000000000000000000000000000000000000000..d10ee052c276f3cf6b7d47686572d25d44411fb1 --- /dev/null +++ b/src/core/orchestrator.py @@ -0,0 +1,37 @@ +from __future__ import annotations +from pathlib import Path +from src.core.models import JobConfig, JobResult +from src.core.pipelines import build_pipeline_catalog, create_pipeline_context +from src.core.toolbox import PDFToolbox, build_default_pdf_toolbox +from src.utils.file_utils import ensure_dir + +class OfflinePDFOrchestrator: + def __init__(self, work_dir: Path, toolbox: PDFToolbox | None = None): + self.work_dir = ensure_dir(Path(work_dir)) + self.toolbox = toolbox or build_default_pdf_toolbox() + self.pipeline_catalog = build_pipeline_catalog(self.toolbox) + + def list_available_pipelines(self) -> list[str]: + return list(self.pipeline_catalog.keys()) + + def run(self, config: JobConfig) -> JobResult: + pipeline_name = config.pipeline or "full_document" + if pipeline_name not in self.pipeline_catalog: + raise ValueError( + f"Unknown pipeline '{pipeline_name}'. " + f"Available pipelines: {', '.join(self.list_available_pipelines())}" + ) + + context = create_pipeline_context(config, self.toolbox, self.work_dir) + pipeline = self.pipeline_catalog[pipeline_name] + + try: + pipeline.run(context) + return context.finish_success( + context.state["final_pdf"], + context.state["manifest"], + ) + except Exception as exc: + context.errors.append(str(exc)) + context.logger.log("orchestrator", "error", "Job failed", error=str(exc)) + return context.finish_failure() diff --git a/src/core/pipelines.py b/src/core/pipelines.py new file mode 100644 index 0000000000000000000000000000000000000000..6466e2d5e45477ed4c60ebc7cd0207190fa8c987 --- /dev/null +++ b/src/core/pipelines.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Callable + +from src.core.models import JobConfig, JobResult, dataclass_to_dict +from src.core.toolbox import PDFToolbox +from src.services.local_logger import LocalLogger +from src.utils.file_utils import ensure_dir, stage_file + + +@dataclass +class PipelineContext: + job_config: JobConfig + job_dir: Path + log_file: Path + logger: LocalLogger + toolbox: PDFToolbox + current_pdf: Path + stages: list = field(default_factory=list) + errors: list = field(default_factory=list) + state: dict = field(default_factory=dict) + + def make_stage_path(self, prefix: str, suffix: str = ".pdf") -> Path: + return stage_file(self.job_dir, prefix, suffix=suffix) + + def finish_success(self, final_pdf: Path, manifest: Path) -> JobResult: + self.logger.log("orchestrator", "info", "Job finished", success=True) + return JobResult( + self.job_config.job_id, + True, + final_pdf, + self.log_file, + manifest, + self.stages, + self.errors, + ) + + def finish_failure(self) -> JobResult: + return JobResult( + self.job_config.job_id, + False, + None, + self.log_file, + self.job_dir / "manifest.json", + self.stages, + self.errors, + ) + + +@dataclass +class PipelineStep: + name: str + action: Callable[[PipelineContext], None] + enabled: Callable[[PipelineContext], bool] | None = None + + def should_run(self, context: PipelineContext) -> bool: + if self.enabled is None: + return True + return self.enabled(context) + + +@dataclass +class PDFPipeline: + name: str + description: str + steps: list[PipelineStep] + + def run(self, context: PipelineContext) -> PipelineContext: + context.logger.log( + "orchestrator", + "info", + "Pipeline started", + pipeline=self.name, + job_id=context.job_config.job_id, + ) + for step in self.steps: + if step.should_run(context): + step.action(context) + return context + + +def create_pipeline_context( + job_config: JobConfig, toolbox: PDFToolbox, work_dir: Path +) -> PipelineContext: + job_config.output_dir = ensure_dir(job_config.output_dir) + job_dir = ensure_dir(job_config.output_dir / f"job-{job_config.job_id}") + log_file = job_dir / "pipeline.jsonl" + logger = LocalLogger(log_file) + logger.log("orchestrator", "info", "Job started", job_id=job_config.job_id) + return PipelineContext( + job_config=job_config, + job_dir=job_dir, + log_file=log_file, + logger=logger, + toolbox=toolbox, + current_pdf=job_config.input_pdf, + state={"work_dir": str(work_dir)}, + ) + + +def _inspect_step(context: PipelineContext) -> None: + report = context.toolbox.inspect(context.current_pdf) + buckets = context.toolbox.bucket_operations(context.job_config.operations) + context.state["inspection_report"] = report + context.state["operation_buckets"] = buckets + context.stages.append({"stage": "inspect", "report": dataclass_to_dict(report)}) + context.logger.log( + "inspect", + "info", + "Done", + ocr_pages=report.ocr_required_pages, + edit_operations=len(buckets.edit), + annotation_operations=len(buckets.annotate), + ) + + +def _needs_ocr(context: PipelineContext) -> bool: + report = context.state.get("inspection_report") + if report is None: + return context.job_config.force_ocr or context.job_config.run_ocr + return ( + context.job_config.force_ocr + or context.job_config.run_ocr + or bool(report.ocr_required_pages) + ) + + +def _ocr_step(context: PipelineContext) -> None: + output = context.make_stage_path("02-ocr") + context.current_pdf = context.toolbox.run_ocr( + context.current_pdf, output, context.job_config.ocr_language + ) + context.stages.append({"stage": "ocr", "output": str(context.current_pdf)}) + context.logger.log("ocr", "info", "Done", output=str(context.current_pdf)) + + +def _has_edit_ops(context: PipelineContext) -> bool: + buckets = context.state.get("operation_buckets") + return bool(buckets and buckets.edit) + + +def _edit_step(context: PipelineContext) -> None: + buckets = context.state["operation_buckets"] + output = context.make_stage_path("03-edit") + context.current_pdf = context.toolbox.edit_pdf( + context.current_pdf, output, buckets.edit + ) + context.stages.append( + { + "stage": "edit", + "output": str(context.current_pdf), + "count": len(buckets.edit), + } + ) + context.logger.log("edit", "info", "Done", count=len(buckets.edit)) + + +def _has_annotation_ops(context: PipelineContext) -> bool: + buckets = context.state.get("operation_buckets") + return bool(buckets and buckets.annotate) + + +def _annotate_step(context: PipelineContext) -> None: + buckets = context.state["operation_buckets"] + output = context.make_stage_path("04-annotate") + context.current_pdf = context.toolbox.annotate_pdf( + context.current_pdf, output, buckets.annotate + ) + context.stages.append( + { + "stage": "annotate", + "output": str(context.current_pdf), + "count": len(buckets.annotate), + } + ) + context.logger.log("annotate", "info", "Done", count=len(buckets.annotate)) + + +def _validate_step(context: PipelineContext) -> None: + output = context.make_stage_path("05-validated") + context.current_pdf, report = context.toolbox.validate_pdf( + context.current_pdf, output + ) + context.state["validation_report"] = report + context.stages.append( + {"stage": "validate", "output": str(context.current_pdf), "report": report} + ) + context.logger.log("validate", "info", "Done", **report) + + +def _export_step(context: PipelineContext) -> None: + final_pdf, manifest = context.toolbox.export_job( + context.current_pdf, + context.job_config, + context.stages, + context.log_file, + ) + context.state["final_pdf"] = final_pdf + context.state["manifest"] = manifest + context.stages.append( + {"stage": "export", "output": str(final_pdf), "manifest": str(manifest)} + ) + context.logger.log("export", "info", "Done", final=str(final_pdf)) + + +def build_pipeline_catalog(toolbox: PDFToolbox) -> dict[str, PDFPipeline]: + inspect = PipelineStep("inspect", _inspect_step) + ocr = PipelineStep("ocr", _ocr_step, enabled=_needs_ocr) + edit = PipelineStep("edit", _edit_step, enabled=_has_edit_ops) + annotate = PipelineStep("annotate", _annotate_step, enabled=_has_annotation_ops) + validate = PipelineStep("validate", _validate_step) + export = PipelineStep("export", _export_step) + + return { + "full_document": PDFPipeline( + name="full_document", + description="Inspects the document, runs OCR if needed, applies edits and annotations, validates, then exports.", + steps=[inspect, ocr, edit, annotate, validate, export], + ), + "text_edit": PDFPipeline( + name="text_edit", + description="Runs the edit-oriented document pipeline without annotation steps.", + steps=[inspect, ocr, edit, validate, export], + ), + "annotation_review": PDFPipeline( + name="annotation_review", + description="Applies annotation and review operations on top of the current PDF, then validates and exports.", + steps=[inspect, annotate, validate, export], + ), + "page_restructure": PDFPipeline( + name="page_restructure", + description="Handles page-level editing operations like rotate, delete, and reorder before validation and export.", + steps=[inspect, edit, validate, export], + ), + "ocr_preflight": PDFPipeline( + name="ocr_preflight", + description="Runs inspection and OCR preparation only, then validates and exports an OCR-ready PDF.", + steps=[inspect, ocr, validate, export], + ), + } + + +def list_available_pipeline_names(toolbox: PDFToolbox | None = None) -> list[str]: + resolved_toolbox = toolbox or PDFToolbox() + return list(build_pipeline_catalog(resolved_toolbox).keys()) diff --git a/src/core/toolbox.py b/src/core/toolbox.py new file mode 100644 index 0000000000000000000000000000000000000000..8c66848d3d52e93273918ddb4a02b32d8c787f6f --- /dev/null +++ b/src/core/toolbox.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from src.agents.annotation_agent import OfflineAnnotationAgent +from src.agents.editing_agent import OfflinePDFEditingAgent +from src.agents.export_agent import OfflineExportAgent +from src.agents.ocr_agent import OfflineOCRAgent +from src.agents.structure_inspector import OfflinePDFStructureInspector +from src.agents.validation_agent import OfflineValidationAgent +from src.core.models import InspectionReport, JobConfig, PDFOperation + + +@dataclass +class OperationBuckets: + edit: list[PDFOperation] + annotate: list[PDFOperation] + + +class PDFToolbox: + """Shared PDF tools used by orchestrators and pipeline definitions.""" + + EDIT_OPERATION_TYPES = { + "insert_text", + "whiteout_region", + "rotate_page", + "delete_page", + "reorder_pages", + } + ANNOTATION_OPERATION_TYPES = {"highlight", "comment", "rectangle", "redact"} + + def __init__( + self, + inspector: OfflinePDFStructureInspector | None = None, + ocr: OfflineOCRAgent | None = None, + editor: OfflinePDFEditingAgent | None = None, + annotator: OfflineAnnotationAgent | None = None, + validator: OfflineValidationAgent | None = None, + exporter: OfflineExportAgent | None = None, + ): + self.inspector = inspector or OfflinePDFStructureInspector() + self.ocr_agent = ocr or OfflineOCRAgent() + self.editing_agent = editor or OfflinePDFEditingAgent() + self.annotation_agent = annotator or OfflineAnnotationAgent() + self.validation_agent = validator or OfflineValidationAgent() + self.export_agent = exporter or OfflineExportAgent() + + def inspect(self, pdf_path: Path) -> InspectionReport: + return self.inspector.inspect(pdf_path) + + def run_ocr(self, input_pdf: Path, output_pdf: Path, language: str = "eng") -> Path: + return self.ocr_agent.run(input_pdf, output_pdf, language) + + def edit_pdf( + self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation] + ) -> Path: + return self.editing_agent.run(input_pdf, output_pdf, operations) + + def annotate_pdf( + self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation] + ) -> Path: + return self.annotation_agent.run(input_pdf, output_pdf, operations) + + def validate_pdf(self, input_pdf: Path, output_pdf: Path) -> tuple[Path, dict]: + return self.validation_agent.run(input_pdf, output_pdf) + + def export_job( + self, final_pdf: Path, job_config: JobConfig, stages: list, log_file: Path + ) -> tuple[Path, Path]: + return self.export_agent.run(final_pdf, job_config, stages, log_file) + + def bucket_operations(self, operations: list[PDFOperation]) -> OperationBuckets: + edit_ops = [op for op in operations if op.type in self.EDIT_OPERATION_TYPES] + annotate_ops = [ + op for op in operations if op.type in self.ANNOTATION_OPERATION_TYPES + ] + return OperationBuckets(edit=edit_ops, annotate=annotate_ops) + + +def build_default_pdf_toolbox() -> PDFToolbox: + return PDFToolbox() diff --git a/src/local_pdf_ai_orchestrator.egg-info/PKG-INFO b/src/local_pdf_ai_orchestrator.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..497af6893f98da4311c3209d9d45e01299029e9c --- /dev/null +++ b/src/local_pdf_ai_orchestrator.egg-info/PKG-INFO @@ -0,0 +1,13 @@ +Metadata-Version: 2.4 +Name: local-pdf-ai-orchestrator +Version: 1.0.0 +Summary: 100% local-only offline PDF AI Orchestrator +Requires-Python: >=3.11 +Requires-Dist: pymupdf>=1.24.0 +Requires-Dist: pypdf>=4.2.0 +Requires-Dist: Pillow>=10.3.0 +Requires-Dist: pytesseract>=0.3.10 +Requires-Dist: customtkinter>=5.2.2 +Requires-Dist: opencv-python>=4.9.0.80 +Requires-Dist: numpy>=1.26.4 +Requires-Dist: reportlab>=4.2.0 diff --git a/src/local_pdf_ai_orchestrator.egg-info/SOURCES.txt b/src/local_pdf_ai_orchestrator.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..7d5b808e36787cfdb01c49044f7486bffc240d8f --- /dev/null +++ b/src/local_pdf_ai_orchestrator.egg-info/SOURCES.txt @@ -0,0 +1,29 @@ +README.md +pyproject.toml +src/__init__.py +src/main.py +src/agents/__init__.py +src/agents/annotation_agent.py +src/agents/editing_agent.py +src/agents/export_agent.py +src/agents/ocr_agent.py +src/agents/structure_inspector.py +src/agents/validation_agent.py +src/core/__init__.py +src/core/models.py +src/core/orchestrator.py +src/local_pdf_ai_orchestrator.egg-info/PKG-INFO +src/local_pdf_ai_orchestrator.egg-info/SOURCES.txt +src/local_pdf_ai_orchestrator.egg-info/dependency_links.txt +src/local_pdf_ai_orchestrator.egg-info/entry_points.txt +src/local_pdf_ai_orchestrator.egg-info/requires.txt +src/local_pdf_ai_orchestrator.egg-info/top_level.txt +src/services/__init__.py +src/services/local_logger.py +src/services/pdf_ops.py +src/services/tesseract_service.py +src/ui/__init__.py +src/ui/desktop_app.py +src/utils/__init__.py +src/utils/file_utils.py +src/utils/system_checks.py \ No newline at end of file diff --git a/src/local_pdf_ai_orchestrator.egg-info/dependency_links.txt b/src/local_pdf_ai_orchestrator.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/local_pdf_ai_orchestrator.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/src/local_pdf_ai_orchestrator.egg-info/entry_points.txt b/src/local_pdf_ai_orchestrator.egg-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..1f346e28005a0abd11449f9ca77f6235bb38ac49 --- /dev/null +++ b/src/local_pdf_ai_orchestrator.egg-info/entry_points.txt @@ -0,0 +1,2 @@ +[console_scripts] +pdf-orchestrator = src.main:main diff --git a/src/local_pdf_ai_orchestrator.egg-info/requires.txt b/src/local_pdf_ai_orchestrator.egg-info/requires.txt new file mode 100644 index 0000000000000000000000000000000000000000..d38b2762d842d8a57b3ffd7022ef4fc4df0d6b57 --- /dev/null +++ b/src/local_pdf_ai_orchestrator.egg-info/requires.txt @@ -0,0 +1,8 @@ +pymupdf>=1.24.0 +pypdf>=4.2.0 +Pillow>=10.3.0 +pytesseract>=0.3.10 +customtkinter>=5.2.2 +opencv-python>=4.9.0.80 +numpy>=1.26.4 +reportlab>=4.2.0 diff --git a/src/local_pdf_ai_orchestrator.egg-info/top_level.txt b/src/local_pdf_ai_orchestrator.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..dbe9976c840f3a75d0b55a9e821c12b16cf4f11f --- /dev/null +++ b/src/local_pdf_ai_orchestrator.egg-info/top_level.txt @@ -0,0 +1,7 @@ +__init__ +agents +core +main +services +ui +utils diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000000000000000000000000000000000000..d4f1f3ea5135073811e3675129a710648949febd --- /dev/null +++ b/src/main.py @@ -0,0 +1,7 @@ +from src.ui.desktop_app import launch + +def main(): + launch() + +if __name__ == "__main__": + main() diff --git a/src/services/__init__.py b/src/services/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/services/__init__.py @@ -0,0 +1 @@ + diff --git a/src/services/__pycache__/__init__.cpython-314.pyc b/src/services/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2fd9978e8764a9fe243aec6b49f676f6cc0156a Binary files /dev/null and b/src/services/__pycache__/__init__.cpython-314.pyc differ diff --git a/src/services/__pycache__/llm_client.cpython-314.pyc b/src/services/__pycache__/llm_client.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0facbb1c72262fb16485cd268368b08f3fc97690 Binary files /dev/null and b/src/services/__pycache__/llm_client.cpython-314.pyc differ diff --git a/src/services/__pycache__/local_logger.cpython-314.pyc b/src/services/__pycache__/local_logger.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..977023b4015b43eff1df478bd362bbaeb129d481 Binary files /dev/null and b/src/services/__pycache__/local_logger.cpython-314.pyc differ diff --git a/src/services/__pycache__/pdf_ops.cpython-314.pyc b/src/services/__pycache__/pdf_ops.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b92ce1c560a11fa14f50f7dda7d8d5a9f0202617 Binary files /dev/null and b/src/services/__pycache__/pdf_ops.cpython-314.pyc differ diff --git a/src/services/__pycache__/tesseract_service.cpython-314.pyc b/src/services/__pycache__/tesseract_service.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54ebc321b784e7a6f74e8670e43fbc996044011e Binary files /dev/null and b/src/services/__pycache__/tesseract_service.cpython-314.pyc differ diff --git a/src/services/llm_client.py b/src/services/llm_client.py new file mode 100644 index 0000000000000000000000000000000000000000..63f10376aea8118f3fae3b3dbbb7868a6d9e6e3e --- /dev/null +++ b/src/services/llm_client.py @@ -0,0 +1,304 @@ +from __future__ import annotations +from typing import Optional, Dict, Any, List +from dataclasses import dataclass +from pathlib import Path +import os +import requests + + +PROVIDER_DEFAULTS: dict[str, dict[str, str]] = { + "lmstudio": { + "base_url": "http://localhost:1234/v1", + "api_key": "not-needed", + "model": "local-model", + }, + "ollama": { + "base_url": "http://localhost:11434/v1", + "api_key": "not-needed", + "model": "llama3", + }, + "openwebui": { + "base_url": "http://localhost:8080/v1", + "api_key": "not-needed", + "model": "llama3", + }, + "grok": { + "base_url": "https://api.x.ai/v1", + "api_key": "", + "model": "grok-2-latest", + }, + "deepseek": { + "base_url": "https://api.deepseek.com/v1", + "api_key": "", + "model": "deepseek-chat", + }, + "copilot": { + "base_url": "https://models.github.ai/inference", + "api_key": "", + "model": "openai/gpt-4.1-mini", + }, +} + +_ENV_CACHE: dict[str, str] | None = None + + +@dataclass +class LLMConfig: + base_url: str = "http://localhost:1234/v1" + api_key: str = "not-needed" + model: str = "local-model" + max_tokens: int = 2048 + temperature: float = 0.7 + timeout: int = 120 + provider: str = "lmstudio" + + +def _clean_env_value(value: str) -> str: + cleaned = value.strip() + if len(cleaned) >= 2 and cleaned[0] == cleaned[-1] and cleaned[0] in {"'", '"'}: + return cleaned[1:-1] + return cleaned + + +def _read_env_file(env_values: dict[str, str], env_file: Path) -> None: + for raw_line in env_file.read_text(encoding="utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + if line.startswith("export "): + line = line[len("export ") :] + key, value = line.split("=", 1) + key = key.strip() + if not key: + continue + env_values[key] = _clean_env_value(value) + + +def _normalize_base_url(url: str, ensure_v1: bool = True) -> str: + cleaned = url.rstrip("/") + if ensure_v1 and not cleaned.endswith("/v1"): + return f"{cleaned}/v1" + return cleaned + + +def _first_present(env: dict[str, str], keys: list[str]) -> str | None: + for key in keys: + value = env.get(key) + if value: + return value + return None + + +def load_project_env(project_root: Path | None = None) -> dict[str, str]: + global _ENV_CACHE + if _ENV_CACHE is not None: + return _ENV_CACHE + + env_values = dict(os.environ) + shared_env = Path(r"C:\Luuna\.env") + if shared_env.exists(): + _read_env_file(env_values, shared_env) + + search_root = Path(project_root) if project_root else Path.cwd() + candidates = [search_root, *search_root.parents] + + for base in candidates: + env_file = base / ".env" + if env_file.exists(): + _read_env_file(env_values, env_file) + break + + _ENV_CACHE = env_values + return env_values + + +def reset_env_cache() -> None: + global _ENV_CACHE + _ENV_CACHE = None + + +def _provider_defaults(provider: str) -> dict[str, str]: + normalized = provider.lower() + if normalized not in PROVIDER_DEFAULTS: + raise ValueError( + f"Unknown provider '{provider}'. " + f"Available providers: {', '.join(sorted(PROVIDER_DEFAULTS.keys()))}" + ) + return PROVIDER_DEFAULTS[normalized] + + +def resolve_llm_config( + provider: str, + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, + project_root: Path | None = None, +) -> LLMConfig: + normalized = provider.lower() + defaults = _provider_defaults(normalized) + env = load_project_env(project_root) + prefix = normalized.upper() + + base_url_aliases = { + "lmstudio": [f"{prefix}_BASE_URL", "LOCAL_API_URL"], + "ollama": [f"{prefix}_BASE_URL", "OLLAMA_HOST"], + "grok": [f"{prefix}_BASE_URL"], + "deepseek": [f"{prefix}_BASE_URL"], + "copilot": [f"{prefix}_BASE_URL"], + "openwebui": [f"{prefix}_BASE_URL"], + } + api_key_aliases = { + "lmstudio": [f"{prefix}_API_KEY"], + "ollama": [f"{prefix}_API_KEY", "OLLAMA_API_KEY"], + "grok": [f"{prefix}_API_KEY", "GROQ_API_KEY", "GROQ_API_KEY_ALT"], + "deepseek": [f"{prefix}_API_KEY"], + "copilot": [f"{prefix}_API_KEY", "GITHUB_TOKEN", "REFINED_GITHUB_TOKEN"], + "openwebui": [f"{prefix}_API_KEY"], + } + model_aliases = { + "lmstudio": [f"{prefix}_MODEL", "ACTIVE_MODEL"], + "ollama": [f"{prefix}_MODEL"], + "grok": [f"{prefix}_MODEL"], + "deepseek": [f"{prefix}_MODEL"], + "copilot": [f"{prefix}_MODEL"], + "openwebui": [f"{prefix}_MODEL"], + } + + resolved_base_url = ( + base_url + or _first_present(env, base_url_aliases.get(normalized, [f"{prefix}_BASE_URL"])) + or defaults["base_url"] + ) + resolved_api_key = ( + api_key + or _first_present(env, api_key_aliases.get(normalized, [f"{prefix}_API_KEY"])) + or defaults["api_key"] + ) + resolved_model = ( + model + or _first_present(env, model_aliases.get(normalized, [f"{prefix}_MODEL"])) + or defaults["model"] + ) + + if normalized in {"lmstudio", "ollama", "openwebui", "grok", "deepseek", "copilot"}: + resolved_base_url = _normalize_base_url(resolved_base_url) + + return LLMConfig( + base_url=resolved_base_url, + api_key=resolved_api_key, + model=resolved_model, + provider=normalized, + ) + + +class LLMClient: + def __init__(self, config: Optional[LLMConfig] = None, provider: str = "lmstudio"): + self.config = config or LLMConfig(provider=provider) + self.provider = self.config.provider or provider + self._session = requests.Session() + if self.config.api_key and self.config.api_key != "not-needed": + self._session.headers.update( + {"Authorization": f"Bearer {self.config.api_key}"} + ) + + def complete(self, prompt: str, system: str = "", **kwargs) -> str: + messages = [] + if system: + messages.append({"role": "system", "content": system}) + messages.append({"role": "user", "content": prompt}) + + payload = { + "model": self.config.model, + "messages": messages, + "max_tokens": kwargs.get("max_tokens", self.config.max_tokens), + "temperature": kwargs.get("temperature", self.config.temperature), + } + + try: + response = self._session.post( + f"{self.config.base_url}/chat/completions", + json=payload, + timeout=self.config.timeout, + ) + response.raise_for_status() + data = response.json() + return data["choices"][0]["message"]["content"] + except Exception as e: + return f"Error: {str(e)}" + + def chat(self, messages: List[Dict[str, str]], **kwargs) -> str: + payload = { + "model": self.config.model, + "messages": messages, + "max_tokens": kwargs.get("max_tokens", self.config.max_tokens), + "temperature": kwargs.get("temperature", self.config.temperature), + } + + try: + response = self._session.post( + f"{self.config.base_url}/chat/completions", + json=payload, + timeout=self.config.timeout, + ) + response.raise_for_status() + data = response.json() + return data["choices"][0]["message"]["content"] + except Exception as e: + return f"Error: {str(e)}" + + @staticmethod + def for_lmstudio( + model: str = "local-model", base_url: str = "http://localhost:1234/v1" + ): + return LLMClient( + LLMConfig( + base_url=base_url, + model=model, + api_key="not-needed", + provider="lmstudio", + ), + "lmstudio", + ) + + @staticmethod + def for_ollama(model: str = "llama3", base_url: str = "http://localhost:11434/v1"): + return LLMClient( + LLMConfig( + base_url=base_url, + model=model, + api_key="not-needed", + provider="ollama", + ), + "ollama", + ) + + @staticmethod + def for_openwebui( + model: str = "llama3", base_url: str = "http://localhost:8080/v1" + ): + return LLMClient( + LLMConfig( + base_url=base_url, + model=model, + api_key="not-needed", + provider="openwebui", + ), + "openwebui", + ) + + @staticmethod + def for_provider( + provider: str, + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, + project_root: Path | None = None, + ): + config = resolve_llm_config( + provider, + model=model, + base_url=base_url, + api_key=api_key, + project_root=project_root, + ) + return LLMClient(config, provider=config.provider) diff --git a/src/services/local_logger.py b/src/services/local_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..80e0a3cc84702d5c65f0da20ad70b5662c43b61b --- /dev/null +++ b/src/services/local_logger.py @@ -0,0 +1,15 @@ +from __future__ import annotations +from pathlib import Path +import json +from datetime import datetime + +class LocalLogger: + def __init__(self, log_file: Path): + self.log_file = log_file + self.log_file.parent.mkdir(parents=True, exist_ok=True) + + def log(self, stage: str, level: str, message: str, **extra): + event = {"timestamp": datetime.utcnow().isoformat()+"Z", + "stage": stage, "level": level, "message": message, **extra} + with self.log_file.open("a", encoding="utf-8") as f: + f.write(json.dumps(event, ensure_ascii=False) + "\n") diff --git a/src/services/pdf_ops.py b/src/services/pdf_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..936ba426d08aa4735b8da8ef40bd11ce733b41e3 --- /dev/null +++ b/src/services/pdf_ops.py @@ -0,0 +1,22 @@ +from __future__ import annotations +from pathlib import Path +import fitz +from pypdf import PdfReader, PdfWriter + +def open_pdf(path: Path) -> fitz.Document: + return fitz.open(path) + +def save_pdf(doc: fitz.Document, out_path: Path) -> Path: + out_path.parent.mkdir(parents=True, exist_ok=True) + doc.save(out_path, garbage=4, deflate=True, clean=True) + return out_path + +def rebuild_pdf(input_path: Path, output_path: Path) -> Path: + reader = PdfReader(str(input_path)) + writer = PdfWriter() + for page in reader.pages: + writer.add_page(page) + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("wb") as f: + writer.write(f) + return output_path diff --git a/src/services/tesseract_service.py b/src/services/tesseract_service.py new file mode 100644 index 0000000000000000000000000000000000000000..7f7ff7d7e678b580f3fc5ebab35be91d57dea7b4 --- /dev/null +++ b/src/services/tesseract_service.py @@ -0,0 +1,25 @@ +from __future__ import annotations +import os, shutil, subprocess +from pathlib import Path +import pytesseract + +class TesseractService: + def __init__(self): + custom = os.environ.get("TESSERACT_CMD") + if custom: + pytesseract.pytesseract.tesseract_cmd = custom + + def ensure_available(self) -> str: + cmd = pytesseract.pytesseract.tesseract_cmd + resolved = shutil.which(cmd) if os.path.basename(cmd) == cmd else cmd + if not resolved or not Path(resolved).exists(): + raise RuntimeError( + "Tesseract not found. Install it locally and add to PATH, " + "or set the TESSERACT_CMD environment variable." + ) + return resolved + + def version(self) -> str: + cmd = self.ensure_available() + out = subprocess.check_output([cmd, "--version"], text=True, errors="ignore") + return out.splitlines()[0].strip() diff --git a/src/ui/__init__.py b/src/ui/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/ui/__init__.py @@ -0,0 +1 @@ + diff --git a/src/ui/__pycache__/__init__.cpython-314.pyc b/src/ui/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cf370680e6a03c21eaf36a27b30e556f0a7b3c6a Binary files /dev/null and b/src/ui/__pycache__/__init__.cpython-314.pyc differ diff --git a/src/ui/__pycache__/desktop_app.cpython-314.pyc b/src/ui/__pycache__/desktop_app.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ebf5fb1ac63010140d452f054d708181d2b46aa Binary files /dev/null and b/src/ui/__pycache__/desktop_app.cpython-314.pyc differ diff --git a/src/ui/desktop_app.py b/src/ui/desktop_app.py new file mode 100644 index 0000000000000000000000000000000000000000..7557749b3679e901469fd196b9f7a52acbc891e4 --- /dev/null +++ b/src/ui/desktop_app.py @@ -0,0 +1,128 @@ +from __future__ import annotations +import sys, os, threading, webbrowser, http.server, json +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from src.core.command_router import CommandRouter +from src.core.models import JobConfig, PDFOperation +from src.core.orchestrator import OfflinePDFOrchestrator + + +UI_HTML = Path(__file__).parent / "web" / "orchestrator.html" +PORT = 18472 +COMMAND_ROUTER = CommandRouter(project_root=Path(__file__).parent.parent.parent) + + +class _Handler(http.server.SimpleHTTPRequestHandler): + def __init__(self, *a, **kw): + super().__init__(*a, directory=str(UI_HTML.parent), **kw) + def log_message(self, *_): pass + + def do_POST(self): + if self.path == "/run": + length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(length) + payload = json.loads(body) + result = _run_job(payload) + resp = json.dumps(result).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(resp))) + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.wfile.write(resp) + elif self.path == "/team-plan": + length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(length) + payload = json.loads(body) + result = COMMAND_ROUTER.plan_command(payload) + resp = json.dumps(result).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(resp))) + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.wfile.write(resp) + else: + self.send_response(404); self.end_headers() + + def do_GET(self): + if self.path == "/team-config": + result = { + "success": True, + "team": COMMAND_ROUTER.describe_team(), + "available_pipelines": COMMAND_ROUTER.orchestrator.list_available_pipelines(), + } + resp = json.dumps(result).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(resp))) + self.send_header("Access-Control-Allow-Origin", "*") + self.end_headers() + self.wfile.write(resp) + return + super().do_GET() + + def do_OPTIONS(self): + self.send_response(204) + self.send_header("Access-Control-Allow-Origin", "*") + self.send_header("Access-Control-Allow-Methods", "GET,POST,OPTIONS") + self.send_header("Access-Control-Allow-Headers", "Content-Type") + self.end_headers() + + +def _run_job(payload: dict) -> dict: + try: + ops = [] + for item in payload.get("operations", []): + ops.append(PDFOperation(**item)) + input_pdf = Path(payload.get("input_pdf", "")) + output_dir = Path(payload.get("output_dir", Path.home() / "pdf-output")) + if not input_pdf.exists(): + return {"success": False, "error": f"Input PDF not found: {input_pdf}"} + cfg = JobConfig( + input_pdf = input_pdf, + output_dir = output_dir, + run_ocr = payload.get("force_ocr", False), + force_ocr = payload.get("force_ocr", False), + ocr_language = payload.get("ocr_language", "eng"), + operations = ops, + ) + orc = OfflinePDFOrchestrator(output_dir / "_jobs") + result = orc.run(cfg) + return { + "success": result.success, + "job_id": result.job_id, + "final_pdf": str(result.final_pdf) if result.final_pdf else "", + "log_file": str(result.log_file), + "stages": result.stages, + "errors": result.errors, + } + except Exception as exc: + return {"success": False, "error": str(exc)} + + +def _start_server(): + server = http.server.HTTPServer(("127.0.0.1", PORT), _Handler) + server.serve_forever() + + +def launch(): + t = threading.Thread(target=_start_server, daemon=True) + t.start() + url = f"http://127.0.0.1:{PORT}/orchestrator.html" + print(f"Opening PDF Orchestrator UI at {url}") + webbrowser.open(url) + + # Keep process alive + import time + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + print("Shutting down.") + + +if __name__ == "__main__": + launch() diff --git a/src/ui/web/orchestrator.html b/src/ui/web/orchestrator.html new file mode 100644 index 0000000000000000000000000000000000000000..8b792acf1ef088c8e98b6ea1d737339bf1a3629b --- /dev/null +++ b/src/ui/web/orchestrator.html @@ -0,0 +1,1256 @@ + + + + + +PDF AI Orchestrator + + + + + + + + +
+ +
+ + + +
+
+ + +
+
+ + +
+ + + + + +
+
+
+
+
+
Agent Pipeline
+
Drag agents from the left to build your processing flow. Configure tasks per agent.
+
+
+
+
+
+
⬡
+

Pipeline is empty

+

Drag agents from the left panel to start building

+
+
+
+ + +
+ + + +
+ + + + + + + diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1 @@ + diff --git a/src/utils/__pycache__/__init__.cpython-314.pyc b/src/utils/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ddd73ba32b695e5ff6b063656f4053d217d1cae Binary files /dev/null and b/src/utils/__pycache__/__init__.cpython-314.pyc differ diff --git a/src/utils/__pycache__/file_utils.cpython-314.pyc b/src/utils/__pycache__/file_utils.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b161686392537849425179ace2f405a4cdec795e Binary files /dev/null and b/src/utils/__pycache__/file_utils.cpython-314.pyc differ diff --git a/src/utils/file_utils.py b/src/utils/file_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3b3dc4ac46cee8b1534ed1424262bbfb6e557fcf --- /dev/null +++ b/src/utils/file_utils.py @@ -0,0 +1,11 @@ +from __future__ import annotations +from pathlib import Path +from datetime import datetime + +def ensure_dir(path: Path) -> Path: + path.mkdir(parents=True, exist_ok=True) + return path + +def stage_file(output_dir: Path, prefix: str, suffix: str = ".pdf") -> Path: + ts = datetime.now().strftime("%Y%m%d-%H%M%S-%f") + return output_dir / f"{prefix}-{ts}{suffix}" diff --git a/src/utils/system_checks.py b/src/utils/system_checks.py new file mode 100644 index 0000000000000000000000000000000000000000..80da53c77ecd31f29a6bcfe71c153f0df591d75b --- /dev/null +++ b/src/utils/system_checks.py @@ -0,0 +1,12 @@ +from __future__ import annotations +from src.services.tesseract_service import TesseractService + +def environment_report() -> dict: + svc = TesseractService() + report: dict = {"tesseract": None, "status": "ok"} + try: + report["tesseract"] = svc.version() + except Exception as exc: + report["status"] = "warning: " + str(exc) + report["tesseract"] = "NOT FOUND" + return report