diff --git a/docs/docs/agent.md b/docs/docs/agent.md
index 68963c0..40d4244 100644
--- a/docs/docs/agent.md
+++ b/docs/docs/agent.md
@@ -1,13 +1,60 @@
 ---
 sidebar_position: 4
 
-title: Agent
+title: Agents
 ---
 
 import CodeEmbed from '@site/src/components/CodeEmbed';
 
-# Agent
+# Agents
+
+Build powerful, flexible agents—backed by LLMs and tools—in a few minutes.
+
+Ollama4j’s agent system lets you bring together the best of LLM reasoning and external tool-use using a simple, declarative YAML configuration. No framework bloat, no complicated setup—just describe your agent, plug in your logic, and go.
+
+---
+
+**Why use agents in Ollama4j?**
+
+- **Effortless Customization:** Instantly adjust your agent’s persona, reasoning strategies, or domain by tweaking YAML. No need to touch your compiled Java code.
+- **Easy Extensibility:** Want new capabilities? Just add or change tools and logic classes—no framework glue or plumbing required.
+- **Fast Experimentation:** Mix-and-match models, instructions, and tools—prototype sophisticated behaviors or orchestrators in minutes.
+- **Clean Separation:** Keep business logic (Java) and agent personality/configuration (YAML) separate for maintainability and clarity.
+
+---
+
+## Define an Agent in YAML
+
+Specify everything about your agent—what LLM it uses, its “personality,” and all callable tools—in a single YAML file.
+
+**Agent YAML keys:**
+
+| Field                   | Description                                                                                                           |
+|-------------------------|-----------------------------------------------------------------------------------------------------------------------|
+| `name`                  | Name of your agent.                                                                                                   |
+| `host`                  | The base URL for your Ollama server (e.g., `http://localhost:11434`).                                                |
+| `model`                 | The LLM backing your agent (e.g., `llama2`, `mistral`, `mixtral`, etc).                                              |
+| `customPrompt`          | _(optional)_ System prompt—instructions or persona for your agent.                                                   |
+| `tools`                 | List of tools the agent can use. Each tool entry describes the name, function, and parameters.                        |
+| `toolFunctionFQCN`      | Fully qualified Java class name implementing the tool logic. Must be present on classpath.                            |
+| `requestTimeoutSeconds` | _(optional)_ How long (seconds) to wait for agent replies.                                                            |
+
+YAML makes it effortless to configure and tweak your agent’s powers and behavior—no code changes needed!
+
+**Example agent YAML:**
+
+<CodeEmbed src="https://raw.githubusercontent.com/ollama4j/ollama4j-examples/refs/heads/main/src/main/resources/agent.yaml" language='yaml'/>
+
+---
+
+## Instantiating and Running Agents in Java
+
+Once your agent is described in YAML, bringing it to life in Java takes only a couple of lines:
+
+<CodeEmbed src="https://raw.githubusercontent.com/ollama4j/ollama4j-examples/refs/heads/main/src/main/java/io/github/ollama4j/examples/AgentExample.java"/>
+
+- **No boilerplate.** Just load and start chatting or calling tools.
+- The API takes care of wiring up LLMs, tool invocation, and instruction handling.
+
+Ready to build your own AI-powered assistant? Just write your YAML, implement the tool logic in Java, and go!
 
