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

lionagi/experimental/tool/schema.py

@@ -0,0 +1,59 @@
+from typing import Any
+from pydantic import field_serializer
+from functools import singledispatchmethod
+from lionagi import logging as _logging
+from lionagi.libs import func_call, AsyncUtil
+from lionagi.core.generic.node import Node
+from .function_calling import FunctionCalling
+
+
+class Tool(Node):
+
+    func: Any
+    schema_: dict | None = None
+    manual: Any | None = None
+    parser: Any | None = None
+    pre_processor: Any | None = None
+    post_processor: Any | None = None
+
+    @property
+    def name(self):
+        return self.schema_["function"]["name"]
+
+    @field_serializer("func")
+    def serialize_func(self, func):
+        return func.__name__
+
+    @singledispatchmethod
+    async def invoke(self, values: Any) -> Any:
+        raise TypeError(f"Unsupported type {type(values)}")
+
+    @invoke.register
+    async def _(self, kwargs: dict):
+
+        out = None
+
+        if self.pre_processor:
+            kwargs = await func_call.call_handler(self.pre_processor, kwargs)
+        try:
+            out = await func_call.call_handler(self.func, **kwargs)
+
+        except Exception as e:
+            _logging.error(f"Error invoking function {self.func_name}: {e}")
+            return None
+
+        if self.post_processor:
+            return await func_call.call_handler(self.post_processor, out)
+
+        return out
+
+    @invoke.register
+    async def _(self, function_calls: FunctionCalling):
+        return await self.invoke(function_calls.kwargs)
+
+    @invoke.register
+    async def _(self, values: list):
+        return await func_call.alcall(self.invoke, values)
+
+
+TOOL_TYPE = bool | Tool | str | list[Tool | str | dict] | dict

lionagi/experimental/tool/tool_manager.py

@@ -0,0 +1,138 @@
+import asyncio
+from functools import singledispatchmethod
+from collections import deque
+from typing import Tuple, Any, TypeVar, Callable
+from lionagi.libs import func_call, convert, ParseUtil
+from lionagi import logging as _logging
+from .schema import Tool, TOOL_TYPE
+from .util import func_to_tool, parse_tool_response
+from .function_calling import FunctionCalling
+
+T = TypeVar("T", bound=Tool)
+
+
+class ToolManager:
+
+    def __init__(
+        self,
+        tool_registry: dict = {},
+        function_calling_tasks: dict[str : deque[FunctionCalling]] = {},
+    ):
+        self.registry = tool_registry
+        self.function_calling_tasks = function_calling_tasks
+
+    @singledispatchmethod
+    def register_tools(self, tools: Any):
+        raise TypeError(f"Unsupported type {type(tools)}")
+
+    @register_tools.register(Tool)
+    def _(self, tools):
+        name = tools.schema_["function"]["name"]
+        if self._has_name(name):
+            err_msg = f"Function {name} is already registered."
+            _logging.error(err_msg)
+            raise ValueError(err_msg)
+        else:
+            self.registry[name] = tools
+            self.function_calling_tasks[name] = deque()
+            return True
+
+    @register_tools.register(Callable)
+    def _(self, tools):
+        tool = func_to_tool(tools)[0]
+        return self.register_tools(tool)
+
+    @register_tools.register(list)
+    def _(self, tools):
+        return func_call.lcall(tools, self.register_tools)
+
+    @singledispatchmethod
+    def register_function_calling(self, func_params: Any):
+        raise TypeError(f"Unsupported type {type(func_params)}")
+
+    @register_function_calling.register(tuple)
+    def _(self, func_params):
+        func = self.registry[func_params[0]].func
+        kwargs = func_params[1]
+        _function_calling = FunctionCalling(func=func, kwargs=kwargs)
+        self.function_calling_tasks[func.__name__].append(_function_calling)
+        return True
+
+    @register_function_calling.register(dict)
+    def _(self, response):
+        tuple_ = parse_tool_response(response)
+        return self.register_function_calling(tuple_)
+
+    @register_function_calling.register(list)
+    def _(self, func_params):
+        return func_call.lcall(func_params, self.register_function_calling)
+
+    async def invoke(self, func_params: Tuple[str, dict[str, Any]]) -> Any:
+        name, kwargs = func_params
+        if not self._has_name(name):
+            raise ValueError(f"Function {name} is not registered.")
+        tool = self.registry[name]
+        func = tool.func
+        parser = tool.parser
+        try:
+            out = await func_call.call_handler(func, **kwargs)
+            return parser(out) if parser else out
+
+        except Exception as e:
+            raise ValueError(
+                f"Error when invoking function {name} with arguments {kwargs} with error message {e}"
+            ) from e
+
+    @property
+    def _schema_list(self) -> list[dict[str, Any]]:
+        return [tool.schema_ for tool in self.registry.values()]
+
+    def get_tool_schema(self, tools: TOOL_TYPE, **kwargs):
+        if isinstance(tools, bool):
+            tool_kwarg = {"tools": self._schema_list}
+            return tool_kwarg | kwargs
+
+        else:
+            if not isinstance(tools, list):
+                tools = [tools]
+            tool_kwarg = {"tools": self._get_tool_schema(tools)}
+            return tool_kwarg | kwargs
+
+    def _has_name(self, name: str) -> bool:
+        return name in self.registry
+
+    @singledispatchmethod
+    def _get_tool_schema(self, tool: Any) -> dict:
+        raise TypeError(f"Unsupported type {type(tool)}")
+
+    @_get_tool_schema.register(dict)
+    def _(self, tool):
+        """
+        assuming that the tool is a schema
+        """
+        return tool
+
+    @_get_tool_schema.register(Tool)
+    def _(self, tool):
+        if self._has_name(tool.name):
+            return self.registry[tool.name].schema_
+        else:
+            err_msg = f"Function {tool.name} is not registered."
+            _logging.error(err_msg)
+            raise ValueError(err_msg)
+
+    @_get_tool_schema.register(str)
+    def _(self, tool):
+        """
+        assuming that the tool is a name
+        """
+        if self._has_name(tool):
+            return self.registry[tool].schema_
+        else:
+            err_msg = f"Function {tool} is not registered."
+            _logging.error(err_msg)
+            raise ValueError(err_msg)
+
+    @_get_tool_schema.register(list)
+    def _(self, tools):
+        return func_call.lcall(tools, self._get_tool_schema)

