Skip to content

Commit

Permalink
FInal revisions and test functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@JexSrs
JexSrs committed May 24, 2024
1 parent 97d1888 commit afa8528
Show file tree
Hide file tree
Showing 15 changed files with 801 additions and 463 deletions.
57 changes: 28 additions & 29 deletions ChatRequestBuilder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,24 +2,23 @@ package ollama

// ChatRequestBuilder represents the chat API request.
type ChatRequestBuilder struct {
Model *string `json:"model"`
Stream *bool `json:"stream"`
StreamBufferSize *int `json:"-"`
StreamFunc func(r *GCResponse, err error) `json:"-"`
Model *string `json:"model"`
Format *string `json:"format"`
Raw *bool `json:"raw"`
Messages []Message `json:"messages"`
KeepAlive *string `json:"keep_alive,omitempty"`
Options *Options `json:"options"`

Format *string `json:"format"`
Images []string `json:"images"`
Raw *bool `json:"raw"`
Messages []Message `json:"messages"`

Options *Options `json:"options"`
Stream *bool `json:"stream"`
StreamBufferSize *int `json:"-"`
StreamFunc func(r *ChatResponse, err error) `json:"-"`
}

// WithModel sets the model used for this request.
//
// Parameters:
// - v: The model name.
func (c ChatFunc) WithModel(v string) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithModel(v string) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
r.Model = &v
}
Expand All @@ -31,39 +30,29 @@ func (c ChatFunc) WithModel(v string) func(*ChatRequestBuilder) {
// - v: A boolean indicating whether to use streaming.
// - bufferSize: The size of the streamed buffer
// - f: The function to handle streaming
func (c ChatFunc) WithStream(v bool, bufferSize int, f func(r *GCResponse, err error)) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithStream(v bool, bufferSize int, fn func(r *ChatResponse, err error)) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
r.Stream = &v
r.StreamBufferSize = &bufferSize
r.StreamFunc = f
r.StreamFunc = fn
}
}

// WithFormat sets the format to return a response in. Currently, the only accepted value is "json".
//
// Parameters:
// - v: The format string.
func (c ChatFunc) WithFormat(v string) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithFormat(v string) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
r.Format = &v
}
}

// WithImage appends an image to the message sent to The image must be base64 encoded.
//
// Parameters:
// - v: The base64 encoded image string.
func (c ChatFunc) WithImage(v string) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
r.Images = append(r.Images, v)
}
}

