Add tests

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Author: mudler

rag/persistency_test.go

@@ -119,4 +119,40 @@ var _ = Describe("PersistentKB", func() {
 			Expect(count).To(Equal(0))
 		})
 	})
+
+	Describe("GetEntryContent", func() {
+		var kb *PersistentKB
+		var testFile string
+
+		BeforeEach(func() {
+			var err error
+			kb, err = NewPersistentCollectionKB(stateFile, assetDir, engine, 1000, openaiClient, "granite-embedding-107m-multilingual")
+			Expect(err).ToNot(HaveOccurred())
+
+			testFile = filepath.Join(tempDir, "getcontent.txt")
+			err = os.WriteFile(testFile, []byte("This is content for GetEntryContent test."), 0644)
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		It("should return entry not found for missing entry", func() {
+			_, err := kb.GetEntryContent("nonexistent.txt")
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("entry not found"))
+		})
+
+		It("should return chunks for stored entry", func() {
+			err := kb.Store(testFile, map[string]string{"type": "test"})
+			Expect(err).ToNot(HaveOccurred())
+
+			results, err := kb.GetEntryContent("getcontent.txt")
+			Expect(err).ToNot(HaveOccurred())
+			Expect(results).ToNot(BeEmpty())
+
+			var fullContent string
+			for _, r := range results {
+				fullContent += r.Content
+			}
+			Expect(fullContent).To(ContainSubstring("This is content for GetEntryContent test"))
+		})
+	})
 })

test/e2e/e2e_test.go

@@ -141,6 +141,26 @@ var _ = Describe("API", func() {
 
 		expectContent(TestCollection, "heist", "the Great Pigeon Heist", localRecall)
 	})
+
+	It("should get entry content", func() {
+		err := localRecall.CreateCollection(TestCollection)
+		Expect(err).ToNot(HaveOccurred())
+
+		fileName := tempContent(story1, localRecall)
+
+		chunks, err := localRecall.GetEntryContent(TestCollection, fileName)
+		if err != nil && err.Error() == "this collection backend does not support listing entry content" {
+			Skip("Backend (e.g. LocalAI) does not support GetEntryContent")
+		}
+		Expect(err).ToNot(HaveOccurred())
+		Expect(chunks).ToNot(BeEmpty())
+
+		var fullContent string
+		for _, c := range chunks {
+			fullContent += c.Content
+		}
+		Expect(fullContent).To(ContainSubstring("Bertie"))
+	})
 })
 
 func tempContent(content string, localRecall *client.Client) string {

test/e2e/persistency_test.go

@@ -113,6 +113,22 @@ var _ = Describe("Persistency", func() {
 			Expect(docs).To(HaveLen(1))
 			Expect(docs[0]).To(Equal("test.txt"))
 		})
+
+		It("should get entry content", func() {
+			metadata := map[string]string{"type": "test"}
+			err := kb.Store(testFile, metadata)
+			Expect(err).To(BeNil())
+
+			results, err := kb.GetEntryContent("test.txt")
+			Expect(err).To(BeNil())
+			Expect(results).ToNot(BeEmpty())
+
+			var fullContent string
+			for _, r := range results {
+				fullContent += r.Content
+			}
+			Expect(fullContent).To(ContainSubstring("This is a test document"))
+		})
 	})
 
 	Context("External Sources", func() {