mirror of
https://github.com/likelovewant/ollama-for-amd.git
synced 2025-12-21 22:33:56 +00:00
embeddings: added embedding command for cl (#12795)
Co-authored-by: A-Akhil <akhilrahul70@gmail.com> This PR introduces a new ollama embed command that allows users to generate embeddings directly from the command line. Added ollama embed MODEL [TEXT...] command for generating text embeddings Supports both direct text arguments and stdin piping for scripted workflows Outputs embeddings as JSON arrays (one per line)
This commit is contained in:
nicole pardal
324
cmd/cmd_test.go
324
cmd/cmd_test.go
@@ -355,6 +355,330 @@ func TestDeleteHandler(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunEmbeddingModel(t *testing.T) {
|
||||
reqCh := make(chan api.EmbedRequest, 1)
|
||||
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/api/show" && r.Method == http.MethodPost {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(api.ShowResponse{
|
||||
Capabilities: []model.Capability{model.CapabilityEmbedding},
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/api/embed" && r.Method == http.MethodPost {
|
||||
var req api.EmbedRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
reqCh <- req
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(api.EmbedResponse{
|
||||
Model: "test-embedding-model",
|
||||
Embeddings: [][]float32{{0.1, 0.2, 0.3}},
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
|
||||
t.Setenv("OLLAMA_HOST", mockServer.URL)
|
||||
t.Cleanup(mockServer.Close)
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(t.Context())
|
||||
cmd.Flags().String("keepalive", "", "")
|
||||
cmd.Flags().Bool("truncate", false, "")
|
||||
cmd.Flags().Int("dimensions", 0, "")
|
||||
cmd.Flags().Bool("verbose", false, "")
|
||||
cmd.Flags().Bool("insecure", false, "")
|
||||
cmd.Flags().Bool("nowordwrap", false, "")
|
||||
cmd.Flags().String("format", "", "")
|
||||
cmd.Flags().String("think", "", "")
|
||||
cmd.Flags().Bool("hidethinking", false, "")
|
||||
|
||||
oldStdout := os.Stdout
|
||||
r, w, _ := os.Pipe()
|
||||
os.Stdout = w
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- RunHandler(cmd, []string{"test-embedding-model", "hello", "world"})
|
||||
}()
|
||||
|
||||
err := <-errCh
|
||||
w.Close()
|
||||
os.Stdout = oldStdout
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("RunHandler returned error: %v", err)
|
||||
}
|
||||
|
||||
var out bytes.Buffer
|
||||
io.Copy(&out, r)
|
||||
|
||||
select {
|
||||
case req := <-reqCh:
|
||||
inputText, _ := req.Input.(string)
|
||||
if diff := cmp.Diff("hello world", inputText); diff != "" {
|
||||
t.Errorf("unexpected input (-want +got):\n%s", diff)
|
||||
}
|
||||
if req.Truncate != nil {
|
||||
t.Errorf("expected truncate to be nil, got %v", *req.Truncate)
|
||||
}
|
||||
if req.KeepAlive != nil {
|
||||
t.Errorf("expected keepalive to be nil, got %v", req.KeepAlive)
|
||||
}
|
||||
if req.Dimensions != 0 {
|
||||
t.Errorf("expected dimensions to be 0, got %d", req.Dimensions)
|
||||
}
|
||||
default:
|
||||
t.Fatal("server did not receive embed request")
|
||||
}
|
||||
|
||||
expectOutput := "[0.1,0.2,0.3]\n"
|
||||
if diff := cmp.Diff(expectOutput, out.String()); diff != "" {
|
||||
t.Errorf("unexpected output (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunEmbeddingModelWithFlags(t *testing.T) {
|
||||
reqCh := make(chan api.EmbedRequest, 1)
|
||||
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/api/show" && r.Method == http.MethodPost {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(api.ShowResponse{
|
||||
Capabilities: []model.Capability{model.CapabilityEmbedding},
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/api/embed" && r.Method == http.MethodPost {
|
||||
var req api.EmbedRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
reqCh <- req
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(api.EmbedResponse{
|
||||
Model: "test-embedding-model",
|
||||
Embeddings: [][]float32{{0.4, 0.5}},
|
||||
LoadDuration: 5 * time.Millisecond,
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
|
||||
t.Setenv("OLLAMA_HOST", mockServer.URL)
|
||||
t.Cleanup(mockServer.Close)
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(t.Context())
|
||||
cmd.Flags().String("keepalive", "", "")
|
||||
cmd.Flags().Bool("truncate", false, "")
|
||||
cmd.Flags().Int("dimensions", 0, "")
|
||||
cmd.Flags().Bool("verbose", false, "")
|
||||
cmd.Flags().Bool("insecure", false, "")
|
||||
cmd.Flags().Bool("nowordwrap", false, "")
|
||||
cmd.Flags().String("format", "", "")
|
||||
cmd.Flags().String("think", "", "")
|
||||
cmd.Flags().Bool("hidethinking", false, "")
|
||||
|
||||
if err := cmd.Flags().Set("truncate", "true"); err != nil {
|
||||
t.Fatalf("failed to set truncate flag: %v", err)
|
||||
}
|
||||
if err := cmd.Flags().Set("dimensions", "2"); err != nil {
|
||||
t.Fatalf("failed to set dimensions flag: %v", err)
|
||||
}
|
||||
if err := cmd.Flags().Set("keepalive", "5m"); err != nil {
|
||||
t.Fatalf("failed to set keepalive flag: %v", err)
|
||||
}
|
||||
|
||||
oldStdout := os.Stdout
|
||||
r, w, _ := os.Pipe()
|
||||
os.Stdout = w
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- RunHandler(cmd, []string{"test-embedding-model", "test", "input"})
|
||||
}()
|
||||
|
||||
err := <-errCh
|
||||
w.Close()
|
||||
os.Stdout = oldStdout
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("RunHandler returned error: %v", err)
|
||||
}
|
||||
|
||||
var out bytes.Buffer
|
||||
io.Copy(&out, r)
|
||||
|
||||
select {
|
||||
case req := <-reqCh:
|
||||
inputText, _ := req.Input.(string)
|
||||
if diff := cmp.Diff("test input", inputText); diff != "" {
|
||||
t.Errorf("unexpected input (-want +got):\n%s", diff)
|
||||
}
|
||||
if req.Truncate == nil || !*req.Truncate {
|
||||
t.Errorf("expected truncate pointer true, got %v", req.Truncate)
|
||||
}
|
||||
if req.Dimensions != 2 {
|
||||
t.Errorf("expected dimensions 2, got %d", req.Dimensions)
|
||||
}
|
||||
if req.KeepAlive == nil || req.KeepAlive.Duration != 5*time.Minute {
|
||||
t.Errorf("unexpected keepalive duration: %v", req.KeepAlive)
|
||||
}
|
||||
default:
|
||||
t.Fatal("server did not receive embed request")
|
||||
}
|
||||
|
||||
expectOutput := "[0.4,0.5]\n"
|
||||
if diff := cmp.Diff(expectOutput, out.String()); diff != "" {
|
||||
t.Errorf("unexpected output (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunEmbeddingModelPipedInput(t *testing.T) {
|
||||
reqCh := make(chan api.EmbedRequest, 1)
|
||||
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/api/show" && r.Method == http.MethodPost {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(api.ShowResponse{
|
||||
Capabilities: []model.Capability{model.CapabilityEmbedding},
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/api/embed" && r.Method == http.MethodPost {
|
||||
var req api.EmbedRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
reqCh <- req
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(api.EmbedResponse{
|
||||
Model: "test-embedding-model",
|
||||
Embeddings: [][]float32{{0.6, 0.7}},
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
|
||||
t.Setenv("OLLAMA_HOST", mockServer.URL)
|
||||
t.Cleanup(mockServer.Close)
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(t.Context())
|
||||
cmd.Flags().String("keepalive", "", "")
|
||||
cmd.Flags().Bool("truncate", false, "")
|
||||
cmd.Flags().Int("dimensions", 0, "")
|
||||
cmd.Flags().Bool("verbose", false, "")
|
||||
cmd.Flags().Bool("insecure", false, "")
|
||||
cmd.Flags().Bool("nowordwrap", false, "")
|
||||
cmd.Flags().String("format", "", "")
|
||||
cmd.Flags().String("think", "", "")
|
||||
cmd.Flags().Bool("hidethinking", false, "")
|
||||
|
||||
// Capture stdin
|
||||
oldStdin := os.Stdin
|
||||
stdinR, stdinW, _ := os.Pipe()
|
||||
os.Stdin = stdinR
|
||||
stdinW.Write([]byte("piped text"))
|
||||
stdinW.Close()
|
||||
|
||||
// Capture stdout
|
||||
oldStdout := os.Stdout
|
||||
stdoutR, stdoutW, _ := os.Pipe()
|
||||
os.Stdout = stdoutW
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- RunHandler(cmd, []string{"test-embedding-model", "additional", "args"})
|
||||
}()
|
||||
|
||||
err := <-errCh
|
||||
stdoutW.Close()
|
||||
os.Stdout = oldStdout
|
||||
os.Stdin = oldStdin
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("RunHandler returned error: %v", err)
|
||||
}
|
||||
|
||||
var out bytes.Buffer
|
||||
io.Copy(&out, stdoutR)
|
||||
|
||||
select {
|
||||
case req := <-reqCh:
|
||||
inputText, _ := req.Input.(string)
|
||||
// Should combine piped input with command line args
|
||||
if diff := cmp.Diff("piped text additional args", inputText); diff != "" {
|
||||
t.Errorf("unexpected input (-want +got):\n%s", diff)
|
||||
}
|
||||
default:
|
||||
t.Fatal("server did not receive embed request")
|
||||
}
|
||||
|
||||
expectOutput := "[0.6,0.7]\n"
|
||||
if diff := cmp.Diff(expectOutput, out.String()); diff != "" {
|
||||
t.Errorf("unexpected output (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunEmbeddingModelNoInput(t *testing.T) {
|
||||
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/api/show" && r.Method == http.MethodPost {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(api.ShowResponse{
|
||||
Capabilities: []model.Capability{model.CapabilityEmbedding},
|
||||
}); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
|
||||
t.Setenv("OLLAMA_HOST", mockServer.URL)
|
||||
t.Cleanup(mockServer.Close)
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(t.Context())
|
||||
cmd.Flags().String("keepalive", "", "")
|
||||
cmd.Flags().Bool("truncate", false, "")
|
||||
cmd.Flags().Int("dimensions", 0, "")
|
||||
cmd.Flags().Bool("verbose", false, "")
|
||||
cmd.Flags().Bool("insecure", false, "")
|
||||
cmd.Flags().Bool("nowordwrap", false, "")
|
||||
cmd.Flags().String("format", "", "")
|
||||
cmd.Flags().String("think", "", "")
|
||||
cmd.Flags().Bool("hidethinking", false, "")
|
||||
|
||||
cmd.SetOut(io.Discard)
|
||||
cmd.SetErr(io.Discard)
|
||||
|
||||
// Test with no input arguments (only model name)
|
||||
err := RunHandler(cmd, []string{"test-embedding-model"})
|
||||
if err == nil || !strings.Contains(err.Error(), "embedding models require input text") {
|
||||
t.Fatalf("expected error about missing input, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetModelfileName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
Reference in New Issue
Block a user