feat: Added search functionality with TF-IDF scoring and inverted index data structure for efficient document retrieval.

pkg/search/index.go

@@ -0,0 +1,63 @@
+package search
+
+import (
+	"bufio"
+	"log"
+	"os"
+	"strings"
+)
+
+// needs to be saved in a file or database
+type InvertedIndex map[string]map[string]int
+type DocumentFrequency map[string]int
+
+var Index InvertedIndex
+var DocFreq DocumentFrequency
+var Files []string
+
+// BuildIndex reads files and builds an inverted index.
+func BuildIndex(files []string) (InvertedIndex, DocumentFrequency, error) {
+	index := make(InvertedIndex)
+	docFreq := make(DocumentFrequency)
+
+	for _, file := range files {
+		f, err := os.Open(file)
+		if err != nil {
+			return nil, nil, err
+		}
+		defer f.Close()
+
+		seenTerms := make(map[string]bool) // Track terms in this document
+		scanner := bufio.NewScanner(f)
+		scanner.Split(bufio.ScanWords)
+
+		for scanner.Scan() {
+			word := strings.ToLower(strings.Trim(scanner.Text(), ",.!?"))
+
+			if index[word] == nil {
+				index[word] = make(map[string]int)
+			}
+			index[word][file]++
+
+			if !seenTerms[word] {
+				docFreq[word]++
+				seenTerms[word] = true
+			}
+		}
+	}
+	return index, docFreq, nil
+}
+
+func TestIndex() {
+	// Index files
+	files := []string{"data/doc1.txt", "data/doc2.txt", "data/doc3.txt"}
+
+	index, docFreq, err := BuildIndex(files)
+	if err != nil {
+		log.Fatalf("Error building index: %v", err)
+	}
+
+	Files = files
+	Index = index
+	DocFreq = docFreq
+}

pkg/search/search.go

@@ -0,0 +1,87 @@
+package search
+
+import (
+	"math"
+	"sort"
+)
+
+// SearchResult stores the document and its relevance score.
+type SearchResult struct {
+	Document string
+	Score    float64
+}
+
+// Search processes different types of queries using TF-IDF scoring.
+func Search(terms []string, queryType string, index InvertedIndex, docFreq DocumentFrequency, numDocs int) []SearchResult {
+	scores := make(map[string]float64)
+
+	if queryType == "AND" {
+		// Ensure all terms appear in the document (AND logic)
+		for _, doc := range intersectDocs(terms, index) {
+			scores[doc] = scoreDoc(terms, doc, index, docFreq, numDocs)
+		}
+	} else if queryType == "OR" {
+		// Include any document that contains at least one of the terms (OR logic)
+		for _, term := range terms {
+			for doc := range index[term] {
+				scores[doc] += scoreDoc([]string{term}, doc, index, docFreq, numDocs)
+			}
+		}
+	} else {
+		// Simple query - score documents based on TF-IDF for any terms
+		for _, term := range terms {
+			for doc := range index[term] {
+				scores[doc] += scoreDoc([]string{term}, doc, index, docFreq, numDocs)
+			}
+		}
+	}
+
+	return rankResults(scores)
+}
+
+// Helper function to score a single document based on terms
+func scoreDoc(terms []string, doc string, index InvertedIndex, docFreq DocumentFrequency, numDocs int) float64 {
+	score := 0.0
+	for _, term := range terms {
+		tf := float64(index[term][doc])
+		idf := math.Log(float64(numDocs) / float64(docFreq[term]))
+		score += tf * idf
+		//fmt.Printf("Score: %f64 %f64 %f64\n", tf, idf, score)
+	}
+	return score
+}
+
+// Helper function to intersect documents for AND logic
+func intersectDocs(terms []string, index InvertedIndex) []string {
+	if len(terms) == 0 {
+		return nil
+	}
+	docs := make(map[string]bool)
+	for doc := range index[terms[0]] {
+		docs[doc] = true
+	}
+	for _, term := range terms[1:] {
+		for doc := range docs {
+			if _, exists := index[term][doc]; !exists {
+				delete(docs, doc)
+			}
+		}
+	}
+	result := []string{}
+	for doc := range docs {
+		result = append(result, doc)
+	}
+	return result
+}
+
+// rankResults sorts the documents by score
+func rankResults(scores map[string]float64) []SearchResult {
+	results := make([]SearchResult, 0, len(scores))
+	for doc, score := range scores {
+		results = append(results, SearchResult{Document: doc, Score: score})
+	}
+	sort.Slice(results, func(i, j int) bool {
+		return results[i].Score > results[j].Score
+	})
+	return results
+}