Added tests for Agent. Use in process plug for mocking.

test/anthropix/agent_test.exs

@@ -0,0 +1,72 @@
+defmodule Anthropix.AgentTest do
+  use ExUnit.Case
+  alias Anthropix.{Agent, APIError, Mock, Tool}
+
+  @tickers %{
+    "General Motors" => "GM",
+  }
+
+  @prices %{
+    "GM" => 39.21,
+  }
+
+  setup_all do
+    ticker_tool = Tool.new([
+      name: "get_ticker_symbol",
+      description: "Gets the stock ticker symbol for a company searched by name. Returns str: The ticker symbol for the company stock. Raises TickerNotFound: if no matching ticker symbol is found.",
+      params: [
+        %{name: "company_name", description: "The name of the company.", type: "string"}
+      ],
+      function: fn name -> @tickers[name] || raise "TickerNotFound" end
+    ])
+
+    price_tool = Tool.new([
+      name: "get_current_stock_price",
+      description: "Gets the current stock price for a company. Returns float: The current stock price. Raises ValueError: if the input symbol is invalid/unknown.",
+      params: [
+        %{name: "symbol", description: "The stock symbol of the company to get the price for.", type: "string"}
+      ],
+      function: fn symbol -> @prices[symbol] || raise "ValueError" end
+    ])
+
+    {:ok, tools: [ticker_tool, price_tool]}
+  end
+
+  describe "chat/2" do
+    test "returns agent with result complete", %{tools: tools} do
+      agent =
+        Mock.client(& Mock.respond(&1, {:agent, :messages}))
+        |> Agent.init(tools)
+
+      assert {:ok, agent} = Agent.chat(agent, [
+        model: "claude-3-sonnet-20240229",
+        system: "Answer like Snoop Dogg.",
+        messages: [
+          %{role: "user", content: "What is the current stock price of General Motors?"}
+        ]
+      ])
+
+      assert length(agent.messages) == 5
+      assert is_map(agent.result)
+      assert is_map(agent.usage)
+
+      expected = "Word, the current stock price for General Motors is $39.21. Representing that big auto money, ya dig? Gotta make them stacks and invest wisely in the motor city players."
+      assert ^expected = hd(agent.result["content"]) |> Map.get("text")
+    end
+
+    test "returns error when model not found", %{tools: tools} do
+      agent =
+        Mock.client(& Mock.respond(&1, 404))
+        |> Agent.init(tools)
+
+      assert {:error, %APIError{type: "not_found"}} = Agent.chat(agent, [
+        model: "claude-3-sonnet-20240229",
+        system: "Answer like Snoop Dogg.",
+        messages: [
+          %{role: "user", content: "What is the current stock price of General Motors?"}
+        ]
+      ])
+    end
+  end
+
+end

test/anthropix_test.exs

@@ -1,12 +1,6 @@
 defmodule AnthropixTest do
   use ExUnit.Case
-  alias Anthropix.APIError
-
-  setup_all do
-    {:ok, pid} = Bandit.start_link(plug: Anthropix.MockServer)
-    on_exit(fn -> Process.exit(pid, :normal) end)
-    {:ok, client: Anthropix.init("test_key", base_url: "http://localhost:4000")}
-  end
+  alias Anthropix.{APIError, Mock}
 
   describe "init without api_key" do
     test "raises if no api_key in config" do
@@ -47,7 +41,8 @@ defmodule AnthropixTest do
   end
 
   describe "chat/2" do
