diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,3 @@ +1.0.0: + +- Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2025, Mercury Technologies + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Gabriella Gonzalez nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,83 @@ +# `pinecone` + +This provides a binding to Pinecone's API using `servant` + +Example usage: + +```haskell +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DuplicateRecordFields #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE OverloadedStrings #-} + +import Pinecone +import Pinecone.Indexes +import Pinecone.Search +import Pinecone.Vectors + +import qualified Control.Exception as Exception +import qualified Data.Text as Text +import qualified System.Environment as Environment + +main :: IO () +main = do + controlEnv <- getClientEnv "https://api.pinecone.io" + + key <- Environment.getEnv "PINECONE_KEY" + + let token = Text.pack key + + let ControlMethods{..} = makeControlMethods controlEnv token + + let open = createIndexWithEmbedding _CreateIndexWithEmbedding + { name = "test" + , cloud = AWS + , region = "us-east-1" + , embed = EmbedRequest + { model = "llama-text-embed-v2" + , metric = Nothing + , read_parameters = Nothing + , write_parameters = Nothing + } + } + + let close IndexModel{ name } = deleteIndex name + + Exception.bracket open close \IndexModel{ name, host } -> do + let waitUntilIndexReady = do + IndexModel{ status } <- describeIndex name + + let Status{ ready } = status + + if ready + then return () + else waitUntilIndexReady + + waitUntilIndexReady + + dataEnv <- getClientEnv host + + let DataMethods{..} = makeDataMethods dataEnv token + + upsertText "test" _Record{ id = "hi", text = "Hello, world!" } + upsertText "test" _Record{ id = "bye", text = "Goodbye, world!" } + + -- Pinecone is eventually consistent, so we have to wait + let waitUntilVectorsReady = do + IndexStats{ totalVectorCount } <- getIndexStats _GetIndexStats + + if totalVectorCount == 2 + then return () + else waitUntilVectorsReady + + waitUntilVectorsReady + + Hits{ hits } <- searchWithText "test" SearchWithText + { query = _Query{ top_k = 1, input = Just "best greeting" } + , fields = Nothing + , rerank = Nothing + } + + print (fmap _id hits) -- ["hi"] +``` diff --git a/pinecone-example/Main.hs b/pinecone-example/Main.hs new file mode 100644 --- /dev/null +++ b/pinecone-example/Main.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DuplicateRecordFields #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE OverloadedStrings #-} + +import Pinecone +import Pinecone.Indexes +import Pinecone.Search +import Pinecone.Vectors + +import qualified Control.Exception as Exception +import qualified Data.Text as Text +import qualified System.Environment as Environment + +main :: IO () +main = do + controlEnv <- getClientEnv "https://api.pinecone.io" + + key <- Environment.getEnv "PINECONE_KEY" + + let token = Text.pack key + + let ControlMethods{..} = makeControlMethods controlEnv token + + let open = createIndexWithEmbedding _CreateIndexWithEmbedding + { name = "test" + , cloud = AWS + , region = "us-east-1" + , embed = EmbedRequest + { model = "llama-text-embed-v2" + , metric = Nothing + , read_parameters = Nothing + , write_parameters = Nothing + } + } + + let close IndexModel{ name } = deleteIndex name + + Exception.bracket open close \IndexModel{ name, host } -> do + let waitUntilIndexReady = do + IndexModel{ status } <- describeIndex name + + let Status{ ready } = status + + if ready + then return () + else waitUntilIndexReady + + waitUntilIndexReady + + dataEnv <- getClientEnv host + + let DataMethods{..} = makeDataMethods dataEnv token + + upsertText "test" _Record{ id = "hi", text = "Hello, world!" } + upsertText "test" _Record{ id = "bye", text = "Goodbye, world!" } + + -- Pinecone is eventually consistent, so we have to wait + let waitUntilVectorsReady = do + IndexStats{ totalVectorCount } <- getIndexStats _GetIndexStats + + if totalVectorCount == 2 + then return () + else waitUntilVectorsReady + + waitUntilVectorsReady + + Hits{ hits } <- searchWithText "test" SearchWithText + { query = _Query{ top_k = 1, input = Just "best greeting" } + , fields = Nothing + , rerank = Nothing + } + + print (fmap _id hits) -- ["hi"] diff --git a/pinecone.cabal b/pinecone.cabal new file mode 100644 --- /dev/null +++ b/pinecone.cabal @@ -0,0 +1,83 @@ +cabal-version: 2.4 +name: pinecone +version: 1.0.0 +synopsis: Servant bindings to Pinecone +description: This package provides comprehensive and type-safe bindings + to Pinecone, providing both a Servant interface and + non-Servant interface for convenience. + . + Read the @README@ below for a fully worked usage example. + . + Otherwise, browse the "Pinecone" module, which is the + intended package entrypoint. +license: BSD-3-Clause +license-file: LICENSE +author: Gabriella Gonzalez +maintainer: GenuineGabriella@gmail.com +copyright: 2025 Mercury Technologies +build-type: Simple +extra-doc-files: CHANGELOG.md README.md + +library + default-language: Haskell2010 + build-depends: base >= 4.15.0.0 && < 5 + , aeson + , containers + , http-api-data + , http-client + , http-client-tls + , scientific + , servant + , servant-client + , servant-client-core + , text + , time + , vector + exposed-modules: Pinecone + , Pinecone.Backups + , Pinecone.Embed + , Pinecone.Imports + , Pinecone.Indexes + , Pinecone.Metadata + , Pinecone.Pagination + , Pinecone.Rerank + , Pinecone.Search + , Pinecone.Vectors + other-modules: Pinecone.Prelude + default-extensions: BlockArguments + , DataKinds + , DeriveGeneric + , DeriveAnyClass + , DerivingStrategies + , DuplicateRecordFields + , GeneralizedNewtypeDeriving + , NamedFieldPuns + , OverloadedLists + , OverloadedStrings + , RecordWildCards + , TypeApplications + , TypeOperators + hs-source-dirs: src + ghc-options: -Wall -Wno-missing-fields + +test-suite tasty + default-language: Haskell2010 + type: exitcode-stdio-1.0 + hs-source-dirs: tasty + main-is: Main.hs + build-depends: base + , pinecone + , tasty + , tasty-hunit + , text + , vector + ghc-options: -Wall + +executable pinecone-example + default-language: Haskell2010 + hs-source-dirs: pinecone-example + main-is: Main.hs + build-depends: base + , pinecone + , text + ghc-options: -Wall diff --git a/src/Pinecone.hs b/src/Pinecone.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone.hs @@ -0,0 +1,240 @@ +-- | 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 baseUrlText) = do + baseUrl <- Client.parseBaseUrl (Text.unpack 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 + { baseUrlScheme = Https + , baseUrlPort = 443 + } + + let managerSettings = TLS.tlsManagerSettings + { HTTP.Client.managerResponseTimeout = + HTTP.Client.responseTimeoutNone + } + + manager <- TLS.newTlsManagerWith managerSettings + + pure (Client.mkClientEnv manager newBaseUrl) + +apiVersion :: Text +apiVersion = "2025-01" + +-- | Get a record of control API methods after providing an API token +makeControlMethods + :: ClientEnv + -- ^ + -> Text + -- ^ API token + -> ControlMethods +makeControlMethods clientEnv token = ControlMethods{..} + where + ( ( listIndexes_ + :<|> createIndex + :<|> createIndexWithEmbedding + :<|> describeIndex + :<|> deleteIndex_ + :<|> configureIndex + ) + :<|> ( listCollections_ + :<|> createCollection + :<|> describeCollection + :<|> deleteCollection_ + ) + :<|> generateVectors + :<|> rerankResults + ) = Client.hoistClient @ControlAPI Proxy (run clientEnv) (Client.client @ControlAPI Proxy) token apiVersion + + listIndexes = fmap indexes listIndexes_ + listCollections = fmap collections listCollections_ + deleteIndex a = void (deleteIndex_ a) + deleteCollection a = void (deleteCollection_ a) + +-- | Get a record of control API methods after providing an API token +makeDataMethods + :: ClientEnv + -- ^ + -> Text + -- ^ API token + -> DataMethods +makeDataMethods clientEnv token = DataMethods{..} + where + ( getIndexStats + :<|> ( ( upsertVectors + :<|> fetchVectors_ + :<|> updateVector_ + :<|> deleteVectors_ + :<|> listVectorIDs + ) + :<|> upsertText_ + ) + :<|> ( searchWithVector + :<|> searchWithText + ) + :<|> ( startImport_ + :<|> listImports + :<|> describeImport + :<|> cancelImport_ + ) + ) = Client.hoistClient @DataAPI Proxy (run clientEnv) (Client.client @DataAPI Proxy) token apiVersion + + fetchVectors a = fetchVectors_ (toList a) + updateVector a = void (updateVector_ a) + deleteVectors a = void (deleteVectors_ a) + upsertText a b = void (upsertText_ a b) + startImport a = do StartImportResponse{ id } <- startImport_ a; return id + cancelImport a = void (cancelImport_ a) + +run :: Client.ClientEnv -> Client.ClientM a -> IO a +run clientEnv clientM = do + result <- Client.runClientM clientM clientEnv + case result of + Left exception -> Exception.throwIO exception + Right a -> return a + +-- | Control plane methods +data ControlMethods = ControlMethods + { listIndexes :: IO (Vector IndexModel) + , createIndex :: CreateIndex -> IO IndexModel + , createIndexWithEmbedding :: CreateIndexWithEmbedding -> IO IndexModel + , describeIndex :: Index -> IO IndexModel + , deleteIndex :: Index -> IO () + , configureIndex :: Index -> ConfigureIndex -> IO IndexModel + , listCollections :: IO (Vector CollectionModel) + , createCollection :: CreateCollection -> IO CollectionModel + , describeCollection :: Collection -> IO CollectionModel + , deleteCollection :: Collection -> IO () + , generateVectors :: GenerateVectors -> IO Embeddings + , rerankResults :: RerankResults -> IO Documents + } + +-- | Data plane methods +data DataMethods = DataMethods + { getIndexStats :: GetIndexStats -> IO IndexStats + , upsertVectors :: UpsertVectors -> IO UpsertStats + , upsertText :: Namespace -> Record -> IO () + , fetchVectors + :: Vector Text + -- ^ IDs + -> Maybe Namespace + -- ^ namespace + -> IO Vectors + , updateVector :: UpdateVector -> IO () + , deleteVectors :: DeleteVectors -> IO () + , listVectorIDs + :: Maybe Text + -- ^ prefix + -> Maybe Natural + -- ^ limit + -> Maybe Text + -- ^ pagination token + -> Maybe Namespace + -- ^ namespace + -> IO VectorIDs + , searchWithVector :: SearchWithVector -> IO Matches + , searchWithText :: Namespace -> SearchWithText -> IO Hits + , startImport :: StartImport -> IO Import + , listImports + :: Maybe Natural + -- ^ limit + -> Maybe Text + -- ^ paginationToken + -> IO Imports + , describeImport :: Import -> IO ImportModel + , 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) diff --git a/src/Pinecone/Backups.hs b/src/Pinecone/Backups.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Backups.hs @@ -0,0 +1,78 @@ +-- | Backups +module Pinecone.Backups + ( -- * Main types + Collection(..) + , Collections(..) + , CreateCollection(..) + , _CreateCollection + , CollectionModel(..) + + -- * Other types + , Status(..) + + -- * Servant + , API + ) where + +import Pinecone.Indexes (Index) +import Pinecone.Prelude + +-- | The name of the collection +newtype Collection = Collection{ text :: Text } + deriving newtype (Eq, FromJSON, IsString, Show, ToHttpApiData, ToJSON) + +-- | The list of collections that exist in the project. +data Collections = Collections + { collections :: Vector CollectionModel + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The configuration needed to create a Pinecone collection. +data CreateCollection = CreateCollection + { name :: Text + , source :: Index + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `CreateCollection` +_CreateCollection :: CreateCollection +_CreateCollection = CreateCollection{ } + +-- | The CollectionModel describes the configuration and status of a Pinecone collection. +data CollectionModel = CollectionModel + { name :: Collection + , status :: Status + , environment :: Text + , size :: Maybe Natural + , dimension :: Maybe Natural + , vector_count :: Maybe Natural + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The status of the operation. +data Status + = Initializing + | Ready + | Terminating + deriving stock (Eq, Generic, Show) + +instance FromJSON Status where + parseJSON = genericParseJSON aesonOptions{ constructorTagModifier = id } + +instance ToJSON Status where + toJSON = genericToJSON aesonOptions{ constructorTagModifier = id } + +-- | Servant API +type API = + "collections" + :> ( Get '[JSON] Collections + + :<|> ReqBody '[JSON] CreateCollection + :> Post '[JSON] CollectionModel + + :<|> Capture "collection_name" Collection + :> Get '[JSON] CollectionModel + + :<|> Capture "collection_name" Collection + :> DeleteAccepted '[JSON] NoContent + ) diff --git a/src/Pinecone/Embed.hs b/src/Pinecone/Embed.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Embed.hs @@ -0,0 +1,117 @@ +-- | Embed +module Pinecone.Embed + ( -- * Main types + GenerateVectors(..) + , _GenerateVectors + , Embeddings(..) + + -- * Other types + , Input(..) + , VectorType(..) + , Embedding(..) + , Usage(..) + + -- * Servant + , API + ) where + +import Pinecone.Metadata (Scalar) +import Pinecone.Prelude + +-- | Generate embeddings for inputs +data GenerateVectors = GenerateVectors + { model :: Text + , inputs :: Vector Text + , parameters :: Map Text Scalar + } deriving stock (Eq, Generic, Show) + +instance FromJSON GenerateVectors where + parseJSON value = do + GenerateVectors_{..} <- parseJSON value + + return GenerateVectors{ inputs = fmap text inputs, ..} + +instance ToJSON GenerateVectors where + toJSON GenerateVectors{..} = + toJSON GenerateVectors_{ inputs = fmap Input inputs, ..} + +data GenerateVectors_ = GenerateVectors_ + { model :: Text + , inputs :: Vector Input + , parameters :: Map Text Scalar + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `GenerateVectors` +_GenerateVectors :: GenerateVectors +_GenerateVectors = GenerateVectors{ } + +-- | Embeddings generated for the input. +data Embeddings = Embeddings + { model :: Text + , vector_type :: VectorType + , data_ :: Vector Embedding + , usage :: Usage + } deriving stock (Eq, Generic, Show) + +instance FromJSON Embeddings where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON Embeddings where + toJSON = genericToJSON aesonOptions + +-- | An embedding generated for an input +data Embedding + = EmbeddingDense + { values :: Vector Double + } + | EmbeddingSparse + { sparse_values :: Vector Double + , sparse_indices :: Vector Natural + , sparse_tokens :: Vector Text + } + deriving stock (Eq, Generic, Show) + +embeddingOptions :: Options +embeddingOptions = aesonOptions + { sumEncoding = + TaggedObject{ tagFieldName = "vector_type", contentsFieldName = "" } + + , constructorTagModifier = stripPrefix "Embedding" + + , tagSingleConstructors = True + } + +instance FromJSON Embedding where + parseJSON = genericParseJSON embeddingOptions + +instance ToJSON Embedding where + toJSON = genericToJSON embeddingOptions + +-- | Input to generate embedding for +data Input = Input + { text :: Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The index vector type +data VectorType = Dense | Sparse + deriving stock (Eq, Generic, Show) + +instance FromJSON VectorType where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON VectorType where + toJSON = genericToJSON aesonOptions + +-- | Usage statistics for the model inference. +data Usage = Usage + { total_tokens :: Natural + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Servant API +type API = + "embed" + :> ReqBody '[JSON] GenerateVectors + :> Post '[JSON] Embeddings diff --git a/src/Pinecone/Imports.hs b/src/Pinecone/Imports.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Imports.hs @@ -0,0 +1,121 @@ +-- | Imports +module Pinecone.Imports + ( -- * Main types + Import(..) + , StartImport(..) + , _StartImport + , StartImportResponse(..) + , Imports(..) + , ImportModel(..) + + -- * Other types + , ErrorMode(..) + , OnError(..) + , Status(..) + + -- * Servant + , API + ) where + +import Pinecone.Pagination +import Pinecone.Prelude +import Prelude hiding (id) + +-- | Unique identifier for the import operation +newtype Import = Import{ text :: Text } + deriving newtype (Eq, FromJSON, IsString, Show, ToHttpApiData, ToJSON) + +-- | Request body for @\/bulk\/imports@ +data StartImport = StartImport + { uri :: Text + , integrationId :: Maybe Text + , errorMode :: Maybe ErrorMode + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `StartImport` +_StartImport :: StartImport +_StartImport = StartImport + { integrationId = Nothing + , errorMode = Nothing + } + +-- | Response body for @\/bulk\/imports@ +data StartImportResponse = StartImportResponse + { id :: Import + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | A list of import operations +data Imports = ListImport + { data_ :: Vector ImportModel + , pagination :: Maybe Pagination + } deriving stock (Eq, Generic, Show) + +instance FromJSON Imports where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON Imports where + toJSON = genericToJSON aesonOptions + +-- | The model for an import operation. +data ImportModel = ImportModel + { id :: Import + , uri :: Text + , status :: Status + , createdAt :: POSIXTime + , finishedAt :: POSIXTime + , percentComplete :: Double + , recordsImported :: Natural + , error :: Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Indicates how to respond to errors during the import process. +data ErrorMode = ErrorMode + { onError :: OnError + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Indicates how to respond to errors during the import process. +data OnError = Abort | Continue + deriving stock (Eq, Generic, Show) + +instance FromJSON OnError where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON OnError where + toJSON = genericToJSON aesonOptions + +-- | The status of the operation. +data Status + = Pending + | InProgress + | Failed + | Completed + | Cancelled + deriving stock (Eq, Generic, Show) + +instance FromJSON Status where + parseJSON = genericParseJSON aesonOptions{ constructorTagModifier = \x -> x } + +instance ToJSON Status where + toJSON = genericToJSON aesonOptions{ constructorTagModifier = \x -> x } + +-- | Servant API +type API = + "bulk" + :> "imports" + :> ( ReqBody '[JSON] StartImport + :> Post '[JSON] StartImportResponse + + :<|> QueryParam "limit" Natural + :> QueryParam "paginationToken" Text + :> Get '[JSON] Imports + + :<|> Capture "id" Import + :> Get '[JSON] ImportModel + + :<|> Capture "id" Import + :> Delete '[JSON] NoContent + ) diff --git a/src/Pinecone/Indexes.hs b/src/Pinecone/Indexes.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Indexes.hs @@ -0,0 +1,367 @@ +-- | Indexes +module Pinecone.Indexes + ( -- * Main types + Index(..) + , Host(..) + , IndexModels(..) + , IndexModel(..) + , CreateIndex(..) + , _CreateIndex + , CreateIndexWithEmbedding(..) + , _CreateIndexWithEmbedding + , ConfigureIndex(..) + , _ConfigureIndex + , GetIndexStats(..) + , _GetIndexStats + , IndexStats(..) + + -- * Other types + , Metric(..) + , Spec(..) + , Pod(..) + , PodType(..) + , Prefix(..) + , Suffix(..) + , MetadataConfig(..) + , Serverless(..) + , Cloud(..) + , Status(..) + , State(..) + , DeletionProtection(..) + , EmbedRequest(..) + , EmbedResponse(..) + , Contents(..) + + -- * API + , ControlAPI + , DataAPI + ) where + +import Pinecone.Embed (VectorType) +import Pinecone.Metadata (Filter) +import Pinecone.Prelude + +import qualified Data.Text as Text + +-- | The name of the index +newtype Index = Index{ text :: Text } + deriving newtype (Eq, FromJSON, IsString, Show, ToHttpApiData, ToJSON) + +-- | The host for the index +newtype Host = Host{ text :: Text } + deriving newtype (Eq, FromJSON, IsString, Show, ToHttpApiData, ToJSON) + +-- | The list of indexes that exist in the project +data IndexModels = IndexModels + { indexes :: Vector IndexModel + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The `IndexModel` describes the configuration and status of a Pinecone +-- index +data IndexModel = IndexModel + { name :: Index + , metric :: Metric + , host :: Host + , spec :: Spec + , status :: Status + , vector_type :: VectorType + , dimension :: Maybe Natural + , deletion_protection :: Maybe DeletionProtection + , tags :: Maybe (Map Text Text) + , embed :: Maybe EmbedResponse + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The desired configuration for the index +data CreateIndex = CreateIndex + { name :: Index + , spec :: Spec + , dimension :: Maybe Natural + , metric :: Maybe Metric + , deletion_protection :: Maybe DeletionProtection + , tags :: Maybe (Map Text Text) + , vector_type :: Maybe VectorType + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `CreateIndex` +_CreateIndex :: CreateIndex +_CreateIndex = CreateIndex + { dimension = Nothing + , metric = Nothing + , deletion_protection = Nothing + , tags = Nothing + , vector_type = Nothing + } + +-- | The desired configuration for the index and associated embedding model +data CreateIndexWithEmbedding = CreateIndexWithEmbedding + { name :: Index + , cloud :: Cloud + , region :: Text + , embed :: EmbedRequest + , deletion_protection :: Maybe DeletionProtection + , tags :: Maybe (Map Text Text) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `CreateIndexWithEmbedding` +_CreateIndexWithEmbedding :: CreateIndexWithEmbedding +_CreateIndexWithEmbedding = CreateIndexWithEmbedding + { deletion_protection = Nothing + , tags = Nothing + } + +-- | The desired pod size and replica configuration for the index +data ConfigureIndex = ConfigureIndex + { spec :: Maybe Spec + , deletion_protection :: Maybe DeletionProtection + , tags :: Maybe (Map Text Text) + , embed :: Maybe EmbedRequest + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `ConfigureIndex` +_ConfigureIndex :: ConfigureIndex +_ConfigureIndex = ConfigureIndex + { spec = Nothing + , deletion_protection = Nothing + , tags = Nothing + , embed = Nothing + } + +-- | Request body for @\/describe_index_stats@ +data GetIndexStats = GetIndexStats + { filter :: Maybe Filter + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `GetIndexStats` +_GetIndexStats :: GetIndexStats +_GetIndexStats = GetIndexStats + { filter = Nothing + } + +-- | Response body for @\/describe_index_stats@ +data IndexStats = IndexStats + { namespaces :: Map Text Contents + , dimension :: Natural + , indexFullness :: Double + , totalVectorCount :: Natural + , metric :: Metric + , vectorType :: VectorType + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The distance metric to be used for similarity search +data Metric = Cosine | Euclidean | DotProduct + deriving stock (Eq, Generic, Show) + +instance FromJSON Metric where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON Metric where + toJSON = genericToJSON aesonOptions + +-- | `Spec` object +data Spec = Spec + { pod :: Maybe Pod + , serverless :: Maybe Serverless + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Configuration needed to deploy a pod-based index. +data Pod = Pod + { environment :: Text + , pod_type :: PodType + , replicas :: Maybe Natural + , shards :: Maybe Natural + , pods :: Maybe Natural + , metadata_config :: Maybe MetadataConfig + , source_collection :: Maybe Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The type of pod to use +data PodType = PodType + { prefix :: Prefix + , suffix :: Suffix + } deriving stock (Eq, Generic, Show) + +instance ToJSON PodType where + toJSON PodType{..} = do + toJSON (Text.concat [ prefixText, ".", suffixText ]) + where + prefixText = case toJSON prefix of + String text -> text + _ -> "" + + suffixText = case toJSON suffix of + String text -> text + _ -> "" + +instance FromJSON PodType where + parseJSON value = do + text <- parseJSON value + + case Text.splitOn "." text of + [ prefixText, suffixText ] -> do + prefix <- parseJSON (String prefixText) + suffix <- parseJSON (String suffixText) + return PodType{..} + _ -> do + fail "Pod type must be of the form PREFIX.SUFFIX" + +-- | The first component of a pod type +data Prefix = S1 | P1 | P2 + deriving stock (Eq, Generic, Show) + +instance FromJSON Prefix where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON Prefix where + toJSON = genericToJSON aesonOptions + +-- | The second component of a pod type +data Suffix = X1 | X2 | X4 | X8 + deriving stock (Eq, Generic, Show) + +instance FromJSON Suffix where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON Suffix where + toJSON = genericToJSON aesonOptions + +-- | Configuration for the behavior of Pinecone's internal metadata index +data MetadataConfig = MetadataConfig + { indexed :: Maybe (Vector Text) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Configuration needed to deploy a serverless index +data Serverless = Serverless + { cloud :: Cloud + , region :: Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The public cloud where you would like your index hosted +data Cloud = GCP | AWS | Azure + deriving stock (Eq, Generic, Show) + +instance FromJSON Cloud where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON Cloud where + toJSON = genericToJSON aesonOptions + +-- | Index status +data Status = Status + { ready :: Bool + , state :: State + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Index state +data State + = Initializing + | InitializationFailed + | ScalingUp + | ScalingDown + | ScalingUpPodSize + | ScalingDownPodSize + | Terminating + | Ready + deriving stock (Eq, Generic, Show) + +instance FromJSON State where + parseJSON = genericParseJSON aesonOptions{ constructorTagModifier = id } + +instance ToJSON State where + toJSON = genericToJSON aesonOptions{ constructorTagModifier = id } + +-- | Whether deletion protection is enabled/disabled for the index. +data DeletionProtection = Disabled | Enabled + deriving stock (Eq, Generic, Show) + +instance FromJSON DeletionProtection where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON DeletionProtection where + toJSON = genericToJSON aesonOptions + +-- | Specify the integrated inference embedding configuration for the index +data EmbedRequest = EmbedRequest + { model :: Text + , metric :: Maybe Metric + , read_parameters :: Maybe (Map Text Value) + , write_parameters :: Maybe (Map Text Value) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON) + +data EmbedRequest_ = EmbedRequest_ + { model :: Text + , field_map :: Map Text Text + , metric :: Maybe Metric + , read_parameters :: Maybe (Map Text Value) + , write_parameters :: Maybe (Map Text Value) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +instance ToJSON EmbedRequest where + toJSON EmbedRequest{..} = toJSON EmbedRequest_{..} + where + field_map = [ ("text", "text") ] + +-- | The embedding model and document fields mapped to embedding inputs +data EmbedResponse = EmbedResponse + { model :: Text + , metric :: Maybe Metric + , dimension :: Maybe Natural + , vector_type :: Maybe VectorType + , field_map :: Maybe (Map Text Text) + , read_parameters :: Maybe (Map Text Value) + , write_parameters :: Maybe (Map Text Value) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | A summary of the contents of a namespace +data Contents = Contents + { vectorCount :: Natural + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Control API +type ControlAPI = + "indexes" + :> ( Get '[JSON] IndexModels + + :<|> ( ReqBody '[JSON] CreateIndex + :> Post '[JSON] IndexModel + ) + + :<|> ( "create-for-model" + :> ReqBody '[JSON] CreateIndexWithEmbedding + :> PostCreated '[JSON] IndexModel + ) + + :<|> ( Capture "index_name" Index + :> Get '[JSON] IndexModel + ) + + :<|> ( Capture "index_name" Index + :> DeleteAccepted '[JSON] NoContent + ) + + :<|> ( Capture "index_name" Index + :> ReqBody '[JSON] ConfigureIndex + :> Patch '[JSON] IndexModel + ) + ) + +-- | Data API +type DataAPI = + "describe_index_stats" + :> ReqBody '[JSON] GetIndexStats + :> Post '[JSON] IndexStats diff --git a/src/Pinecone/Metadata.hs b/src/Pinecone/Metadata.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Metadata.hs @@ -0,0 +1,140 @@ +-- | Pinecone's filtering query language +module Pinecone.Metadata + ( Scalar(..) + , Filter(..) + ) where + +import Pinecone.Prelude + +-- | A scalar value used for metadata filters +data Scalar + = ScalarNumber Scientific + | ScalarString Text + | ScalarBoolean Bool + deriving stock (Eq, Generic, Show) + +instance IsString Scalar where + fromString string = ScalarString (fromString string) + +instance Num Scalar where + fromInteger n = ScalarNumber (fromInteger n) + + ScalarNumber l + ScalarNumber r = ScalarNumber (l + r) + _ + ScalarNumber r = ScalarNumber r + l + _ = l + + ScalarNumber l * ScalarNumber r = ScalarNumber (l * r) + _ * ScalarNumber r = ScalarNumber r + l * _ = l + + abs (ScalarNumber n) = ScalarNumber (abs n) + abs n = n + + signum (ScalarNumber n) = ScalarNumber (signum n) + signum n = n + + negate (ScalarNumber n) = ScalarNumber (negate n) + negate n = n + +instance ToJSON Scalar where + toJSON (ScalarNumber scientific) = Number scientific + toJSON (ScalarString text) = String text + toJSON (ScalarBoolean bool) = Bool bool + +instance FromJSON Scalar where + parseJSON (Number scientific) = pure (ScalarNumber scientific) + parseJSON (String text) = pure (ScalarString text) + parseJSON (Bool bool) = pure (ScalarBoolean bool) + parseJSON object = typeMismatch "Number/String/Boolean" object + +-- | Metadata query language +data Filter + = Equal Text Scalar + | NotEqual Text Scalar + | GreaterThan Text Scientific + | GreaterThanOrEqual Text Scientific + | LessThan Text Scientific + | LessThanOrEqual Text Scientific + | In Text Scalar + | NotIn Text Scalar + | Exists Text Bool + | And (Vector Filter) + | Or (Vector Filter) + deriving stock (Eq, Generic, Show) + +instance ToJSON Filter where + toJSON (Equal field literal) = + toJSON @(Map Text (Map Text Scalar)) [ (field, [ ("$eq", literal) ]) ] + toJSON (NotEqual field literal) = + toJSON @(Map Text (Map Text Scalar)) [ (field, [ ("$ne", literal) ]) ] + toJSON (GreaterThan field literal) = + toJSON @(Map Text (Map Text Scientific)) [ (field, [ ("$gt", literal) ]) ] + toJSON (GreaterThanOrEqual field literal) = + toJSON @(Map Text (Map Text Scientific)) [ (field, [ ("$gte", literal) ]) ] + toJSON (LessThan field literal) = + toJSON @(Map Text (Map Text Scientific)) [ (field, [ ("$lt", literal) ]) ] + toJSON (LessThanOrEqual field literal) = + toJSON @(Map Text (Map Text Scientific)) [ (field, [ ("$lte", literal) ]) ] + toJSON (In field literal) = + toJSON @(Map Text (Map Text Scalar)) [ (field, [ ("$in", literal) ]) ] + toJSON (NotIn field literal) = + toJSON @(Map Text (Map Text Scalar)) [ (field, [ ("$nin", literal) ]) ] + toJSON (Exists field literal) = + toJSON @(Map Text (Map Text Bool)) [ (field, [ ("$exists", literal) ]) ] + toJSON (And clauses) = + toJSON @(Map Text (Vector Filter)) [ ("$and", clauses) ] + toJSON (Or clauses) = + toJSON @(Map Text (Vector Filter)) [ ("$or", clauses) ] + +instance FromJSON Filter where + parseJSON v = + parseEqual v + <|> parseNotEqual v + <|> parseGreaterThan v + <|> parseGreaterThanOrEqual v + <|> parseLessThan v + <|> parseLessThanOrEqual v + <|> parseIn v + <|> parseNotIn v + <|> parseAnd v + <|> parseOr v + where + parseEqual value = do + [ (field, [ ("$eq", literal) ]) ] <- parseJSON @(Map Text (Map Text Scalar)) value + return (Equal field literal) + + parseNotEqual value = do + [ (field, [ ("$ne", literal) ]) ] <- parseJSON @(Map Text (Map Text Scalar)) value + return (NotEqual field literal) + + parseGreaterThan value = do + [ (field, [ ("$gt", literal) ]) ] <- parseJSON @(Map Text (Map Text Scientific)) value + return (GreaterThan field literal) + + parseGreaterThanOrEqual value = do + [ (field, [ ("$gte", literal) ]) ] <- parseJSON @(Map Text (Map Text Scientific)) value + return (GreaterThanOrEqual field literal) + + parseLessThan value = do + [ (field, [ ("$lt", literal) ]) ] <- parseJSON @(Map Text (Map Text Scientific)) value + return (LessThan field literal) + + parseLessThanOrEqual value = do + [ (field, [ ("$lte", literal) ]) ] <- parseJSON @(Map Text (Map Text Scientific)) value + return (LessThanOrEqual field literal) + + parseIn value = do + [ (field, [ ("$in", literal) ]) ] <- parseJSON @(Map Text (Map Text Scalar)) value + return (In field literal) + + parseNotIn value = do + [ (field, [ ("$nin", literal) ]) ] <- parseJSON @(Map Text (Map Text Scalar)) value + return (NotIn field literal) + + parseAnd value = do + [ ("$and", clauses) ] <- parseJSON @(Map Text (Vector Filter)) value + return (And clauses) + + parseOr value = do + [ ("$or", clauses) ] <- parseJSON @(Map Text (Vector Filter)) value + return (Or clauses) diff --git a/src/Pinecone/Pagination.hs b/src/Pinecone/Pagination.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Pagination.hs @@ -0,0 +1,13 @@ +-- | Pagination +module Pinecone.Pagination + ( -- * Types + Pagination(..) + ) where + +import Pinecone.Prelude + +-- | Pagination +data Pagination = Pagination + { next :: Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) diff --git a/src/Pinecone/Prelude.hs b/src/Pinecone/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Prelude.hs @@ -0,0 +1,90 @@ +module Pinecone.Prelude + ( -- * JSON + aesonOptions + , labelModifier + , stripPrefix + + -- * Re-exports + , module Control.Applicative + , module Data.Aeson + , module Data.Aeson.Types + , module Data.Map + , module Data.Scientific + , module Data.String + , module Data.Text + , module Data.Time.Clock.POSIX + , module Data.Vector + , module GHC.Generics + , module Numeric.Natural + , module Servant.API + , module Web.HttpApiData + ) where + +import Control.Applicative (Alternative(..)) +import Data.Aeson.Types (typeMismatch) +import Data.Map (Map) +import Data.Scientific (Scientific) +import Data.String (IsString(..)) +import Data.Text (Text) +import Data.Time.Clock.POSIX (POSIXTime) +import Data.Vector (Vector) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) +import Web.HttpApiData (ToHttpApiData(..)) + +import Data.Aeson + ( genericToJSON + , genericParseJSON + , FromJSON(..) + , Object + , Options(..) + , SumEncoding(..) + , ToJSON(..) + , Value(..) + ) +import Servant.API + ( Capture + , Delete + , DeleteAccepted + , Get + , Header' + , JSON + , NoContent + , Patch + , Post + , PostCreated + , QueryParam + , QueryParam' + , QueryParams + , ReqBody + , Required + , Strict + , (:<|>)(..) + , (:>) + ) + +import qualified Data.Aeson as Aeson +import qualified Data.List as List +import qualified Data.Char as Char + +dropTrailingUnderscore :: String -> String +dropTrailingUnderscore "_" = "" +dropTrailingUnderscore "" = "" +dropTrailingUnderscore (c : cs) = c : dropTrailingUnderscore cs + +labelModifier :: String -> String +labelModifier = map Char.toLower . dropTrailingUnderscore + +stripPrefix :: String -> String -> String +stripPrefix prefix string = labelModifier suffix + where + suffix = case List.stripPrefix prefix string of + Nothing -> string + Just x -> x + +aesonOptions :: Options +aesonOptions = Aeson.defaultOptions + { fieldLabelModifier = labelModifier + , constructorTagModifier = labelModifier + , omitNothingFields = True + } diff --git a/src/Pinecone/Rerank.hs b/src/Pinecone/Rerank.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Rerank.hs @@ -0,0 +1,82 @@ +-- | Rerank +module Pinecone.Rerank + ( -- * Main types + RerankResults(..) + , _RerankResults + , Documents(..) + + -- * Other types + , Document(..) + , Usage(..) + + -- * Servant + , API + ) where + +import Pinecone.Metadata (Scalar) +import Pinecone.Prelude +import Pinecone.Vectors (Record) + +-- | Rerank documents for the given query +data RerankResults = RerankResults + { model :: Text + , query :: Text + , documents :: Vector Record + , top_n :: Maybe Natural + , return_documents :: Bool + , parameters :: Map Text Scalar + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON) + +data RerankResults_ = RerankResults_ + { model :: Text + , query :: Text + , documents :: Vector Record + , top_n :: Maybe Natural + , return_documents :: Bool + , rank_fields :: Vector Text + , parameters :: Map Text Scalar + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +instance ToJSON RerankResults where + toJSON RerankResults{..} = toJSON RerankResults_{..} + where + rank_fields = [ "text" ] + +-- | Default `RerankResults` +_RerankResults :: RerankResults +_RerankResults = RerankResults + { top_n = Nothing + } + +-- | The result of a reranking request. +data Documents = Documents + { model :: Text + , data_ :: Vector Document + , usage :: Usage + } deriving stock (Eq, Generic, Show) + +instance FromJSON Documents where + parseJSON = genericParseJSON aesonOptions + +instance ToJSON Documents where + toJSON = genericToJSON aesonOptions + +-- | A ranked document with a relevance score and an index position. +data Document = Document + { index :: Natural + , score :: Double + , document :: Maybe Record + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Usage statistics for the model inference. +data Usage = Usage + { rerank_units :: Natural + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Servant API +type API = + "rerank" :> ReqBody '[JSON] RerankResults :> Post '[JSON] Documents diff --git a/src/Pinecone/Search.hs b/src/Pinecone/Search.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Search.hs @@ -0,0 +1,207 @@ +-- | Search +module Pinecone.Search + ( -- * Main types + SearchWithVector(..) + , _SearchWithVector + , Matches(..) + , SearchWithText(..) + , _SearchWithText + , Hits(..) + + -- * Other types + , Match(..) + , Query(..) + , _Query + , VectorQuery(..) + , Rerank(..) + , Hit(..) + , Usage(..) + + -- * Servant + , API + ) where + +import Prelude hiding (filter) +import Pinecone.Embed (Input(..)) +import Pinecone.Metadata (Filter, Scalar) +import Pinecone.Prelude +import Pinecone.Vectors (Namespace, SparseValues) + +-- | Request body for @\/query@ +data SearchWithVector = SearchWithVector + { topK :: Natural + , namespace :: Maybe Namespace + , filter :: Maybe Filter + , includeValues :: Maybe Bool + , includeMetadata :: Maybe Bool + , vector :: Maybe (Vector Double) + , sparseVector :: Maybe SparseValues + , id :: Maybe Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `SearchWithVector` +_SearchWithVector :: SearchWithVector +_SearchWithVector = SearchWithVector + { namespace = Nothing + , filter = Nothing + , includeValues = Nothing + , includeMetadata = Nothing + , vector = Nothing + , sparseVector = Nothing + , id = Nothing + } + +-- | Response body for @\/query@ +data Matches = Matches + { matches :: Vector Match + , namespace :: Namespace + , usage :: Usage + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | A search request for records in a specific namespace. +data SearchWithText = SearchWithText + { query :: Query + , fields :: Maybe (Vector Text) + , rerank :: Maybe Rerank + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `SearchWithText` +_SearchWithText :: SearchWithText +_SearchWithText = SearchWithText + { fields = Nothing + , rerank = Nothing + } + +-- | A successful search namespace response. +data Hits = Hits + { usage :: Usage + , hits :: Vector Hit + } deriving stock (Eq, Generic, Show) + +instance FromJSON Hits where + parseJSON value = do + Hits_{..} <- parseJSON value + + let Result{..} = result + + return Hits{..} + +instance ToJSON Hits where + toJSON Hits{..} = toJSON Hits_{..} + where + result = Result{..} + +data Hits_ = Hits_ + { result :: Result + , usage :: Usage + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Match for a vector +data Match = Match + { id :: Text + , score :: Maybe Double + , values :: Maybe (Vector Double) + , sparseValues :: Maybe SparseValues + , metadata :: Maybe (Map Text Scalar) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | The query inputs to search with +data Query = Query + { top_k :: Natural + , filter :: Maybe Filter + , input :: Maybe Text + , vector :: Maybe VectorQuery + } deriving stock (Eq, Generic, Show) + +data Query_ = Query_ + { top_k :: Natural + , filter :: Maybe Filter + , inputs :: Maybe Input + , vector :: Maybe VectorQuery + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +instance FromJSON Query where + parseJSON value = do + Query_{..} <- parseJSON value + + let input = do + Input{..} <- inputs + return text + + return Query{..} + +instance ToJSON Query where + toJSON Query{..} = toJSON Query_{..} + where + inputs = do + text <- input + + return Input{..} + +-- | Default `Query` +_Query :: Query +_Query = Query + { filter = Nothing + , input = Nothing + , vector = Nothing + } + +-- | Vector query +data VectorQuery = VectorQuery + { values :: Maybe (Vector Double) + , sparse_values :: Maybe (Vector Double) + , sparse_indices :: Maybe (Vector Natural) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Parameters for reranking the initial search results +data Rerank = Rerank + { model :: Text + , rank_fields :: Vector Text + , top_n :: Maybe Natural + , parameters :: Maybe (Map Text Value) + , query :: Maybe Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Search result +data Result = Result + { hits :: Vector Hit + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Hit for the search document request +data Hit = Hit + { _id :: Text + , _score :: Double + , fields :: Value + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Usage +data Usage = Usage + { read_units :: Maybe Natural + , embed_total_tokens :: Maybe Natural + , rerank_units :: Maybe Natural + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Servant API +type API = + ( "query" + :> ReqBody '[JSON] SearchWithVector + :> Post '[JSON] Matches + ) + :<|> ( "records" + :> "namespaces" + :> Capture "namespace" Namespace + :> "search" + :> ReqBody '[JSON] SearchWithText + :> Post '[JSON] Hits + ) diff --git a/src/Pinecone/Vectors.hs b/src/Pinecone/Vectors.hs new file mode 100644 --- /dev/null +++ b/src/Pinecone/Vectors.hs @@ -0,0 +1,234 @@ +-- | Vectors +module Pinecone.Vectors + ( -- * Main types + Namespace(..) + , UpsertVectors(..) + , _UpsertVectors + , UpsertStats(..) + , Vectors(..) + , UpdateVector(..) + , _UpdateVector + , DeleteVectors(..) + , _DeleteVectors + , VectorIDs(..) + , Record(..) + , _Record + + -- * Other types + , VectorObject(..) + , SparseValues(..) + , Usage(..) + + -- * Servant + , API + ) where + +import Pinecone.Metadata (Filter, Scalar) +import Pinecone.Pagination (Pagination) +import Pinecone.Prelude +import Prelude hiding (id) + +import qualified Data.Map as Map + +-- | The namespace of a record +newtype Namespace = Namespace{ text :: Text } + deriving newtype (Eq, FromJSON, IsString, Show, ToHttpApiData, ToJSON) + +-- | Request body for @\/vectors\/upsert@ +data UpsertVectors = UpsertVectors + { vectors :: Vector VectorObject + , namespace :: Maybe Namespace + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `UpsertVectors` +_UpsertVectors :: UpsertVectors +_UpsertVectors = UpsertVectors + { namespace = Nothing + } + +-- | Response body for @\/vectors\/upsert@ +data UpsertStats = UpsertStats + { upsertedCount :: Natural + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Response body for @\/vectors\/fetch@ +data Vectors = Vectors + { vectors :: Map Text VectorObject + , namespace :: Namespace + , usage :: Maybe Usage + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Request body for @\/vectors\/update@ +data UpdateVector = UpdateVector + { id :: Text + , values :: Maybe (Vector Double) + , sparseValues :: Maybe SparseValues + , setMetadata :: Maybe (Map Text Scalar) + , namespace :: Maybe Namespace + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `UpdateVector` +_UpdateVector :: UpdateVector +_UpdateVector = UpdateVector + { values = Nothing + , sparseValues = Nothing + , setMetadata = Nothing + , namespace = Nothing + } + +-- | Request body for @\/vectors\/delete@ +data DeleteVectors = DeleteVectors + { ids :: Maybe (Vector Text) + , deleteAll :: Maybe Bool + , namespace :: Maybe Namespace + , filter :: Maybe Filter + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Default `DeleteVectors` +_DeleteVectors :: DeleteVectors +_DeleteVectors = DeleteVectors + { deleteAll = Nothing + , namespace = Nothing + , filter = Nothing + } + +-- | Response body for @\/vectors\/list@ +data VectorIDs = VectorIDs + { vectors :: Vector Text + , pagination :: Maybe Pagination + , namespace :: Namespace + , usage :: Usage + } deriving stock (Eq, Generic, Show) + +instance FromJSON VectorIDs where + parseJSON value = do + VectorIDs_{..} <- parseJSON value + + return VectorIDs{ vectors = do VectorID{ id } <- vectors; return id, ..} + +data VectorIDs_ = VectorIDs_ + { vectors :: Vector VectorID + , pagination :: Maybe Pagination + , namespace :: Namespace + , usage :: Usage + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Record to upsert +data Record = Record + { id :: Text + , text :: Text + , metadata :: Maybe (Map Text Scalar) + } deriving stock (Eq, Generic, Show) + +instance FromJSON Record where + parseJSON value = do + m <- parseJSON value + + text <- case Map.lookup "text" m of + Nothing -> do + fail "Missing text field" + Just v -> do + parseJSON v + + id <- case Map.lookup "_id" m of + Nothing -> do + fail "Missing id field" + Just v -> do + parseJSON v + + metadata <- fmap Just (traverse parseJSON (Map.delete "_id" (Map.delete "text" m))) + + return Record{..} + +instance ToJSON Record where + toJSON Record{..} = toJSON (Map.union reserved nonReserved) + where + reserved = + [ ("_id", toJSON id) + , ("text", toJSON text) + ] + + nonReserved = case metadata of + Nothing -> Map.empty + Just m -> Map.map toJSON m + +-- | Default `Record` +_Record :: Record +_Record = Record + { metadata = Nothing + } + +-- | A vector +data VectorObject = VectorObject + { id :: Text + , values :: Maybe (Vector Double) + , sparseValues :: Maybe SparseValues + , metadata :: Maybe (Map Text Scalar) + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Vector sparse data +data SparseValues = SparseValues + { indices :: Vector Natural + , values :: Vector Double + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Vector ID +data VectorID = VectorID + { id :: Text + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Usage +data Usage = Usage + { readUnits :: Natural + } deriving stock (Eq, Generic, Show) + deriving anyclass (FromJSON, ToJSON) + +-- | Servant API +type API = + ( "vectors" + :> ( ( "upsert" + :> ReqBody '[JSON] UpsertVectors + :> Post '[JSON] UpsertStats + ) + + :<|> ( "fetch" + :> QueryParams "ids" Text + :> QueryParam "namespace" Namespace + :> Get '[JSON] Vectors + ) + + :<|> ( "update" + :> ReqBody '[JSON] UpdateVector + :> Post '[JSON] NoContent + ) + + :<|> ( "delete" + :> ReqBody '[JSON] DeleteVectors + :> Post '[JSON] NoContent + ) + + :<|> ( "list" + :> QueryParam "prefix" Text + :> QueryParam "limit" Natural + :> QueryParam "paginationToken" Text + :> QueryParam "namespace" Namespace + :> Get '[JSON] VectorIDs + ) + ) + ) + :<|> ( "records" + :> "namespaces" + :> Capture "namespace" Namespace + :> "upsert" + :> ReqBody '[JSON] Record + :> PostCreated '[JSON] NoContent + ) diff --git a/tasty/Main.hs b/tasty/Main.hs new file mode 100644 --- /dev/null +++ b/tasty/Main.hs @@ -0,0 +1,223 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE DuplicateRecordFields #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE OverloadedLists #-} + +module Main where + +import Pinecone (DataMethods(..), ControlMethods(..)) +import Pinecone.Embed (GenerateVectors(..)) +import Pinecone.Rerank (Document(..), Documents(..), RerankResults(..)) +import Prelude hiding (id) + +import Pinecone.Indexes + ( Cloud(..) + , CreateIndexWithEmbedding(..) + , ConfigureIndex(..) + , EmbedRequest(..) + , GetIndexStats(..) + , IndexModel(..) + , IndexStats(..) + , Status(..) + ) +import Pinecone.Search + ( Hit(..) + , Hits(..) + , Matches(..) + , Query(..) + , SearchWithText(..) + , SearchWithVector(..) + ) +import Pinecone.Vectors + ( DeleteVectors(..) + , Record(..) + , UpdateVector(..) + , UpsertVectors(..) + , UpsertStats(..) + , VectorIDs(..) + , VectorObject(..) + , Vectors(..) + ) + +import qualified Control.Exception as Exception +import qualified Data.Text as Text +import qualified Data.Vector as Vector +import qualified Pinecone +import qualified System.Environment as Environment +import qualified Test.Tasty as Tasty +import qualified Test.Tasty.HUnit as HUnit + +main :: IO () +main = do + controlEnv <- Pinecone.getClientEnv "https://api.pinecone.io" + + key <- Environment.getEnv "PINECONE_KEY" + + let token = Text.pack key + + let ControlMethods{..} = Pinecone.makeControlMethods controlEnv token + + let namespace = "test" + + Tasty.defaultMain do + HUnit.testCase "Vectors" do + let open = do + createIndexWithEmbedding CreateIndexWithEmbedding + { name = "vectors-test" + , cloud = AWS + , region = "us-east-1" + , embed = EmbedRequest + { model = "llama-text-embed-v2" + , metric = Nothing + , read_parameters = Nothing + , write_parameters = Nothing + } + , deletion_protection = Nothing + , tags = Nothing + } + + let close IndexModel{ name } = deleteIndex name + + Exception.bracket open close \IndexModel{ name, host } -> do + _ <- generateVectors GenerateVectors + { model = "llama-text-embed-v2" + , inputs = [ "Hello, world!" ] + , parameters = [ ("input_type", "query") ] + } + + let hello = Record + { id = "hi" + , text = "Hello, world" + , metadata = Nothing + } + + let goodbye = Record + { id = "bye" + , text = "Goodbye, world" + , metadata = Nothing + } + + Documents{ data_ = [ Document{ document = Just Record{ id = "hi" } } ] } <- rerankResults RerankResults + { model = "bge-reranker-v2-m3" + , query = "best greeting" + , documents = [ hello, goodbye ] + , top_n = Just 1 + , return_documents = True + , parameters = [ ] + } + + dataEnv <- Pinecone.getClientEnv host + + let DataMethods{..} = Pinecone.makeDataMethods dataEnv token + + let waitUntilIndexReady = do + indexModel <- describeIndex name + + let IndexModel{ status } = indexModel + + let Status{ ready } = status + + if ready + then return indexModel + else waitUntilIndexReady + + indexModel <- waitUntilIndexReady + + indexes <- listIndexes + + case indexes of + [ indexModel₀ ] + | indexModel == indexModel₀ -> return () + _ -> HUnit.assertFailure "GET /indexes - wrong models" + + _ <- configureIndex name ConfigureIndex + { spec = Nothing + , deletion_protection = Nothing + , tags = Just [ ("foo", "bar") ] + , embed = Nothing + } + + IndexModel{ tags = Just [ ("foo", "bar") ] } <- describeIndex name + + UpsertStats{..} <- upsertVectors UpsertVectors + { vectors = + [ VectorObject + { id = "vector-0" + , values = Just (Vector.replicate 1024 0.1) + , sparseValues = Nothing + , metadata = Nothing + } + ] + , namespace = Just namespace + } + + HUnit.assertEqual "" 1 upsertedCount + + deleteVectors DeleteVectors + { ids = Just [ "vector-0" ] + , deleteAll = Nothing + , namespace = Just namespace + , filter = Nothing + } + + upsertText namespace hello + + upsertText namespace goodbye + + updateVector UpdateVector + { id = "hi" + , values = Nothing + , sparseValues = Nothing + , setMetadata = Just [ ("category", "greeting") ] + , namespace = Just namespace + } + + let waitUntilVectorsReady = do + IndexStats{..} <- getIndexStats GetIndexStats + { filter = Nothing + } + + if totalVectorCount == 2 + then return () + else waitUntilVectorsReady + + waitUntilVectorsReady + + Vectors{ vectors = [ ("bye", _), ("hi", _) ] } <- fetchVectors [ "hi", "bye" ] (Just namespace) + + VectorIDs{ vectors = [ "bye", "hi" ] } <- listVectorIDs Nothing Nothing Nothing (Just namespace) + + Hits{ hits = [ Hit{ _id = "hi" } ] } <- searchWithText namespace SearchWithText + { query = Query + { top_k = 1 + , filter = Nothing + , input = Just "Hi!" + , vector = Nothing + } + , fields = Nothing + , rerank = Nothing + } + + Matches{ matches } <- searchWithVector SearchWithVector + { topK = 1 + , namespace = Just namespace + , filter = Nothing + , includeValues = Nothing + , includeMetadata = Nothing + , vector = Just (Vector.replicate 1024 0.1) + , sparseVector = Nothing + , id = Nothing + } + + HUnit.assertEqual "" 1 (Vector.length matches) + + deleteVectors DeleteVectors + { ids = Nothing + , deleteAll = Just True + , namespace = Just namespace + , filter = Nothing + } + + return ()