lionagi 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

lionagi/core/work/worker.py

@@ -0,0 +1,195 @@
+"""
+Copyright 2024 HaiyangLi
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from abc import ABC
+from functools import wraps
+import asyncio
+from lionagi import logging as _logging
+from lionagi.libs.ln_func_call import pcall
+from lionagi.core.work.work_function import WorkFunction
+from lionagi.core.work.work import Work
+
+
+class Worker(ABC):
+    """
+    This class represents a worker that handles multiple work functions.
+
+    Attributes:
+        name (str): The name of the worker.
+        work_functions (dict[str, WorkFunction]): Dictionary mapping assignments to WorkFunction objects.
+    """
+
+    name: str = "Worker"
+    work_functions: dict[str, WorkFunction] = {}
+
+    def __init__(self) -> None:
+        self.stopped = False
+
+    async def stop(self):
+        """
+        Stops the worker and all associated work functions.
+        """
+        self.stopped = True
+        _logging.info(f"Stopping worker {self.name}")
+        non_stopped_ = []
+
+        for func in self.work_functions.values():
+            worklog = func.worklog
+            await worklog.stop()
+            if not worklog.stopped:
+                non_stopped_.append(func.name)
+
+        if len(non_stopped_) > 0:
+            _logging.error(f"Could not stop worklogs: {non_stopped_}")
+        _logging.info(f"Stopped worker {self.name}")
+
+    async def is_progressable(self):
+        """
+        Checks if any work function is progressable and the worker is not stopped.
+
+        Returns:
+            bool: True if any work function is progressable and the worker is not stopped, else False.
+        """
+
+        return (
+            any([await i.is_progressable() for i in self.work_functions.values()])
+            and not self.stopped
+        )
+
+    async def process(self, refresh_time=1):
+        """
+        Processes all work functions periodically.
+
+        Args:
+            refresh_time (int): Time interval between each process cycle.
+        """
+        while await self.is_progressable():
+            await pcall([i.process(refresh_time) for i in self.work_functions.values()])
+            await asyncio.sleep(refresh_time)
+
+    # TODO: Implement process method
+
+    # async def process(self, refresh_time=1):
+    #     while not self.stopped:
+    #         tasks = [
+    #             asyncio.create_task(func.process(refresh_time=refresh_time))
+    #             for func in self.work_functions.values()
+    #         ]
+    #         await asyncio.wait(tasks)
+    #         await asyncio.sleep(refresh_time)
+
+    async def _wrapper(
+        self,
+        *args,
+        func=None,
+        assignment=None,
+        capacity=None,
+        retry_kwargs=None,
+        guidance=None,
+        **kwargs,
+    ):
+        """
+        Internal wrapper to handle work function execution.
+
+        Args:
+            func (Callable): The function to be executed.
+            assignment (str): The assignment description.
+            capacity (int): Capacity for the work log.
+            retry_kwargs (dict): Retry arguments for the function.
+            guidance (str): Guidance or documentation for the function.
+        """
+        if getattr(self, "work_functions", None) is None:
+            self.work_functions = {}
+
+        if func.__name__ not in self.work_functions:
+            self.work_functions[func.__name__] = WorkFunction(
+                assignment=assignment,
+                function=func,
+                retry_kwargs=retry_kwargs or {},
+                guidance=guidance or func.__doc__,
+                capacity=capacity,
+            )
+
+        work_func: WorkFunction = self.work_functions[func.__name__]
+        task = asyncio.create_task(work_func.perform(self, *args, **kwargs))
+        work = Work(async_task=task)
+        await work_func.worklog.append(work)
+        return True
+
+
+def work(
+    assignment=None,
+    capacity=10,
+    guidance=None,
+    retry_kwargs=None,
+    refresh_time=1,
+    timeout=10,
+):
+    """
+    Decorator to mark a method as a work function.
+
+    Args:
+        assignment (str): The assignment description of the work function.
+        capacity (int): Capacity for the work log.
+        guidance (str): Guidance or documentation for the work function.
+        retry_kwargs (dict): Retry arguments for the work function.
+        refresh_time (int): Time interval between each process cycle.
+        timeout (int): Timeout for the work function.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(
+            self: Worker,
+            *args,
+            func=func,
+            assignment=assignment,
+            capacity=capacity,
+            retry_kwargs=retry_kwargs,
+            guidance=guidance,
+            **kwargs,
+        ):
+            retry_kwargs = retry_kwargs or {}
+            retry_kwargs["timeout"] = retry_kwargs.get("timeout", timeout)
+            return await self._wrapper(
+                *args,
+                func=func,
+                assignment=assignment,
+                capacity=capacity,
+                retry_kwargs=retry_kwargs,
+                guidance=guidance,
+                **kwargs,
+            )
+
+        return wrapper
+
+    return decorator
+
+
+# # Example
+# from lionagi import Session
+# from lionagi.experimental.work.work_function import work
+
+
+# class MyWorker(Worker):
+
+#     @work(assignment="instruction, context -> response")
+#     async def chat(instruction=None, context=None):
+#         session = Session()
+#         return await session.chat(instruction=instruction, context=context)
+
+
+# await a.chat(instruction="Hello", context={})

