docai-toolkit 0.1.0__py3-none-any.whl

docai_toolkit/__init__.py

@@ -0,0 +1,3 @@
+"""DocAI Toolkit package (renamed to avoid PyPI conflicts)."""
+
+__all__ = ["config"]

docai_toolkit/config.py

@@ -0,0 +1,79 @@
+import json
+import os
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+from typing import Optional
+
+CONFIG_PATH = Path.home() / ".docai" / "config.json"
+
+
+@dataclass
+class OcrConfig:
+    provider: str = "deepseek"  # deepseek | tesseract
+    api_key: Optional[str] = None
+    model: Optional[str] = None  # provider-specific
+    endpoint: Optional[str] = None  # user-defined OCR API endpoint
+
+
+@dataclass
+class EmbeddingConfig:
+    backend: str = "sentence-transformers"  # sentence-transformers | huggingface-hub
+    model: str = "all-mpnet-base-v2"
+    device: str = "auto"
+    endpoint: Optional[str] = None  # user-defined embedding API endpoint
+    api_key: Optional[str] = None  # for hosted endpoints
+
+
+@dataclass
+class LlmConfig:
+    backend: str = "huggingface-hub"  # huggingface-hub | local-gguf | openai-compatible
+    model: str = "mistralai/Mistral-7B-Instruct-v0.1"
+    api_key: Optional[str] = None
+    max_new_tokens: int = 256
+    endpoint: Optional[str] = None  # user-defined generation endpoint
+
+
+@dataclass
+class AppConfig:
+    output_dir: Path = field(default_factory=lambda: Path("./outputs"))
+    ocr: OcrConfig = field(default_factory=OcrConfig)
+    embeddings: EmbeddingConfig = field(default_factory=EmbeddingConfig)
+    llm: LlmConfig = field(default_factory=LlmConfig)
+
+    @classmethod
+    def from_env(cls) -> "AppConfig":
+        cfg = cls.load_from_file(CONFIG_PATH) or cls()
+        hf_token = (
+            os.getenv("HF_TOKEN")
+            or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+            or os.getenv("DOC_AI_HF_TOKEN")
+        )
+        if hf_token:
+            cfg.llm.api_key = hf_token
+            cfg.embeddings.api_key = hf_token
+
+        output_dir_env = os.getenv("DOC_AI_OUTPUT_DIR")
+        if output_dir_env:
+            cfg.output_dir = Path(output_dir_env)
+        return cfg
+
+    @classmethod
+    def load_from_file(cls, path: Path | None) -> "AppConfig | None":
+        if not path:
+            return None
+        if not path.exists():
+            return None
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return cls(
+            output_dir=Path(data.get("output_dir", "./outputs")),
+            ocr=OcrConfig(**data.get("ocr", {})),
+            embeddings=EmbeddingConfig(**data.get("embeddings", {})),
+            llm=LlmConfig(**data.get("llm", {})),
+        )
+
+    def save(self, path: Path | None = None) -> None:
+        path = path or CONFIG_PATH
+        path.parent.mkdir(parents=True, exist_ok=True)
+        data = asdict(self)
+        data["output_dir"] = str(self.output_dir)
+        path.write_text(json.dumps(data, indent=2), encoding="utf-8")

docai_toolkit/hf_client.py