-    test "generates a response for a given prompt", %{client: client} do
+    test "generates a response for a given prompt" do
+      client = Mock.client(& Mock.respond(&1, :messages))
       assert {:ok, res} = Anthropix.chat(client, [
         model: "claude-3-sonnet-20240229",
         messages: [
@@ -60,7 +55,8 @@ defmodule AnthropixTest do
       assert Enum.all?(res["content"], &is_map/1)
     end
 
-    test "streams a response for a given prompt", %{client: client} do
+    test "streams a response for a given prompt" do
+      client = Mock.client(& Mock.stream(&1, :messages))
       assert {:ok, stream} = Anthropix.chat(client, [
         model: "claude-3-sonnet-20240229",
         messages: [
@@ -75,7 +71,8 @@ defmodule AnthropixTest do
       assert get_in(last, ["usage", "output_tokens"]) == 34
     end
 
-    test "returns error when model not found", %{client: client} do
+    test "returns error when model not found" do
+      client = Mock.client(& Mock.respond(&1, 404))
       assert {:error, %APIError{type: "not_found"}} = Anthropix.chat(client, [
         model: "not-found",
         messages: [
@@ -86,7 +83,8 @@ defmodule AnthropixTest do
   end
 
   describe "streaming methods" do
-    test "with stream: true, returns a lazy enumerable", %{client: client} do
+    test "with stream: true, returns a lazy enumerable" do
+      client = Mock.client(& Mock.stream(&1, :messages))
       assert {:ok, stream} = Anthropix.chat(client, [
         model: "claude-3-sonnet-20240229",
         messages: [
@@ -99,8 +97,9 @@ defmodule AnthropixTest do
       assert Enum.to_list(stream) |> length() == 31
     end
 
-    test "with stream: pid, returns a task and sends messages to pid", %{client: client} do
+    test "with stream: pid, returns a task and sends messages to pid" do
       {:ok, pid} = Anthropix.StreamCatcher.start_link()
+      client = Mock.client(& Mock.stream(&1, :messages))
       assert {:ok, task} = Anthropix.chat(client, [
         model: "claude-3-sonnet-20240229",
         messages: [

test/support/mock_server.ex → test/support/mock.ex

@@ -1,24 +1,71 @@
-defmodule Anthropix.MockServer do
-  use Plug.Router
+defmodule Anthropix.Mock do
+  alias Plug.Conn.Status
+  import Plug.Conn
 
   @mocks %{
-    messages: """
-    {
-      "content": [
-        {
-          "type": "text",
-          "text": "Here's a haiku about the color of the sky:\\n\\nBlue canvas stretched wide,\\nClouds drift lazily across,\\nSky's endless expanse."
+    :messages => %{
+      "content" => [
+        %{
+          "type" => "text",
+          "text" => "Here's a haiku about the color of the sky:\\n\\nBlue canvas stretched wide,\\nClouds drift lazily across,\\nSky's endless expanse."
         }
       ],
-      "id": "msg_test",
-      "model": "claude-3-sonnet-20240229",
-      "role": "assistant",
-      "stop_reason": "end_turn",
-      "stop_sequence": null,
-      "type": "message",
-      "usage": {"input_tokens": 18, "output_tokens": 36}
+      "id" => "msg_test",
+      "model" => "claude-3-sonnet-20240229",
+      "role" => "assistant",
+      "stop_reason" => "end_turn",
+      "stop_sequence" => nil,
+      "type" => "message",
+      "usage" => %{"input_tokens" => 18, "output_tokens" => 36}
+    },
+
+    {:agent, :messages, 1} => %{
+      "content" => [
+        %{
+          "text" => "Here's how I'll find the current stock price for General Motors, yo:\n\n<function_calls>\n<invoke>\n<tool_name>get_ticker_symbol</tool_name>\n<parameters>\n<company_name>General Motors</company_name>\n</parameters>\n</invoke>\n",
+          "type" => "text"
+        }
+      ],
+      "id" => "msg_01FJNLQSEpck6s1cimMJ9Ru1",
+      "model" => "claude-3-sonnet-20240229",
+      "role" => "assistant",
+      "stop_reason" => "stop_sequence",
+      "stop_sequence" => "</function_calls>",
+      "type" => "message",
+      "usage" => %{"input_tokens" => 333, "output_tokens" => 73}
+    },
+
+    {:agent, :messages, 3} => %{
+      "content" => [
+        %{
+          "text" => "Aight, got the ticker symbol GM for General Motors. Now let me look up that current stock price:\n\n<function_calls>\n<invoke>\n<tool_name>get_current_stock_price</tool_name>\n<parameters>\n<symbol>GM</symbol>\n</parameters>\n</invoke>\n",
+          "type" => "text"
+        }
+      ],
+      "id" => "msg_01QjBCr47TNUzpPvqnr5Yfn5",
+      "model" => "claude-3-sonnet-20240229",
+      "role" => "assistant",
+      "stop_reason" => "stop_sequence",
+      "stop_sequence" => "</function_calls>",
+      "type" => "message",
+      "usage" => %{"input_tokens" => 439, "output_tokens" => 77}
+    },
+
+    {:agent, :messages, 5} => %{
+      "content" => [
+        %{
+          "text" => "Word, the current stock price for General Motors is $39.21. Representing that big auto money, ya dig? Gotta make them stacks and invest wisely in the motor city players.",
+          "type" => "text"
+        }
+      ],
+      "id" => "msg_01Nq9eY6wnrSSs5QDKYD9yHE",
+      "model" => "claude-3-sonnet-20240229",
+      "role" => "assistant",
+      "stop_reason" => "end_turn",
+      "stop_sequence" => nil,
+      "type" => "message",
+      "usage" => %{"input_tokens" => 553, "output_tokens" => 45}
     }
-    """
   }
 
   @stream_mocks %{
@@ -66,38 +113,42 @@ defmodule Anthropix.MockServer do
     ]
   }
 
-  plug :match
-  plug Plug.Parsers, parsers: [:json], json_decoder: Jason
-  plug :dispatch
-
-  post "/messages", do: handle_request(conn, :messages)
-
-  defp handle_request(conn, name) do
-    case conn.body_params do
-      %{"model" => "not-found"} -> respond(conn, 404)
-      %{"stream" => true} -> stream(conn, name)
-      _ -> respond(conn, name)
-    end
+  @spec client(function()) :: Anthropix.t()
+  def client(plug) when is_function(plug, 1) do
+    struct(Anthropix, req: Req.new(plug: plug))
   end
 
-  defp respond(conn, name) when is_atom(name) do
+  @spec respond(Plug.Conn.t(), term()) :: Plug.Conn.t()
+  def respond(conn, name) when is_atom(name) do
     conn
     |> put_resp_header("content-type", "application/json")
-    |> send_resp(200, @mocks[name])
+    |> send_resp(200, Jason.encode!(@mocks[name]))
   end
 
-  defp respond(conn, status) when is_number(status) do
+  def respond(conn, status) when is_number(status) do
     conn
     |> put_resp_header("content-type", "application/json")
     |> send_resp(status, Jason.encode!(%{
       error: %{
-        type: Plug.Conn.Status.reason_atom(status),
-        message: Plug.Conn.Status.reason_phrase(status),
+        type: Status.reason_atom(status),
+        message: Status.reason_phrase(status),
       }
     }))
   end
 
-  defp stream(conn, name) do
+  def respond(conn, {:agent, :messages}) do
+    # Manually run through the json parser plug
+    opts = Plug.Parsers.init(parsers: [:json], json_decoder: Jason)
+    conn = Plug.Parsers.call(conn, opts)
+
+    m = conn.body_params["messages"]
+    conn
+    |> put_resp_header("content-type", "application/json")
+    |> send_resp(200, Jason.encode!(@mocks[{:agent, :messages, length(m)}]))
+  end
+
+  @spec stream(Plug.Conn.t(), term()) :: Plug.Conn.t()
+  def stream(conn, name) when is_atom(name) do
     Enum.reduce(@stream_mocks[name], send_chunked(conn, 200), fn chunk, conn ->
       {:ok, conn} = chunk(conn, to_sse_event(chunk))
       conn