ollama client

This commit is contained in:
Erik Winter 2024-06-11 07:29:31 +02:00
parent ca4411b3a0
commit 7cc563ab05
2 changed files with 82 additions and 1 deletions

View File

@ -3,14 +3,31 @@ package main
import ( import (
"flag" "flag"
"fmt" "fmt"
"os"
) )
var ( var (
inputFile = flag.String("i", "", "input file") inputFile = flag.String("i", "", "input file")
model = flag.String("m", "llama3", "model file")
) )
func main() { func main() {
flag.Parse() flag.Parse()
if *inputFile == "" {
fmt.Println("Hello World") fmt.Println("No input file specified")
os.Exit(1)
}
ollamaHost := os.Getenv("OLLAMA_HOST")
if ollamaHost == "" {
fmt.Println("OLLAMA_HOST environment variable not set")
os.Exit(1)
}
ollama := NewOllama(ollamaHost)
res, err := ollama.Generate(*model, "Could you translate the following text into English? 'Mijn fietsband is lek. Wat moet ik nu doen'")
if err != nil {
fmt.Println(err)
os.Exit(1)
}
fmt.Println(res)
} }

View File

@ -0,0 +1,64 @@
package main
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
)
type Ollama struct {
baseURL string
c *http.Client
}
func NewOllama(baseURL string) *Ollama {
return &Ollama{
baseURL: baseURL,
c: &http.Client{},
}
}
func (o *Ollama) Generate(model, prompt string) (string, error) {
url := fmt.Sprintf("%s/api/generate", o.baseURL)
reqBody := struct {
Model string
Prompt string
Format string
Stream bool
}{
Model: model,
Prompt: prompt,
Format: "json",
Stream: false,
}
reqBodyJSON, err := json.Marshal(reqBody)
if err != nil {
return "", err
}
req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(reqBodyJSON))
if err != nil {
return "", err
}
res, err := o.c.Do(req)
if err != nil {
return "", err
}
body, err := io.ReadAll(res.Body)
if err != nil {
return "", err
}
defer res.Body.Close()
resBody := struct {
Response string
}{}
if err := json.Unmarshal(body, &resBody); err != nil {
return "", err
}
return resBody.Response, nil
}