main.go

Documentation: gitlab.hexacode.org/go-libs/chromem-go/wasm

     1  //go:build js
     2  
     3  package main
     4  
     5  import (
     6  	"context"
     7  	"errors"
     8  	"syscall/js"
     9  
    10  	"gitlab.hexacode.org/go-libs/chromem-go"
    11  )
    12  
    13  var c *chromem.Collection
    14  
    15  func main() {
    16  	js.Global().Set("initDB", js.FuncOf(initDB))
    17  	js.Global().Set("addDocument", js.FuncOf(addDocument))
    18  	js.Global().Set("query", js.FuncOf(query))
    19  
    20  	select {} // prevent main from exiting
    21  }
    22  
    23  // Exported function to initialize the database and collection.
    24  // Takes an OpenAI API key as argument.
    25  func initDB(this js.Value, args []js.Value) interface{} {
    26  	if len(args) != 1 {
    27  		return "expected 1 argument with the OpenAI API key"
    28  	}
    29  
    30  	openAIAPIKey := args[0].String()
    31  	embeddingFunc := chromem.NewEmbeddingFuncOpenAI(openAIAPIKey, chromem.EmbeddingModelOpenAI3Small)
    32  
    33  	db := chromem.NewDB()
    34  	var err error
    35  	c, err = db.CreateCollection("chromem", nil, embeddingFunc)
    36  	if err != nil {
    37  		return err.Error()
    38  	}
    39  
    40  	return nil
    41  }
    42  
    43  // Exported function to add documents to the collection.
    44  // Takes the document ID and content as arguments.
    45  func addDocument(this js.Value, args []js.Value) interface{} {
    46  	ctx := context.Background()
    47  
    48  	var id string
    49  	var content string
    50  	var err error
    51  	if len(args) != 2 {
    52  		err = errors.New("expected 2 arguments with the document ID and content")
    53  	} else {
    54  		id = args[0].String()
    55  		content = args[1].String()
    56  	}
    57  
    58  	handler := js.FuncOf(func(this js.Value, args []js.Value) interface{} {
    59  		resolve := args[0]
    60  		reject := args[1]
    61  		go func() {
    62  			if err != nil {
    63  				handleErr(err, reject)
    64  				return
    65  			}
    66  
    67  			err = c.AddDocument(ctx, chromem.Document{
    68  				ID:      id,
    69  				Content: content,
    70  			})
    71  			if err != nil {
    72  				handleErr(err, reject)
    73  				return
    74  			}
    75  			resolve.Invoke()
    76  		}()
    77  		return nil
    78  	})
    79  
    80  	promiseConstructor := js.Global().Get("Promise")
    81  	return promiseConstructor.New(handler)
    82  }
    83  
    84  // Exported function to query the collection
    85  // Takes the query string and the number of documents to return as argument.
    86  func query(this js.Value, args []js.Value) interface{} {
    87  	ctx := context.Background()
    88  
    89  	var q string
    90  	var err error
    91  	if len(args) != 1 {
    92  		err = errors.New("expected 1 argument with the query string")
    93  	} else {
    94  		q = args[0].String()
    95  	}
    96  
    97  	handler := js.FuncOf(func(this js.Value, args []js.Value) interface{} {
    98  		resolve := args[0]
    99  		reject := args[1]
   100  		go func() {
   101  			if err != nil {
   102  				handleErr(err, reject)
   103  				return
   104  			}
   105  
   106  			res, err := c.Query(ctx, q, 1, nil, nil)
   107  			if err != nil {
   108  				handleErr(err, reject)
   109  				return
   110  			}
   111  
   112  			// Convert response to JS values
   113  			// TODO: Return more than one result
   114  			o := js.Global().Get("Object").New()
   115  			o.Set("ID", res[0].ID)
   116  			o.Set("Similarity", res[0].Similarity)
   117  			o.Set("Content", res[0].Content)
   118  
   119  			resolve.Invoke(o)
   120  		}()
   121  		return nil
   122  	})
   123  
   124  	promiseConstructor := js.Global().Get("Promise")
   125  	return promiseConstructor.New(handler)
   126  }
   127  
   128  func handleErr(err error, reject js.Value) {
   129  	errorConstructor := js.Global().Get("Error")
   130  	errorObject := errorConstructor.New(err.Error())
   131  	reject.Invoke(errorObject)
   132  }
   133
View as plain text