From 55c8e43fb37ec108eb3f1fe69f147fd37f3deea3 Mon Sep 17 00:00:00 2001
From: extio-de <github@extio.de>
Date: Mon, 8 Dec 2025 22:31:42 +0100
Subject: [PATCH] feat: Update responseHandler method to accept AgentContext

---
 README.md                                         |  2 +-
 pom.xml                                           |  2 +-
 src/main/java/de/extio/lmlib/agent/Agent.java     |  8 ++++++++
 src/main/java/de/extio/lmlib/agent/BaseAgent.java | 12 ++++++++----
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 0637184..6ed8412 100644
--- a/README.md
+++ b/README.md
@@ -555,7 +555,7 @@ public class BlogPostAgent implements BaseAgent {
     
     // OPTIONAL: How to parse/store the response (default: TextAgentResponseHandler("response"))
     @Override 
-    public AgentResponseHandler responseHandler() { 
+    public AgentResponseHandler responseHandler(final AgentContext ctx) { 
         return new TextAgentResponseHandler("blogPost"); 
     }
     
diff --git a/pom.xml b/pom.xml
index 646faa1..6f65991 100755
--- a/pom.xml
+++ b/pom.xml
@@ -4,7 +4,7 @@
 	<artifactId>lm-lib</artifactId>
 	<groupId>de.extio</groupId>
 	<name>lm-lib</name>
-	<version>1.9.0</version>
+	<version>1.9.2</version>
 	
 	<properties>
 		<java.version>21</java.version>
diff --git a/src/main/java/de/extio/lmlib/agent/Agent.java b/src/main/java/de/extio/lmlib/agent/Agent.java
index 1d6ff9d..62c8908 100755
--- a/src/main/java/de/extio/lmlib/agent/Agent.java
+++ b/src/main/java/de/extio/lmlib/agent/Agent.java
@@ -54,6 +54,14 @@ public record Agent(String name,
         return BaseAgent.super.modelProfile(context);
     }
     
+    @Override
+    public AgentResponseHandler responseHandler(final AgentContext context) {
+        if (responseHandler != null) {
+            return responseHandler;
+        }
+        return BaseAgent.super.responseHandler(context);
+    }
+
     @Override
     public void preProcess(final AgentContext context) {
         if (preProcessor != null) {
diff --git a/src/main/java/de/extio/lmlib/agent/BaseAgent.java b/src/main/java/de/extio/lmlib/agent/BaseAgent.java
index 095cdca..3441ea6 100644
--- a/src/main/java/de/extio/lmlib/agent/BaseAgent.java
+++ b/src/main/java/de/extio/lmlib/agent/BaseAgent.java
@@ -56,7 +56,7 @@ public interface BaseAgent {
 		return null;
 	}
 	
-	default AgentResponseHandler responseHandler() {
+	default AgentResponseHandler responseHandler(final AgentContext context) {
 		return new TextAgentResponseHandler("response");
 	}
 	
@@ -107,12 +107,12 @@ public interface BaseAgent {
 						
 						boolean parseable = false;
 						for (int i = 0; i < 2; i++) {
-							LOGGER.debug("Conversation: {}", conversation);
+							LOGGER.debug("Conversation: {}", conversation);	
 							
-							final var responseHandler = this.responseHandler();
 							final var modelProfile = this.modelProfile(split.context());
+							final var responseHandler = this.responseHandler(split.context());
+							long requestStart = 0;
 							Completion completion = null;
-							final var requestStart = System.currentTimeMillis();
 							if (split.context().isStreaming()) {
 								if (responseHandler instanceof final StreamedAgentResponseHandler streamedResponseHandler) {
 									streamedResponseHandler.beforeStream(split.context());
@@ -125,17 +125,21 @@ public interface BaseAgent {
 									}
 								};
 								if (modelProfile != null) {
+									requestStart = System.currentTimeMillis();
 									completion = client.streamConversation(modelProfile, conversation, chunkConsumer);
 								}
 								else {
+									requestStart = System.currentTimeMillis();
 									completion = client.streamConversation(this.modelCategory(split.context()), conversation, chunkConsumer);
 								}
 							}
 							else {
 								if (modelProfile != null) {
+									requestStart = System.currentTimeMillis();
 									completion = client.conversation(modelProfile, conversation);
 								}
 								else {
+									requestStart = System.currentTimeMillis();
 									completion = client.conversation(this.modelCategory(split.context()), conversation);
 								}
 							}