chatterer 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

chatterer/examples/transcription_api.py

@@ -0,0 +1,127 @@
+# pyright: reportUnknownVariableType=false, reportUnknownMemberType=false, reportArgumentType=false, reportMissingTypeStubs=false
+
+from io import BytesIO
+from pathlib import Path
+from typing import cast
+
+from openai import OpenAI
+from pydub import AudioSegment
+from spargear import ArgumentSpec, BaseArguments
+
+# Maximum chunk length in seconds
+MAX_CHUNK_DURATION = 600
+
+
+class TranscriptionApiArguments(BaseArguments):
+    in_path = ArgumentSpec(
+        ["in-path"],
+        type=Path,
+        help="The audio file to transcribe.",
+    )
+    out_path = ArgumentSpec(
+        ["--out-path"],
+        type=Path,
+        default=None,
+        help="Path to save the transcription output.",
+    )
+    model: ArgumentSpec[str] = ArgumentSpec(
+        ["--model"],
+        default="gpt-4o-transcribe",
+        help="The model to use for transcription.",
+    )
+    api_key: ArgumentSpec[str] = ArgumentSpec(
+        ["--api-key"],
+        default=None,
+        help="The API key for authentication.",
+    )
+    base_url: ArgumentSpec[str] = ArgumentSpec(
+        ["--base-url"],
+        default="https://api.openai.com/v1",
+        help="The base URL for the API.",
+    )
+
+    def run(self) -> None:
+        audio_path = self.in_path.unwrap()
+        model = self.model.unwrap()
+
+        client = OpenAI(api_key=self.api_key.value, base_url=self.base_url.value)
+
+        audio = load_audio_segment(audio_path)
+
+        segments = split_audio(audio, MAX_CHUNK_DURATION)
+        print(f"[i] Audio duration: {len(audio) / 1000:.1f}s; splitting into {len(segments)} segment(s)")
+
+        transcripts: list[str] = []
+        for idx, seg in enumerate(segments, start=1):
+            print(f"[i] Transcribing segment {idx}/{len(segments)}...")
+            transcripts.append(transcribe_segment(seg, client, model))
+
+        full_transcript = "\n\n".join(transcripts)
+        output_path: Path = self.out_path.value or audio_path.with_suffix(".txt")
+        output_path.write_text(full_transcript, encoding="utf-8")
+        print(f"[✓] Transcription saved to: {output_path}")
+
+
+def load_audio_segment(file_path: Path) -> AudioSegment:
+    """
+    Load an audio file as an AudioSegment. Convert to mp3 format in-memory if needed.
+    """
+    ext = file_path.suffix.lower()[1:]
+    audio = AudioSegment.from_file(file_path.as_posix(), format=ext if ext != "mp3" else None)
+    if ext != "mp3":
+        buffer = BytesIO()
+        audio.export(buffer, format="mp3")
+        buffer.seek(0)
+        audio = AudioSegment.from_file(buffer, format="mp3")
+    return audio
+
+
+def split_audio(audio: AudioSegment, max_duration_s: int) -> list[AudioSegment]:
+    """
+    Split the AudioSegment into chunks no longer than max_duration_s seconds.
+    """
+    chunk_length_ms = (max_duration_s - 1) * 1000
+    duration_ms = len(audio)
+    segments: list[AudioSegment] = []
+    segment_idx: int = 0
+    for start_ms in range(0, duration_ms, chunk_length_ms):
+        end_ms = min(start_ms + chunk_length_ms, duration_ms)
+        segment = cast(AudioSegment, audio[start_ms:end_ms])
+        segments.append(segment)
+        # with open(f"segment_{segment_idx}.mp3", "wb") as f:
+        #     segment.export(f, format="mp3")
+        segment_idx += 1
+    return segments
+
+
+def transcribe_segment(segment: AudioSegment, client: OpenAI, model: str) -> str:
+    """
+    Transcribe a single AudioSegment chunk and return its text.
+    """
+    buffer = BytesIO()
+    segment.export(buffer, format="mp3")
+    buffer.seek(0)
+    mp3_bytes = buffer.read()
+    response = client.audio.transcriptions.create(
+        model=model,
+        prompt="Transcribe whole text from audio.",
+        file=("audio.mp3", mp3_bytes),
+        response_format="text",
+        stream=True,
+    )
+    for res in response:
+        if res.type == "transcript.text.delta":
+            print(res.delta, end="", flush=True)
+        if res.type == "transcript.text.done":
+            print()
+            return res.text
+    else:
+        raise RuntimeError("No transcription result found.")
+
+
+def main() -> None:
+    TranscriptionApiArguments().run()
+
+
+if __name__ == "__main__":
+    main()

