planai 0.1.0__py3-none-any.whl

planai/__init__.py

@@ -0,0 +1,33 @@
+# Copyright 2024 Niels Provos
+#
+# 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.
+# src/planai/__init__.py
+
+from .cached_task import CachedTaskWorker
+from .graph import Graph
+from .joined_task import JoinedTaskWorker
+from .llm_interface import llm_from_config
+from .llm_task import CachedLLMTaskWorker, LLMTaskWorker
+from .task import TaskWorker, TaskWorkItem
+
+# If you want to control what gets imported with "from planai import *"
+__all__ = [
+    "Graph",
+    "llm_from_config",
+    "LLMTaskWorker",
+    "TaskWorker",
+    "TaskWorkItem",
+    "CachedTaskWorker",
+    "CachedLLMTaskWorker",
+    "JoinedTaskWorker",
+]

planai/cached_task.py

@@ -0,0 +1,113 @@
+# Copyright 2024 Niels Provos
+#
+# 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 hashlib
+import logging
+import threading
+from typing import List, Type
+
+from diskcache import Cache
+from pydantic import Field, PrivateAttr
+
+from .task import TaskWorker, TaskWorkItem
+
+
+class CachedTaskWorker(TaskWorker):
+    cache_dir: str = Field("./cache", description="Directory to store the cache")
+    cache_size_limit: int = Field(
+        25_000_000_000, description="Cache size limit in bytes"
+    )
+    _cache: Cache = PrivateAttr()
+    _lock: threading.Lock = PrivateAttr(default_factory=threading.Lock)
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self._cache = Cache(self.cache_dir, size_limit=self.cache_size_limit)
+
+    def _pre_consume_work(self, task: TaskWorkItem):
+        self.pre_consume_work(task)
+
+        cache_key = self._get_cache_key(task)
+        cached_results = self._cache.get(cache_key)
+
+        if cached_results is not None:
+            logging.info("Cache hit for %s with key: %s", self.name, cache_key)
+            self._publish_cached_results(cached_results, task)
+        else:
+            logging.info("Cache miss for %s with key: %s", self.name, cache_key)
+            self.consume_work(task)
+            
+        self.post_consume_work(task)
+
+    def pre_consume_work(self, task: TaskWorkItem):
+        """
+        This method is called before consuming the work item. It will be called even if the task has been cached.
+        It can be used for state manipulation, e.g. changing state for a class specific cache key.
+
+        Args:
+            task (TaskWorkItem): The work item to be consumed.
+
+        Returns:
+            None
+        """
+        pass
+
+    def post_consume_work(self, task: TaskWorkItem):
+        """
+        This method is called after consuming a work item in the task. It will be called even if the task has been cached.
+        It can be used for state manipulation, e.g. changing state for a class specific cache key.
+
+        Args:
+            task (TaskWorkItem): The work item that was consumed.
+
+        Returns:
+            None
+        """
+        pass
+
+    def _get_cache_key(self, task: TaskWorkItem) -> str:
+        """Generate a unique cache key for the input task."""
+        task_dict = task.model_dump()
+        task_str = str(sorted(task_dict.items()))  # Ensure consistent ordering
+        task_str += f" - {self.name}"  # Include the task name to avoid collisions
+        for output_type in self.output_types:
+            task_str += f" - {output_type.__name__}"
+        if extra_key := self.extra_cache_key(task):
+            task_str += f" - {extra_key}"
+        return hashlib.sha1(task_str.encode()).hexdigest()
+
+    def extra_cache_key(self, task: TaskWorkItem) -> str:
+        """Can be implemented by subclasses to provide additional cache key information."""
+        return ""
+
+    def _publish_cached_results(
+        self, cached_results: List[TaskWorkItem], input_task: TaskWorkItem
+    ):
+        """Publish cached results."""
+        for result in cached_results:
+            super().publish_work(result, input_task=input_task)
+
+    def publish_work(self, task: TaskWorkItem, input_task: TaskWorkItem):
+        """Publish work and cache the results."""
+        super().publish_work(task, input_task=input_task)
+
+        cache_key = self._get_cache_key(input_task)
+        try:
+            # since the cache key includes the name of this task worker, we can use the lock in this class
+            with self._lock:
+                cached_results = self._cache.get(cache_key, default=[])
+                cached_results.append(task)
+                self._cache.set(cache_key, cached_results)
+            logging.info("Task %s cached results for key: %s", self.name, cache_key)
+        except Exception as e:
+            logging.error("Error caching results for key %s: %s", cache_key, str(e))