-:::warning[Note]
-This is work in progress
-:::
\ No newline at end of file
diff --git a/docs/src/components/CodeEmbed/index.js b/docs/src/components/CodeEmbed/index.js
index 1a269d1..5a38aac 100644
--- a/docs/src/components/CodeEmbed/index.js
+++ b/docs/src/components/CodeEmbed/index.js
@@ -1,84 +1,14 @@
-// import React, { useState, useEffect } from 'react';
-// import CodeBlock from '@theme/CodeBlock';
-// import Icon from '@site/src/components/Icon';
-
-
-// const CodeEmbed = ({ src }) => {
-//     const [code, setCode] = useState('');
-//     const [loading, setLoading] = useState(true);
-//     const [error, setError] = useState(null);
-
-//     useEffect(() => {
-//         let isMounted = true;
-
-//         const fetchCodeFromUrl = async (url) => {
-//             if (!isMounted) return;
-
-//             setLoading(true);
-//             setError(null);
-
-//             try {
-//                 const response = await fetch(url);
-//                 if (!response.ok) {
-//                     throw new Error(`HTTP error! status: ${response.status}`);
-//                 }
-//                 const data = await response.text();
-//                 if (isMounted) {
-//                     setCode(data);
-//                 }
-//             } catch (err) {
-//                 console.error('Failed to fetch code:', err);
-//                 if (isMounted) {
-//                     setError(err);
-//                     setCode(`// Failed to load code from ${url}\n// ${err.message}`);
-//                 }
-//             } finally {
-//                 if (isMounted) {
-//                     setLoading(false);
-//                 }
-//             }
-//         };
-
-//         if (src) {
-//             fetchCodeFromUrl(src);
-//         }
-
-//         return () => {
-//             isMounted = false;
-//         };
-//     }, [src]);
-
-//     const githubUrl = src ? src.replace('https://raw.githubusercontent.com', 'https://github.com').replace('/refs/heads/', '/blob/') : null;
-//     const fileName = src ? src.substring(src.lastIndexOf('/') + 1) : null;
-
-//     return (
-//         loading ? (
-//             <div>Loading code...</div>
-//         ) : error ? (
-//             <div>Error: {error.message}</div>
-//         ) : (
-//             <div style={{ backgroundColor: 'transparent', padding: '0px', borderRadius: '5px' }}>
-//                 <div style={{ textAlign: 'right' }}>
-//                     {githubUrl && (
-//                         <a href={githubUrl} target="_blank" rel="noopener noreferrer" style={{ paddingRight: '15px', color: 'gray', fontSize: '0.8em', fontStyle: 'italic', display: 'inline-flex', alignItems: 'center' }}>
-//                             View on GitHub
-//                             <Icon icon="mdi:github" height="48" />
-//                         </a>
-//                     )}
-//                 </div>
-//                 <CodeBlock title={fileName} className="language-java">{code}</CodeBlock>
-//             </div>
-//         )
-//     );
-// };
-
-// export default CodeEmbed;
-import React, { useState, useEffect } from 'react';
+import React, {useState, useEffect} from 'react';
 import CodeBlock from '@theme/CodeBlock';
 import Icon from '@site/src/components/Icon';
 
