import langchain
from langchain.chains.summarize import load_summarize_chain
from langchain.docstore.document import Document
from langchain.text_splitter import CharacterTextSplitter
from steamship import File, Task
from steamship.invocable import PackageService, post

from steamship_langchain.cache import SteamshipCache
from steamship_langchain.llms import OpenAI


class SummarizeAudioPackage(PackageService):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        langchain.llm_cache = SteamshipCache(client=self.client)
        self.llm = OpenAI(client=self.client, cache=True)

    @post("summarize_file")
    def summarize_file(self, file_handle: str) -> str:
        file = File.get(self.client, handle=file_handle)
        text_splitter = CharacterTextSplitter()
        texts = []
        for block in file.blocks:
            texts.extend(text_splitter.split_text(block.text))
        docs = [Document(page_content=t) for t in texts]
        chain = load_summarize_chain(self.llm, chain_type="map_reduce")
        return chain.run(docs)

    @post("summarize_audio_file")
    def summarize_audio_file(self, file_handle: str) -> Task[str]:
        transcriber = self.client.use_plugin("whisper-s2t-blockifier")
        audio_file = File.get(self.client, handle=file_handle)
        transcribe_task = audio_file.blockify(plugin_instance=transcriber.handle)
        return self.invoke_later(
            "summarize_file",
            wait_on_tasks=[transcribe_task],
            arguments={"file_handle": audio_file.handle},
        )

langchain.text_splitter.CharacterTextSplitter langchain.chains.summarize.load_summarize_chain langchain.docstore.document.Document
[((613, 635), 'steamship.invocable.post', 'post', (['"""summarize_file"""'], {}), "('summarize_file')\n", (617, 635), False, 'from steamship.invocable import PackageService, post\n'), ((1078, 1106), 'steamship.invocable.post', 'post', (['"""summarize_audio_file"""'], {}), "('summarize_audio_file')\n", (1082, 1106), False, 'from steamship.invocable import PackageService, post\n'), ((514, 548), 'steamship_langchain.cache.SteamshipCache', 'SteamshipCache', ([], {'client': 'self.client'}), '(client=self.client)\n', (528, 548), False, 'from steamship_langchain.cache import SteamshipCache\n'), ((568, 606), 'steamship_langchain.llms.OpenAI', 'OpenAI', ([], {'client': 'self.client', 'cache': '(True)'}), '(client=self.client, cache=True)\n', (574, 606), False, 'from steamship_langchain.llms import OpenAI\n'), ((706, 747), 'steamship.File.get', 'File.get', (['self.client'], {'handle': 'file_handle'}), '(self.client, handle=file_handle)\n', (714, 747), False, 'from steamship import File, Task\n'), ((772, 795), 'langchain.text_splitter.CharacterTextSplitter', 'CharacterTextSplitter', ([], {}), '()\n', (793, 795), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((985, 1040), 'langchain.chains.summarize.load_summarize_chain', 'load_summarize_chain', (['self.llm'], {'chain_type': '"""map_reduce"""'}), "(self.llm, chain_type='map_reduce')\n", (1005, 1040), False, 'from langchain.chains.summarize import load_summarize_chain\n'), ((1266, 1307), 'steamship.File.get', 'File.get', (['self.client'], {'handle': 'file_handle'}), '(self.client, handle=file_handle)\n', (1274, 1307), False, 'from steamship import File, Task\n'), ((928, 952), 'langchain.docstore.document.Document', 'Document', ([], {'page_content': 't'}), '(page_content=t)\n', (936, 952), False, 'from langchain.docstore.document import Document\n')]