chatterer/examples/upstage_parser.py

@@ -0,0 +1,95 @@
+def resolve_import_path_and_get_logger():
+    # ruff: noqa: E402
+    import logging
+    import sys
+
+    if __name__ == "__main__" and "." not in sys.path:
+        sys.path.append(".")
+
+    logger = logging.getLogger(__name__)
+    return logger
+
+
+logger = resolve_import_path_and_get_logger()
+from pathlib import Path
+
+from langchain_core.documents.base import Blob
+from spargear import ArgumentSpec, BaseArguments
+
+from chatterer import Chatterer, UpstageDocumentParseParser
+from chatterer.tools.upstage_document_parser import (
+    DEFAULT_IMAGE_DIR,
+    DOCUMENT_PARSE_BASE_URL,
+    DOCUMENT_PARSE_DEFAULT_MODEL,
+    OCR,
+    Category,
+    OutputFormat,
+    SplitType,
+)
+
+
+class UpstageParserArguments(BaseArguments):
+    in_path: ArgumentSpec[Path] = ArgumentSpec(["in-path"], help="Path to the input file.")
+    out_path: ArgumentSpec[Path] = ArgumentSpec(["--out-path"], default=None, help="Output file path.")
+    api_key: ArgumentSpec[str] = ArgumentSpec(["--api-key"], default=None, help="API key for the Upstage API.")
+    base_url: ArgumentSpec[str] = ArgumentSpec(["--base-url"], default=DOCUMENT_PARSE_BASE_URL, help="Base URL for the Upstage API.")
+    model: ArgumentSpec[str] = ArgumentSpec(["--model"], default=DOCUMENT_PARSE_DEFAULT_MODEL, help="Model to use for parsing.")
+    split: ArgumentSpec[SplitType] = ArgumentSpec(["--split"], default="none", help="Split type for parsing.")
+    ocr: ArgumentSpec[OCR] = ArgumentSpec(["--ocr"], default="auto", help="OCR type for parsing.")
+    output_format: ArgumentSpec[OutputFormat] = ArgumentSpec(["--output-format"], default="markdown", help="Output format.")
+    coordinates: ArgumentSpec[bool] = ArgumentSpec(["--coordinates"], action="store_true", help="Include coordinates.")
+    base64_encoding: ArgumentSpec[list[Category]] = ArgumentSpec(["--base64-encoding"], default=["figure"], help="Base64 encoding for specific categories.")
+    image_description_instruction: ArgumentSpec[str] = ArgumentSpec(
+        ["--image-description-instruction"],
+        default="Describe the image in detail.",
+        help="Instruction for image description.",
+    )
+    image_dir: ArgumentSpec[str] = ArgumentSpec(
+        ["--image-dir"],
+        default=DEFAULT_IMAGE_DIR,
+        help="Directory for image paths.",
+    )
+    chatterer: ArgumentSpec[Chatterer] = ArgumentSpec(
+        ["--chatterer"],
+        default=None,
+        help="Chatterer instance for communication.",
+        type=Chatterer.from_provider,
+    )
+
+    def run(self) -> None:
+        input = UpstageParserArguments.in_path.unwrap().resolve()
+        out = UpstageParserArguments.out_path.value or input.with_suffix(".md")
+
+        parser = UpstageDocumentParseParser(
+            api_key=UpstageParserArguments.api_key.value,
+            base_url=UpstageParserArguments.base_url.unwrap(),
+            model=UpstageParserArguments.model.unwrap(),
+            split=UpstageParserArguments.split.unwrap(),
+            ocr=UpstageParserArguments.ocr.unwrap(),
+            output_format=UpstageParserArguments.output_format.unwrap(),
+            coordinates=UpstageParserArguments.coordinates.unwrap(),
+            base64_encoding=UpstageParserArguments.base64_encoding.unwrap(),
+            image_description_instruction=UpstageParserArguments.image_description_instruction.unwrap(),
+            image_dir=UpstageParserArguments.image_dir.value,
+            chatterer=UpstageParserArguments.chatterer.value,
+        )
+
+        docs = parser.parse(Blob.from_path(input))  # pyright: ignore[reportUnknownMemberType]
+
+        if UpstageParserArguments.image_dir.value:
+            for path, image in parser.image_data.items():
+                (path := Path(path)).parent.mkdir(parents=True, exist_ok=True)
+                path.write_bytes(image)
+                logger.info(f"Saved image to `{path}`")
+
+        markdown: str = "\n\n".join(f"<!--- page {i} -->\n{doc.page_content}" for i, doc in enumerate(docs, 1))
+        out.write_text(markdown, encoding="utf-8")
+        logger.info(f"Parsed `{input}` to `{out}`")
+
+
+def main() -> None:
+    UpstageParserArguments().run()
+
+
+if __name__ == "__main__":
+    main()

