From fa193172e054789e925a7f457867335a79693079 Mon Sep 17 00:00:00 2001
From: hupe1980 <frankhuebner1980@gmail.com>
Date: Sat, 9 Sep 2023 00:09:56 +0200
Subject: [PATCH] Add pinecone vectorstore example

---
 examples/pinecone_vectorstore/main.go | 84 +++++++++++++++++++++++++++
 integration/pinecone/rest.go          |  9 +++
 integration/pinecone/types.go         |  5 ++
 3 files changed, 98 insertions(+)
 create mode 100644 examples/pinecone_vectorstore/main.go

diff --git a/examples/pinecone_vectorstore/main.go b/examples/pinecone_vectorstore/main.go
new file mode 100644
index 0000000..ae9c52b
--- /dev/null
+++ b/examples/pinecone_vectorstore/main.go
@@ -0,0 +1,84 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/hupe1980/golc/embedding"
+	"github.com/hupe1980/golc/integration/pinecone"
+	"github.com/hupe1980/golc/schema"
+	"github.com/hupe1980/golc/vectorstore"
+)
+
+func main() {
+	openai, err := embedding.NewOpenAI(os.Getenv("OPENAI_API_KEY"))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	client, err := pinecone.NewRestClient(os.Getenv("PINECONE_API_KEY"), pinecone.Endpoint{
+		IndexName:   "golc",
+		ProjectName: os.Getenv("PINECONE_PROJECT"),
+		Environment: "us-west1-gcp-free",
+	})
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	vs, err := vectorstore.NewPinecone(client, openai, "textKey", func(po *vectorstore.PineconeOptions) {
+		po.TopK = 1
+	})
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	if err = vs.AddDocuments(context.Background(), []schema.Document{
+		{
+			PageContent: "Pizza is an Italian dish consisting of a flat, round base of dough topped with various ingredients, including tomato sauce, cheese, and various toppings.",
+			Metadata: map[string]any{
+				"cousine": "Italian",
+			},
+		},
+		{
+			PageContent: "Sushi is a Japanese dish consisting of vinegared rice combined with various ingredients, such as seafood, vegetables, and occasionally tropical fruits.",
+			Metadata: map[string]any{
+				"cousine": "Japanese",
+			},
+		},
+		{
+			PageContent: "A burger is a sandwich consisting of a cooked ground meat patty, usually beef, placed in a sliced bread roll or bun.",
+			Metadata: map[string]any{
+				"cousine": "American",
+			},
+		},
+		{
+			PageContent: "Pad Thai is a popular Thai stir-fried noodle dish made with rice noodles, eggs, tofu or shrimp, bean sprouts, peanuts, and lime.",
+			Metadata: map[string]any{
+				"cousine": "Thai",
+			},
+		},
+		{
+			PageContent: "Sashimi is a Japanese delicacy consisting of thinly sliced raw seafood, served with soy sauce and wasabi.",
+			Metadata: map[string]any{
+				"cousine": "Japanese",
+			},
+		},
+		{
+			PageContent: "Lasagna is an Italian pasta dish made with layers of pasta sheets, meat sauce, cheese, and bechamel sauce.",
+			Metadata: map[string]any{
+				"cousine": "Italian",
+			},
+		},
+	}); err != nil {
+		log.Fatal(err)
+	}
+
+	docs, err := vs.SimilaritySearch(context.Background(), "Cheeseburger")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Println(docs)
+}
diff --git a/integration/pinecone/rest.go b/integration/pinecone/rest.go
index c64d557..cdf8021 100644
--- a/integration/pinecone/rest.go
+++ b/integration/pinecone/rest.go
@@ -40,6 +40,15 @@ func (p *RestClient) Upsert(ctx context.Context, req *UpsertRequest) (*UpsertRes
 		return nil, err
 	}
 
+	if res.StatusCode != 200 {
+		errorResponse := ErrorResponse{}
+		if err := json.Unmarshal(body, &errorResponse); err != nil {
+			return nil, err
+		}
+
+		return nil, fmt.Errorf("pincoe error: %s", errorResponse.Message)
+	}
+
 	upsertResponse := UpsertResponse{}
 	if err := json.Unmarshal(body, &upsertResponse); err != nil {
 		return nil, err
diff --git a/integration/pinecone/types.go b/integration/pinecone/types.go
index 9c7984f..9e2366a 100644
--- a/integration/pinecone/types.go
+++ b/integration/pinecone/types.go
@@ -55,3 +55,8 @@ type QueryResponse struct {
 	Matches   []*Match `json:"matches"`
 	Namespace string   `json:"namespace"`
 }
+
+type ErrorResponse struct {
+	Code    int    `json:"code"`
+	Message string `json:"message"`
+}