planai/dispatcher.py

@@ -0,0 +1,236 @@
+# Copyright 2024 Niels Provos
+#
+# 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 logging
+import threading
+import time
+from collections import defaultdict, deque
+from queue import Empty, Queue
+from threading import Event, Lock
+from typing import TYPE_CHECKING, DefaultDict, Dict, List, Tuple, Type
+
+from .task import TaskWorker, TaskWorkItem
+from .web_interface import is_quit_requested, run_web_interface
+
+if TYPE_CHECKING:
+    from .graph import Graph
+
+
+# Type aliases
+TaskID = int
+TaskName = str
+ProvenanceChain = Tuple[Tuple[TaskName, TaskID], ...]
+
+
+class Dispatcher:
+    def __init__(self, graph: "Graph", web_port=5000):
+        self.graph = graph
+        self.work_queue = Queue()
+        self.provenance: DefaultDict[ProvenanceChain, int] = defaultdict(int)
+        self.notifiers: DefaultDict[ProvenanceChain, List[TaskWorker]] = defaultdict(
+            list
+        )
+        self.provenance_lock = Lock()
+        self.notifiers_lock = Lock()
+        self.stop_event = Event()
+        self.active_tasks = 0
+        self.task_completion_event = threading.Event()
+        self.web_port = web_port
+        self.debug_active_tasks: Dict[str, Tuple[TaskWorker, TaskWorkItem]] = {}
+        self.completed_tasks: deque = deque(maxlen=100)  # Keep last 100 completed tasks
+        self.total_completed_tasks = 0
+        self.task_id_counter = 0
+        self.task_lock = threading.Lock()
+
+    def _add_provenance(self, task: TaskWorkItem):
+        with self.provenance_lock:
+            for i in range(1, len(task._provenance) + 1):
+                prefix = tuple(task._provenance[:i])
+                self.provenance[prefix] = self.provenance.get(prefix, 0) + 1
+
+    def _remove_provenance(self, task: TaskWorkItem):
+        to_notify = set()
+        with self.provenance_lock:
+            for i in range(1, len(task._provenance) + 1):
+                prefix = tuple(task._provenance[:i])
+                self.provenance[prefix] -= 1
+                if self.provenance[prefix] == 0:
+                    del self.provenance[prefix]
+                    to_notify.add(prefix)
+
+        for prefix in to_notify:
+            self._notify_task_completion(prefix)
+
+    def watch(self, prefix: ProvenanceChain, notifier: TaskWorker) -> bool:
+        if not isinstance(prefix, tuple):
+            raise ValueError("Prefix must be a tuple")
+        with self.notifiers_lock:
+            if notifier not in self.notifiers[prefix]:
+                self.notifiers[prefix].append(notifier)
+                return True
+        return False
+
+    def unwatch(self, prefix: ProvenanceChain, notifier: TaskWorker) -> bool:
+        if not isinstance(prefix, tuple):
+            raise ValueError("Prefix must be a tuple")
+        with self.notifiers_lock:
+            if prefix in self.notifiers and notifier in self.notifiers[prefix]:
+                self.notifiers[prefix].remove(notifier)
+                if not self.notifiers[prefix]:
+                    del self.notifiers[prefix]
+                return True
+        return False
+
+    def _notify_task_completion(self, prefix: tuple):
+        to_notify = []
+        with self.notifiers_lock:
+            for notifier in self.notifiers[prefix]:
+                to_notify.append((notifier, prefix))
+
+        for notifier, prefix in to_notify:
+            self.active_tasks += 1
+            future = self.graph._thread_pool.submit(notifier.notify, prefix)
+            future.add_done_callback(self._notify_completed)
+
+    def _dispatch_once(self) -> bool:
+        try:
+            worker, task = self.work_queue.get(timeout=1)
+            self.active_tasks += 1
+            future = self.graph._thread_pool.submit(self._execute_task, worker, task)
+
+            # Use a named function instead of a lambda to avoid closure issues
+            def task_completed_callback(future, task=task):
+                self._task_completed(task, future)
+
+            future.add_done_callback(task_completed_callback)
+            return True
+
+        except Empty:
+            return False
+
+    def dispatch(self):
+        while (
+            not self.stop_event.is_set()
+            or not self.work_queue.empty()
+            or self.active_tasks > 0
+        ):
+            self._dispatch_once()
+
+    def _execute_task(self, worker: TaskWorker, task: TaskWorkItem):
+        task_id = self._get_next_task_id()
+        with self.task_lock:
+            self.debug_active_tasks[task_id] = (worker, task)
+
+        try:
+            worker._pre_consume_work(task)
+        finally:
+            with self.task_lock:
+                if task_id in self.debug_active_tasks:
+                    del self.debug_active_tasks[task_id]
+                    self.completed_tasks.appendleft(
+                        (task_id, worker, task)
+                    )  # may need to move this to _task_completed
+
+    def _get_next_task_id(self):
+        with self.task_lock:
+            self.task_id_counter += 1
+            return self.task_id_counter
+
+    def _task_to_dict(self, worker: TaskWorker, task: TaskWorkItem) -> Dict:
+        return {
+            "id": self._get_task_id(task),
+            "type": type(task).__name__,
+            "worker": worker.name,
+            "provenance": [f"{worker}_{id}" for worker, id in task._provenance],
+            "input_provenance": [
+                input_task.model_dump() for input_task in task._input_provenance
+            ],
+        }
+
+    def get_queued_tasks(self) -> List[Dict]:
+        return [
+            self._task_to_dict(worker, task) for worker, task in self.work_queue.queue
+        ]
+
+    def get_active_tasks(self) -> List[Dict]:
+        with self.task_lock:
+            return [
+                self._task_to_dict(worker, task)
+                for task_id, (worker, task) in self.debug_active_tasks.items()
+            ]
+
+    def get_completed_tasks(self) -> List[Dict]:
+        with self.task_lock:
+            return [
+                self._task_to_dict(worker, task)
+                for task_id, worker, task in self.completed_tasks
+            ]
+
+    def _get_task_id(self, task: TaskWorkItem) -> str:
+        # Use the last entry in the _provenance list as the task ID
+        if task._provenance:
+            return f"{task._provenance[-1][0]}_{task._provenance[-1][1]}"
+        else:
+            # Fallback in case _provenance is empty
+            return f"unknown_{id(task)}"
+
+    def _notify_completed(self, future):
+        self.active_tasks -= 1
+        if self.active_tasks == 0 and self.work_queue.empty():
+            self.task_completion_event.set()
+
+    def _task_completed(self, task: TaskWorkItem, future):
+        try:
+            # This will raise an exception if the task failed
+            result = future.result()
+
+            # Handle successful task completion
+            logging.info(f"Task {task} completed successfully")
+            self.total_completed_tasks += 1
+
+        except Exception as e:
+            # Handle task failure
+            logging.exception(f"Task {task} failed with exception: {str(e)}")
+
+            # Anything else that needs to be done when a task fails?
+
+        finally:
+            # This code will run whether the task succeeded or failed
+            self._remove_provenance(task)
+            self.active_tasks -= 1
+            if self.active_tasks == 0 and self.work_queue.empty():
+                self.task_completion_event.set()
+
+    def add_work(self, worker: TaskWorker, task: TaskWorkItem):
+        self._add_provenance(task)
+        self.work_queue.put((worker, task))
+
+    def stop(self):
+        self.stop_event.set()
+
+    def wait_for_completion(self, wait_for_quit=False):
+        self.task_completion_event.wait()
+
+        if wait_for_quit:
+            while not is_quit_requested():
+                # Sleep for a short time to avoid busy waiting
+                time.sleep(0.1)
+
+    def start_web_interface(self):
+        web_thread = threading.Thread(
+            target=run_web_interface, args=(self, self.web_port)
+        )
+        web_thread.daemon = (
+            True  # This ensures the web thread will exit when the main thread exits
+        )
+        web_thread.start()

