From 397aeec3303015eafb5d87c406655bb41687f893 Mon Sep 17 00:00:00 2001
From: Jesse Drelick <info@jessedrelick.com>
Date: Tue, 20 Aug 2024 21:54:42 -0400
Subject: [PATCH] prompt test

---
 mix.exs                 | 1 +
 test/agens/job_test.exs | 8 +++++---
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/mix.exs b/mix.exs
index 91abbe9..9c894c5 100644
--- a/mix.exs
+++ b/mix.exs
@@ -40,6 +40,7 @@ defmodule Agens.MixProject do
     [
       {:ex_doc, ">= 0.0.0", only: :dev, runtime: false},
       {:bumblebee, "~> 0.5.3"},
+      {:nx, "~> 0.7.3"},
       {:exla, "~> 0.7.0"},
       {:dialyxir, "~> 1.4", only: [:dev, :test], runtime: false},
       {:excoveralls, "~> 0.17.1", only: :test}
diff --git a/test/agens/job_test.exs b/test/agens/job_test.exs
index 7f9e54a..8d0056a 100644
--- a/test/agens/job_test.exs
+++ b/test/agens/job_test.exs
@@ -4,6 +4,7 @@ defmodule Agens.JobTest do
   import ExUnit.CaptureLog
 
   alias Agens.{Agent, Job}
+  alias Test.Support.Tools.NoopTool
 
   @lm_result_timeout 100_000
 
@@ -253,7 +254,7 @@ defmodule Agens.JobTest do
   describe "prompt" do
     setup [:start_agens, :start_serving]
 
-    @tag :skip
+    @tag capture_log: true
     test "full prompt" do
       job_name = :test_prompt_job
       agent_name = :test_prompt_agent
@@ -264,6 +265,7 @@ defmodule Agens.JobTest do
         constraints: "test agent constraints",
         context: "test agent context",
         reflection: "test agent reflection"
+        # TODO
         # examples: [
         #   %{input: "A", output: "C"},
         #   %{input: "F", output: "H"},
@@ -298,8 +300,8 @@ defmodule Agens.JobTest do
 
       assert_receive {:job_started, ^job_name}
 
-      assert_receive {:step_started, {^job_name, 0}, "test input"}
-      assert_receive {:tool_started, {^job_name, 0}, "STUB RUN"}
+      assert_receive {:step_started, {^job_name, 0}, ^input}
+      assert_receive {:tool_started, {^job_name, 0}, "sent 'test input' to: test_prompt_agent"}
       assert_receive {:tool_raw, {^job_name, 0}, %{}}
       assert_receive {:tool_result, {^job_name, 0}, "TRUE"}
       assert_receive {:step_result, {^job_name, 0}, "TRUE"}