-- | Main entrypoint to the API
module Pinecone
    ( -- * Methods
      getClientEnv
    , makeControlMethods
    , ControlMethods(..)
    , makeDataMethods
    , DataMethods(..)

      -- * Servant
    , ControlAPI
    , DataAPI
    ) where

import Data.Foldable (toList)
import Data.Functor (void)
import Data.Proxy (Proxy(..))
import Pinecone.Embed (GenerateVectors, Embeddings)
import Pinecone.Prelude
import Pinecone.Rerank (Documents(..), RerankResults(..))
import Prelude hiding (id)
import Servant.Client (ClientEnv)
import Servant.Client.Core (BaseUrl(..), Scheme(..))

import Pinecone.Backups
    (CreateCollection, Collection, CollectionModel, Collections(..))
import Pinecone.Imports
    (Import, ImportModel, Imports, StartImport, StartImportResponse(..))
import Pinecone.Indexes
    ( ConfigureIndex
    , CreateIndex
    , CreateIndexWithEmbedding
    , GetIndexStats
    , Host(..)
    , Index
    , IndexModel
    , IndexModels(..)
    , IndexStats
    )
import Pinecone.Search
    ( Hits
    , Matches
    , SearchWithText
    , SearchWithVector
    )
import Pinecone.Vectors
    ( DeleteVectors
    , Namespace
    , Record
    , UpdateVector
    , UpsertVectors
    , UpsertStats
    , VectorIDs
    , Vectors
    )

import qualified Control.Exception as Exception
import qualified Data.Text as Text
import qualified Pinecone.Backups as Backups
import qualified Pinecone.Embed as Embed
import qualified Pinecone.Imports as Imports
import qualified Pinecone.Indexes as Indexes
import qualified Pinecone.Rerank as Rerank
import qualified Pinecone.Search as Search
import qualified Pinecone.Vectors as Vectors
import qualified Network.HTTP.Client as HTTP.Client
import qualified Network.HTTP.Client.TLS as TLS
import qualified Servant.Client as Client

-- | Convenient utility to get a `ClientEnv` for the most common use case
getClientEnv
    :: Host
    -- ^ Base URL for API
    --
    -- __CAREFULLY NOTE:__ This should be @\"app.pinecone.io\"@ for
    -- `makeControlMethods` and should be /the index host/ for `makeDataMethods`
    -- (i.e. the `Pinecone.Indexes.host` field of a returned `IndexModel`).
    -> IO ClientEnv
getClientEnv :: Host -> IO ClientEnv
getClientEnv (Host Text
baseUrlText) = do
    BaseUrl
baseUrl <- String -> IO BaseUrl
forall (m :: * -> *). MonadThrow m => String -> m BaseUrl
Client.parseBaseUrl (Text -> String
Text.unpack Text
baseUrlText)

    -- The hosts returned by the Pinecone API don't include the `https://`
    -- prefix, which means that by default they're parsed as insecure URLs,
    -- which we fix here.
    --
    -- If you don't like this behavior or you specifically want to use an
    -- insecure URL then just use the `mkClientEnv` function yourself instead
    -- of using this `getClientEnv` utility function.
    let newBaseUrl :: BaseUrl
newBaseUrl = BaseUrl
baseUrl
          { baseUrlScheme = Https
          , baseUrlPort = 443
          }

    let managerSettings :: ManagerSettings
managerSettings = ManagerSettings
TLS.tlsManagerSettings
            { HTTP.Client.managerResponseTimeout =
                HTTP.Client.responseTimeoutNone
            }

    Manager
manager <- ManagerSettings -> IO Manager
forall (m :: * -> *). MonadIO m => ManagerSettings -> m Manager
TLS.newTlsManagerWith ManagerSettings
managerSettings

    ClientEnv -> IO ClientEnv
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Manager -> BaseUrl -> ClientEnv
Client.mkClientEnv Manager
manager BaseUrl
newBaseUrl)