planai/graph.py

@@ -0,0 +1,172 @@
+# Copyright 2024 Niels Provos
+#
+# 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 concurrent.futures import ThreadPoolExecutor
+from threading import Thread
+from typing import Dict, List, Optional, Set, Tuple, Type
+
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+
+from .dispatcher import Dispatcher
+from .task import TaskWorker, TaskWorkItem
+
+
+class Graph(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    name: str
+    workers: Set[TaskWorker] = Field(default_factory=set)
+    dependencies: Dict[TaskWorker, List[TaskWorker]] = Field(default_factory=dict)
+
+    _dispatcher: Optional[Dispatcher] = PrivateAttr(default=None)
+    _thread_pool: Optional[ThreadPoolExecutor] = PrivateAttr(default=None)
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        if self._thread_pool is None:
+            self._thread_pool = ThreadPoolExecutor()
+
+    def add_worker(self, task: TaskWorker) -> "Graph":
+        """Add a task to the Graph."""
+        if task in self.workers:
+            raise ValueError(f"Task {task} already exists in the Graph.")
+        self.workers.add(task)
+        self.dependencies[task] = []
+        task.set_graph(self)
+        return self
+
+    def add_workers(self, *workers: TaskWorker) -> "Graph":
+        """Add multiple tasks to the Graph."""
+        for worker in workers:
+            self.add_worker(worker)
+        return self
+
+    def set_dependency(
+        self, upstream: TaskWorker, downstream: TaskWorker
+    ) -> TaskWorker:
+        """Set a dependency between two tasks."""
+        if upstream not in self.workers or downstream not in self.workers:
+            raise ValueError(
+                "Both tasks must be added to the Graph before setting dependencies."
+            )
+
+        if downstream not in self.dependencies[upstream]:
+            self.dependencies[upstream].append(downstream)
+            upstream.register_consumer(
+                task_cls=downstream.get_taskworkitem_class(),
+                consumer=downstream,
+            )
+
+        return downstream
+
+    def validate_graph(self) -> None:
+        """Return the execution order of tasks based on dependencies."""
+        pass
+
+    def run(self, initial_tasks: List[Tuple[TaskWorker, TaskWorkItem]]) -> None:
+        """Execute the Graph by initiating source tasks."""
+        dispatcher = Dispatcher(self)
+        dispatch_thread = Thread(target=dispatcher.dispatch)
+        dispatch_thread.start()
+        dispatcher.start_web_interface()
+        self._dispatcher = dispatcher
+
+        # let the workers now that we are about to start
+        for worker in self.workers:
+            worker.init()
+
+        for worker, task in initial_tasks:
+            success, error = worker.validate_taskworkitem(type(task), worker)
+            if not success:
+                raise error
+
+            dispatcher.add_work(worker, task)
+
+        # Wait for all tasks to complete
+        dispatcher.wait_for_completion(wait_for_quit=True)
+        dispatcher.stop()
+        dispatch_thread.join()
+        self._thread_pool.shutdown(wait=True)
+
+        for worker in self.workers:
+            worker.completed()
+
+    def __str__(self) -> str:
+        return f"Graph: {self.name} with {len(self.workers)} tasks"
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+
+def main():
+    # Define custom TaskWorkItem classes
+    class Task1WorkItem(TaskWorkItem):
+        data: str
+
+    class Task2WorkItem(TaskWorkItem):
+        processed_data: str
+
+    class Task3WorkItem(TaskWorkItem):
+        final_result: str
+
+    # Define custom TaskWorker classes
+    class Task1Worker(TaskWorker):
+        output_types: Set[Type[TaskWorkItem]] = {Task2WorkItem}
+
+        def consume_work(self, task: Task1WorkItem):
+            print(f"Task1 consuming: {task.data}")
+            processed = f"Processed: {task.data.upper()}"
+            self.publish_work(Task2WorkItem(processed_data=processed), input_task=task)
+
+    class Task2Worker(TaskWorker):
+        output_types: Set[Type[TaskWorkItem]] = {Task3WorkItem}
+
+        def consume_work(self, task: Task2WorkItem):
+            print(f"Task2 consuming: {task.processed_data}")
+            final = f"Final: {task.processed_data}!"
+            self.publish_work(Task3WorkItem(final_result=final), input_task=task)
+
+    class Task3Worker(TaskWorker):
+        output_types: Set[Type[TaskWorkItem]] = set()
+
+        def consume_work(self, task: Task3WorkItem):
+            print(f"Task3 consuming: {task.final_result}")
+            print("Workflow complete!")
+
+    # Create Graph
+    graph = Graph(name="Simple Workflow")
+
+    # Create tasks
+    task1 = Task1Worker()
+    task2 = Task2Worker()
+    task3 = Task3Worker()
+
+    # Add tasks to Graph
+    graph.add_worker(task1).add_worker(task2).add_worker(task3)
+
+    # Set dependencies
+    graph.set_dependency(task1, task2).next(task3)
+
+    # Validate Graph
+    execution_order = graph.validate_graph()
+    print(f"Execution order: {execution_order}")
+
+    # Prepare initial work item
+    initial_work = [(task1, Task1WorkItem(data="Hello, Graph!"))]
+
+    # Run the Graph
+    graph.run(initial_work)
+
+
+if __name__ == "__main__":
+    main()

planai/joined_task.py

@@ -0,0 +1,56 @@
+# Copyright 2024 Niels Provos
+#
+# 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 abstractmethod
+from typing import Dict, List, Type
+
+from pydantic import PrivateAttr
+
+from .cached_task import CachedTaskWorker
+from .task import TaskWorker, TaskWorkItem
+
+
+class JoinedTaskWorker(TaskWorker):
+    join_type: Type[TaskWorker]
+    _joined_results: Dict[tuple, List[TaskWorkItem]] = PrivateAttr(default_factory=dict)
+
+    @abstractmethod
+    def consume_work_joined(self, task: List[TaskWorkItem]):
+        """
+        A subclass needs to implement consume_work only for the type hint.
+        It still needs to call the super() method in consume_work.
+        All accumulated results will be delivered to this method.
+        """
+        pass
+
+    def consume_work(self, task: TaskWorkItem):
+        prefix = task.prefix_for_input_task(self.join_type)
+        if prefix is None:
+            raise ValueError(
+                f"Task {task} does not have a prefix for {self.join_type.__name__} in provenance."
+            )
+        if not prefix in self._joined_results:
+            # we will register the watch for the prefix when we see it for the first time.
+            if not self.watch(prefix):
+                raise ValueError(
+                    f"Task {task} does not have a watch for {self.join_type} in provenance."
+                )
+
+        # we accumulate the results by the prefix described by join_type
+        # we will deliver them to the sub-class when we get notified
+        self._joined_results.setdefault(prefix, []).append(task)
+
+    def notify(self, prefix: str):
+        if prefix not in self._joined_results:
+            raise ValueError(f"Task {prefix} does not have any results to join.")
+        self.consume_work_joined(self._joined_results.pop(prefix))