diff --git a/README.md b/README.md index f6b159a..4fdb45d 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [![Go Reference](https://pkg.go.dev/badge/github.com/philippgille/chromem-go.svg)](https://pkg.go.dev/github.com/philippgille/chromem-go) -In-memory vector database for Go with Chroma-like interface. +In-memory vector database for Go with Chroma-like interface and zero third-party dependencies. It's not a library to connect to the Chroma database. It's an in-memory database on its own, meant to enable retrieval augmented generation (RAG) applications in Go *without having to run a separate database*. As such, the focus is not scale or performance, but simplicity. @@ -74,6 +74,7 @@ Initially, only a minimal subset of all of Chroma's interface is implemented or ## Features +- [X] Zero dependencies on third party libraries - Embedding creators: - [X] [OpenAI ada v2](https://platform.openai.com/docs/guides/embeddings/embedding-models) (default) - [X] Bring your own diff --git a/embedding.go b/embedding.go index 53619f5..bde1e84 100644 --- a/embedding.go +++ b/embedding.go @@ -1,14 +1,30 @@ package chromem import ( + "bytes" "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" "os" +) - "github.com/sashabaranov/go-openai" +const ( + baseURLOpenAI = "https://api.openai.com/v1" + embeddingModelOpenAI = "text-embedding-ada-002" ) +type openAIResponse struct { + Data []struct { + Embedding []float32 `json:"embedding"` + } `json:"data"` +} + // CreateEmbeddingsDefault returns a function that creates embeddings for a document using using // OpenAI`s ada v2 model via their API. +// The model supports a maximum document length of 8192 tokens. // The API key is read from the environment variable "OPENAI_API_KEY". func CreateEmbeddingsDefault() EmbeddingFunc { apiKey := os.Getenv("OPENAI_API_KEY") @@ -17,19 +33,65 @@ func CreateEmbeddingsDefault() EmbeddingFunc { // CreateEmbeddingsOpenAI returns a function that creates the embeddings for a document // using OpenAI`s ada v2 model via their API. +// The model supports a maximum document length of 8192 tokens. func CreateEmbeddingsOpenAI(apiKey string) EmbeddingFunc { - client := openai.NewClient(apiKey) + // We don't set a default timeout here, although it's usually a good idea. + // In our case though, the library user can set the timeout on the context, + // and it might have to be a long timeout, depending on the document size. + client := &http.Client{} + return func(ctx context.Context, document string) ([]float32, error) { - req := openai.EmbeddingRequest{ - Input: document, - Model: openai.AdaEmbeddingV2, + // Prepare the request body. + reqBody, err := json.Marshal(map[string]string{ + "input": document, + "model": embeddingModelOpenAI, + }) + if err != nil { + return nil, fmt.Errorf("couldn't marshal request body: %w", err) } - res, err := client.CreateEmbeddings(ctx, req) + // Create the request. Creating it with context is important for a timeout + // to be possible, because the client is configured without a timeout. + req, err := http.NewRequestWithContext(ctx, "POST", baseURLOpenAI+"/embeddings", bytes.NewBuffer(reqBody)) if err != nil { - return nil, err + return nil, fmt.Errorf("couldn't create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+apiKey) + + // Send the request. + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("couldn't send request: %w", err) + } + defer resp.Body.Close() + + // Check the response status. + if resp.StatusCode != http.StatusOK { + body, err := io.ReadAll(resp.Body) + if err != nil { + panic(err) + } + fmt.Println("========", string(body)) + return nil, errors.New("error response from the OpenAI API: " + resp.Status) + } + + // Read and decode the response body. + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("couldn't read response body: %w", err) + } + var embeddingResponse openAIResponse + err = json.Unmarshal(body, &embeddingResponse) + if err != nil { + return nil, fmt.Errorf("couldn't unmarshal response body: %w", err) + } + + // Check if the response contains embeddings. + if len(embeddingResponse.Data) == 0 || len(embeddingResponse.Data[0].Embedding) == 0 { + return nil, errors.New("no embeddings found in the response") } - return res.Data[0].Embedding, nil + return embeddingResponse.Data[0].Embedding, nil } } diff --git a/example/go.mod b/example/go.mod new file mode 100644 index 0000000..e9238e8 --- /dev/null +++ b/example/go.mod @@ -0,0 +1,10 @@ +module github.com/philippgille/chromem-go/example + +go 1.21 + +require ( + github.com/philippgille/chromem-go v0.0.0 + github.com/sashabaranov/go-openai v1.17.9 +) + +replace github.com/philippgille/chromem-go => ./.. diff --git a/example/go.sum b/example/go.sum new file mode 100644 index 0000000..45f7040 --- /dev/null +++ b/example/go.sum @@ -0,0 +1,2 @@ +github.com/sashabaranov/go-openai v1.17.9 h1:QEoBiGKWW68W79YIfXWEFZ7l5cEgZBV4/Ow3uy+5hNY= +github.com/sashabaranov/go-openai v1.17.9/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= diff --git a/go.mod b/go.mod index 60b1bc8..c135bd3 100644 --- a/go.mod +++ b/go.mod @@ -1,5 +1,3 @@ module github.com/philippgille/chromem-go go 1.21 - -require github.com/sashabaranov/go-openai v1.17.9 diff --git a/go.sum b/go.sum index 45f7040..e69de29 100644 --- a/go.sum +++ b/go.sum @@ -1,2 +0,0 @@ -github.com/sashabaranov/go-openai v1.17.9 h1:QEoBiGKWW68W79YIfXWEFZ7l5cEgZBV4/Ow3uy+5hNY= -github.com/sashabaranov/go-openai v1.17.9/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=