apiVersion :: Text
apiVersion :: Text
apiVersion = Text
"2025-01"

-- | Get a record of control API methods after providing an API token
makeControlMethods
    :: ClientEnv
    -- ^
    -> Text
    -- ^ API token
    -> ControlMethods
makeControlMethods :: ClientEnv -> Text -> ControlMethods
makeControlMethods ClientEnv
clientEnv Text
token = ControlMethods{IO (Vector IndexModel)
IO (Vector CollectionModel)
GenerateVectors -> IO Embeddings
CreateIndexWithEmbedding -> IO IndexModel
CreateIndex -> IO IndexModel
Index -> IO ()
Index -> IO IndexModel
Index -> ConfigureIndex -> IO IndexModel
CreateCollection -> IO CollectionModel
Collection -> IO ()
Collection -> IO CollectionModel
RerankResults -> IO Documents
createIndex :: CreateIndex -> IO IndexModel
createIndexWithEmbedding :: CreateIndexWithEmbedding -> IO IndexModel
describeIndex :: Index -> IO IndexModel
configureIndex :: Index -> ConfigureIndex -> IO IndexModel
createCollection :: CreateCollection -> IO CollectionModel
describeCollection :: Collection -> IO CollectionModel
generateVectors :: GenerateVectors -> IO Embeddings
rerankResults :: RerankResults -> IO Documents
listIndexes :: IO (Vector IndexModel)
listCollections :: IO (Vector CollectionModel)
deleteIndex :: Index -> IO ()
deleteCollection :: Collection -> IO ()
$sel:listIndexes:ControlMethods :: IO (Vector IndexModel)
$sel:createIndex:ControlMethods :: CreateIndex -> IO IndexModel
$sel:createIndexWithEmbedding:ControlMethods :: CreateIndexWithEmbedding -> IO IndexModel
$sel:describeIndex:ControlMethods :: Index -> IO IndexModel
$sel:deleteIndex:ControlMethods :: Index -> IO ()
$sel:configureIndex:ControlMethods :: Index -> ConfigureIndex -> IO IndexModel
$sel:listCollections:ControlMethods :: IO (Vector CollectionModel)
$sel:createCollection:ControlMethods :: CreateCollection -> IO CollectionModel
$sel:describeCollection:ControlMethods :: Collection -> IO CollectionModel
$sel:deleteCollection:ControlMethods :: Collection -> IO ()
$sel:generateVectors:ControlMethods :: GenerateVectors -> IO Embeddings
$sel:rerankResults:ControlMethods :: RerankResults -> IO Documents
..}
  where
    (       (     IO IndexModels
listIndexes_
            :<|>  CreateIndex -> IO IndexModel
createIndex
            :<|>  CreateIndexWithEmbedding -> IO IndexModel
createIndexWithEmbedding
            :<|>  Index -> IO IndexModel
describeIndex
            :<|>  Index -> IO NoContent
deleteIndex_
            :<|>  Index -> ConfigureIndex -> IO IndexModel
configureIndex
            )
      :<|>  (     IO Collections
listCollections_
            :<|>  CreateCollection -> IO CollectionModel
createCollection
            :<|>  Collection -> IO CollectionModel
describeCollection
            :<|>  Collection -> IO NoContent
deleteCollection_
            )
      :<|>  GenerateVectors -> IO Embeddings
generateVectors
      :<|>  RerankResults -> IO Documents
rerankResults
      ) = forall api (m :: * -> *) (n :: * -> *).
HasClient ClientM api =>
Proxy api -> (forall a. m a -> n a) -> Client m api -> Client n api
Client.hoistClient @ControlAPI Proxy ControlAPI
forall {k} (t :: k). Proxy t
Proxy (ClientEnv -> ClientM a -> IO a
forall a. ClientEnv -> ClientM a -> IO a
run ClientEnv
clientEnv) (forall api.
HasClient ClientM api =>
Proxy api -> Client ClientM api
Client.client @ControlAPI Proxy ControlAPI
forall {k} (t :: k). Proxy t
Proxy) Text
token Text
apiVersion

    listIndexes :: IO (Vector IndexModel)