// WithRaw bypasses the templating system and provides a full prompt.
//
// Parameters:
// - v: A boolean indicating whether to use raw mode.
func (c ChatFunc) WithRaw(v bool) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithRaw(v bool) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
r.Raw = &v
}
Expand All @@ -73,7 +62,7 @@ func (c ChatFunc) WithRaw(v bool) func(*ChatRequestBuilder) {
//
// Parameters:
// - v: The temperature value.
func (c ChatFunc) WithTemperature(v float64) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithTemperature(v float64) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
if r.Options == nil {
r.Options = &Options{}
Expand All @@ -87,7 +76,7 @@ func (c ChatFunc) WithTemperature(v float64) func(*ChatRequestBuilder) {
//
// Parameters:
// - v: The seed value.
func (c ChatFunc) WithSeed(v int) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithSeed(v int) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
if r.Options == nil {
r.Options = &Options{}
Expand All @@ -101,7 +90,7 @@ func (c ChatFunc) WithSeed(v int) func(*ChatRequestBuilder) {
//
// Parameters:
// - v: The message to append.
func (c ChatFunc) WithMessage(v Message) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithMessage(v Message) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
if v.Role == nil {
v.Role = pointer("user")
Expand All @@ -111,11 +100,21 @@ func (c ChatFunc) WithMessage(v Message) func(*ChatRequestBuilder) {
}
}

// WithKeepAlive controls how long the model will stay loaded into memory following the request.
//
// Parameters:
// - v: The keep alive duration.
func (f *ChatFunc) WithKeepAlive(v string) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
r.KeepAlive = &v
}
}

// WithOptions sets the options for this request. It will override any settings set before, such as temperature and seed.
//
// Parameters:
// - v: The options to set.
func (c ChatFunc) WithOptions(v Options) func(*ChatRequestBuilder) {
func (f *ChatFunc) WithOptions(v Options) func(*ChatRequestBuilder) {
return func(r *ChatRequestBuilder) {
r.Options = &v
}
Expand Down
20 changes: 10 additions & 10 deletions GenerateEmbeddingsBuilder.go → GenerateEmbeddingsRequestBuilder.go
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
package ollama

// GenerateEmbeddingsBuilder represents the generate embeddings API request.
type GenerateEmbeddingsBuilder struct {
// GenerateEmbeddingsRequestBuilder represents the generate embeddings API request.
type GenerateEmbeddingsRequestBuilder struct {
Model *string `json:"model"`
Prompt *string `json:"prompt"`
KeepAlive *string `json:"keep_alive"`
Expand All @@ -12,8 +12,8 @@ type GenerateEmbeddingsBuilder struct {
//
// Parameters:
// - v: The model name.
func (c GenerateEmbeddingsFunc) WithModel(v string) func(*GenerateEmbeddingsBuilder) {
return func(r *GenerateEmbeddingsBuilder) {
func (c GenerateEmbeddingsFunc) WithModel(v string) func(*GenerateEmbeddingsRequestBuilder) {
return func(r *GenerateEmbeddingsRequestBuilder) {
r.Model = &v
}
}
Expand All @@ -22,8 +22,8 @@ func (c GenerateEmbeddingsFunc) WithModel(v string) func(*GenerateEmbeddingsBuil
//
// Parameters:
// - v: The prompt string.
func (c GenerateEmbeddingsFunc) WithPrompt(v string) func(*GenerateEmbeddingsBuilder) {
return func(r *GenerateEmbeddingsBuilder) {
func (c GenerateEmbeddingsFunc) WithPrompt(v string) func(*GenerateEmbeddingsRequestBuilder) {
return func(r *GenerateEmbeddingsRequestBuilder) {
r.Prompt = &v
}
}
Expand All @@ -32,8 +32,8 @@ func (c GenerateEmbeddingsFunc) WithPrompt(v string) func(*GenerateEmbeddingsBui
//
// Parameters:
// - v: The keep alive string.
func (c GenerateEmbeddingsFunc) WithKeepAlive(v string) func(*GenerateEmbeddingsBuilder) {
return func(r *GenerateEmbeddingsBuilder) {
func (c GenerateEmbeddingsFunc) WithKeepAlive(v string) func(*GenerateEmbeddingsRequestBuilder) {
return func(r *GenerateEmbeddingsRequestBuilder) {
r.KeepAlive = &v
}
}
Expand All @@ -42,8 +42,8 @@ func (c GenerateEmbeddingsFunc) WithKeepAlive(v string) func(*GenerateEmbeddings
//
// Parameters:
// - v: The options to set.
func (c GenerateEmbeddingsFunc) WithOptions(v Options) func(*GenerateEmbeddingsBuilder) {
return func(r *GenerateEmbeddingsBuilder) {
func (c GenerateEmbeddingsFunc) WithOptions(v Options) func(*GenerateEmbeddingsRequestBuilder) {
return func(r *GenerateEmbeddingsRequestBuilder) {
r.Options = &v
}
}
68 changes: 55 additions & 13 deletions GenerateRequestBuilder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,20 @@ package ollama

// GenerateRequestBuilder represents the generate API request.
type GenerateRequestBuilder struct {
Model *string `json:"model"`
Prompt *string `json:"prompt"`
Stream *bool `json:"stream"`
StreamBufferSize *int `json:"-"`
StreamFunc func(r *GCResponse, err error) `json:"-"`

// Format the format to return a response in. Currently, the only accepted value is json
Format *string `json:"format"`
Images []string `json:"images"`
Raw *bool `json:"raw"`

Options *Options `json:"options"`
Model *string `json:"model"`
Prompt *string `json:"prompt"`
System *string `json:"system"`
Template *string `json:"template"`
Format *string `json:"format"`
Images []string `json:"images"`
Raw *bool `json:"raw"`
Context []int `json:"context,omitempty"`
KeepAlive *string `json:"keep_alive,omitempty"`
Options *Options `json:"options"`

Stream *bool `json:"stream"`
StreamBufferSize *int `json:"-"`
StreamFunc func(r *GenerateResponse, err error) `json:"-"`
}

// WithModel sets the model used for this request.
Expand All @@ -36,13 +38,53 @@ func (c GenerateFunc) WithPrompt(v string) func(*GenerateRequestBuilder) {
}
}

// WithSystem overrides the model's default system message/prompt.
//
// Parameters:
// - v: The system string.
func (c GenerateFunc) WithSystem(v string) func(*GenerateRequestBuilder) {
return func(r *GenerateRequestBuilder) {
r.System = &v
}
}

// WithTemplate overrides the model's default prompt template.
//
// Parameters:
// - v: The template string.
func (c GenerateFunc) WithTemplate(v string) func(*GenerateRequestBuilder) {
return func(r *GenerateRequestBuilder) {
r.Template = &v
}
}

// WithContext overrides the model's default prompt template.
//
// Parameters:
// - v: The content int array.
func (c GenerateFunc) WithContext(v []int) func(*GenerateRequestBuilder) {
return func(r *GenerateRequestBuilder) {
r.Context = v
}
}

// WithKeepAlive controls how long the model will stay loaded in memory following this request.
//
// Parameters:
// - v: The keep alive duration.
func (c GenerateFunc) WithKeepAlive(v string) func(*GenerateRequestBuilder) {
return func(r *GenerateRequestBuilder) {
r.KeepAlive = &v
}
}

// WithStream passes a function to allow reading stream types.
//
// Parameters:
// - v: A boolean indicating whether to use streaming.
// - bufferSize: The size of the streamed buffer
// - f: The function to handle streaming types.
func (c GenerateFunc) WithStream(v bool, bufferSize int, f func(r *GCResponse, err error)) func(*GenerateRequestBuilder) {
func (c GenerateFunc) WithStream(v bool, bufferSize int, f func(r *GenerateResponse, err error)) func(*GenerateRequestBuilder) {
return func(r *GenerateRequestBuilder) {
r.Stream = &v
r.StreamBufferSize = &bufferSize
Expand Down
Loading

0 comments on commit afa8528

Please sign in to comment.