lionagi 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

lionagi/core/work/work_queue.py

@@ -0,0 +1,103 @@
+"""
+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.
+"""
+
+import asyncio
+from lionagi.core.work.work import WorkStatus
+
+
+class WorkQueue:
+    """
+    A class representing a queue for managing work.
+
+    Attributes:
+        capacity (int): The maximum number of tasks the queue can handle.
+        queue (asyncio.Queue): The queue holding the tasks.
+        _stop_event (asyncio.Event): Event to signal stopping `execute` of the queue.
+        available_capacity (int): The remaining number of tasks the queue can handle.
+        execution_mode (bool): If `execute` is running.
+    """
+
+    def __init__(self, capacity=5):
+        if capacity < 0:
+            raise ValueError("initial capacity must be >= 0")
+        self.capacity = capacity
+        self.queue = asyncio.Queue()
+        self._stop_event = asyncio.Event()
+        self.available_capacity = capacity
+        self.execution_mode = False
+
+    async def enqueue(self, work) -> None:
+        """Enqueue a work item."""
+        await self.queue.put(work)
+
+    async def dequeue(self):
+        """Dequeue a work item."""
+        return await self.queue.get()
+
+    async def join(self) -> None:
+        """Block until all items in the queue have been processed."""
+        await self.queue.join()
+
+    async def stop(self) -> None:
+        """Signal the queue to stop processing."""
+        self._stop_event.set()
+
+    @property
+    def stopped(self) -> bool:
+        """Return whether the queue has been stopped."""
+        return self._stop_event.is_set()
+
+    # async def process(self):
+        # async def _parse_work(work):
+        #     async with self.semaphore:
+        #         await work.perform()
+        #
+        # tasks = set()
+        # while self.queue.qsize() > 0:
+        #     next = await self.dequeue()
+        #     next.status = WorkStatus.IN_PROGRESS
+        #     task = asyncio.create_task(_parse_work(next))
+        #     tasks.add(task)
+        #
+        # await asyncio.wait(tasks)
+
+    async def process(self) -> None:
+        """Process the work items in the queue."""
+        tasks = set()
+        while self.available_capacity > 0 and self.queue.qsize() > 0:
+            next = await self.dequeue()
+            next.status = WorkStatus.IN_PROGRESS
+            task = asyncio.create_task(next.perform())
+            tasks.add(task)
+            self.available_capacity -= 1
+
+        if tasks:
+            await asyncio.wait(tasks)
+            self.available_capacity = self.capacity
+
+    async def execute(self, refresh_time=1):
+        """
+            Continuously executes the process method at a specified refresh interval.
+
+            Args:
+                refresh_time (int, optional): The time in seconds to wait between
+                    successive calls to `process`. Defaults to 1.
+        """
+        self.execution_mode = True
+        while not self.stopped:
+            await self.process()
+            await asyncio.sleep(refresh_time)
+        self.execution_mode = False

lionagi/core/work/worker.py

@@ -0,0 +1,258 @@
+"""
+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
+from typing import Callable
+import asyncio
+import inspect
+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
+from lionagi.core.report.form import Form
+from lionagi.core.collections.abc import get_lion_id
+
+
+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.
+        forms (dict[str, Form]): Dictionary mapping form identifier to Form objects.
+        default_form (str|None): The default form to be used by the worker.
+    """
+
+    name: str = "Worker"
+    work_functions: dict[str, WorkFunction] = {}
+
+    def __init__(self, forms=None, default_form=None) -> None:
+        # self.stopped = False
+        self.forms: dict[str, Form] = forms or {}
+        self.default_form = default_form
+
+    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 change_default_form(self, form_key):
+        """
+        Changes the default form to the specified form key.
+
+        Args:
+            form_key (str): The key of the form to set as the default.
+
+        Raises:
+            ValueError: If the form key does not exist in the forms dictionary.
+
+        """
+        if form_key not in self.forms.keys():
+            raise ValueError(f"Unable to change default form. Key {form_key} does not exist.")
+        self.default_form = self.forms[form_key]
+
+    # 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,
+        form_param_key=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.
+            form_param_key (str): The key to identify the form parameter in
+                the function's signature. This parameter is used to locate and fill
+                the appropriate form according to the assignment. Raises an error
+                if the form parameter key is not found in the function's signature.
+            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,
+                capacity=capacity,
+            )
+
+        work_func: WorkFunction = self.work_functions[func.__name__]
+
+        # locate form that should be filled according to the assignment
+        if form_param_key:
+            func_signature = inspect.signature(func)
+            if form_param_key not in func_signature.parameters:
+                raise KeyError(f"Failed to locate form. \"{form_param_key}\" is not defined in the function.")
+            if "self" in func_signature.parameters:
+                bound_args = func_signature.bind(None, *args, **kwargs)
+            else:
+                bound_args = func_signature.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            arguments = bound_args.arguments
+
+            form_key = arguments.get(form_param_key)
+            try:
+                form_key = get_lion_id(form_key)
+            except:
+                pass
+            form = self.forms.get(form_key) or self.default_form
+
+            if form:
+                subform = form.__class__(assignment=work_func.assignment, task=work_func.guidance)
+                for k in subform.input_fields:
+                    v = getattr(form, k, None)
+                    setattr(subform, k, v)
+                subform.origin = form
+                kwargs = {"form": subform} | kwargs
+            else:
+                raise ValueError(f"Cannot locate form in Worker's forms and default_form is not available.")
+
+        task = work_func.perform(self, *args, **kwargs)
+        work = Work(async_task=task, async_task_name=work_func.name)
+        await work_func.worklog.append(work)
+        return work
+
+
+def work(
+    assignment=None,
+    form_param_key=None,
+    capacity=10,
+    guidance=None,
+    retry_kwargs=None,
+    timeout=10,
+):
+    """
+    Decorator to mark a method as a work function.
+
+    Args:
+        assignment (str): The assignment description of the work function.
+        form_param_key (str): The key to identify the form parameter in
+                the function's signature. This parameter is used to locate and fill
+                the appropriate form according to the assignment. Raises an error
+                if the form parameter key is not found in the function's signature.
+        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.
+        timeout (int): Timeout for the work function.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(
+            self: Worker,
+            *args,
+            func=func,
+            assignment=assignment,
+            form_param_key=form_param_key,
+            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,
+                form_param_key=form_param_key,
+                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,120 @@
+"""
+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.
+        """
+        while len(self.pending) > 0:
+            work: Work = self.pile[self.pending.popleft()]
+            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