listIndexes = (IndexModels -> Vector IndexModel)
-> IO IndexModels -> IO (Vector IndexModel)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap IndexModels -> Vector IndexModel
indexes IO IndexModels
listIndexes_
    listCollections :: IO (Vector CollectionModel)
listCollections = (Collections -> Vector CollectionModel)
-> IO Collections -> IO (Vector CollectionModel)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Collections -> Vector CollectionModel
collections IO Collections
listCollections_
    deleteIndex :: Index -> IO ()
deleteIndex Index
a = IO NoContent -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Index -> IO NoContent
deleteIndex_ Index
a)
    deleteCollection :: Collection -> IO ()
deleteCollection Collection
a = IO NoContent -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Collection -> IO NoContent
deleteCollection_ Collection
a)

-- | Get a record of control API methods after providing an API token
makeDataMethods
    :: ClientEnv
    -- ^
    -> Text
    -- ^ API token
    -> DataMethods
makeDataMethods :: ClientEnv -> Text -> DataMethods
makeDataMethods ClientEnv
clientEnv Text
token = DataMethods{Maybe Nat -> Maybe Text -> IO Imports
Maybe Text
-> Maybe Nat -> Maybe Text -> Maybe Namespace -> IO VectorIDs
Vector Text -> Maybe Namespace -> IO Vectors
StartImport -> IO Import
Import -> IO ()
Import -> IO ImportModel
GetIndexStats -> IO IndexStats
DeleteVectors -> IO ()
UpdateVector -> IO ()
UpsertVectors -> IO UpsertStats
Namespace -> Record -> IO ()
Namespace -> SearchWithText -> IO Hits
SearchWithVector -> IO Matches
forall {t :: * -> *}.
Foldable t =>
t Text -> Maybe Namespace -> IO Vectors
getIndexStats :: GetIndexStats -> IO IndexStats
upsertVectors :: UpsertVectors -> IO UpsertStats
listVectorIDs :: Maybe Text
-> Maybe Nat -> Maybe Text -> Maybe Namespace -> IO VectorIDs
searchWithVector :: SearchWithVector -> IO Matches
searchWithText :: Namespace -> SearchWithText -> IO Hits
listImports :: Maybe Nat -> Maybe Text -> IO Imports
describeImport :: Import -> IO ImportModel
fetchVectors :: forall {t :: * -> *}.
Foldable t =>
t Text -> Maybe Namespace -> IO Vectors
updateVector :: UpdateVector -> IO ()
deleteVectors :: DeleteVectors -> IO ()
upsertText :: Namespace -> Record -> IO ()
startImport :: StartImport -> IO Import
cancelImport :: Import -> IO ()
$sel:getIndexStats:DataMethods :: GetIndexStats -> IO IndexStats
$sel:upsertVectors:DataMethods :: UpsertVectors -> IO UpsertStats
$sel:upsertText:DataMethods :: Namespace -> Record -> IO ()
$sel:fetchVectors:DataMethods :: Vector Text -> Maybe Namespace -> IO Vectors
$sel:updateVector:DataMethods :: UpdateVector -> IO ()
$sel:deleteVectors:DataMethods :: DeleteVectors -> IO ()
$sel:listVectorIDs:DataMethods :: Maybe Text
-> Maybe Nat -> Maybe Text -> Maybe Namespace -> IO VectorIDs
$sel:searchWithVector:DataMethods :: SearchWithVector -> IO Matches
$sel:searchWithText:DataMethods :: Namespace -> SearchWithText -> IO Hits
$sel:startImport:DataMethods :: StartImport -> IO Import
$sel:listImports:DataMethods :: Maybe Nat -> Maybe Text -> IO Imports
$sel:describeImport:DataMethods :: Import -> IO ImportModel
$sel:cancelImport:DataMethods :: Import -> IO ()
..}
  where
    (      GetIndexStats -> IO IndexStats
getIndexStats
      :<|> (      (     UpsertVectors -> IO UpsertStats
upsertVectors
                  :<|>  [Text] -> Maybe Namespace -> IO Vectors
fetchVectors_
                  :<|>  UpdateVector -> IO NoContent
updateVector_
                  :<|>  DeleteVectors -> IO NoContent
deleteVectors_
                  :<|>  Maybe Text
-> Maybe Nat -> Maybe Text -> Maybe Namespace -> IO VectorIDs
listVectorIDs
                  )
            :<|>  Namespace -> Record -> IO NoContent
upsertText_
            )
      :<|>  (     SearchWithVector -> IO Matches
searchWithVector
            :<|>  Namespace -> SearchWithText -> IO Hits
searchWithText
            )
      :<|>  (     StartImport -> IO StartImportResponse
startImport_
            :<|>  Maybe Nat -> Maybe Text -> IO Imports
listImports
            :<|>  Import -> IO ImportModel
describeImport
            :<|>  Import -> IO NoContent
cancelImport_
            )
      ) = forall api (m :: * -> *) (n :: * -> *).
