52 lines
1.2 KiB
Go
52 lines
1.2 KiB
Go
package ollama
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
)
|
|
|
|
type OllamaClient struct {
|
|
client api.Client
|
|
}
|
|
|
|
func New(ctx context.Context, host *url.URL, http *http.Client) *OllamaClient {
|
|
client := api.NewClient(host, http)
|
|
return &OllamaClient{client: *client}
|
|
}
|
|
|
|
func (o *OllamaClient) EmbedTexts(ctx context.Context, posts []string) ([][]float64, error) {
|
|
vectors := make([][]float64, 0, len(posts))
|
|
for _, t := range posts {
|
|
resp, err := o.client.Embeddings(ctx, &api.EmbeddingRequest{
|
|
Model: "nomic-embed-text",
|
|
Prompt: t,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
vectors = append(vectors, resp.Embedding)
|
|
}
|
|
return vectors, nil
|
|
}
|
|
|
|
func (o *OllamaClient) SummarizeStyle(ctx context.Context, posts []string) (string, error) {
|
|
prompt := "Summarize writing style, recurring themes, tone, personality traits of this user:\n\n" +
|
|
strings.Join(posts, "\n") + "\n\nStyle summary:"
|
|
var fullResponse strings.Builder
|
|
err := o.client.Generate(ctx, &api.GenerateRequest{
|
|
Model: "mistral",
|
|
Prompt: prompt,
|
|
}, func(resp api.GenerateResponse) error {
|
|
fullResponse.WriteString(resp.Response)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return fullResponse.String(), nil
|
|
}
|