close pg.store connection

lib/embeddings/load.go

@@ -11,6 +11,7 @@ import (
 	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/textsplitter"
 	"github.com/tmc/langchaingo/vectorstores"
+	"github.com/tmc/langchaingo/vectorstores/pgvector"
 )
 
 
@@ -19,26 +20,23 @@ import (
 func LoadDocsToStore(docs []schema.Document, store vectorstores.VectorStore)  {
 	fmt.Println("loading data from")
 
-	/*
-	store, err := GetVectorStore()
-
-	if err != nil {
-		log.Panic(err)
-	}
-	*/
-	//docs := getSampleDocs()
-
 	fmt.Println("no. of documents to be loaded", len(docs))
 
 	_, err := store.AddDocuments(context.Background(), docs)
-
 	if err != nil {
 		log.Panic(err)
 	}
-
 	fmt.Println("data successfully loaded into vector store")
-
 	log.Println(err)
+
+	defer func() {
+		var pgvStore pgvector.Store
+		pgvStore, ok := store.(pgvector.Store)
+		if !ok {
+		  log.Fatalf("store does not implement pgvector.Store")
+		}
+		pgvStore.Close()
+	  }()
 }
 
 /*

lib/embeddings/query.go

@@ -3,11 +3,13 @@ package embeddings
 import (
 	"context"
 	"fmt"
+	"log"
 
 	"github.com/tmc/langchaingo/chains"
 	"github.com/tmc/langchaingo/llms/openai"
 	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/vectorstores"
+	"github.com/tmc/langchaingo/vectorstores/pgvector"
 )
 
 
@@ -43,6 +45,18 @@ func Rag(ai_url string,api_token string,question string, numOfResults int,store
 	}
 
 	fmt.Println("====final answer====\n", result)
+
+
+	defer func() {
+		var pgvStore pgvector.Store
+		pgvStore, ok := store.(pgvector.Store)
+		if !ok {
+		  log.Fatalf("store does not implement pgvector.Store")
+		}
+		pgvStore.Close()
+	  }()
+
+
 	return result,nil
 }
 
@@ -60,6 +74,15 @@ func SemanticSearch(searchQuery string, maxResults int, store vectorstores.Vecto
 		fmt.Println("============================")
 
 	}
+
+	defer func() {
+		var pgvStore pgvector.Store
+		pgvStore, ok := store.(pgvector.Store)
+		if !ok {
+		  log.Fatalf("store does not implement pgvector.Store")
+		}
+		pgvStore.Close()
+	  }()
 	return searchResults,nil
 }