lionagi/experimental/tool/util.py

@@ -0,0 +1,16 @@
+from typing import Tuple
+from lionagi.libs import convert
+
+
+def parse_tool_response(response: dict) -> Tuple[str, dict]:
+    try:
+        func = response["action"][7:]
+        args = convert.to_dict(response["arguments"])
+        return func, args
+    except Exception:
+        try:
+            func = response["recipient_name"].split(".")[-1]
+            args = response["parameters"]
+            return func, args
+        except:
+            raise ValueError("response is not a valid function call")

lionagi/experimental/validator/rule.py

@@ -0,0 +1,139 @@
+from lionagi.libs import validation_funcs
+from abc import abstractmethod
+
+
+class Rule:
+
+    def __init__(self, **kwargs):
+        self.validation_kwargs = kwargs
+        self.fix = kwargs.get("fix", False)
+
+    @abstractmethod
+    def condition(self, **kwargs):
+        pass
+
+    @abstractmethod
+    async def validate(self, value, **kwargs):
+        pass
+
+
+class ChoiceRule(Rule):
+
+    def condition(self, choices=None):
+        return choices is not None
+
+    def check(self, choices=None):
+        if choices and not isinstance(choices, list):
+            try:
+                choices = [i.value for i in choices]
+            except Exception as e:
+                raise ValueError(f"failed to get choices") from e
+        return choices
+
+    def fix(self, value, choices=None, **kwargs):
+        v_ = validation_funcs["enum"](value, choices=choices, fix_=True, **kwargs)
+        return v_
+
+    async def validate(self, value, choices=None, **kwargs):
+        if self.condition(choices):
+            if value in self.check(choices):
+                return value
+            if self.fix:
+                kwargs = {**self.validation_kwargs, **kwargs}
+                return self.fix(value, choices, **kwargs)
+            raise ValueError(f"{value} is not in chocies {choices}")
+
+
+class ActionRequestRule(Rule):
+
+    def condition(self, annotation=None):
+        return any("actionrequest" in i for i in annotation)
+
+    async def validate(self, value, annotation=None):
+        if self.condition(annotation):
+            try:
+                return validation_funcs["action"](value)
+            except Exception as e:
+                raise ValueError(f"failed to validate field") from e
+
+
+class BooleanRule(Rule):
+
+    def condition(self, annotation=None):
+        return "bool" in annotation and "str" not in annotation
+
+    async def validate(self, value, annotation=None):
+        if self.condition(annotation):
+            try:
+                return validation_funcs["bool"](
+                    value, fix_=self.fix, **self.validation_kwargs
+                )
+            except Exception as e:
+                raise ValueError(f"failed to validate field") from e
+
+
+class NumberRule(Rule):
+
+    def condition(self, annotation=None):
+        return (
+            any([i in annotation for i in ["int", "float", "number"]])
+            and "str" not in annotation
+        )
+
+    async def validate(self, value, annotation=None):
+        if self.condition(annotation):
+            if "float" in annotation:
+                self.validation_kwargs["num_type"] = float
+                if "precision" not in self.validation_kwargs:
+                    self.validation_kwargs["precision"] = 32
+
+            try:
+                return validation_funcs["number"](
+                    value, fix_=self.fix, **self.validation_kwargs
+                )
+            except Exception as e:
+                raise ValueError(f"failed to validate field") from e
+
+
+class DictRule(Rule):
+
+    def condition(self, annotation=None):
+        return "dict" in annotation
+
+    async def validate(self, value, annotation=None, keys=None):
+        if self.condition(annotation):
+            if "str" not in annotation or keys:
+                try:
+                    return validation_funcs["dict"](
+                        value, keys=keys, fix_=self.fix, **self.validation_kwargs
+                    )
+                except Exception as e:
+                    raise ValueError(f"failed to validate field") from e
+            raise ValueError(f"failed to validate field")
+
+
+class StringRule(Rule):
+
+    def condition(self, annotation=None):
+        return "str" in annotation
+
+    async def validate(self, value, annotation=None):
+        if self.condition(annotation):
+            try:
+                return validation_funcs["str"](
+                    value, fix_=self.fix, **self.validation_kwargs
+                )
+            except Exception as e:
+                raise ValueError(f"failed to validate field") from e
+
+
+from enum import Enum
+
+
+class DEFAULT_RULES(Enum):
+    CHOICE = ChoiceRule
+    ACTION_REQUEST = ActionRequestRule
+    BOOL = BooleanRule
+    NUMBER = NumberRule
+    DICT = DictRule
+    STR = StringRule

