baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

simplify api client

Browse source
This commit is contained in:
Michael Yang 2023-07-06 15:02:10 -07:00
parent c4b9e84945
commit b0e63bfb4c
2 changed files with 21 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
api/client.go
View file

@ -5,14 +5,24 @@ import (
"bytes" "bytes"
"context" "context"
"encoding/json" "encoding/json"
"fmt"
"io" "io"
"net/http" "net/http"
"net/url"
) )
type Client struct { type Client struct {
URL string base url.URL
HTTP http.Client }
func NewClient(hosts ...string) *Client {
host := "127.0.0.1:11434"
if len(hosts) > 0 {
host = hosts[0]
}
return &Client{
base: url.URL{Scheme: "http", Host: host},
}
} }
func (c *Client) stream(ctx context.Context, method string, path string, reqData any, fn func(bts []byte) error) error { func (c *Client) stream(ctx context.Context, method string, path string, reqData any, fn func(bts []byte) error) error {
@ -27,23 +37,21 @@ func (c *Client) stream(ctx context.Context, method string, path string, reqData
reqBody = bytes.NewReader(data) reqBody = bytes.NewReader(data)
} }
url := fmt.Sprintf("%s%s", c.URL, path) request, err := http.NewRequestWithContext(ctx, method, c.base.JoinPath(path).String(), reqBody)
req, err := http.NewRequestWithContext(ctx, method, url, reqBody)
if err != nil { if err != nil {
return err return err
} }
req.Header.Set("Content-Type", "application/json") request.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json") request.Header.Set("Accept", "application/json")
res, err := c.HTTP.Do(req) response, err := http.DefaultClient.Do(request)
if err != nil { if err != nil {
return err return err
} }
defer res.Body.Close() defer response.Body.Close()
scanner := bufio.NewScanner(res.Body) scanner := bufio.NewScanner(response.Body)
for scanner.Scan() { for scanner.Scan() {
if err := fn(scanner.Bytes()); err != nil { if err := fn(scanner.Bytes()); err != nil {
return err return err

16
cmd/cmd.go
View file

@ -36,10 +36,7 @@ func RunRun(cmd *cobra.Command, args []string) error {
} }
func pull(model string) error { func pull(model string) error {
client, err := NewAPIClient() client := api.NewClient()
if err != nil {
return err
}
var bar *progressbar.ProgressBar var bar *progressbar.ProgressBar
return client.Pull( return client.Pull(
@ -68,10 +65,7 @@ func RunGenerate(_ *cobra.Command, args []string) error {
} }
func generate(model string, prompts ...string) error { func generate(model string, prompts ...string) error {
client, err := NewAPIClient() client := api.NewClient()
if err != nil {
return err
}
for _, prompt := range prompts { for _, prompt := range prompts {
client.Generate(context.Background(), &api.GenerateRequest{Model: model, Prompt: prompt}, func(resp api.GenerateResponse) error { client.Generate(context.Background(), &api.GenerateRequest{Model: model, Prompt: prompt}, func(resp api.GenerateResponse) error {
@ -121,12 +115,6 @@ func RunServer(_ *cobra.Command, _ []string) error {
return server.Serve(ln) return server.Serve(ln)
} }
func NewAPIClient() (*api.Client, error) {
return &api.Client{
URL: "http://localhost:11434",
}, nil
}
func NewCLI() *cobra.Command { func NewCLI() *cobra.Command {
log.SetFlags(log.LstdFlags | log.Lshortfile) log.SetFlags(log.LstdFlags | log.Lshortfile)