From 98090ea8638ff969b45da562a2450df7a7a713c7 Mon Sep 17 00:00:00 2001
From: lebrunel <124721263+lebrunel@users.noreply.github.com>
Date: Wed, 21 Feb 2024 14:58:07 +0000
Subject: [PATCH] user agent header. a few tests covering Ollama.init/1

---
 lib/ollama.ex        |  8 +++++++-
 test/ollama_test.exs | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/lib/ollama.ex b/lib/ollama.ex
index 885a910..21c071b 100644
--- a/lib/ollama.ex
+++ b/lib/ollama.ex
@@ -1,4 +1,5 @@
 defmodule Ollama do
+  @version Keyword.fetch!(Mix.Project.config(), :version)
   @moduledoc """
   ![Ollama-ex](https://raw.githubusercontent.com/lebrunel/ollama-ex/main/media/poster.webp)
 
@@ -21,7 +22,7 @@ defmodule Ollama do
   ```elixir
   def deps do
     [
-      {:ollama, "#{Keyword.fetch!(Mix.Project.config(), :version)}"}
+      {:ollama, "#{@version}"}
     ]
   end
   ```
@@ -201,6 +202,9 @@ defmodule Ollama do
 
   @default_req_opts [
     base_url: "http://localhost:11434/api",
+    headers: [
+      {"user-agent", "ollama-ex/#{@version}"}
+    ],
     receive_timeout: 60_000,
   ]
 
@@ -240,9 +244,11 @@ defmodule Ollama do
 
   @spec init_req(keyword()) :: Req.Request.t()
   defp init_req(opts) do
+    {headers, opts} = Keyword.pop(opts, :headers, [])
     @default_req_opts
     |> Keyword.merge(opts)
     |> Req.new()
+    |> Req.update(headers: headers)
   end
 
 
diff --git a/test/ollama_test.exs b/test/ollama_test.exs
index 99465e9..516bc27 100644
--- a/test/ollama_test.exs
+++ b/test/ollama_test.exs
@@ -8,6 +8,39 @@ defmodule OllamaTest do
     {:ok, client: Ollama.init("http://localhost:4000")}
   end
 
+  describe "init/2" do
+    test "default client" do
+      client = Ollama.init()
+      assert "http://localhost:11434/api" = client.req.options.base_url
+      assert %{"user-agent" => _val} = client.req.headers
+    end
+
+    test "client with custom base url" do
+      client = Ollama.init("https://ollama.my.site/api")
+      assert "https://ollama.my.site/api" = client.req.options.base_url
+    end
+
+    test "client with custom req opts" do
+      client = Ollama.init(receive_timeout: :infinity)
+      assert "http://localhost:11434/api" = client.req.options.base_url
+      assert :infinity = client.req.options.receive_timeout
+    end
+
+    test "client with custom req struct" do
+      client = Ollama.init(Req.new(base_url: "https://ollama.my.site/api"))
+      assert "https://ollama.my.site/api" = client.req.options.base_url
+    end
+
+    test "client with merged headers" do
+      client = Ollama.init(headers: [
+        {"User-Agent", "testing"},
+        {"X-Test", "testing"},
+      ])
+      assert "http://localhost:11434/api" = client.req.options.base_url
+      assert %{"user-agent" => ["testing"], "x-test" => ["testing"]} = client.req.headers
+    end
+  end
+
   describe "chat2" do
     test "generates a response for a given prompt", %{client: client} do
       assert {:ok, res} = Ollama.chat(client, [