@@ -0,0 +1,43 @@
+"""Minimal Hugging Face / custom endpoint client."""
+
+import json
+import urllib.error
+import urllib.request
+from typing import Any, Dict, Optional
+
+
+class HuggingFaceClient:
+    def __init__(self, api_token: str | None, default_endpoint: Optional[str] = None, timeout: float = 30.0) -> None:
+        self.api_token = api_token
+        self.default_endpoint = default_endpoint
+        self.timeout = timeout
+
+    def post_json(
+        self,
+        payload: Dict[str, Any] | bytes,
+        endpoint: Optional[str] = None,
+        content_type: str = "application/json",
+    ) -> Dict[str, Any] | Any:
+        url = endpoint or self.default_endpoint
+        if not url:
+            raise ValueError("Endpoint is required for HuggingFaceClient.")
+
+        if isinstance(payload, bytes):
+            data = payload
+        else:
+            data = json.dumps(payload).encode("utf-8")
+
+        req = urllib.request.Request(url, data=data)
+        req.add_header("Content-Type", content_type)
+        if self.api_token:
+            req.add_header("Authorization", f"Bearer {self.api_token}")
+
+        try:
+            with urllib.request.urlopen(req, timeout=self.timeout) as resp:  # nosec B310 - user-configured endpoints
+                body = resp.read().decode("utf-8")
+                try:
+                    return json.loads(body)
+                except json.JSONDecodeError:
+                    return body
+        except urllib.error.URLError as exc:
+            raise RuntimeError(f"Request to {url} failed: {exc}") from exc

docai_toolkit/ocr/__init__.py

@@ -0,0 +1,4 @@
+from .clients import OcrClient, DeepSeekOcrClient, RemoteOcrClient, TesseractOcrClient
+from .pipeline import run_ocr_to_markdown
+
+__all__ = ["OcrClient", "DeepSeekOcrClient", "RemoteOcrClient", "TesseractOcrClient", "run_ocr_to_markdown"]

docai_toolkit/ocr/clients.py

@@ -0,0 +1,83 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List
+
+from docai_toolkit.hf_client import HuggingFaceClient
+
+
+@dataclass
+class PageResult:
+    page_number: int
+    text: str
+
+
+class OcrClient(ABC):
+    @abstractmethod
+    def recognize(self, pdf_path: Path) -> List[PageResult]:
+        raise NotImplementedError
+
+
+class DeepSeekOcrClient(OcrClient):
+    """Placeholder for DeepSeek OCR API integration."""
+
+    def __init__(self, api_key: str, model: str | None = None) -> None:
+        self.api_key = api_key
+        self.model = model or "deepseek-ocr-default"
+
+    def recognize(self, pdf_path: Path) -> List[PageResult]:
+        raise NotImplementedError("DeepSeek OCR integration not yet implemented.")
+
+
+class RemoteOcrClient(OcrClient):
+    """Generic OCR via a Hugging Face Inference or custom endpoint."""
+
+    def __init__(self, api_key: str | None, endpoint: str, model: str | None = None) -> None:
+        self.client = HuggingFaceClient(api_key, default_endpoint=endpoint)
+        self.model = model
+
+    def recognize(self, pdf_path: Path) -> List[PageResult]:
+        with open(pdf_path, "rb") as handle:
+            pdf_bytes = handle.read()
+
+        payload = pdf_bytes
+        response = self.client.post_json(payload, content_type="application/pdf")
+
+        # Accept a few possible response shapes.
+        if isinstance(response, str):
+            pages = [response]
+        elif isinstance(response, list):
+            pages = [str(item) for item in response]
+        elif isinstance(response, dict):
+            if "pages" in response:
+                pages = [p.get("text", "") if isinstance(p, dict) else str(p) for p in response["pages"]]
+            elif "text" in response:
+                pages = [str(response["text"])]
+            else:
+                pages = [str(response)]
+        else:
+            pages = [""]
+
+        return [PageResult(page_number=i + 1, text=pages[i]) for i in range(len(pages))]
+
+
+class TesseractOcrClient(OcrClient):
+    """Local fallback using pytesseract + pdf2image."""
+
+    def __init__(self) -> None:
+        try:
+            import pytesseract  # type: ignore
+            from pdf2image import convert_from_path  # type: ignore
+        except ImportError as exc:  # pragma: no cover - optional path
+            raise RuntimeError("pytesseract and pdf2image are required for TesseractOcrClient.") from exc
+
+        self._pytesseract = pytesseract
+        self._convert_from_path = convert_from_path
+
+    def recognize(self, pdf_path: Path) -> List[PageResult]:
+        images = self._convert_from_path(str(pdf_path))
+        results: List[PageResult] = []
+        for idx, image in enumerate(images):
+            text = self._pytesseract.image_to_string(image)
+            results.append(PageResult(page_number=idx + 1, text=text))
+        return results

