lexoid 0.1.6__py3-none-any.whl

lexoid/api.py

@@ -0,0 +1,200 @@
+import os
+import re
+import tempfile
+from concurrent.futures import ProcessPoolExecutor
+from enum import Enum
+from glob import glob
+from time import time
+from typing import Union, Dict, List
+
+from loguru import logger
+
+from lexoid.core.parse_type.llm_parser import parse_llm_doc
+from lexoid.core.parse_type.static_parser import parse_static_doc
+from lexoid.core.utils import (
+    convert_to_pdf,
+    download_file,
+    is_supported_url_file_type,
+    is_supported_file_type,
+    recursive_read_html,
+    router,
+    split_pdf,
+)
+
+
+class ParserType(Enum):
+    LLM_PARSE = "LLM_PARSE"
+    STATIC_PARSE = "STATIC_PARSE"
+    AUTO = "AUTO"
+
+
+def parse_chunk(
+    path: str, parser_type: ParserType, raw: bool, **kwargs
+) -> List[Dict] | str:
+    """
+    Parses a file using the specified parser type.
+
+    Args:
+        path (str): The file path or URL.
+        parser_type (ParserType): The type of parser to use (LLM_PARSE, STATIC_PARSE, or AUTO).
+        raw (bool): Whether to return raw text or structured data.
+        **kwargs: Additional arguments for the parser.
+
+    Returns:
+        List[Dict] | str: Parsed document data as a list of dictionaries or raw text.
+    """
+    if parser_type == ParserType.AUTO:
+        parser_type = ParserType[router(path)]
+        logger.debug(f"Auto-detected parser type: {parser_type}")
+
+    kwargs["start"] = (
+        int(os.path.basename(path).split("_")[1]) - 1 if kwargs.get("split") else 0
+    )
+    if parser_type == ParserType.STATIC_PARSE:
+        logger.debug("Using static parser")
+        return parse_static_doc(path, raw, **kwargs)
+    else:
+        logger.debug("Using LLM parser")
+        return parse_llm_doc(path, raw, **kwargs)
+
+
+def parse_chunk_list(
+    file_paths: List[str], parser_type: ParserType, raw: bool, kwargs: Dict
+) -> List[Dict | str]:
+    """
+    Parses a list of files using the specified parser type.
+
+    Args:
+        file_paths (list): List of file paths.
+        parser_type (ParserType): The type of parser to use.
+        raw (bool): Whether to return raw text or structured data.
+        kwargs (dict): Additional arguments for the parser.
+
+    Returns:
+        List[Dict | str]: List of parsed documents with raw text and/or metadata.
+    """
+    local_docs = []
+    for file_path in file_paths:
+        result = parse_chunk(file_path, parser_type, raw, **kwargs)
+        if isinstance(result, list):
+            local_docs.extend(result)
+        else:
+            local_docs.append(result.replace("<page break>", "\n\n"))
+    return local_docs
+
+
+def parse(
+    path: str,
+    parser_type: Union[str, ParserType] = "LLM_PARSE",
+    raw: bool = False,
+    pages_per_split: int = 4,
+    max_processes: int = 4,
+    **kwargs,
+) -> Union[List[Dict], str]:
+    """
+    Parses a document or URL, optionally splitting it into chunks and using multiprocessing.
+
+    Args:
+        path (str): The file path or URL.
+        parser_type (Union[str, ParserType], optional): The type of parser to use ("LLM_PARSE", "STATIC_PARSE", or "AUTO"). Defaults to "LLM_PARSE".
+        raw (bool, optional): Whether to return raw text or structured data. Defaults to False.
+        pages_per_split (int, optional): Number of pages per split for chunking. Defaults to 4.
+        max_processes (int, optional): Maximum number of processes for parallel processing. Defaults to 4.
+        **kwargs: Additional arguments for the parser.
+
+    Returns:
+        Union[List[Dict], str]: Parsed document data as a list of dictionaries or raw text.
+    """
+    kwargs["title"] = os.path.basename(path)
+    kwargs["pages_per_split_"] = pages_per_split
+    as_pdf = kwargs.get("as_pdf", False)
+    depth = kwargs.get("depth", 1)
+    if type(parser_type) == str:
+        parser_type = ParserType[parser_type]
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        if (
+            path.lower().endswith((".doc", ".docx"))
+            and parser_type != ParserType.STATIC_PARSE
+        ):
+            as_pdf = True
+
+        if path.startswith(("http://", "https://")):
+            download_dir = os.path.join(temp_dir, "downloads/")
+            os.makedirs(download_dir, exist_ok=True)
+            if is_supported_url_file_type(path):
+                path = download_file(path, download_dir)
+            elif as_pdf:
+                pdf_path = os.path.join(download_dir, f"webpage_{int(time())}.pdf")
+                path = convert_to_pdf(path, pdf_path)
+            else:
+                return recursive_read_html(path, depth, raw)
+
+        assert is_supported_file_type(
+            path
+        ), f"Unsupported file type {os.path.splitext(path)[1]}"
+
+        if as_pdf and not path.lower().endswith(".pdf"):
+            pdf_path = os.path.join(temp_dir, "converted.pdf")
+            path = convert_to_pdf(path, pdf_path)
+
+        if not path.lower().endswith(".pdf") or parser_type == ParserType.STATIC_PARSE:
+            kwargs["split"] = False
+            all_docs = parse_chunk(path, parser_type, raw, **kwargs)
+            if raw:
+                all_docs = [all_docs]
+        else:
+            kwargs["split"] = True
+            split_dir = os.path.join(temp_dir, "splits/")
+            os.makedirs(split_dir, exist_ok=True)
+            split_pdf(path, split_dir, pages_per_split)
+            split_files = sorted(glob(os.path.join(split_dir, "*.pdf")))
+
+            chunk_size = max(1, len(split_files) // max_processes)
+            file_chunks = [
+                split_files[i : i + chunk_size]
+                for i in range(0, len(split_files), chunk_size)
+            ]
+
+            process_args = [(chunk, parser_type, raw, kwargs) for chunk in file_chunks]
+
+            if max_processes == 1 or len(file_chunks) == 1:
+                all_docs = [parse_chunk_list(*args) for args in process_args]
+            else:
+                with ProcessPoolExecutor(max_workers=max_processes) as executor:
+                    all_docs = list(executor.map(parse_chunk_list, *zip(*process_args)))
+
+            all_docs = [item for sublist in all_docs for item in sublist]
+
+    if depth > 1:
+        new_docs = all_docs.copy()
+        for doc in all_docs:
+            urls = re.findall(
+                r'https?://[^\s<>"\']+|www\.[^\s<>"\']+(?:\.[^\s<>"\']+)*',
+                doc if raw else doc["content"],
+            )
+            for url in urls:
+                if "](" in url:
+                    url = url.split("](")[-1]
+                logger.debug(f"Reading content from {url}")
+                if not url.startswith("http"):
+                    url = "https://" + url
+
+                kwargs_cp = kwargs.copy()
+                kwargs_cp["depth"] = depth - 1
+                res = parse(
+                    url,
+                    parser_type=parser_type,
+                    raw=raw,
+                    pages_per_split=pages_per_split,
+                    max_processes=max_processes,
+                    **kwargs_cp,
+                )
+
+                if raw:
+                    new_docs.append(res)
+                else:
+                    new_docs.extend(res)
+        all_docs = new_docs
+
+    return "\n".join(all_docs) if raw else all_docs

lexoid/core/parse_type/llm_parser.py

@@ -0,0 +1,200 @@
+import base64
+import io
+import mimetypes
+import os
+from typing import Dict, List
+
+import pypdfium2 as pdfium
+import requests
+from lexoid.core.prompt_templates import (
+    INSTRUCTIONS_ADD_PG_BREAK,
+    OPENAI_USER_PROMPT,
+    PARSER_PROMPT,
+)
+from lexoid.core.utils import convert_image_to_pdf
+from loguru import logger
+from openai import OpenAI
+
+
+def parse_llm_doc(path: str, raw: bool, **kwargs) -> List[Dict] | str:
+    if "model" not in kwargs:
+        kwargs["model"] = "gemini-1.5-flash"
+    model = kwargs.get("model")
+    if model.startswith("gemini"):
+        return parse_with_gemini(path, raw, **kwargs)
+    elif model.startswith("gpt"):
+        return parse_with_gpt(path, raw, **kwargs)
+    else:
+        raise ValueError(f"Unsupported model: {model}")
+
+
+def parse_with_gemini(path: str, raw: bool, **kwargs) -> List[Dict] | str:
+    api_key = os.environ.get("GOOGLE_API_KEY")
+    if not api_key:
+        raise ValueError("GOOGLE_API_KEY environment variable is not set")
+
+    url = f"https://generativelanguage.googleapis.com/v1beta/models/{kwargs['model']}:generateContent?key={api_key}"
+
+    # Check if the file is an image and convert to PDF if necessary
+    mime_type, _ = mimetypes.guess_type(path)
+    if mime_type and mime_type.startswith("image"):
+        pdf_content = convert_image_to_pdf(path)
+        mime_type = "application/pdf"
+        base64_file = base64.b64encode(pdf_content).decode("utf-8")
+    else:
+        with open(path, "rb") as file:
+            file_content = file.read()
+        base64_file = base64.b64encode(file_content).decode("utf-8")
+
+    # Ideally, we do this ourselves. But, for now this might be a good enough.
+    custom_instruction = f"""- Total number of pages: {kwargs["pages_per_split_"]}. {INSTRUCTIONS_ADD_PG_BREAK}"""
+    if kwargs["pages_per_split_"] == 1:
+        custom_instruction = ""
+
+    payload = {
+        "contents": [
+            {
+                "parts": [
+                    {
+                        "text": PARSER_PROMPT.format(
+                            custom_instructions=custom_instruction
+                        )
+                    },
+                    {"inline_data": {"mime_type": mime_type, "data": base64_file}},
+                ]
+            }
+        ],
+        "generationConfig": {
+            "temperature": kwargs.get("temperature", 0.7),
+        },
+    }
+
+    headers = {"Content-Type": "application/json"}
+
+    response = requests.post(url, json=payload, headers=headers)
+    response.raise_for_status()
+
+    result = response.json()
+
+    raw_text = "".join(
+        part["text"]
+        for candidate in result.get("candidates", [])
+        for part in candidate.get("content", {}).get("parts", [])
+        if "text" in part
+    )
+
+    result = ""
+    if "<output>" in raw_text:
+        result = raw_text.split("<output>")[1].strip()
+    if "</output>" in result:
+        result = result.split("</output>")[0].strip()
+
+    if raw:
+        return result
+
+    return [
+        {
+            "metadata": {
+                "title": kwargs["title"],
+                "page": kwargs.get("start", 0) + page_no,
+            },
+            "content": page,
+        }
+        for page_no, page in enumerate(result.split("<page-break>"), start=1)
+        if page.strip()
+    ]
+
+
+def convert_pdf_page_to_base64(
+    pdf_document: pdfium.PdfDocument, page_number: int
+) -> str:
+    """Convert a PDF page to a base64-encoded PNG string."""
+    page = pdf_document[page_number]
+    # Render with 4x scaling for better quality
+    pil_image = page.render(scale=4).to_pil()
+
+    # Convert to base64
+    img_byte_arr = io.BytesIO()
+    pil_image.save(img_byte_arr, format="PNG")
+    img_byte_arr.seek(0)
+    return base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
+
+
+def parse_with_gpt(path: str, raw: bool, **kwargs) -> List[Dict] | str:
+    client = OpenAI()
+
+    # Handle different input types
+    mime_type, _ = mimetypes.guess_type(path)
+    if mime_type and mime_type.startswith("image"):
+        # Single image processing
+        with open(path, "rb") as img_file:
+            image_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+            images = [(0, image_base64)]
+    else:
+        # PDF processing
+        pdf_document = pdfium.PdfDocument(path)
+        images = [
+            (page_num, convert_pdf_page_to_base64(pdf_document, page_num))
+            for page_num in range(len(pdf_document))
+        ]
+
+    # Process each page/image
+    all_results = []
+    for page_num, image_base64 in images:
+        messages = [
+            {
+                "role": "system",
+                "content": PARSER_PROMPT,
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": f"{OPENAI_USER_PROMPT} (Page {page_num + 1})",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{image_base64}"},
+                    },
+                ],
+            },
+        ]
+
+        # Get completion from GPT-4 Vision
+        response = client.chat.completions.create(
+            model=kwargs["model"],
+            temperature=kwargs.get("temperature", 0.7),
+            messages=messages,
+        )
+
+        # Extract the response text
+        page_text = response.choices[0].message.content
+        if kwargs.get("verbose", None):
+            logger.debug(f"Page {page_num + 1} response: {page_text}")
+        result = ""
+        if "<output>" in page_text:
+            result = page_text.split("<output>")[1].strip()
+        if "</output>" in result:
+            result = result.split("</output>")[0].strip()
+        all_results.append((page_num, result))
+
+    # Sort results by page number and combine
+    all_results.sort(key=lambda x: x[0])
+    all_texts = [text for _, text in all_results]
+    combined_text = "<page-break>".join(all_texts)
+
+    if raw:
+        return combined_text
+
+    return [
+        {
+            "metadata": {
+                "title": kwargs["title"],
+                "page": kwargs.get("start", 0) + page_no,
+            },
+            "content": page,
+        }
+        for page_no, page in enumerate(all_texts, start=1)
+        if page.strip()
+    ]