{-# LANGUAGE TypeFamilies #-}

{- |
Module      : Langchain.Retriever.Core
Description : Retrieval mechanism implementation for LangChain Haskell
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental

Haskell implementation of LangChain's retrieval abstraction, providing:

- Document retrieval based on semantic similarity
- Integration with vector stores
- Runnable interface for workflow composition

Example usage:

@
-- Hypothetical vector store instance
vectorStore :: MyVectorStore
vectorStore = ...

-- Create retriever
retriever :: VectorStoreRetriever MyVectorStore
retriever = VectorStoreRetriever vectorStore

-- Retrieve relevant documents
docs <- invoke retriever "Haskell programming"
-- Right [Document {pageContent = "...", ...}, ...]
@
-}
module Langchain.Retriever.Core
  ( Retriever (..)
  , VectorStoreRetriever (..)
  ) where

import Langchain.DocumentLoader.Core (Document)
import Langchain.Runnable.Core
import Langchain.VectorStore.Core

import Data.Text (Text)

{- | Typeclass for document retrieval systems
Implementations should return documents relevant to a given query.

Example instance for a custom retriever:

@
data CustomRetriever = CustomRetriever

instance Retriever CustomRetriever where
  _get_relevant_documents _ query = do
    -- Custom retrieval logic
    return $ Right [Document ("Result for: " <> query) mempty]
@
-}
class Retriever a where
  -- | Retrieve documents relevant to the query
  --
  --   Example:
  --
  --   >>> _get_relevant_documents (VectorStoreRetriever myStore) "AI"
  --   Right [Document "AI definition...", ...]
  _get_relevant_documents :: a -> Text -> IO (Either String [Document])

{- | Vector store-backed retriever implementation
Wraps any 'VectorStore' instance to provide similarity-based retrieval.

Example usage:

@
-- Using a hypothetical FAISS vector store
faissStore :: FAISSStore
faissStore = ...

-- Create vector store retriever
vsRetriever = VectorStoreRetriever faissStore

-- Get similar documents
docs <- _get_relevant_documents vsRetriever "machine learning"
-- Returns top 5 relevant documents by default
@
-}
newtype VectorStore a => VectorStoreRetriever a = VectorStoreRetriever {forall a. VectorStore a => VectorStoreRetriever a -> a
vs :: a}
  deriving (VectorStoreRetriever a -> VectorStoreRetriever a -> Bool
(VectorStoreRetriever a -> VectorStoreRetriever a -> Bool)
-> (VectorStoreRetriever a -> VectorStoreRetriever a -> Bool)
-> Eq (VectorStoreRetriever a)
forall a.
Eq a =>
VectorStoreRetriever a -> VectorStoreRetriever a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a.
Eq a =>
VectorStoreRetriever a -> VectorStoreRetriever a -> Bool
== :: VectorStoreRetriever a -> VectorStoreRetriever a -> Bool
$c/= :: forall a.
Eq a =>
VectorStoreRetriever a -> VectorStoreRetriever a -> Bool
/= :: VectorStoreRetriever a -> VectorStoreRetriever a -> Bool
Eq, Int -> VectorStoreRetriever a -> ShowS
[VectorStoreRetriever a] -> ShowS
VectorStoreRetriever a -> String
(Int -> VectorStoreRetriever a -> ShowS)
-> (VectorStoreRetriever a -> String)
-> ([VectorStoreRetriever a] -> ShowS)
-> Show (VectorStoreRetriever a)
forall a.
(VectorStore a, Show a) =>
Int -> VectorStoreRetriever a -> ShowS
forall a.
(VectorStore a, Show a) =>
[VectorStoreRetriever a] -> ShowS
forall a.
(VectorStore a, Show a) =>
VectorStoreRetriever a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a.
(VectorStore a, Show a) =>
Int -> VectorStoreRetriever a -> ShowS
showsPrec :: Int -> VectorStoreRetriever a -> ShowS
$cshow :: forall a.
(VectorStore a, Show a) =>
VectorStoreRetriever a -> String
show :: VectorStoreRetriever a -> String
$cshowList :: forall a.
(VectorStore a, Show a) =>
[VectorStoreRetriever a] -> ShowS
showList :: [VectorStoreRetriever a] -> ShowS
Show)

{- | Runnable interface for vector store retrievers
Allows integration with LangChain workflows and expressions.

Example:

>>> invoke (VectorStoreRetriever store) "Quantum computing"
Right [Document "Quantum theory...", ...]
-}
instance VectorStore a => Retriever (VectorStoreRetriever a) where
  _get_relevant_documents :: VectorStoreRetriever a -> Text -> IO (Either String [Document])
_get_relevant_documents (VectorStoreRetriever a
v) Text
query = a -> Text -> Int -> IO (Either String [Document])
forall m.
VectorStore m =>
m -> Text -> Int -> IO (Either String [Document])
similaritySearch a
v Text
query Int
5

{- | Runnable interface for vector store retrievers
Allows integration with LangChain workflows and expressions.

Example:

>>> invoke (VectorStoreRetriever store) "Quantum computing"
Right [Document "Quantum theory...", ...]
-}
instance VectorStore a => Runnable (VectorStoreRetriever a) where
  type RunnableInput (VectorStoreRetriever a) = Text
  type RunnableOutput (VectorStoreRetriever a) = [Document]

  invoke :: VectorStoreRetriever a
-> RunnableInput (VectorStoreRetriever a)
-> IO (Either String (RunnableOutput (VectorStoreRetriever a)))
invoke VectorStoreRetriever a
retriever RunnableInput (VectorStoreRetriever a)
query = VectorStoreRetriever a -> Text -> IO (Either String [Document])
forall a. Retriever a => a -> Text -> IO (Either String [Document])
_get_relevant_documents VectorStoreRetriever a
retriever Text
RunnableInput (VectorStoreRetriever a)
query

{- $examples
Test case patterns:
1. Basic retrieval
   >>> let retriever = VectorStoreRetriever mockStore
   >>> _get_relevant_documents retriever "Test"
   Right [Document "Test content" ...]

2. Runnable integration
   >>> run retriever "Hello"
   Right [Document "Greeting response" ...]

3. Error handling
   >>> _get_relevant_documents (VectorStoreRetriever invalidStore) "Query"
   Left "Vector store error"
-}