docai_toolkit/ocr/pipeline.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+from typing import Iterable
+
+from .clients import OcrClient, PageResult
+
+
+def run_ocr_to_markdown(pdf_path: Path, output_dir: Path, client: OcrClient) -> Path:
+    """Run OCR on a PDF and save combined Markdown output."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+    md_path = output_dir / f"{pdf_path.stem}.md"
+    counter = 1
+    while md_path.exists():
+        md_path = output_dir / f"{pdf_path.stem}-{counter}.md"
+        counter += 1
+
+    pages: Iterable[PageResult] = client.recognize(pdf_path)
+    lines: list[str] = []
+    for page in pages:
+        lines.append(f"# Page {page.page_number}")
+        lines.append("")
+        lines.append(page.text.strip())
+        lines.append("\n")
+
+    md_path.write_text("\n".join(lines).strip() + "\n", encoding="utf-8")
+    return md_path

docai_toolkit/rag/__init__.py

@@ -0,0 +1,4 @@
+from .index import build_index_from_markdown, load_index
+from .chat import chat_over_corpus
+
+__all__ = ["build_index_from_markdown", "load_index", "chat_over_corpus"]

docai_toolkit/rag/chat.py

@@ -0,0 +1,55 @@
+from typing import List, Optional
+
+from docai_toolkit.hf_client import HuggingFaceClient
+
+try:
+    from langchain_community.llms import HuggingFacePipeline
+    from transformers import pipeline
+except ImportError as _chat_import_error:  # pragma: no cover - optional dependency
+    HuggingFacePipeline = None  # type: ignore[assignment]
+    pipeline = None  # type: ignore[assignment]
+    _CHAT_IMPORT_ERROR = _chat_import_error
+else:
+    _CHAT_IMPORT_ERROR = None
+
+
+def chat_over_corpus(
+    db,
+    query: str,
+    model_id: str = "mistralai/Mistral-7B-Instruct-v0.1",
+    endpoint: Optional[str] = None,
+    api_key: Optional[str] = None,
+    max_new_tokens: int = 256,
+) -> str:
+    """Simple retrieve-then-generate over a FAISS db."""
+    docs = db.similarity_search(query, k=4)
+    context_blocks: List[str] = []
+    for doc in docs:
+        src = doc.metadata.get("source", "unknown")
+        context_blocks.append(f"Source: {src}\n{doc.page_content}")
+    context = "\n\n".join(context_blocks)
+
+    prompt = f"Answer the question using only the provided context.\n\nContext:\n{context}\n\nQuestion:\n{query}\n\nAnswer:"
+
+    if endpoint:
+        client = HuggingFaceClient(api_key, default_endpoint=endpoint)
+        resp = client.post_json(
+            {
+                "inputs": prompt,
+                "parameters": {"max_new_tokens": max_new_tokens},
+            }
+        )
+        if isinstance(resp, list) and resp and isinstance(resp[0], dict) and "generated_text" in resp[0]:
+            return resp[0]["generated_text"]
+        if isinstance(resp, dict) and "generated_text" in resp:
+            return resp["generated_text"]
+        if isinstance(resp, str):
+            return resp
+        raise ValueError("Unexpected response from generation endpoint.")
+
+    if pipeline is None or HuggingFacePipeline is None:
+        raise RuntimeError("transformers/langchain-community required for local HF generation.")
+
+    pipe = pipeline("text-generation", model=model_id, device_map="auto")
+    llm = HuggingFacePipeline(pipeline=pipe)
+    return llm(prompt)

docai_toolkit/rag/index.py

@@ -0,0 +1,111 @@
+from pathlib import Path
+from typing import Iterable, List, Optional
+
+try:
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
+    from langchain_community.vectorstores import FAISS
+    from langchain.embeddings.base import Embeddings
+    from docai_toolkit.hf_client import HuggingFaceClient
+except ImportError as _langchain_exc:  # pragma: no cover - optional dependency
+    RecursiveCharacterTextSplitter = None  # type: ignore[assignment]
+    FAISS = None  # type: ignore[assignment]
+    Embeddings = object  # type: ignore[assignment]
+    HuggingFaceClient = None  # type: ignore[assignment]
+    _LANGCHAIN_IMPORT_ERROR = _langchain_exc
+else:
+    _LANGCHAIN_IMPORT_ERROR = None
+
+try:
+    from sentence_transformers import SentenceTransformer  # type: ignore
+except ImportError as exc:  # pragma: no cover - optional dependency
+    SentenceTransformer = None  # type: ignore[misc,assignment]
+    _IMPORT_ERROR = exc
+else:
+    _IMPORT_ERROR = None
+
+
+def _load_embedding_model(model_name: str):
+    if SentenceTransformer is None:
+        raise RuntimeError("sentence-transformers not installed") from _IMPORT_ERROR
+    return SentenceTransformer(model_name)
+
+
+class RemoteEmbeddings(Embeddings):
+    """Call a remote embedding endpoint that accepts JSON {\"inputs\": text}."""
+
+    def __init__(self, endpoint: str, api_key: Optional[str] = None):
+        self.client = HuggingFaceClient(api_key, default_endpoint=endpoint)
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        try:
+            resp = self.client.post_json({"inputs": texts})
+            return self._extract_batch(resp, len(texts))
+        except Exception:
+            vectors: List[List[float]] = []
+            for text in texts:
+                single = self.client.post_json({"inputs": text})
+                vectors.append(self._extract_vector(single))
+            return vectors
+
+    def embed_query(self, text: str) -> List[float]:
+        resp = self.client.post_json({"inputs": text})
+        return self._extract_vector(resp)
+
+    @staticmethod
+    def _extract_vector(response):
+        if isinstance(response, list) and response and isinstance(response[0], list):
+            return response[0]
+        if isinstance(response, list):
+            return response
+        raise ValueError("Unexpected embedding response format.")
+
+    @staticmethod
+    def _extract_batch(response, expected: int) -> List[List[float]]:
+        if isinstance(response, list) and response and isinstance(response[0], list):
+            return response
+        if isinstance(response, list) and len(response) == expected and isinstance(response[0], dict) and "embedding" in response[0]:
+            return [item["embedding"] for item in response]
+        raise ValueError("Unexpected batch embedding response format.")
+
+
+def build_index_from_markdown(
+    markdown_files: Iterable[Path],
+    embedding_model: str = "all-mpnet-base-v2",
+    chunk_size: int = 1000,
+    chunk_overlap: int = 200,
+    persist_path: Optional[Path] = None,
+    embedding_endpoint: Optional[str] = None,
+    embedding_api_key: Optional[str] = None,
+):
+    if RecursiveCharacterTextSplitter is None or FAISS is None:
+        raise RuntimeError("langchain is required for RAG. Install langchain and langchain-community.")
+
+    texts: List[str] = []
+    metadatas: List[dict] = []
+
+    for path in markdown_files:
+        content = path.read_text(encoding="utf-8")
+        texts.append(content)
+        metadatas.append({"source": str(path)})
+
+    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    docs = splitter.create_documents(texts, metadatas=metadatas)
+
+    if embedding_endpoint:
+        embeddings = RemoteEmbeddings(embedding_endpoint, api_key=embedding_api_key)
+    else:
+        embeddings = _load_embedding_model(embedding_model)
+    db = FAISS.from_documents(docs, embeddings)
+    if persist_path:
+        persist_path.parent.mkdir(parents=True, exist_ok=True)
+        db.save_local(str(persist_path))
+    return db
+
+
+def load_index(persist_path: Path):
+    if not persist_path.exists():
+        raise FileNotFoundError(f"Persisted index not found at {persist_path}")
+    if FAISS is None:
+        raise RuntimeError("langchain-community is required to load indexes.")
+    embeddings = None  # embeddings are restored from disk
+    return FAISS.load_local(str(persist_path), embeddings, allow_dangerous_deserialization=False)