lionagi/core/work/worklog.py

@@ -0,0 +1,124 @@
+"""
+Copyright 2024 HaiyangLi
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from lionagi.core.collections.abc import Progressable
+from lionagi.core.collections import pile, progression, Pile
+from lionagi.core.work.work import Work, WorkStatus
+from lionagi.core.work.work_queue import WorkQueue
+
+
+class WorkLog(Progressable):
+    """
+    A class representing a log of work items.
+
+    Attributes:
+        pile (Pile): A pile containing work items.
+        pending (Progression): A progression of pending work items.
+        queue (WorkQueue): A queue to manage the execution of work items.
+    """
+
+    def __init__(self, capacity=10, workpile=None):
+        """
+        Initializes a new instance of WorkLog.
+
+        Args:
+            capacity (int): The capacity of the work queue.
+            workpile (Pile, optional): An optional pile of initial work items.
+        """
+        self.pile = (
+            workpile if workpile and isinstance(workpile, Pile) else pile({}, Work)
+        )
+        self.pending = progression(workpile) if workpile else progression()
+        self.queue = WorkQueue(capacity=capacity)
+
+    async def append(self, work: Work):
+        """
+        Appends a new work item to the log.
+
+        Args:
+            work (Work): The work item to append.
+        """
+        self.pile.append(work)
+        self.pending.append(work)
+
+    async def forward(self):
+        """
+        Forwards pending work items to the queue if capacity allows.
+        """
+        if not self.queue.available_capacity:
+            return
+        else:
+            while len(self.pending) > 0 and self.queue.available_capacity:
+                work: Work = self.pile[self.pending.popleft()]
+                work.status = WorkStatus.IN_PROGRESS
+                await self.queue.enqueue(work)
+
+    async def stop(self):
+        """
+        Stops the work queue.
+        """
+        await self.queue.stop()
+
+    @property
+    def pending_work(self):
+        """
+        Retrieves the pile of pending work items.
+
+        Returns:
+            Pile: A pile of pending work items.
+        """
+        return pile([i for i in self.pile if i.status == WorkStatus.PENDING])
+
+    @property
+    def stopped(self):
+        """
+        Checks if the work queue is stopped.
+
+        Returns:
+            bool: True if the work queue is stopped, else False.
+        """
+        return self.queue.stopped
+
+    @property
+    def completed_work(self):
+        """
+        Retrieves the pile of completed work items.
+
+        Returns:
+            Pile: A pile of completed work items.
+        """
+        return pile([i for i in self.pile if i.status == WorkStatus.COMPLETED])
+
+    def __contains__(self, work):
+        """
+        Checks if a work item is in the pile.
+
+        Args:
+            work (Work): The work item to check.
+
+        Returns:
+            bool: True if the work item is in the pile, else False.
+        """
+        return work in self.pile
+
+    def __iter__(self):
+        """
+        Returns an iterator over the work pile.
+
+        Returns:
+            Iterator: An iterator over the work pile.
+        """
+        return iter(self.pile)

lionagi/experimental/compressor/base.py

@@ -0,0 +1,46 @@
+from abc import ABC
+from lionagi.core.collections import iModel
+
+
+class TokenCompressor(ABC):
+    """
+    NOTICE:
+        The token compressor system is inspired by LLMLingua.
+        https://github.com/microsoft/LLMLingua
+
+        MIT License
+        Copyright (c) Microsoft Corporation.
+
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+
+    Authors:
+        Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, Lili Qiu
+        @inproceedings{jiang-etal-2023-llmlingua,
+            title = "{LLML}ingua: Compressing Prompts for Accelerated Inference of Large Language Models",
+            author = "Huiqiang Jiang and Qianhui Wu and Chin-Yew Lin and Yuqing Yang and Lili Qiu",
+            booktitle = "Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing",
+            month = dec,
+            year = "2023",
+            publisher = "Association for Computational Linguistics",
+            url = "https://aclanthology.org/2023.emnlp-main.825",
+            doi = "10.18653/v1/2023.emnlp-main.825",
+            pages = "13358--13376",
+        }
+
+    LionAGI Modifications:
+        - Only borrowed the concept of token compression via perplexity
+        - Removed the dependency on the LLMLingua library
+        - use logprobs from GPT model to calculate perplexity
+        - added async ability to the functions
+        - used lionagi existing iModel class for API calls
+    """
+
+    def __init__(self, imodel: iModel, tokenizer=None, splitter=None):
+        self.imodel = imodel
+        self.tokenizer = tokenizer
+        self.splitter = splitter

lionagi/experimental/compressor/llm_compressor.py

@@ -0,0 +1,247 @@
+import asyncio
+from lionagi import alcall
+from lionagi.libs.ln_convert import to_list
+import numpy as np
+from lionagi.core.collections import iModel
+from .base import TokenCompressor
+from lionagi.libs.ln_tokenize import TokenizeUtil
+from time import time
+
+# inspired by LLMLingua, MIT License, Copyright (c) Microsoft Corporation.
+# https://github.com/microsoft/LLMLingua
+
+
+class LLMCompressor(TokenCompressor):
+
+    def __init__(
+        self,
+        imodel: iModel = None,
+        system_msg=None,
+        tokenizer=None,  # must be a callable or object with a tokenize method
+        splitter=None,  # must be a callable or object with a split/chunk/segment method
+        target_ratio=0.2,
+        n_samples=5,  # the cumulative samples to take in each perplexity calculation
+        chunk_size=64,
+        max_tokens_per_sample=80,
+        min_compression_score=0,  # (0-1) the minimum score to consider for compression, 0 means all
+        split_overlap=0,
+        split_threshold=0,
+        verbose=True,
+    ):
+        imodel = imodel or iModel(model="gpt-3.5-turbo", temperature=0.3)
+        super().__init__(imodel=imodel, tokenizer=tokenizer, splitter=splitter)
+        self.system_msg = (
+            system_msg
+            or "Concisely summarize and compress the information for storage:"
+        )
+        self.target_ratio = target_ratio
+        self.n_samples = n_samples
+        self.chunk_size = chunk_size
+        self.max_tokens_per_sample = max_tokens_per_sample
+        self.min_compression_score = min_compression_score
+        self.verbose = verbose
+        self.split_overlap = split_overlap
+        self.split_threshold = split_threshold
+
+    def tokenize(self, text, encoding_name=None, return_byte=False, **kwargs):
+        """
+        by default you can use `encoding_name` to be one of,
+        ['gpt2', 'r50k_base', 'p50k_base', 'p50k_edit', 'cl100k_base', 'o200k_base']
+
+        or you can use `encoding_model` that tiktoken supports in their mapping such as "gpt-4o"
+        """
+        if not self.tokenizer:
+            return TokenizeUtil.tokenize(
+                text,
+                encoding_model=self.imodel.iModel_name,
+                encoding_name=encoding_name,
+                return_byte=return_byte,
+            )
+
+        if hasattr(self.tokenizer, "tokenize"):
+            return self.tokenizer.tokenize(text, **kwargs)
+
+        return self.tokenizer(text, **kwargs)
+
+    def split(
+        self,
+        text,
+        chunk_size=None,
+        overlap=None,
+        threshold=None,
+        by_chars=False,
+        return_tokens=False,
+        return_byte=False,
+        **kwargs,
+    ):
+        if not self.splitter:
+            splitter = (
+                TokenizeUtil.chunk_by_chars
+                if by_chars
+                else TokenizeUtil.chunk_by_tokens
+            )
+            return splitter(
+                text,
+                chunk_size or self.chunk_size,
+                overlap or self.split_overlap,
+                threshold or self.split_threshold,
+                return_tokens=return_tokens,
+                return_byte=return_byte,
+            )
+
+        a = [
+            getattr(self.splitter, i, None)
+            for i in ["split", "chunk", "segment"]
+            if i is not None
+        ][0]
+        a = getattr(self.splitter, a)
+        return a(text, **kwargs)
+
+    async def rank_by_pplex(
+        self, items: list, initial_text=None, cumulative=False, n_samples=None, **kwargs
+    ):
+        """
+        rank a list of items according to their perplexity
+        an item can be a single token or a list of tokens
+
+        kwargs: additional arguments to pass to the model
+        """
+
+        async def _get_item_perplexity(item):
+            item = item if isinstance(item, list) else [item]
+            item = (
+                item[: self.max_tokens_per_sample]
+                if len(item) > self.max_tokens_per_sample
+                else item
+            )
+            return await self.imodel.compute_perplexity(
+                initial_context=initial_text,
+                tokens=item,
+                n_samples=n_samples or self.n_samples,
+                system_msg=self.system_msg,
+                **kwargs,
+            )
+
+        if not isinstance(items, list):
+            items = self.tokenize(items)
+
+        if len(items) == 1:
+            return [items]  # no need to rank a single item
+
+        _segments = []
+        _context = initial_text or ""
+        _task = []
+
+        if cumulative:
+            for i in items:
+                if isinstance(i, list):
+                    _context += " " + " ".join(i).strip()
+                else:
+                    _context += " " + i.strip()
+
+                _segments.append(_context)
+        else:
+            _segments = items
+
+        for i in _segments:
+            _task.append(asyncio.create_task(_get_item_perplexity(i)))
+
+        results = await asyncio.gather(*_task)
+        results = [(item, pplex) for item, pplex in zip(items, results)]
+        return sorted(results, key=lambda x: x[1]["logprobs"], reverse=True)
+
+    async def compress(
+        self,
+        text,
+        target_ratio=None,
+        initial_text=None,
+        cumulative=False,
+        split_kwargs=None,
+        split_overlap=None,
+        split_threshold=None,
+        rank_by="perplexity",
+        min_compression_score=None,
+        verbose=True,
+        **kwargs,
+    ):
+        start = time()
+        if split_kwargs is None:
+            split_kwargs = {}
+            split_kwargs["chunk_size"] = self.max_tokens_per_sample
+            split_kwargs["overlap"] = split_overlap or 0
+            split_kwargs["threshold"] = split_threshold or 0
+
+        len_tokens = len(self.tokenize(text))
+
+        items = self.split(text, return_tokens=True, **split_kwargs)
+
+        if rank_by == "perplexity":
+            ranked_items = await self.rank_by_pplex(
+                items=items, initial_text=initial_text, cumulative=cumulative, **kwargs
+            )
+
+            prompt_tokens = sum([i[1]["num_prompt_tokens"] for i in ranked_items])
+
+            num_completion_tokens = sum(
+                [i[1]["num_completion_tokens"] for i in ranked_items]
+            )
+
+            price = (
+                prompt_tokens * 0.5 / 1000000 + num_completion_tokens * 1.5 / 1000000
+            )
+
+            selected_items = self.select_by_pplex(
+                ranked_items=ranked_items,
+                target_compression_ratio=target_ratio or self.target_ratio,
+                original_length=len_tokens,
+                min_pplex=min_compression_score or self.min_compression_score,
+            )
+
+            if verbose:
+                msg = ""
+                msg += f"Original Token number: {len_tokens}\n"
+
+                def _f(i):
+                    if isinstance(i, str):
+                        i = self.tokenize(i)
+
+                    if isinstance(i, list):
+                        return len(to_list(i, dropna=True, flatten=True))
+
+                len_ = sum([_f(i) for i in selected_items])
+                msg += f"Selected Token number: {len_}\n"
+                msg += f"Token Compression Ratio: {len_ / len_tokens:.03f}\n"
+                msg += f"Compression Time: {time() - start:.04f} seconds\n"
+                msg += f"Compression Model: {self.imodel.iModel_name}\n"
+                msg += f"Compression Method: {rank_by}\n"
+                msg += f"Compression Usage: ${price:.05f}\n"
+                print(msg)
+
+            a = "".join([i.strip() for i in selected_items]).strip()
+            a = a.replace("\n\n", "")
+            return a
+
+        raise ValueError(f"Ranking method {rank_by} is not supported")
+
+    def select_by_pplex(
+        self, ranked_items, target_compression_ratio, original_length, min_pplex=None
+    ):
+        min_pplex = min_pplex or 0
+
+        desired_length = int(original_length * target_compression_ratio)
+
+        items = []
+        current_length = 0
+
+        for item, info in ranked_items:
+            if info["perplexity"] > min_pplex:
+                item = self.tokenize(item) if isinstance(item, str) else item
+                item = item if isinstance(item, list) else [item]
+                item = to_list(item, dropna=True, flatten=True)
+                if current_length + len(item) > desired_length:
+                    break
+                else:
+                    current_length += len(item)
+                    items.append("".join(item))
+
+        return items

lionagi/experimental/compressor/llm_summarizer.py

@@ -0,0 +1,61 @@
+# from lionagi.core.collections import iModel
+# from .base import TokenCompressor
+
+
+# class LLMSummarizer(TokenCompressor):
+
+#     def __init__(
+#         self, imodel: iModel = None, system_msg=None, tokenizer=None, splitter=None,
+#         max_tokens=25, target_ratio=0.3
+#     ):
+#         imodel = imodel or iModel(model="gpt-3.5-turbo", max_tokens=max_tokens)
+#         super().__init__(imodel=imodel, tokenizer=tokenizer, splitter=splitter)
+#         self.system_msg = (
+#             system_msg
+#             or "Summarize the following sentence to be concise and informative:"
+#         )
+#         self.target_ratio = target_ratio
+
+#     async def summarize_sentence(self, sentence, **kwargs):
+#         messages = [
+#             {"role": "system", "content": self.system_msg},
+#             {"role": "user", "content": sentence},
+#         ]
+#         response = await self.imodel.call_chat_completion(messages, **kwargs)
+#         return response["choices"][0]["message"]["content"]
+
+#     def tokenize(self, text):
+#         tokenize_func = self.tokenizer or tokenize
+#         return tokenize_func(text)
+
+#     def split(self, text):
+#         split_func = self.splitter or split_into_segments
+#         return split_func(text)
+
+#     # Function to enforce maximum sentence length
+#     def enforce_max_sentence_length(self, sentence, max_words=25):
+#         words = self.tokenize(sentence)
+#         if len(words) > max_words:
+#             sentence = ' '.join(words[:max_words])
+#         return sentence
+
+#     async def summarize_text(self, text, max_length_per_sentence=25, target_ratio=None, **kwargs):
+#         sentences = self.split(text)
+#         summarized = await alcall(
+#             sentences, self.summarize_sentence, **kwargs
+#         )
+#         summarized = [
+#             self.enforce_max_sentence_length(sentence, max_length_per_sentence)
+#             for sentence in summarized
+#         ]
+
+#         original_length = len(self.tokenize(text))
+#         summarized_length = len(self.tokenize(' '.join(summarized)))
+#         current_ratio = summarized_length / original_length
+
+#         target_ratio = target_ratio or self.target_ratio
+#         if current_ratio > target_ratio:
+#             words_to_remove = int((current_ratio - target_ratio) * original_length)
+#             return ' '.join(summarized[:-words_to_remove])
+
+#         return ' '.join(summarized)