Add ollama provider (#31)

Author: joelpaulkoch

lib/rag/ai/ollama.ex

@@ -0,0 +1,91 @@
+defmodule Rag.Ai.Ollama do
+  @moduledoc """
+  Implementation of `Rag.Ai.Provider` using Ollama.
+  """
+
+  @behaviour Rag.Ai.Provider
+
+  @type t :: %__MODULE__{
+          embeddings_url: String.t() | nil,
+          embeddings_model: String.t() | nil,
+          text_url: String.t() | nil,
+          text_model: String.t() | nil
+        }
+  defstruct embeddings_url: "http://localhost:11434/api/embed",
+            embeddings_model: nil,
+            text_url: "http://localhost:11434/api/chat",
+            text_model: nil
+
+  @impl Rag.Ai.Provider
+  def new(attrs) do
+    struct!(__MODULE__, attrs)
+  end
+
+  @impl Rag.Ai.Provider
+  def generate_embeddings(%__MODULE__{} = provider, texts, _opts \\ []) do
+    req_params =
+      [
+        json: %{"model" => provider.embeddings_model, "input" => texts}
+      ]
+
+    with {:ok, %Req.Response{status: 200} = response} <-
+           Req.post(provider.embeddings_url, req_params),
+         {:ok, embeddings} <- get_embeddings(response) do
+      {:ok, embeddings}
+    else
+      {:ok, %Req.Response{status: status}} ->
+        {:error, "HTTP request failed with status code #{status}"}
+
+      {:error, reason} ->
+        {:error, reason}
+    end
+  end
+
+  defp get_embeddings(response) do
+    path = ["embeddings"]
+
+    case get_in(response.body, path) do
+      nil ->
+        {:error,
+         "failed to access embeddings from path embeddings in response #{inspect(response.body)}"}
+
+      embeddings ->
+        {:ok, embeddings}
+    end
+  end
+
+  @impl Rag.Ai.Provider
+  def generate_text(%__MODULE__{} = provider, prompt, _opts \\ []) do
+    req_params =
+      [
+        json: %{
+          "model" => provider.text_model,
+          "messages" => [%{role: :user, content: prompt}],
+          "stream" => false
+        }
+      ]
+
+    with {:ok, %Req.Response{status: 200} = response} <- Req.post(provider.text_url, req_params),
+         {:ok, text} <- get_text(response) do
+      {:ok, text}
+    else
+      {:ok, %Req.Response{status: status}} ->
+        {:error, "HTTP request failed with status code #{status}"}
+
+      {:error, reason} ->
+        {:error, reason}
+    end
+  end
+
+  defp get_text(response) do
+    path = ["message", "content"]
+
+    case get_in(response.body, path) do
+      nil ->
+        {:error, "failed to access text from path response in response #{inspect(response.body)}"}
+
+      text ->
+        {:ok, text}
+    end
+  end
+end

test/rag/embedding/http_test.exs

@@ -208,5 +208,13 @@ defmodule Rag.Embedding.HttpTest do
       [%{text: "hello", embedding: _embedding}] =
         Embedding.generate_embeddings_batch([%{text: "hello"}], provider, [])
     end
+
+    @tag :integration_test
+    test "ollama embeddings" do
+      provider = Ai.Ollama.new(%{embeddings_model: "unclemusclez/jina-embeddings-v2-base-code"})
+
+      assert [%{text: "hello", embedding: _embedding}] =
+               Embedding.generate_embeddings_batch([%{text: "hello"}], provider, [])
+    end
   end
 end

test/rag/generation/http_test.exs

@@ -40,5 +40,16 @@ defmodule Rag.Generation.HttpTest do
       %Generation{query: "test?", response: _response} =
         Generation.generate_response(%Generation{query: "test?", prompt: "prompt"}, provider)
     end
+
+    @tag :integration_test
+    test "ollama generation" do
+      provider = Ai.Ollama.new(%{text_model: "llama3.2:latest"})
+
+      assert %Generation{query: "test?", response: _response} =
+               Generation.generate_response(
+                 %Generation{query: "test?", prompt: "prompt"},
+                 provider
+               )
+    end
   end
 end