HasClient ClientM api =>
Proxy api -> (forall a. m a -> n a) -> Client m api -> Client n api
Client.hoistClient @DataAPI Proxy DataAPI
forall {k} (t :: k). Proxy t
Proxy (ClientEnv -> ClientM a -> IO a
forall a. ClientEnv -> ClientM a -> IO a
run ClientEnv
clientEnv) (forall api.
HasClient ClientM api =>
Proxy api -> Client ClientM api
Client.client @DataAPI Proxy DataAPI
forall {k} (t :: k). Proxy t
Proxy) Text
token Text
apiVersion

    fetchVectors :: t Text -> Maybe Namespace -> IO Vectors
fetchVectors t Text
a = [Text] -> Maybe Namespace -> IO Vectors
fetchVectors_ (t Text -> [Text]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t Text
a)
    updateVector :: UpdateVector -> IO ()
updateVector UpdateVector
a = IO NoContent -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (UpdateVector -> IO NoContent
updateVector_ UpdateVector
a)
    deleteVectors :: DeleteVectors -> IO ()
deleteVectors DeleteVectors
a = IO NoContent -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (DeleteVectors -> IO NoContent
deleteVectors_ DeleteVectors
a)
    upsertText :: Namespace -> Record -> IO ()
upsertText Namespace
a Record
b = IO NoContent -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Namespace -> Record -> IO NoContent
upsertText_ Namespace
a Record
b)
    startImport :: StartImport -> IO Import
startImport StartImport
a = do StartImportResponse{ Import
id :: Import
$sel:id:StartImportResponse :: StartImportResponse -> Import
id } <- StartImport -> IO StartImportResponse
startImport_ StartImport
a; Import -> IO Import
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Import
id
    cancelImport :: Import -> IO ()
cancelImport Import
a = IO NoContent -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Import -> IO NoContent
cancelImport_ Import
a)

run :: Client.ClientEnv -> Client.ClientM a -> IO a
run :: forall a. ClientEnv -> ClientM a -> IO a
run ClientEnv
clientEnv ClientM a
clientM = do
    Either ClientError a
result <- ClientM a -> ClientEnv -> IO (Either ClientError a)
forall a. ClientM a -> ClientEnv -> IO (Either ClientError a)
Client.runClientM ClientM a
clientM ClientEnv
clientEnv
    case Either ClientError a
result of
        Left ClientError
exception -> ClientError -> IO a
forall e a. Exception e => e -> IO a
Exception.throwIO ClientError
exception
        Right a
a -> a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

