feat(structured-outputs): typed runAs[T] in agent abstraction

Author: kubinio123

claude/src/main/scala/sttp/ai/claude/ClaudeSyncClient.scala

@@ -18,14 +18,6 @@ class ClaudeSyncClient(config: ClaudeConfig, backend: SyncBackend = DefaultSyncB
       case Right(response) => response
     }
 
-  /** Creates a typed message response. The response schema is derived from `T` via Tapir and set on the request as a structured-output
-    * format (unless one is already set), and the response's text content is parsed back into `T` via uPickle.
-    *
-    * @param request
-    *   Message request. If [[MessageRequest.usesStructuredOutput]] is false, a JSON schema for `T` is set automatically.
-    * @tparam T
-    *   The return type, which must have both a [[TapirSchema]] and a [[SnakePickle.Reader]] available.
-    */
   def createMessageAs[T: TapirSchema: SnakePickle.Reader](request: MessageRequest): T = {
     val withSchema =
       if (request.usesStructuredOutput) request

claude/src/test/scala/sttp/ai/claude/integration/ClaudeAgentIntegrationSpec.scala

@@ -17,7 +17,7 @@ class ClaudeAgentIntegrationSpec extends AgentIntegrationSpecBase {
     val config = ClaudeConfig.fromEnv
     val client = ClaudeClient(config)
     val agentConfig = AgentConfig(maxIterations = maxIterations, userTools = tools).right.get
-    val allTools = agentConfig.userTools ++ AgentConfig.systemTools
+    val allTools = agentConfig.userTools ++ AgentConfig.systemTools(agentConfig)
     val agentBackend = new ClaudeAgentBackend[Identity](
       client,
       "claude-haiku-4-5-20251001",
@@ -26,4 +26,17 @@ class ClaudeAgentIntegrationSpec extends AgentIntegrationSpecBase {
     )(IdentityMonad)
     Agent(agentBackend, agentConfig)(IdentityMonad)
   }
+
+  override def createTypedAgent[T](
+      maxIterations: Int,
+      tools: Seq[AgentTool[_]],
+      responseSchema: ResponseSchema[T]
+  ): Agent[Identity] = {
+    val agentConfig = AgentConfig(
+      maxIterations = maxIterations,
+      userTools = tools,
+      responseSchema = Some(responseSchema)
+    ).right.get
+    ClaudeAgent.synchronous(ClaudeConfig.fromEnv, "claude-haiku-4-5-20251001", agentConfig)
+  }
 }

core/src/main/scala/sttp/ai/core/agent/Agent.scala

@@ -1,8 +1,8 @@
 package sttp.ai.core.agent
 
+import sttp.ai.core.json.SnakePickle
 import sttp.client4.Backend
 import sttp.monad.MonadError
-import sttp.ai.core.json.SnakePickle
 
 class Agent[F[_]](
     agentBackend: AgentBackend[F],
@@ -96,6 +96,16 @@ class Agent[F[_]](
     loop(initialHistory, 0, Seq.empty)
   }
 
+  def runAs[T](
+      initialPrompt: String
+  )(backend: Backend[F])(implicit r: SnakePickle.Reader[T]): F[AgentResult[Either[AgentParseError, T]]] =
+    monad.map(run(initialPrompt)(backend)) { res =>
+      val parsed: Either[AgentParseError, T] =
+        try Right(SnakePickle.read[T](res.finalAnswer))
+        catch { case e: Exception => Left(AgentParseError(res.finalAnswer, e)) }
+      AgentResult(parsed, res.iterations, res.toolCalls, res.finishReason)
+    }
+
   private def executeTool[T](tool: AgentTool[T], toolCall: ToolCall): Either[String, String] = {
     val parseResult: Either[Exception, T] =
       try

core/src/main/scala/sttp/ai/core/agent/AgentResult.scala

@@ -10,16 +10,21 @@ object FinishReason {
   case class Error(message: String) extends FinishReason
 }
 
-case class ToolCallRecord(
+final case class ToolCallRecord(
     toolName: String,
     input: String,
     output: String,
     iteration: Int
 )
 
-case class AgentResult[T](
+final case class AgentResult[T](
     finalAnswer: T,
     iterations: Int,
     toolCalls: Seq[ToolCallRecord],
     finishReason: FinishReason
 )
+
+final case class AgentParseError(
+    rawAnswer: String,
+    cause: Throwable
+)

core/src/test/scala/sttp/ai/core/agent/AgentSpec.scala

@@ -604,6 +604,63 @@ class AgentSpec extends AnyFlatSpec with Matchers {
     cfg.systemPrompt.getOrElse("") should include("JSON object matching the provided input schema")
   }
 
+  "Agent.runAs[T]" should "return Right(T) when the model emits a well-formed structured payload" in {
+    val cfg = AgentConfig(
+      maxIterations = 3,
+      responseSchema = Some(ResponseSchema.derived[WeatherSummary]())
+    ).toOption.get
+
+    val (loop, _) = createLoop(
+      Seq(
+        AgentResponse(
+          "",
+          Seq(
+            ToolCall(
+              id = "call_1",
+              toolName = "finish",
+              input = """{"city":"Krakow","temp_c":12.0,"conditions":"sunny"}"""
+            )
+          ),
+          StopReason.ToolUse
+        )
+      ),
+      cfg
+    )
+
+    val result = loop.runAs[WeatherSummary]("What's the weather?")(backend)
+
+    result.finishReason shouldBe FinishReason.ToolFinish: Unit
+    result.iterations shouldBe 1: Unit
+    result.finalAnswer shouldBe Right(WeatherSummary("Krakow", 12.0, "sunny"))
+  }
+
+  it should "return Left(AgentParseError) preserving the trace when the answer can't be parsed as T" in {
+    val cfg = AgentConfig(
+      maxIterations = 3,
+      responseSchema = Some(ResponseSchema.derived[WeatherSummary]())
+    ).toOption.get
+
+    val (loop, _) = createLoop(
+      Seq(
+        AgentResponse(
+          "",
+          Seq(ToolCall(id = "call_1", toolName = "finish", input = """{"wrong":"shape"}""")),
+          StopReason.ToolUse
+        )
+      ),
+      cfg
+    )
+
+    val result = loop.runAs[WeatherSummary]("What's the weather?")(backend)
+
+    result.iterations shouldBe 1: Unit
+    result.finishReason shouldBe FinishReason.ToolFinish: Unit
+    result.finalAnswer.isLeft shouldBe true: Unit
+    val err = result.finalAnswer.left.toOption.get
+    err.rawAnswer should include("Invalid arguments"): Unit
+    err.cause should not be null
+  }
+
   it should "terminate with the parse-error message as finalAnswer when finish is called with a malformed payload" in {
     val cfg = AgentConfig(
       maxIterations = 3,

core/src/test/scala/sttp/ai/core/agent/integration/AgentIntegrationSpecBase.scala

@@ -14,6 +14,13 @@ abstract class AgentIntegrationSpecBase extends AnyFlatSpec with Matchers {
   def apiKeyEnvVar: String
   def createAgent(maxIterations: Int, tools: Seq[AgentTool[_]]): Agent[Identity]
 
+  def createTypedAgent[T](
+      maxIterations: Int,
+      tools: Seq[AgentTool[_]],
+      responseSchema: ResponseSchema[T]
+  ): Agent[Identity] =
+    cancel(s"$providerName typed agent factory not implemented for this spec")
+
   protected val maybeApiKey: Option[String] = sys.env.get(apiKeyEnvVar)
 
   case class CalculatorInput(operation: String, a: Double, b: Double)
@@ -151,4 +158,28 @@ abstract class AgentIntegrationSpecBase extends AnyFlatSpec with Matchers {
     result.iterations shouldBe 1
     ()
   }
+
+  case class TripSummary(weather: String, calculation: String, conclusion: String)
+  implicit val tripSummaryRW: SnakePickle.ReadWriter[TripSummary] = SnakePickle.macroRW
+  implicit val tripSummarySchema: Schema[TripSummary] = Schema.derived
+
+  it should "return a typed structured answer via runAs[T]" in {
+    if (maybeApiKey.isEmpty) cancel(s"$apiKeyEnvVar not defined - skipping integration test")
+    val backend = DefaultSyncBackend()
+    try {
+      val rs = ResponseSchema.derived[TripSummary]()
+      val agent = createTypedAgent(maxIterations = 6, tools = Seq(weatherTool, calculatorTool), responseSchema = rs)
+
+      val result = agent.runAs[TripSummary](
+        "What's the weather in Paris? Also, what is 15 multiplied by 3? Then summarise both into the structured response."
+      )(backend)
+
+      result.finishReason shouldBe FinishReason.ToolFinish: Unit
+      result.finalAnswer.isRight shouldBe true: Unit
+      val summary = result.finalAnswer.toOption.get
+      summary.weather should not be empty: Unit
+      summary.calculation should not be empty: Unit
+      summary.conclusion should not be empty
+    } finally backend.close()
+  }
 }

examples/src/main/scala/examples/AgentLoopExample.scala

@@ -74,3 +74,51 @@ object AgentLoopExample extends App {
     }
   finally backend.close()
 }
+
+object TypedAgentLoopExample extends App {
+
+  case class TripSummary(weatherSummary: String, calculation: String, conclusion: String) derives SnakePickle.ReadWriter, Schema
+
+  case class WeatherInput(location: String, unit: Option[String]) derives SnakePickle.ReadWriter, Schema
+
+  val weatherTool = AgentTool.fromFunction("get_weather", "Get the current weather for a location") { (input: WeatherInput) =>
+    val unit = input.unit.getOrElse("celsius")
+    s"The weather in ${input.location} is 22°${if (unit == "celsius") "C" else "F"}, sunny"
+  }
+
+  case class CalculatorInput(operation: String, a: Double, b: Double) derives SnakePickle.ReadWriter, Schema
+
+  val calculatorTool = AgentTool.fromFunction("calculate", "Perform a mathematical calculation") { (input: CalculatorInput) =>
+    val result = input.operation match {
+      case "add"      => input.a + input.b
+      case "subtract" => input.a - input.b
+      case "multiply" => input.a * input.b
+      case "divide"   => if (input.b != 0) input.a / input.b else Double.NaN
+      case _          => 0.0
+    }
+    s"${input.a} ${input.operation} ${input.b} = $result"
+  }
+
+  val cfg = AgentConfig(
+    maxIterations = 8,
+    userTools = Seq(weatherTool, calculatorTool),
+    responseSchema = Some(ResponseSchema.derived[TripSummary]())
+  ).toOption.get
+
+  val openai = OpenAI.fromEnv
+  val backend = DefaultSyncBackend()
+  try {
+    val agent = OpenAIAgent.synchronous(openai, "gpt-4o-mini", cfg)
+    val prompt = "What's the weather in Paris? Also, what is 15 multiplied by 23? Provide a complete answer."
+
+    agent.runAs[TripSummary](prompt)(backend).finalAnswer match {
+      case Right(summary) =>
+        println(s"weather:     ${summary.weatherSummary}")
+        println(s"calculation: ${summary.calculation}")
+        println(s"conclusion:  ${summary.conclusion}")
+      case Left(err) =>
+        println(s"Failed to parse structured answer: ${err.cause.getMessage}")
+        println(s"Raw answer was: ${err.rawAnswer}")
+    }
+  } finally backend.close()
+}

openai/src/test/scala/sttp/ai/openai/integration/OpenAIAgentIntegrationSpec.scala

@@ -15,7 +15,7 @@ class OpenAIAgentIntegrationSpec extends AgentIntegrationSpecBase {
   override def createAgent(maxIterations: Int, tools: Seq[AgentTool[_]]): Agent[Identity] = {
     val openai = OpenAI.fromEnv
     val agentConfig = AgentConfig(maxIterations = maxIterations, userTools = tools).right.get
-    val allTools = agentConfig.userTools ++ AgentConfig.systemTools
+    val allTools = agentConfig.userTools ++ AgentConfig.systemTools(agentConfig)
     val agentBackend = new OpenAIAgentBackend[Identity](
       openai,
       "gpt-4o-mini",
@@ -24,4 +24,18 @@ class OpenAIAgentIntegrationSpec extends AgentIntegrationSpecBase {
     )(IdentityMonad)
     Agent(agentBackend, agentConfig)(IdentityMonad)
   }
+
+  override def createTypedAgent[T](
+      maxIterations: Int,
+      tools: Seq[AgentTool[_]],
+      responseSchema: ResponseSchema[T]
+  ): Agent[Identity] = {
+    val openai = OpenAI.fromEnv
+    val agentConfig = AgentConfig(
+      maxIterations = maxIterations,
+      userTools = tools,
+      responseSchema = Some(responseSchema)
+    ).right.get
+    OpenAIAgent.synchronous(openai, "gpt-4o-mini", agentConfig)
+  }
 }