lionagi/experimental/validator/validator.py

@@ -0,0 +1,56 @@
+from pydantic import BaseModel, Field
+from .rule import DEFAULT_RULES, Rule
+
+
+rules_ = {
+    "choice": DEFAULT_RULES.CHOICE.value,
+    "actionrequest": DEFAULT_RULES.ACTION_REQUEST.value,
+    "bool": DEFAULT_RULES.BOOL.value,
+    "number": DEFAULT_RULES.NUMBER.value,
+    "dict": DEFAULT_RULES.DICT.value,
+    "str": DEFAULT_RULES.STR.value,
+}
+
+order_ = [
+    "choice",
+    "actionrequest",
+    "bool",
+    "number",
+    "dict",
+    "str",
+]
+
+
+class Validator(BaseModel):
+    """
+    rules contain all rules that this validator can apply to data
+    the order determines which rule gets applied in what sequence.
+    notice, if a rule is not present in the orders, it will not be applied.
+    """
+
+    rules: dict[str, Rule] = Field(
+        default=rules_,
+        description="The rules to be used for validation.",
+    )
+
+    order: list[str] = Field(
+        default=order_,
+        description="The order in which the rules should be applied.",
+    )
+
+    async def validate(self, value, *args, strict=False, **kwargs):
+
+        for i in self.order:
+            if i in self.rules:
+                try:
+                    if (
+                        a := await self.rules[i].validate(value, *args, **kwargs)
+                        is not None
+                    ):
+                        return a
+                except Exception as e:
+                    raise ValueError(f"failed to validate field") from e
+        if strict:
+            raise ValueError(f"failed to validate field")
+
+        return value

lionagi/experimental/work/__init__.py

@@ -0,0 +1,10 @@
+from .schema import WorkLog
+from .work_function import WorkFunction, work
+from .worker import Worker
+
+__all__ = [
+    "WorkLog",
+    "WorkFunction",
+    "work",
+    "Worker",
+]

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

@@ -0,0 +1,73 @@
+from collections import deque
+from enum import Enum
+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."""
+
+    PENDING = "PENDING"
+    IN_PROGRESS = "IN_PROGRESS"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    
+    
+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()
+
+
+    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

@@ -0,0 +1,67 @@
+import asyncio
+from typing import Callable, Any
+from lionagi.libs import func_call
+from functools import wraps
+from pydantic import Field
+
+from lionagi.core.generic import BaseComponent
+
+from .schema import Work, WorkLog
+
+
+
+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):
+        return self.function.__name__
+
+    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
+                )
+                
+            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

@@ -0,0 +1,56 @@
+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={})