[Refactor] Rename core to agent (#55)

* rename core to agent

* add templates module and templates attribute to TaskHandler

* changelog

Author: nerdai

CHANGELOG.md

@@ -10,12 +10,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ### Changed
 
+- Rename `llm_agents_from_scratch.core` to `llm_agents_from_scratch.agent` (#55)
 - Revised `TaskHandler.get_next_step()` to return `TaskStep | TaskResult` (#54)
 - Fixed bug in `OllamaLLM.chat()` where chat history was coming after user message (#51)
 - Fixed bug in `TaskHandler.run_step()` where tool names were passed to `llm.chat()` (#46)
 
 ### Added
 
+- Add `~agent.templates` module and add `TaskHandler.templates` attribute (#55)
 - Add `enable_console_logging` and `disable_console_logging` to not stream logs as a library by default (#54)
 - Add first working cookbook for a simple `LLMAgent` and task (#54)
 - Add `data_structures.task_handler.GetNextStep` (#54)

src/llm_agents_from_scratch/__init__.py

@@ -4,12 +4,12 @@
 
 # Disable the F403 warning for wildcard imports
 # ruff: noqa: F403, F401
-from .core import *
-from .core import __all__ as _core_all
+from .agent import *
+from .agent import __all__ as _agent_all
 from .tools import *
 from .tools import __all__ as _tool_all
 
 __version__ = VERSION
 
 
-__all__ = sorted(_core_all + _tool_all)  # noqa: PLE0605
+__all__ = sorted(_agent_all + _tool_all)  # noqa: PLE0605

src/llm_agents_from_scratch/core/__init__.py renamed to src/llm_agents_from_scratch/agent/__init__.py

@@ -1,4 +1,4 @@
-from .agent import LLMAgent
+from .core import LLMAgent
 from .task_handler import TaskHandler
 
 __all__ = ["LLMAgent", "TaskHandler"]

src/llm_agents_from_scratch/core/agent.py renamed to src/llm_agents_from_scratch/agent/core.py

@@ -18,43 +18,7 @@
 from llm_agents_from_scratch.errors import TaskHandlerError
 from llm_agents_from_scratch.logger import get_logger
 
-DEFAULT_GET_NEXT_INSTRUCTION_PROMPT = """You are overseeing an assistant's
-progress in accomplishing a user instruction. Provided below is the assistant's
-current response to the original task instruction. Also provided, is an
-internal 'thinking' process of the assistant that the user has not seen.
-
-Determine if the current the response is sufficient to answer the original task
-instruction. In the case that it is not, provide a new instruction to the
-assistant in order to help them improve upon their current response.
-
-<user-instruction>
-{instruction}
-</user-instruction>
-
-<current-response>
-{current_response}
-</current-response>
-
-<thinking-process>
-{current_rollout}
-</thinking-process>
-"""
-
-DEFAULT_USER_MESSAGE = "{instruction}"
-
-DEFAULT_ROLLOUT_BLOCK_FROM_CHAT_MESSAGE = "{role}: {content}"
-
-DEFAULT_SYSTEM_MESSAGE_WITHOUT_ROLLOUT = """You are a helpful assistant."""
-
-DEFAULT_SYSTEM_MESSAGE = """You are a helpful assistant.
-
-Here is some past dialogue and context, where another assistant was working
-towards completing the task.
-
-<history>
-{current_rollout}
-</history>
-"""
+from .templates import TaskHandlerTemplates, default_task_handler_templates
 
 
 class TaskHandler(asyncio.Future):
@@ -64,6 +28,7 @@ class TaskHandler(asyncio.Future):
         task: The task to execute.
         llm: The backbone LLM.
         tools_registry: The tools the LLM agent can use represented as a dict.
+        templates: Associated prompt templates.
         rollout: The execution log of the task.
         logger: TaskHandler logger.
     """
@@ -73,6 +38,7 @@ def __init__(
         task: Task,
         llm: BaseLLM,
         tools: list[BaseTool | AsyncBaseTool],
+        templates: TaskHandlerTemplates = default_task_handler_templates,
         *args: Any,
         **kwargs: Any,
     ) -> None:
@@ -82,7 +48,7 @@ def __init__(
             task (Task): The task to process.
             llm (BaseLLM): The backbone LLM.
             tools (list[BaseTool]): The tools the LLM can use.
-            logger
+            templates (TaskHandlerTemplates): Associated prompt templates.
             *args: Additional positional arguments.
             **kwargs: Additional keyword arguments.
         """
@@ -91,6 +57,7 @@ def __init__(
         self.llm = llm
         self.tools_registry = {t.name: t for t in tools}
         self.rollout = ""
+        self.templates = templates
         self._background_task: asyncio.Task | None = None
         self._lock: asyncio.Lock = asyncio.Lock()
         self.logger = get_logger(self.__class__.__name__)
@@ -135,7 +102,7 @@ def _rollout_contribution_from_single_run_step(
                 )
 
             rollout_contributions.append(
-                DEFAULT_ROLLOUT_BLOCK_FROM_CHAT_MESSAGE.format(
+                self.templates["rollout_block_from_chat_message"].format(
                     role=role.value,
                     content=content,
                 ),
@@ -161,7 +128,7 @@ async def get_next_step(
                 instruction=self.task.instruction,
                 last_step=False,
             )
-        prompt = DEFAULT_GET_NEXT_INSTRUCTION_PROMPT.format(
+        prompt = self.templates["get_next_step"].format(
             instruction=self.task.instruction,
             current_rollout=rollout,
             current_response=previous_step_result.content,
@@ -217,17 +184,17 @@ async def run_step(self, step: TaskStep) -> TaskStepResult:
         # include rollout as context in the system message
         system_message = ChatMessage(
             role=ChatRole.SYSTEM,
-            content=DEFAULT_SYSTEM_MESSAGE.format(
+            content=self.templates["system_message"].format(
                 original_instruction=self.task.instruction,
                 current_rollout=rollout,
             )
             if rollout
-            else DEFAULT_SYSTEM_MESSAGE_WITHOUT_ROLLOUT,
+            else self.templates["system_message_without_rollout"],
         )
         self.logger.debug(f"💬 SYSTEM: {system_message.content}")
         user_message = ChatMessage(
             role=ChatRole.USER,
-            content=DEFAULT_USER_MESSAGE.format(
+            content=self.templates["user_message"].format(
                 instruction=step.instruction,
             ),
         )

src/llm_agents_from_scratch/core/task_handler.py renamed to src/llm_agents_from_scratch/agent/task_handler.py

@@ -0,0 +1,60 @@
+"""Prompt templates for LLMAgent (TaskHandler)."""
+
+from typing import TypedDict
+
+DEFAULT_GET_NEXT_INSTRUCTION_PROMPT = """You are overseeing an assistant's
+progress in accomplishing a user instruction. Provided below is the assistant's
+current response to the original task instruction. Also provided, is an
+internal 'thinking' process of the assistant that the user has not seen.
+
+Determine if the current the response is sufficient to answer the original task
+instruction. In the case that it is not, provide a new instruction to the
+assistant in order to help them improve upon their current response.
+
+<user-instruction>
+{instruction}
+</user-instruction>
+
+<current-response>
+{current_response}
+</current-response>
+
+<thinking-process>
+{current_rollout}
+</thinking-process>
+"""
+
+DEFAULT_USER_MESSAGE = "{instruction}"
+
+DEFAULT_ROLLOUT_BLOCK_FROM_CHAT_MESSAGE = "{role}: {content}"
+
+DEFAULT_SYSTEM_MESSAGE_WITHOUT_ROLLOUT = """You are a helpful assistant."""
+
+DEFAULT_SYSTEM_MESSAGE = """You are a helpful assistant.
+
+Here is some past dialogue and context, where another assistant was working
+towards completing the task.
+
+<history>
+{current_rollout}
+</history>
+"""
+
+
+class TaskHandlerTemplates(TypedDict):
+    """Prompt templates dict for TaskHandler."""
+
+    get_next_step: str
+    rollout_block_from_chat_message: str
+    system_message_without_rollout: str
+    system_message: str
+    user_message: str
+
+
+default_task_handler_templates = TaskHandlerTemplates(
+    get_next_step=DEFAULT_GET_NEXT_INSTRUCTION_PROMPT,
+    rollout_block_from_chat_message=DEFAULT_ROLLOUT_BLOCK_FROM_CHAT_MESSAGE,
+    system_message_without_rollout=DEFAULT_SYSTEM_MESSAGE_WITHOUT_ROLLOUT,
+    system_message=DEFAULT_SYSTEM_MESSAGE,
+    user_message=DEFAULT_USER_MESSAGE,
+)

src/llm_agents_from_scratch/agent/templates.py

@@ -4,8 +4,8 @@
 
 import pytest
 
+from llm_agents_from_scratch.agent import LLMAgent, TaskHandler
 from llm_agents_from_scratch.base.llm import BaseLLM
-from llm_agents_from_scratch.core import LLMAgent, TaskHandler
 from llm_agents_from_scratch.data_structures.agent import (
     Task,
     TaskResult,

tests/test_agent.py

@@ -4,11 +4,11 @@
 
 import pytest
 
-from llm_agents_from_scratch.base.llm import BaseLLM
-from llm_agents_from_scratch.core import TaskHandler
-from llm_agents_from_scratch.core.task_handler import (
-    DEFAULT_SYSTEM_MESSAGE_WITHOUT_ROLLOUT,
+from llm_agents_from_scratch.agent import TaskHandler
+from llm_agents_from_scratch.agent.templates import (
+    default_task_handler_templates,
 )
+from llm_agents_from_scratch.base.llm import BaseLLM
 from llm_agents_from_scratch.data_structures import (
     ChatMessage,
     ChatRole,
@@ -367,7 +367,9 @@ async def plus_two(arg1: int) -> int:
         chat_messages=[
             ChatMessage(
                 role=ChatRole.SYSTEM,
-                content=DEFAULT_SYSTEM_MESSAGE_WITHOUT_ROLLOUT.format(
+                content=default_task_handler_templates[
+                    "system_message_without_rollout"
+                ].format(
                     original_instruction="mock instruction",
                     current_rollout="",
                 ),
@@ -410,7 +412,9 @@ async def test_run_step_without_tool_calls() -> None:
         chat_messages=[
             ChatMessage(
                 role=ChatRole.SYSTEM,
-                content=DEFAULT_SYSTEM_MESSAGE_WITHOUT_ROLLOUT.format(
+                content=default_task_handler_templates[
+                    "system_message_without_rollout"
+                ].format(
                     original_instruction="mock instruction",
                     current_rollout="",
                 ),