{-# LANGUAGE RecordWildCards #-}

{- |
Module      : Langchain.Embeddings.Ollama
Description : Ollama integration for text embeddings in LangChain Haskell
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental

Ollama implementation of LangChain's embedding interface. Supports document and query
embedding generation through Ollama's API.

Example usage:

@
-- Create Ollama embeddings configuration
ollamaEmb = OllamaEmbeddings
  { model = "llama3"
  , defaultTruncate = Just True
  , defaultKeepAlive = Just "5m"
  }

-- Embed query text
queryVec <- embedQuery ollamaEmb "What is Haskell?"
-- Right [0.12, 0.34, ...]

-- Embed document collection
doc <- Document "Haskell is a functional programming language" mempty
docsVec <- embedDocuments ollamaEmb [doc]
-- Right [[0.56, 0.78, ...]]
@
-}
module Langchain.Embeddings.Ollama
  ( OllamaEmbeddings (..)
  ) where

import Data.Maybe
import Data.Ollama.Embeddings
import Data.Text (Text)
import Langchain.DocumentLoader.Core
import Langchain.Embeddings.Core

{- | Ollama-specific embedding configuration
Contains parameters for controlling:

- Model selection
- Input truncation behavior
- Model caching via keep-alive

Example configuration:

>>> OllamaEmbeddings "nomic-embed" (Just False) (Just "1h")
OllamaEmbeddings {model = "nomic-embed", ...}
-}
data OllamaEmbeddings = OllamaEmbeddings
  { OllamaEmbeddings -> Text
model :: Text
  -- ^ The name of the Ollama model to use for embeddings
  , OllamaEmbeddings -> Maybe Bool
defaultTruncate :: Maybe Bool
  -- ^ Optional flag to truncate input if supported by the API
  , OllamaEmbeddings -> Maybe Text
defaultKeepAlive :: Maybe Text
  -- ^ Keep model loaded for specified duration (e.g., "5m")
  }
  deriving (Int -> OllamaEmbeddings -> ShowS
[OllamaEmbeddings] -> ShowS
OllamaEmbeddings -> String
(Int -> OllamaEmbeddings -> ShowS)
-> (OllamaEmbeddings -> String)
-> ([OllamaEmbeddings] -> ShowS)
-> Show OllamaEmbeddings
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> OllamaEmbeddings -> ShowS
showsPrec :: Int -> OllamaEmbeddings -> ShowS
$cshow :: OllamaEmbeddings -> String
show :: OllamaEmbeddings -> String
$cshowList :: [OllamaEmbeddings] -> ShowS
showList :: [OllamaEmbeddings] -> ShowS
Show, OllamaEmbeddings -> OllamaEmbeddings -> Bool
(OllamaEmbeddings -> OllamaEmbeddings -> Bool)
-> (OllamaEmbeddings -> OllamaEmbeddings -> Bool)
-> Eq OllamaEmbeddings
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: OllamaEmbeddings -> OllamaEmbeddings -> Bool
== :: OllamaEmbeddings -> OllamaEmbeddings -> Bool
$c/= :: OllamaEmbeddings -> OllamaEmbeddings -> Bool
/= :: OllamaEmbeddings -> OllamaEmbeddings -> Bool
Eq)

go :: EmbeddingResp -> Either String [Float]
go :: EmbeddingResp -> Either String [Float]
go EmbeddingResp
embResp =
  case [[Float]] -> Maybe [Float]
forall a. [a] -> Maybe a
listToMaybe (EmbeddingResp -> [[Float]]
embedding_ EmbeddingResp
embResp) of
    Maybe [Float]
Nothing -> String -> Either String [Float]
forall a b. a -> Either a b
Left String
"Embeddings are empty"
    Just [Float]
x -> [Float] -> Either String [Float]
forall a b. b -> Either a b
Right [Float]
x

{- | Ollama implementation of the 'Embeddings' interface [[6]].
Uses Ollama's embedding API for vector generation. Handles:
- Multiple document embedding via batch processing
- Query embedding for similarity searches
- Error propagation from API responses

Example instance usage:
@
-- Embed multiple documents
docs <- load (FileLoader "data.txt")
case docs of
  Right documents -> do
    vecs <- embedDocuments ollamaEmb documents
    -- Use vectors for semantic search
  Left err -> print err
@
-}
instance Embeddings OllamaEmbeddings where
  -- \| Document embedding implementation [[3]]:
  --  Processes each document individually through Ollama's API.
  --
  --  Example:
  --  >>> let doc = Document "Test content" mempty
  --  >>> embedDocuments ollamaEmb [doc]
  --  Right [[0.1, 0.2, ...], ...]
  --
  embedDocuments :: OllamaEmbeddings -> [Document] -> IO (Either String [[Float]])
embedDocuments (OllamaEmbeddings {Maybe Bool
Maybe Text
Text
model :: OllamaEmbeddings -> Text
defaultTruncate :: OllamaEmbeddings -> Maybe Bool
defaultKeepAlive :: OllamaEmbeddings -> Maybe Text
model :: Text
defaultTruncate :: Maybe Bool
defaultKeepAlive :: Maybe Text
..}) [Document]
docs = do
    -- For each input text, make an individual API call
    [Either String EmbeddingResp]
results <- (Document -> IO (Either String EmbeddingResp))
-> [Document] -> IO [Either String EmbeddingResp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Document
doc -> Text
-> Text
-> Maybe Bool
-> Maybe Text
-> IO (Either String EmbeddingResp)
embeddingOps Text
model (Document -> Text
pageContent Document
doc) Maybe Bool
defaultTruncate Maybe Text
defaultKeepAlive) [Document]
docs
    -- Combine the results, handling errors appropriately
    Either String [[Float]] -> IO (Either String [[Float]])
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String [[Float]] -> IO (Either String [[Float]]))
-> Either String [[Float]] -> IO (Either String [[Float]])
forall a b. (a -> b) -> a -> b
$
      [Either String EmbeddingResp] -> Either String [EmbeddingResp]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [Either String EmbeddingResp]
