go-ollama: first release v0.1.0

generate_embeddings.go

@@ -0,0 +1,70 @@
+package ollama
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+)
+
+type GenerateEmbeddingsRequest struct {
+	Model      string                            `json:"model"`
+	Input      []string                          `json:"input"`
+	Truncate   *bool                             `json:"truncate,omitempty"`
+	Dimensions *int                              `json:"dimensions,omitempty"`
+	KeepAlive  string                            `json:"keep_alive,omitempty"`
+	Options    *GenerateEmbeddingsRequestOptions `json:"options,omitempty"`
+}
+
+type GenerateEmbeddingsRequestOptions struct {
+	Seed        *int     `json:"seed,omitempty"`
+	Temperature *float32 `json:"temperature,omitempty"`
+	TopK        *int     `json:"top_k,omitempty"`
+	TopP        *float32 `json:"top_p,omitempty"`
+	MinP        *float32 `json:"min_p,omitempty"`
+	Stop        []string `json:"stop,omitempty"`
+	NumCtx      *int     `json:"num_ctx,omitempty"`
+	NumPredict  *int     `json:"num_predict,omitempty"`
+}
+
+type GenerateEmbeddingsResponse struct {
+	Model           string      `json:"model"`
+	Embeddings      [][]float32 `json:"embeddings"`
+	TotalDuration   int         `json:"total_duration"`
+	LoadDuration    int         `json:"load_duration"`
+	PromptEvalCount int         `json:"prompt_eval_count"`
+}
+
+func (o Ollama) GenerateEmbeddings(reqBody GenerateEmbeddingsRequest) (GenerateEmbeddingsResponse, int, error) {
+	reqBodyBytes, err := json.Marshal(reqBody)
+	if err != nil {
+		return GenerateEmbeddingsResponse{}, -1, err
+	}
+
+	req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/embed", o.baseUrl), bytes.NewReader(reqBodyBytes))
+	if err != nil {
+		return GenerateEmbeddingsResponse{}, -1, err
+	}
+
+	for key, val := range o.customHeaders {
+		req.Header.Set(key, val)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return GenerateEmbeddingsResponse{}, -1, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		return GenerateEmbeddingsResponse{}, resp.StatusCode, errors.New("status code is not 200")
+	}
+
+	var respBody GenerateEmbeddingsResponse
+	if err := json.NewDecoder(resp.Body).Decode(&respBody); err != nil {
+		return GenerateEmbeddingsResponse{}, -1, err
+	}
+	return respBody, resp.StatusCode, nil
+}