Improved agent processing loop and logger module (#54)

* improved agent

* more coverage

* nit docstring

* more coverage

* wip

* changelog

Author: nerdai

CHANGELOG.md

@@ -10,11 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ### Changed
 
+- 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 `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)
 - Add `logger.set_log_level()` and logger attribute to `TaskHandler` (#51)
 - Added library logger `llm_agents_from_scratch.logger` (#50)
 - Remove `OllamaLLM` from root import -- too slow! (#45)

docs/notebooks/first_llm_agent.ipynb

@@ -15,13 +15,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from logging import DEBUG\n",
-    "\n",
-    "from llm_agents_from_scratch import LLMAgent\n",
-    "from llm_agents_from_scratch.logger import get_logger, set_log_level\n",
-    "\n",
-    "set_log_level(DEBUG)\n",
-    "nb_logger = get_logger(\"notebook\")"
+    "from llm_agents_from_scratch import LLMAgent"
    ]
   },
   {
@@ -172,16 +166,41 @@
   {
    "cell_type": "code",
    "execution_count": 9,
-   "id": "4d09a484-9eb7-4525-923e-d00ebfe0434a",
+   "id": "c3a940e0-3fb0-43e9-8858-f5277a79df37",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from llm_agents_from_scratch.data_structures.agent import Task"
+    "LOGGING_ENABLED = True"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 10,
+   "id": "bfb51e69-ea36-495d-bcc1-fc7d164b6d40",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import logging\n",
+    "\n",
+    "from llm_agents_from_scratch.logger import enable_console_logging\n",
+    "\n",
+    "if LOGGING_ENABLED:\n",
+    "    enable_console_logging(logging.INFO)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "4d09a484-9eb7-4525-923e-d00ebfe0434a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llm_agents_from_scratch.data_structures import Task"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
    "id": "1d5d512e-d400-4398-980f-a4142d367add",
    "metadata": {},
    "outputs": [],
@@ -193,7 +212,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 13,
    "id": "e99cca54-0d21-4c9d-9a17-bf595d110fe5",
    "metadata": {},
    "outputs": [
@@ -202,74 +221,12 @@
      "output_type": "stream",
      "text": [
       "INFO (llm_agents_fs.LLMAgent) :      🚀 Starting task: Add one to fifty-five point three. Use an appropriate tool!\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      🧵 Rollout: \n",
-      "INFO (llm_agents_fs.TaskHandler) :      🧠 New Step: Add one to fifty-five point three. Use an appropriate tool!\n",
       "INFO (llm_agents_fs.TaskHandler) :      ⚙️ Processing Step: Add one to fifty-five point three. Use an appropriate tool!\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      🧵 Rollout: \n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      💬 SYSTEM: You are a helpful assistant.\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      💬 USER: Add one to fifty-five point three. Use an appropriate tool!\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      💬 ASSISTANT: \n",
-      "INFO (llm_agents_fs.TaskHandler) :      🛠️ Executing Tool Call: add_one\n",
-      "INFO (llm_agents_fs.TaskHandler) :      ✅ Successful Tool Call: 56.3\n",
-      "INFO (llm_agents_fs.TaskHandler) :      ✅ Step Result: After adding one to fifty-five point three, the result is 56.3.\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      🧵 Rollout: user: Add one to fifty-five point three. Use an appropriate tool!\n",
-      "assistant: I need to make a tool call(s) to add_one\n",
-      "tool: \n",
-      "The below is a tool call response for a given tool call.\n",
-      "<tool-call>\n",
-      "tool name: add_one\n",
-      "arguments: {'x': '55.3'}\n",
-      "</tool-call>\n",
-      "\n",
-      "<result>\n",
-      "56.3\n",
-      "</result>\n",
-      "\n",
-      "assistant: After adding one to fifty-five point three, the result is 56.3.\n",
-      "INFO (llm_agents_fs.TaskHandler) :      🧠 New Step: There is sufficient context to complete the task within the provided history — another tool call may not even be necessary! Proceed with completion and mark this as the last step. NOTE: that the real user has not seen this dialogue, and so just provide the final task result without referencing it.\n",
-      "INFO (llm_agents_fs.TaskHandler) :      ⚙️ Processing Step: There is sufficient context to complete the task within the provided history — another tool call may not even be necessary! Proceed with completion and mark this as the last step. NOTE: that the real user has not seen this dialogue, and so just provide the final task result without referencing it.\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      🧵 Rollout: user: Add one to fifty-five point three. Use an appropriate tool!\n",
-      "assistant: I need to make a tool call(s) to add_one\n",
-      "tool: \n",
-      "The below is a tool call response for a given tool call.\n",
-      "<tool-call>\n",
-      "tool name: add_one\n",
-      "arguments: {'x': '55.3'}\n",
-      "</tool-call>\n",
-      "\n",
-      "<result>\n",
-      "56.3\n",
-      "</result>\n",
-      "\n",
-      "assistant: After adding one to fifty-five point three, the result is 56.3.\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      💬 SYSTEM: You are a helpful assistant.\n",
-      "\n",
-      "Here is some past dialogue and context, where another assistant was working\n",
-      "towards completing the task.\n",
-      "\n",
-      "<history>\n",
-      "user: Add one to fifty-five point three. Use an appropriate tool!\n",
-      "assistant: I need to make a tool call(s) to add_one\n",
-      "tool: \n",
-      "The below is a tool call response for a given tool call.\n",
-      "<tool-call>\n",
-      "tool name: add_one\n",
-      "arguments: {'x': '55.3'}\n",
-      "</tool-call>\n",
-      "\n",
-      "<result>\n",
-      "56.3\n",
-      "</result>\n",
-      "\n",
-      "assistant: After adding one to fifty-five point three, the result is 56.3.\n",
-      "</history>\n",
-      "\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      💬 USER: There is sufficient context to complete the task within the provided history — another tool call may not even be necessary! Proceed with completion and mark this as the last step. NOTE: that the real user has not seen this dialogue, and so just provide the final task result without referencing it.\n",
-      "DEBUG (llm_agents_fs.TaskHandler) :      💬 ASSISTANT: \n",
       "INFO (llm_agents_fs.TaskHandler) :      🛠️ Executing Tool Call: add_one\n",
       "INFO (llm_agents_fs.TaskHandler) :      ✅ Successful Tool Call: 56.3\n",
-      "INFO (llm_agents_fs.TaskHandler) :      ✅ Step Result: After adding one to fifty-five point three, the result is 56.3.\n",
-      "INFO (llm_agents_fs.LLMAgent) :      🏁 Task completed: After adding one to fifty-five point three, the result is 56.3.\n"
+      "INFO (llm_agents_fs.TaskHandler) :      ✅ Step Result: Adding one to fifty-five point three gives us fifty-six point three.\n",
+      "INFO (llm_agents_fs.TaskHandler) :      No new step required.\n",
+      "INFO (llm_agents_fs.LLMAgent) :      🏁 Task completed: Adding one to fifty-five point three gives us fifty-six point three.\n"
      ]
     }
    ],
@@ -279,30 +236,57 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 14,
    "id": "0fde7c61-f0d9-4ef9-b0b6-90386667bbb8",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "TaskResult(task=Task(instruction='Add one to fifty-five point three. Use an appropriate tool!'), content='After adding one to fifty-five point three, the result is 56.3.', rollout=\"user: Add one to fifty-five point three. Use an appropriate tool!\\nassistant: I need to make a tool call(s) to add_one\\ntool: \\nThe below is a tool call response for a given tool call.\\n<tool-call>\\ntool name: add_one\\narguments: {'x': '55.3'}\\n</tool-call>\\n\\n<result>\\n56.3\\n</result>\\n\\nassistant: After adding one to fifty-five point three, the result is 56.3.user: There is sufficient context to complete the task within the provided history — another tool call may not even be necessary! Proceed with completion and mark this as the last step. NOTE: that the real user has not seen this dialogue, and so just provide the final task result without referencing it.\\nassistant: I need to make a tool call(s) to add_one\\ntool: \\nThe below is a tool call response for a given tool call.\\n<tool-call>\\ntool name: add_one\\narguments: {'x': '55.3'}\\n</tool-call>\\n\\n<result>\\n56.3\\n</result>\\n\\nassistant: After adding one to fifty-five point three, the result is 56.3.\", error=False)"
-      ]
-     },
-     "execution_count": 12,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The result is fifty-six point three.\n"
+     ]
     }
    ],
    "source": [
-    "result"
+    "print(result.task_result)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "id": "93b1cfa1-1ac0-420d-b36e-007de5cdfc7f",
    "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "assistant: Add one to fifty-five point three. Use an appropriate tool!\n",
+      "assistant: I need to make a tool call(s) to add_one\n",
+      "tool: {\n",
+      "    \"tool_call\": {\n",
+      "        \"tool_name\": \"add_one\",\n",
+      "        \"arguments\": {\n",
+      "            \"x\": \"55.3\"\n",
+      "        }\n",
+      "    },\n",
+      "    \"content\": \"56.3\",\n",
+      "    \"error\": false\n",
+      "}\n",
+      "assistant: Adding one to fifty-five point three gives us fifty-six point three.\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(result.rollout)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1a01698-dcdc-4796-a41b-6029eafa9ab5",
+   "metadata": {},
    "outputs": [],
    "source": []
   }

src/llm_agents_from_scratch/core/agent.py

@@ -6,7 +6,12 @@
 
 from llm_agents_from_scratch.base.llm import BaseLLM
 from llm_agents_from_scratch.base.tool import BaseTool
-from llm_agents_from_scratch.data_structures import Task, TaskResult
+from llm_agents_from_scratch.data_structures import (
+    Task,
+    TaskHandlerResult,
+    TaskResult,
+    TaskStep,
+)
 from llm_agents_from_scratch.logger import get_logger
 
 from .task_handler import TaskHandler
@@ -53,23 +58,29 @@ def run(self, task: Task) -> TaskHandler:
         async def _run() -> None:
             """Asynchronously process the task."""
             self.logger.info(f"🚀 Starting task: {task.instruction}")
+            step_result = None
             while not task_handler.done():
                 try:
-                    step = await task_handler.get_next_step()
-                    step_result = await task_handler.run_step(step)
-                    if step.last_step:
-                        async with task_handler._lock:
-                            rollout = task_handler.rollout
-
-                        task_result = TaskResult(
-                            task=task,
-                            content=step_result.content,
-                            rollout=rollout,
-                        )
-                        task_handler.set_result(task_result)
-                        self.logger.info(
-                            f"🏁 Task completed: {task_result.content}",
-                        )
+                    next_step = await task_handler.get_next_step(step_result)
+
+                    match next_step:
+                        case TaskStep():
+                            step_result = await task_handler.run_step(
+                                next_step,
+                            )
+                        case TaskResult():
+                            async with task_handler._lock:
+                                rollout = task_handler.rollout
+
+                            task_handler.set_result(
+                                TaskHandlerResult(
+                                    task_result=next_step,
+                                    rollout=rollout,
+                                ),
+                            )
+                            self.logger.info(
+                                f"🏁 Task completed: {next_step.content}",
+                            )
 
                 except Exception as e:
                     task_handler.set_exception(e)