results Either String [EmbeddingResp]
-> ([EmbeddingResp] -> Either String [[Float]])
-> Either String [[Float]]
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \[EmbeddingResp]
resps ->
        (EmbeddingResp -> Either String [Float])
-> [EmbeddingResp] -> Either String [[Float]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM EmbeddingResp -> Either String [Float]
go [EmbeddingResp]
resps

  -- \| Query embedding implementation:
  --  Generates vector representation for search queries.
  --
  --  Example:
  --  >>> embedQuery ollamaEmb "Explain monads"
  --  Right [0.3, 0.4, ...]
  --
  embedQuery :: OllamaEmbeddings -> Text -> IO (Either String [Float])
embedQuery (OllamaEmbeddings {Maybe Bool
Maybe Text
Text
model :: OllamaEmbeddings -> Text
defaultTruncate :: OllamaEmbeddings -> Maybe Bool
defaultKeepAlive :: OllamaEmbeddings -> Maybe Text
model :: Text
defaultTruncate :: Maybe Bool
defaultKeepAlive :: Maybe Text
..}) Text
query = do
    Either String EmbeddingResp
res <- Text
-> Text
-> Maybe Bool
-> Maybe Text
-> IO (Either String EmbeddingResp)
embeddingOps Text
model Text
query Maybe Bool
defaultTruncate Maybe Text
defaultKeepAlive
    case (EmbeddingResp -> [[Float]])
-> Either String EmbeddingResp -> Either String [[Float]]
forall a b. (a -> b) -> Either String a -> Either String b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap EmbeddingResp -> [[Float]]
embedding_ Either String EmbeddingResp
res of
      Left String
err -> Either String [Float] -> IO (Either String [Float])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String [Float] -> IO (Either String [Float]))
-> Either String [Float] -> IO (Either String [Float])
forall a b. (a -> b) -> a -> b
$ String -> Either String [Float]
forall a b. a -> Either a b
Left String
err
      Right [[Float]]
lst ->
        case [[Float]] -> Maybe [Float]
forall a. [a] -> Maybe a
listToMaybe [[Float]]
lst of
          Maybe [Float]
Nothing -> Either String [Float] -> IO (Either String [Float])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String [Float] -> IO (Either String [Float]))
-> Either String [Float] -> IO (Either String [Float])
forall a b. (a -> b) -> a -> b
$ String -> Either String [Float]
forall a b. a -> Either a b
Left String
"Embeddings are empty"
          Just [Float]
x -> Either String [Float] -> IO (Either String [Float])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String [Float] -> IO (Either String [Float]))
-> Either String [Float] -> IO (Either String [Float])
forall a b. (a -> b) -> a -> b
$ [Float] -> Either String [Float]
forall a b. b -> Either a b
Right [Float]
x