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

lionagi/experimental/work/async_queue.py

@@ -0,0 +1,54 @@
+import asyncio
+
+
+class WorkQueue:
+
+    def __init__(self, capacity=5):
+
+        self.queue = asyncio.Queue()
+        self._stop_event = asyncio.Event()
+        self.capacity = capacity
+        self.semaphore = asyncio.Semaphore(capacity)
+
+    async def enqueue(self, work) -> None:
+        await self.queue.put(work)
+
+    async def dequeue(self):
+        return await self.queue.get()
+
+    async def join(self) -> None:
+        await self.queue.join()
+
+    async def stop(self) -> None:
+        self._stop_event.set()
+
+    @property
+    def available_capacity(self):
+        if (a:= self.capacity - self.queue.qsize()) > 0:
+            return a    
+        return None
+    
+    @property
+    def stopped(self) -> bool:
+        return self._stop_event.is_set()
+
+
+    async def process(self, refresh_time=1) -> None:
+        tasks = set()
+        while self.queue.qsize() > 0 and not self.stopped:
+            if not self.available_capacity and tasks:
+                _, done = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+                tasks.difference_update(done)
+
+            async with self.semaphore:
+                next = await self.dequeue()
+                if next is None:
+                    break
+                task = asyncio.create_task(next.perform())
+                tasks.add(task)
+
+            if tasks:
+                await asyncio.wait(tasks)
+        await asyncio.sleep(refresh_time)
+
+

lionagi/experimental/work/schema.py

@@ -1,8 +1,12 @@
+from collections import deque
 from enum import Enum
-from typing import Any, Dict, List
-from pydantic import Field
+import asyncio
+from typing import Any
+
+from lionagi.libs import SysUtil
 from lionagi.core.generic import BaseComponent
 
+from .async_queue import WorkQueue
 
 class WorkStatus(str, Enum):
     """Enum to represent different statuses of work."""
@@ -11,20 +15,59 @@ class WorkStatus(str, Enum):
     IN_PROGRESS = "IN_PROGRESS"
     COMPLETED = "COMPLETED"
     FAILED = "FAILED"
-    CANCELLED = "CANCELLED"
+    
+    
+class Work(BaseComponent):
+    status: WorkStatus = WorkStatus.PENDING
+    result: Any = None
+    error: Any = None
+    async_task: asyncio.Task | None = None
+    completion_timestamp: str | None = None
+    
+    async def perform(self):
+        try:        
+            result = await self.async_task
+            self.result = result
+            self.status = WorkStatus.COMPLETED
+            self.async_task = None
+        except Exception as e:
+            self.error = e
+            self.status = WorkStatus.FAILED
+        finally:
+            self.completion_timestamp = SysUtil.get_timestamp()
 
 
-class Work(BaseComponent):
-    """Base component for handling individual units of work."""
-
-    form_id: str = Field(..., description="ID of the form for this work")
-    priority: int = Field(default=0, description="Priority of the work")
-    status: WorkStatus = Field(
-        default=WorkStatus.PENDING, description="Current status of the work"
-    )
-    deliverables: Dict[str, Any] | list = Field(
-        default={}, description="Deliverables produced by the work"
-    )
-    dependencies: List["Work"] = Field(
-        default_factory=list, description="List of work items this work depends on"
-    )
+    def __str__(self):
+        return f"Work(id={self.id_}, status={self.status}, created_at={self.timestamp}, completed_at={self.completion_timestamp})"
+
+class WorkLog:
+    
+    def __init__(self, capacity=5, pile=None):
+        self.pile = pile or {}
+        self.pending_sequence = deque()
+        self.queue = WorkQueue(capacity=capacity)
+
+    async def append(self, work: Work):
+        self.pile[work.id_] = work
+        self.pending_sequence.append(work.id_)
+        
+    async def forward(self):
+        if not self.queue.available_capacity:
+            return
+        else:
+            while self.pending_sequence and self.queue.available_capacity:
+                work = self.pile[self.pending_sequence.popleft()]
+                work.status = WorkStatus.IN_PROGRESS
+                await self.queue.enqueue(work)
+                
+
+    async def stop(self):
+        await self.queue.stop()
+
+    @property
+    def stopped(self):
+        return self.queue.stopped
+
+    @property
+    def completed_work(self):
+        return {k: v for k, v in self.pile.items() if v.status == WorkStatus.COMPLETED}

lionagi/experimental/work/work_function.py

@@ -1,89 +1,67 @@
 import asyncio
-from typing import Any, Callable, Dict, List
-from pydantic import Field
-from functools import wraps
-from lionagi import logging as _logging
+from typing import Callable, Any
 from lionagi.libs import func_call
-from lionagi.core.generic import BaseComponent
+from functools import wraps
+from pydantic import Field
 
-from .schema import Work, WorkStatus
-from ._logger import WorkLog
-from .worker import Worker
+from lionagi.core.generic import BaseComponent
 