chatterer/examples/webpage_to_markdown.py

@@ -0,0 +1,79 @@
+def resolve_import_path_and_get_logger():
+    # ruff: noqa: E402
+    import logging
+    import sys
+
+    if __name__ == "__main__" and "." not in sys.path:
+        sys.path.append(".")
+
+    logger = logging.getLogger(__name__)
+    return logger
+
+
+logger = resolve_import_path_and_get_logger()
+from pathlib import Path
+from typing import Literal
+
+from spargear import ArgumentSpec, BaseArguments
+
+from chatterer import Chatterer, MarkdownLink, PlayWrightBot
+
+
+class WebpageToMarkdownArgs(BaseArguments):
+    url: ArgumentSpec[str] = ArgumentSpec(["url"], help="The URL to crawl.")
+    out_path: str = Path(__file__).with_suffix(".md").as_posix()
+    """The output file path for the markdown file."""
+    chatterer: ArgumentSpec[Chatterer] = ArgumentSpec(
+        ["--llm"],
+        default=None,
+        type=Chatterer.from_provider,
+        help="The Chatterer backend and model to use for filtering the markdown.",
+    )
+    engine: Literal["firefox", "chromium", "webkit"] = "firefox"
+    """The browser engine to use."""
+
+    def run(self) -> None:
+        chatterer = self.chatterer.value
+        url: str = self.url.unwrap().strip()
+        out_path: Path = Path(self.out_path).resolve()
+        with PlayWrightBot(chatterer=chatterer, engine=self.engine) as bot:
+            md = bot.url_to_md(url)
+            out_path.write_text(md, encoding="utf-8")
+            if chatterer is not None:
+                md_llm = bot.url_to_md_with_llm(url.strip())
+                out_path.write_text(md_llm, encoding="utf-8")
+            links = MarkdownLink.from_markdown(md, referer_url=url)
+            for link in links:
+                if link.type == "link":
+                    print(f"- [{truncate_string(link.url)}] {truncate_string(link.inline_text)} ({truncate_string(link.inline_title)})")
+                elif link.type == "image":
+                    print(f"- ![{truncate_string(link.url)}] ({truncate_string(link.inline_text)})")
+
+    async def arun(self) -> None:
+        chatterer = self.chatterer.value
+        url: str = self.url.unwrap().strip()
+        out_path: Path = Path(self.out_path).resolve()
+        async with PlayWrightBot(chatterer=chatterer, engine=self.engine) as bot:
+            md = await bot.aurl_to_md(url)
+            out_path.write_text(md, encoding="utf-8")
+            if chatterer is not None:
+                md_llm = await bot.aurl_to_md_with_llm(url.strip())
+                out_path.write_text(md_llm, encoding="utf-8")
+            links = MarkdownLink.from_markdown(md, referer_url=url)
+            for link in links:
+                if link.type == "link":
+                    print(f"- [{truncate_string(link.url)}] {truncate_string(link.inline_text)} ({truncate_string(link.inline_title)})")
+                elif link.type == "image":
+                    print(f"- ![{truncate_string(link.url)}] ({truncate_string(link.inline_text)})")
+
+
+def truncate_string(s: str) -> str:
+    return s[:50] + "..." if len(s) > 50 else s
+
+
+def main() -> None:
+    WebpageToMarkdownArgs().run()
+
+
+if __name__ == "__main__":
+    main()