-- | Control plane methods
data ControlMethods = ControlMethods
    { ControlMethods -> IO (Vector IndexModel)
listIndexes :: IO (Vector IndexModel)
    , ControlMethods -> CreateIndex -> IO IndexModel
createIndex :: CreateIndex -> IO IndexModel
    , ControlMethods -> CreateIndexWithEmbedding -> IO IndexModel
createIndexWithEmbedding :: CreateIndexWithEmbedding -> IO IndexModel
    , ControlMethods -> Index -> IO IndexModel
describeIndex :: Index -> IO IndexModel
    , ControlMethods -> Index -> IO ()
deleteIndex :: Index -> IO ()
    , ControlMethods -> Index -> ConfigureIndex -> IO IndexModel
configureIndex :: Index -> ConfigureIndex -> IO IndexModel
    , ControlMethods -> IO (Vector CollectionModel)
listCollections :: IO (Vector CollectionModel)
    , ControlMethods -> CreateCollection -> IO CollectionModel
createCollection :: CreateCollection -> IO CollectionModel
    , ControlMethods -> Collection -> IO CollectionModel
describeCollection :: Collection -> IO CollectionModel
    , ControlMethods -> Collection -> IO ()
deleteCollection :: Collection -> IO ()
    , ControlMethods -> GenerateVectors -> IO Embeddings
generateVectors :: GenerateVectors -> IO Embeddings
    , ControlMethods -> RerankResults -> IO Documents
rerankResults :: RerankResults -> IO Documents
    }

-- | Data plane methods
data DataMethods = DataMethods
    { DataMethods -> GetIndexStats -> IO IndexStats
getIndexStats :: GetIndexStats -> IO IndexStats
    , DataMethods -> UpsertVectors -> IO UpsertStats
upsertVectors :: UpsertVectors -> IO UpsertStats
    , DataMethods -> Namespace -> Record -> IO ()
upsertText :: Namespace -> Record -> IO ()
    , DataMethods -> Vector Text -> Maybe Namespace -> IO Vectors
fetchVectors
        :: Vector Text
        -- ^ IDs
        -> Maybe Namespace
        -- ^ namespace
        -> IO Vectors
    , DataMethods -> UpdateVector -> IO ()
updateVector :: UpdateVector -> IO ()
    , DataMethods -> DeleteVectors -> IO ()
deleteVectors :: DeleteVectors -> IO ()
    , DataMethods
-> Maybe Text
-> Maybe Nat
-> Maybe Text
-> Maybe Namespace
-> IO VectorIDs
listVectorIDs
        :: Maybe Text
        -- ^ prefix
        -> Maybe Natural
        -- ^ limit
        -> Maybe Text
        -- ^ pagination token
        -> Maybe Namespace
        -- ^ namespace
        -> IO VectorIDs
    , DataMethods -> SearchWithVector -> IO Matches
searchWithVector :: SearchWithVector -> IO Matches
    , DataMethods -> Namespace -> SearchWithText -> IO Hits
searchWithText :: Namespace -> SearchWithText -> IO Hits
    , DataMethods -> StartImport -> IO Import
startImport :: StartImport -> IO Import
    , DataMethods -> Maybe Nat -> Maybe Text -> IO Imports
listImports
        :: Maybe Natural
        -- ^ limit
        -> Maybe Text
        -- ^ paginationToken
        -> IO Imports
    , DataMethods -> Import -> IO ImportModel
describeImport :: Import -> IO ImportModel
    , DataMethods -> Import -> IO ()
cancelImport :: Import -> IO ()
    }

-- | Index management
type ControlAPI =
        Header' [ Required, Strict ] "Api-Key" Text
    :>  Header' [ Required, Strict ] "X-Pinecone-API-Version" Text
    :>  (Indexes.ControlAPI :<|> Backups.API :<|> Embed.API :<|> Rerank.API)

-- | Index operations
type DataAPI =
        Header' [ Required, Strict ] "Api-Key" Text
    :>  Header' [ Required, Strict ] "X-Pinecone-API-Version" Text
    :>  (Indexes.DataAPI :<|> Vectors.API :<|> Search.API :<|> Imports.API)