+from .schema import Work, WorkLog
 
-class WorkFunction(BaseComponent):
-    """Work function management and execution."""
 
-    function: Callable
-    args: List[Any] = Field(default_factory=list)
-    kwargs: Dict[str, Any] = Field(default_factory=dict)
-    retry_kwargs: Dict[str, Any] = Field(default_factory=dict)
-    worklog: WorkLog = Field(default_factory=WorkLog)
-    instruction: str = Field(
-        default="", description="Instruction for the work function"
-    )
-    refresh_time: float = Field(
-        default=0.5, description="Time to wait before checking for pending work"
-    )
 
+class WorkFunction:
+    
+    def __init__(
+        self, assignment, function, retry_kwargs=None, 
+        instruction = None, capacity=5
+    ):
+        
+        self.assignment = assignment
+        self.function = function
+        self.retry_kwargs = retry_kwargs or {}
+        self.instruction = instruction or function.__doc__
+        self.worklog = WorkLog(capacity=capacity)
+    
+    
     @property
     def name(self):
-        """Get the name of the work function."""
         return self.function.__name__
 
-    async def execute(self):
-        """Execute pending work items."""
-        while self.worklog.pending:
-            work_id = self.worklog.pending.popleft()
-            work = self.worklog.logs[work_id]
-            if work.status == WorkStatus.PENDING:
-                try:
-                    await func_call.rcall(self._execute, work, **work.retry_kwargs)
-                except Exception as e:
-                    work.status = WorkStatus.FAILED
-                    _logging.error(f"Work {work.id_} failed with error: {e}")
-                    self.worklog.errored.append(work.id_)
-            else:
-                _logging.warning(
-                    f"Work {work.id_} is in {work.status} state "
-                    "and cannot be executed."
+    async def perform(self, *args, **kwargs):
+        kwargs = {**self.retry_kwargs, **kwargs}
+        return await func_call.rcall(self.function, *args, **kwargs)
+    
+    async def process(self, refresh_time=1):
+        await self.worklog.process(refresh_time=refresh_time)
+    
+    async def stop(self):
+        await self.worklog.queue.stop()
+    
+    
+    
+def work(assignment, capacity=5):
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(self, *args, retry_kwargs=None, instruction=None, **kwargs):
+            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 {},
+                    instruction=instruction or func.__doc__,
+                    capacity=capacity
                 )
-            await asyncio.sleep(self.refresh_time)
-
-    async def _execute(self, work: Work):
-        """Execute a single work item."""
-        work.status = WorkStatus.IN_PROGRESS
-        result = await self.function(*self.args, **self.kwargs)
-        work.deliverables = result
-        work.status = WorkStatus.COMPLETED
-        return result
-
-
-def workfunc(func):
-
-    @wraps(func)
-    async def wrapper(self: Worker, *args, **kwargs):
-        # Retrieve the worker instance ('self')
-        if not hasattr(self, "work_functions"):
-            self.work_functions = {}
-
-        if func.__name__ not in self.work_functions:
-            # Create WorkFunction with the function and its docstring as instruction
-            self.work_functions[func.__name__] = WorkFunction(
-                function=func,
-                instruction=func.__doc__,
-                args=args,
-                kwargs=kwargs,
-                retry_kwargs=kwargs.pop("retry_kwargs", {}),
-            )
-
-        # Retrieve the existing WorkFunction
-        work_function: WorkFunction = self.work_functions[func.__name__]
-        # Update args and kwargs for this call
-        work_function.args = args
-        work_function.kwargs = kwargs
-
-        # Execute the function using WorkFunction's managed execution process
-        return await work_function.execute()
-
-    return wrapper
+                
+            work_func: WorkFunction = self.work_functions[func.__name__]
+            task = asyncio.create_task(work_func.perform(*args, **kwargs))
+            work = Work(async_task=task)
+            work_func: WorkFunction = self.work_functions[func.__name__]
+            await work_func.worklog.append(work)
+            return True
+        
+        return wrapper
+    return decorator

lionagi/experimental/work/worker.py

@@ -1,12 +1,56 @@
-from abc import ABC
-from pydantic import Field
-from lionagi.core.generic import BaseComponent
-
-
-class Worker(BaseComponent, ABC):
-    form_templates: dict = Field(
-        default={}, description="The form templates of the worker"
-    )
-    work_functions: dict = Field(
-        default={}, description="The work functions of the worker"
-    )
+from abc import ABC, abstractmethod
+from lionagi import logging as _logging
+from .work_function import WorkFunction
+import asyncio
+
+class Worker(ABC):
+    # This is a class that will be used to create a worker object
+    # work_functions are keyed by assignment {assignment: WorkFunction}
+    
+    name: str = "Worker"
+    work_functions: dict[str, WorkFunction] = {}
+    
+    def __init__(self) -> None:
+        self.stopped = False
+        
+    async def stop(self):
+        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 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)
+            
+
+# # 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={})