-
-const CodeEmbed = ({ src }) => {
+/**
+ * CodeEmbed component to display code fetched from a URL in a CodeBlock.
+ * @param {object} props
+ * @param {string} props.src - Source URL to fetch the code from.
+ * @param {string} [props.language='java'] - Language for syntax highlighting in CodeBlock.
+ */
+const CodeEmbed = ({src, language = 'java'}) => {
     const [code, setCode] = useState('');
     const [loading, setLoading] = useState(true);
     const [error, setError] = useState(null);
@@ -127,7 +57,7 @@ const CodeEmbed = ({ src }) => {
     const fileName = src ? src.substring(src.lastIndexOf('/') + 1) : null;
 
     const title = (
-        <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
+        <div style={{display: 'flex', justifyContent: 'space-between', alignItems: 'center'}}>
             <a
                 href={githubUrl}
                 target="_blank"
@@ -146,9 +76,15 @@ const CodeEmbed = ({ src }) => {
                 <span>{fileName}</span>
             </a>
             {githubUrl && (
-                <a href={githubUrl} target="_blank" rel="noopener noreferrer" style={{ color: 'gray', fontSize: '0.9em', fontStyle: 'italic', display: 'inline-flex', alignItems: 'center' }}>
+                <a href={githubUrl} target="_blank" rel="noopener noreferrer" style={{
+                    color: 'gray',
+                    fontSize: '0.9em',
+                    fontStyle: 'italic',
+                    display: 'inline-flex',
+                    alignItems: 'center'
+                }}>
                     View on GitHub
-                    <Icon icon="mdi:github" height="1em" />
+                    <Icon icon="mdi:github" height="1em"/>
                 </a>
             )}
         </div>
@@ -160,8 +96,8 @@ const CodeEmbed = ({ src }) => {
         ) : error ? (
             <div>Error: {error.message}</div>
         ) : (
-            <div style={{ backgroundColor: 'transparent', padding: '0px', borderRadius: '5px' }}>
-                <CodeBlock title={title} className="language-java">{code}</CodeBlock>
+            <div style={{backgroundColor: 'transparent', padding: '0px', borderRadius: '5px'}}>
+                <CodeBlock title={title} language={language}>{code}</CodeBlock>
             </div>
         )
     );
diff --git a/src/main/java/io/github/ollama4j/agent/Agent.java b/src/main/java/io/github/ollama4j/agent/Agent.java
index 5712cdf..bd17232 100644
--- a/src/main/java/io/github/ollama4j/agent/Agent.java
+++ b/src/main/java/io/github/ollama4j/agent/Agent.java
@@ -38,32 +38,44 @@ import lombok.*;
  * </ul>
  */
 public class Agent {
-    /** The agent's display name */
+    /**
+     * The agent's display name
+     */
     private final String name;
 
-    /** List of supported tools for this agent */
+    /**
+     * List of supported tools for this agent
+     */
     private final List<Tools.Tool> tools;
 
-    /** Ollama client instance for communication with the API */
+    /**
+     * Ollama client instance for communication with the API
+     */
     private final Ollama ollamaClient;
 
-    /** The model name used for chat completions */
+    /**
+     * The model name used for chat completions
+     */
     private final String model;
 
-    /** Persists chat message history across rounds */
+    /**
+     * Persists chat message history across rounds
+     */
     private final List<OllamaChatMessage> chatHistory;
 
-    /** Optional custom system prompt for the agent */
+    /**
+     * Optional custom system prompt for the agent
+     */
     private final String customPrompt;
 
     /**
      * Constructs a new Agent.
      *
-     * @param name The agent's given name.
+     * @param name         The agent's given name.
      * @param ollamaClient The Ollama API client instance to use.
-     * @param model The model name to use for chat completion.
+     * @param model        The model name to use for chat completion.
      * @param customPrompt A custom prompt to prepend to all conversations (may be null).
-     * @param tools List of available tools for function calling.
+     * @param tools        List of available tools for function calling.
      */
     public Agent(
             String name,
@@ -162,7 +174,8 @@ public class Agent {
      * @return The model's response as a string.
      * @throws OllamaException If there is a problem with the Ollama API.
      */
-    public String interact(String userInput) throws OllamaException {
+    public String interact(String userInput, OllamaChatStreamObserver chatTokenHandler)
+            throws OllamaException {
         // Build a concise and readable description of available tools
         String availableToolsDescription =
                 tools.isEmpty()
@@ -202,11 +215,6 @@ public class Agent {
                         .withModel(model)
                         .withMessages(chatHistory)
                         .build();
-
-        OllamaChatStreamObserver chatTokenHandler =
-                new OllamaChatStreamObserver(
-                        new ConsoleOutputGenerateTokenHandler(),
-                        new ConsoleOutputGenerateTokenHandler());
         OllamaChatResult response = ollamaClient.chat(request, chatTokenHandler);
 
         // Update chat history for continuity
@@ -230,7 +238,11 @@ public class Agent {
             System.out.print("\n[You]: ");
             String input = sc.nextLine();
             if ("exit".equalsIgnoreCase(input)) break;
-            this.interact(input);
+            this.interact(
+                    input,
+                    new OllamaChatStreamObserver(
+                            new ConsoleOutputGenerateTokenHandler(),
+                            new ConsoleOutputGenerateTokenHandler()));
         }
     }
 
@@ -267,23 +279,35 @@ public class Agent {
     @Getter
     @EqualsAndHashCode(callSuper = false)
     private static class AgentToolSpec extends Tools.ToolSpec {
-        /** Fully qualified class name of the tool's {@link ToolFunction} implementation */
+        /**
+         * Fully qualified class name of the tool's {@link ToolFunction} implementation
+         */
         private String toolFunctionFQCN = null;
 
-        /** Instance of the {@link ToolFunction} to invoke */
+        /**
+         * Instance of the {@link ToolFunction} to invoke
+         */
         private ToolFunction toolFunctionInstance = null;
     }
 
-    /** Bean for describing a tool function parameter for use in agent YAML definitions. */
+    /**
+     * Bean for describing a tool function parameter for use in agent YAML definitions.
+     */
     @Data
     public class AgentToolParameter {
-        /** The parameter's type (e.g., string, number, etc.) */
+        /**
+         * The parameter's type (e.g., string, number, etc.)
+         */
         private String type;
 
-        /** Description of the parameter */
+        /**
+         * Description of the parameter
+         */
         private String description;
 
-        /** Whether this parameter is required */
+        /**
+         * Whether this parameter is required
+         */
         private boolean required;
 
         /**