├── .github ├── semantic.yml └── workflows │ └── build.yml ├── .gitignore ├── LICENSE ├── README.md ├── chat.go ├── embedding.go ├── embedding ├── embedding.go ├── option.go ├── request.go └── response.go ├── go.mod ├── go.sum ├── prompt ├── message.go ├── option.go └── prompt.go └── provider ├── ollama ├── chat_model.go ├── client.go ├── client_test.go ├── embedding_model.go └── types.go └── openai └── chat_model.go /.github/semantic.yml: -------------------------------------------------------------------------------- 1 | # Always validate the PR title AND all the commits 2 | titleAndCommits: true -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | test: 7 | runs-on: ubuntu-latest 8 | strategy: 9 | matrix: 10 | go: ['1.23', '1.24'] 11 | steps: 12 | - name: Set up Go 13 | uses: actions/setup-go@v4 14 | with: 15 | go-version: ${{ matrix.go }} 16 | 17 | - name: Check out code 18 | uses: actions/checkout@v2 19 | 20 | - name: Install dependencies 21 | run: | 22 | go mod download 23 | - name: Run Unit tests 24 | run: | 25 | go test -race -covermode atomic -coverprofile=covprofile ./... 26 | - name: Install goveralls 27 | run: go install github.com/mattn/goveralls@latest 28 | - name: Send coverage 29 | env: 30 | COVERALLS_TOKEN: ${{ secrets.GITHUB_TOKEN }} 31 | run: goveralls -coverprofile=covprofile -service=github 32 | 33 | # semantic-release: 34 | # needs: [test] 35 | # runs-on: ubuntu-latest 36 | # steps: 37 | # - uses: actions/checkout@v2 38 | # - uses: actions/setup-node@v2 39 | # with: 40 | # node-version: 'lts/*' 41 | # 42 | # - name: Run semantic-release 43 | # if: github.repository == 'tech1024/goai' && github.event_name == 'push' 44 | # env: 45 | # GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 46 | # run: npx semantic-release -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # If you prefer the allow list template instead of the deny list, see community template: 2 | # https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore 3 | # 4 | # Binaries for programs and plugins 5 | *.exe 6 | *.exe~ 7 | *.dll 8 | *.so 9 | *.dylib 10 | 11 | # ide 12 | .vscode 13 | .idea 14 | 15 | # Test binary, built with `go test -c` 16 | *.test 17 | 18 | # Output of the go coverage tool, specifically when used with LiteIDE 19 | *.out 20 | 21 | # Dependency directories (remove the comment below to include it) 22 | # vendor/ 23 | 24 | # Go workspace file 25 | go.work 26 | go.work.sum 27 | 28 | # env file 29 | .env 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Go AI 2 | == 3 | 4 | [![build](https://github.com/tech1024/goai/actions/workflows/build.yml/badge.svg)](https://github.com/tech1024/goai/actions/workflows/build.yml) 5 | [![Coverage Status](https://coveralls.io/repos/github/tech1024/goai/badge.svg?branch=main)](https://coveralls.io/github/tech1024/goai?branch=main) 6 | [![Go Report Card](https://goreportcard.com/badge/github.com/tech1024/goai)](https://goreportcard.com/report/github.com/tech1024/goai) 7 | [![Godoc](https://godoc.org/github.com/tech1024/goai?status.svg)](https://pkg.go.dev/github.com/tech1024/goai) 8 | [![Release](https://img.shields.io/github/release/tech1024/goai.svg)](https://github.com/tech1024/goai/releases/latest) 9 | 10 | A golang API library for AI Engineering. 11 | 12 | This is a high level feature overview. 13 | 14 | - Chat Completion 15 | - Embedding 16 | 17 | ## Installation 18 | 19 | ```shell 20 | go get -u 'github.com/tech1024/goai' 21 | ``` 22 | 23 | ## Getting Started 24 | 25 | ```go 26 | package main 27 | 28 | import ( 29 | "context" 30 | "log" 31 | 32 | "github.com/tech1024/goai" 33 | "github.com/tech1024/goai/provider/ollama" 34 | ) 35 | 36 | func main() { 37 | ollamaClient, _ := ollama.NewClient("http://127.0.0.1:11434") 38 | chat := goai.NewChat(ollama.NewNewChatModel(ollamaClient, "deepseek-r1")) 39 | result, err := chat.Chat(context.Background(), "What can you do for me ?") 40 | if err != nil { 41 | log.Fatal(err) 42 | } 43 | 44 | log.Println(result) 45 | } 46 | ``` 47 | 48 | ## License 49 | 50 | This project is licensed under the [Apache 2.0 license](LICENSE). 51 | 52 | ## Contact 53 | 54 | If you have any issues or feature requests, please contact us. PR is welcomed. 55 | - https://github.com/tech1024/goai/issues -------------------------------------------------------------------------------- /chat.go: -------------------------------------------------------------------------------- 1 | package goai 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/tech1024/goai/prompt" 7 | ) 8 | 9 | type ChatModel interface { 10 | Call(ctx context.Context, prompt prompt.Prompt) (string, error) 11 | Stream(ctx context.Context, prompt prompt.Prompt, receive func([]byte) error) error 12 | } 13 | 14 | func NewChat(chatModel ChatModel) *Chat { 15 | return &Chat{ 16 | chatModel: chatModel, 17 | } 18 | } 19 | 20 | type Chat struct { 21 | chatModel ChatModel 22 | } 23 | 24 | // Chat send a message, it returns string 25 | func (c *Chat) Chat(ctx context.Context, content string) (string, error) { 26 | return c.Prompt(ctx, prompt.NewPrompt( 27 | prompt.UserMessage(content), 28 | )) 29 | } 30 | 31 | // ChatStream send a message, need to receive its returns 32 | func (c *Chat) ChatStream(ctx context.Context, content string, receive func([]byte) error) error { 33 | return c.Stream(ctx, prompt.NewPrompt( 34 | prompt.UserMessage(content), 35 | ), receive) 36 | } 37 | 38 | func (c *Chat) Prompt(ctx context.Context, p prompt.Prompt) (string, error) { 39 | return c.chatModel.Call(ctx, p) 40 | } 41 | 42 | func (c *Chat) Stream(ctx context.Context, p prompt.Prompt, fn func([]byte) error) error { 43 | return c.chatModel.Stream(ctx, p, fn) 44 | } 45 | -------------------------------------------------------------------------------- /embedding.go: -------------------------------------------------------------------------------- 1 | package goai 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/tech1024/goai/embedding" 7 | ) 8 | 9 | type EmbeddingModel interface { 10 | Call(context.Context, embedding.Request) (embedding.Response, error) 11 | } 12 | 13 | type Embedding struct { 14 | embeddingModel EmbeddingModel 15 | } 16 | 17 | func (e *Embedding) Embed(ctx context.Context, text string) ([]float32, error) { 18 | es, err := e.Embeds(ctx, text) 19 | if err != nil { 20 | return nil, err 21 | } 22 | 23 | return es[0], nil 24 | } 25 | 26 | func (e *Embedding) Embeds(ctx context.Context, texts ...string) ([][]float32, error) { 27 | response, err := e.embeddingModel.Call( 28 | ctx, 29 | embedding.NewRequest(texts, embedding.Option{}), 30 | ) 31 | 32 | if err != nil { 33 | return nil, err 34 | } 35 | 36 | return response.List(), nil 37 | } 38 | -------------------------------------------------------------------------------- /embedding/embedding.go: -------------------------------------------------------------------------------- 1 | package embedding 2 | 3 | type Embedding struct { 4 | Embedding []float32 5 | Index int 6 | } 7 | -------------------------------------------------------------------------------- /embedding/option.go: -------------------------------------------------------------------------------- 1 | package embedding 2 | 3 | type Option struct { 4 | // Model the model to use for the chat. 5 | Model string 6 | 7 | Dimensions int 8 | } 9 | -------------------------------------------------------------------------------- /embedding/request.go: -------------------------------------------------------------------------------- 1 | package embedding 2 | 3 | func NewRequest(inputs []string, option Option) Request { 4 | return Request{ 5 | Inputs: inputs, 6 | Option: option, 7 | } 8 | } 9 | 10 | type Request struct { 11 | Inputs []string 12 | Option Option 13 | } 14 | -------------------------------------------------------------------------------- /embedding/response.go: -------------------------------------------------------------------------------- 1 | package embedding 2 | 3 | type Response struct { 4 | Embeddings []Embedding 5 | } 6 | 7 | func (r *Response) List() [][]float32 { 8 | fs := make([][]float32, len(r.Embeddings)) 9 | for i, embedding := range r.Embeddings { 10 | fs[i] = embedding.Embedding 11 | } 12 | 13 | return fs 14 | } 15 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/tech1024/goai 2 | 3 | go 1.23 4 | 5 | require github.com/sashabaranov/go-openai v1.38.0 // indirect 6 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/sashabaranov/go-openai v1.38.0 h1:hNN5uolKwdbpiqOn7l+Z2alch/0n0rSFyg4n+GZxR5k= 2 | github.com/sashabaranov/go-openai v1.38.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= 3 | -------------------------------------------------------------------------------- /prompt/message.go: -------------------------------------------------------------------------------- 1 | package prompt 2 | 3 | const ( 4 | MessageTypeUser MessageType = "user" 5 | MessageTypeAssistant MessageType = "assistant" 6 | MessageTypeSystem MessageType = "system" 7 | MessageTypeTool MessageType = "tool" 8 | ) 9 | 10 | // MessageType Enumeration representing types of Message in a chat application. 11 | type MessageType string 12 | 13 | func (mt MessageType) String() string { 14 | return string(mt) 15 | } 16 | 17 | type Message interface { 18 | // Type the message type. 19 | Type() MessageType 20 | 21 | // Text the content of the message 22 | Text() string 23 | 24 | // Metadata the metadata associated with the content. 25 | Metadata() map[string]any 26 | } 27 | 28 | type defaultMessage struct { 29 | _type MessageType 30 | text string 31 | metadata map[string]any 32 | } 33 | 34 | func (m *defaultMessage) Type() MessageType { 35 | return m._type 36 | } 37 | 38 | func (m *defaultMessage) Text() string { 39 | return m.text 40 | } 41 | 42 | func (m *defaultMessage) Metadata() map[string]any { 43 | return m.metadata 44 | } 45 | 46 | // UserMessage a message of the type 'user' 47 | func UserMessage(message string) *defaultMessage { 48 | return &defaultMessage{ 49 | _type: MessageTypeUser, 50 | text: message, 51 | } 52 | } 53 | 54 | // AssistantMessage a message of the type 'assistant' 55 | func AssistantMessage(message string) *defaultMessage { 56 | return &defaultMessage{ 57 | _type: MessageTypeAssistant, 58 | text: message, 59 | } 60 | } 61 | 62 | // SystemMessage a message of the type 'system' 63 | func SystemMessage(message string) *defaultMessage { 64 | return &defaultMessage{ 65 | _type: MessageTypeSystem, 66 | text: message, 67 | } 68 | } 69 | 70 | // ToolMessage a message of the type 'tool' 71 | func ToolMessage(message string) *defaultMessage { 72 | return &defaultMessage{ 73 | _type: MessageTypeTool, 74 | text: message, 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /prompt/option.go: -------------------------------------------------------------------------------- 1 | package prompt 2 | 3 | type Option struct { 4 | // Model the model to use for the chat. 5 | Model string 6 | } 7 | -------------------------------------------------------------------------------- /prompt/prompt.go: -------------------------------------------------------------------------------- 1 | package prompt 2 | 3 | type Prompt struct { 4 | Messages []Message 5 | ChatOption Option 6 | } 7 | 8 | func NewPrompt(messages ...Message) Prompt { 9 | return Prompt{ 10 | Messages: messages, 11 | } 12 | } 13 | -------------------------------------------------------------------------------- /provider/ollama/chat_model.go: -------------------------------------------------------------------------------- 1 | package ollama 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/tech1024/goai/prompt" 7 | ) 8 | 9 | func NewNewChatModel(client *Client, model string) *ChatModel { 10 | return &ChatModel{ 11 | client: client, 12 | model: model, 13 | } 14 | } 15 | 16 | type ChatModel struct { 17 | client *Client 18 | model string 19 | } 20 | 21 | func (chatModel *ChatModel) Call(ctx context.Context, prompt prompt.Prompt) (string, error) { 22 | req, err := chatModel.buildChatRequest(prompt) 23 | if err != nil { 24 | return "", err 25 | } 26 | 27 | resp, err := chatModel.client.Chat(ctx, req) 28 | if err != nil { 29 | return "", err 30 | } 31 | 32 | return resp.Message.Content, nil 33 | } 34 | 35 | func (chatModel *ChatModel) Stream(ctx context.Context, prompt prompt.Prompt, fn func([]byte) error) error { 36 | req, err := chatModel.buildChatRequest(prompt) 37 | if err != nil { 38 | return err 39 | } 40 | 41 | err = chatModel.client.ChatStream(ctx, req, fn) 42 | if err != nil { 43 | return err 44 | } 45 | 46 | return nil 47 | } 48 | 49 | func (chatModel *ChatModel) buildChatRequest(prompt prompt.Prompt) (*ChatRequest, error) { 50 | request := ChatRequest{ 51 | Model: chatModel.model, 52 | Messages: make([]Message, len(prompt.Messages)), 53 | } 54 | for i, message := range prompt.Messages { 55 | request.Messages[i] = Message{ 56 | Role: message.Type().String(), 57 | Content: message.Text(), 58 | } 59 | } 60 | if prompt.ChatOption.Model != "" { 61 | request.Model = prompt.ChatOption.Model 62 | } 63 | 64 | return &request, nil 65 | } 66 | -------------------------------------------------------------------------------- /provider/ollama/client.go: -------------------------------------------------------------------------------- 1 | package ollama 2 | 3 | import ( 4 | "bufio" 5 | "bytes" 6 | "context" 7 | "encoding/json" 8 | "errors" 9 | "fmt" 10 | "io" 11 | "net/http" 12 | "net/url" 13 | "runtime" 14 | "time" 15 | ) 16 | 17 | func NewClient(baseUrl string) (*Client, error) { 18 | var err error 19 | 20 | client := Client{} 21 | 22 | client.baseUrl, err = url.Parse(baseUrl) 23 | if err != nil { 24 | return nil, err 25 | } 26 | 27 | client.httpClient = http.DefaultClient 28 | 29 | return &client, nil 30 | } 31 | 32 | type Client struct { 33 | baseUrl *url.URL // baseUrl The base url of the Client server. 34 | httpClient *http.Client 35 | } 36 | 37 | // Chat part 38 | 39 | // Message is a single message in a chat sequence. The message contains the 40 | // role ("system", "user", or "assistant"), the content and an optional list 41 | // of images. 42 | type Message struct { 43 | Role string `json:"role"` 44 | Content string `json:"content"` 45 | Images []ImageData `json:"images,omitempty"` 46 | ToolCalls []ToolCall `json:"tool_calls,omitempty"` 47 | } 48 | 49 | // ImageData represents the raw binary data of an image file. 50 | type ImageData []byte 51 | 52 | // ChatRequest describes a request sent by [Client.Chat]. 53 | type ChatRequest struct { 54 | // Model is the model name, as in [GenerateRequest]. 55 | Model string `json:"model"` 56 | 57 | // Messages is the messages of the chat - can be used to keep a chat memory. 58 | Messages []Message `json:"messages"` 59 | 60 | // Stream enables streaming of returned responses; true by default. 61 | Stream bool `json:"stream"` 62 | 63 | // Format is the format to return the response in (e.g. "json"). 64 | Format json.RawMessage `json:"format,omitempty"` 65 | 66 | // KeepAlive controls how long the model will stay loaded into memory 67 | // following the request. 68 | KeepAlive *Duration `json:"keep_alive,omitempty"` 69 | 70 | // Tools is an optional list of tools the model has access to. 71 | Tools `json:"tools,omitempty"` 72 | 73 | // Options lists model-specific options. 74 | Options map[string]interface{} `json:"options"` 75 | } 76 | 77 | type ChatResponse struct { 78 | Model string `json:"model"` 79 | CreatedAt time.Time `json:"created_at"` 80 | Message Message `json:"message"` 81 | DoneReason string `json:"done_reason,omitempty"` 82 | 83 | Done bool `json:"done"` 84 | 85 | Metrics 86 | } 87 | 88 | type Metrics struct { 89 | TotalDuration time.Duration `json:"total_duration,omitempty"` 90 | LoadDuration time.Duration `json:"load_duration,omitempty"` 91 | PromptEvalCount int `json:"prompt_eval_count,omitempty"` 92 | PromptEvalDuration time.Duration `json:"prompt_eval_duration,omitempty"` 93 | EvalCount int `json:"eval_count,omitempty"` 94 | EvalDuration time.Duration `json:"eval_duration,omitempty"` 95 | } 96 | 97 | func (c *Client) Chat(ctx context.Context, request *ChatRequest) (*ChatResponse, error) { 98 | var response ChatResponse 99 | request.Stream = false 100 | err := c.Post(ctx, "/api/chat", request, &response) 101 | if err != nil { 102 | return nil, err 103 | } 104 | 105 | return &response, err 106 | } 107 | 108 | func (c *Client) ChatStream(ctx context.Context, request *ChatRequest, fn func([]byte) error) error { 109 | request.Stream = true 110 | err := c.stream(ctx, http.MethodPost, "/api/chat", request, fn) 111 | if err != nil { 112 | return err 113 | } 114 | 115 | return err 116 | } 117 | 118 | // Embedding part 119 | 120 | // EmbedRequest is the request passed to [Client.Embed]. 121 | type EmbedRequest struct { 122 | // Model is the model name. 123 | Model string `json:"model"` 124 | 125 | // Input is the input to embed. 126 | Input any `json:"input"` 127 | 128 | // KeepAlive controls how long the model will stay loaded in memory following 129 | // this request. 130 | KeepAlive *Duration `json:"keep_alive,omitempty"` 131 | 132 | Truncate *bool `json:"truncate,omitempty"` 133 | 134 | // Options lists model-specific options. 135 | Options map[string]interface{} `json:"options"` 136 | } 137 | 138 | // EmbedResponse is the response from [Client.Embed]. 139 | type EmbedResponse struct { 140 | Model string `json:"model"` 141 | Embeddings [][]float32 `json:"embeddings"` 142 | 143 | TotalDuration time.Duration `json:"total_duration,omitempty"` 144 | LoadDuration time.Duration `json:"load_duration,omitempty"` 145 | PromptEvalCount int `json:"prompt_eval_count,omitempty"` 146 | } 147 | 148 | // Embed generates embeddings from a model. 149 | func (c *Client) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { 150 | var resp EmbedResponse 151 | err := c.Post(ctx, "/api/embed", req, &resp) 152 | 153 | if err != nil { 154 | return nil, err 155 | } 156 | 157 | return &resp, nil 158 | } 159 | 160 | // EmbeddingRequest is the request passed to [Client.Embeddings]. 161 | type EmbeddingRequest struct { 162 | // Model is the model name. 163 | Model string `json:"model"` 164 | 165 | // Prompt is the textual prompt to embed. 166 | Prompt string `json:"prompt"` 167 | 168 | // KeepAlive controls how long the model will stay loaded in memory following 169 | // this request. 170 | KeepAlive *Duration `json:"keep_alive,omitempty"` 171 | 172 | // Options lists model-specific options. 173 | Options map[string]interface{} `json:"options"` 174 | } 175 | 176 | // EmbeddingResponse is the response from [Client.Embeddings]. 177 | type EmbeddingResponse struct { 178 | Embedding []float64 `json:"embedding"` 179 | } 180 | 181 | // Embeddings generates an embedding from a model. 182 | func (c *Client) Embeddings(ctx context.Context, req *EmbeddingRequest) (*EmbeddingResponse, error) { 183 | var resp EmbeddingResponse 184 | if err := c.Post(ctx, "/api/embeddings", req, &resp); err != nil { 185 | return nil, err 186 | } 187 | return &resp, nil 188 | } 189 | 190 | // Request part 191 | 192 | func (c *Client) Post(ctx context.Context, path string, data, response any) error { 193 | httpResp, err := c.do(ctx, http.MethodPost, path, data) 194 | if err != nil { 195 | return err 196 | } 197 | 198 | return c.processResponse(httpResp, response) 199 | } 200 | 201 | func (c *Client) do(ctx context.Context, method, path string, data any) (*http.Response, error) { 202 | var body io.Reader 203 | var err error 204 | switch reqData := data.(type) { 205 | case io.Reader: 206 | body = reqData 207 | case nil: 208 | default: 209 | var jsonData []byte 210 | jsonData, err = c.marshalJSON(data) 211 | if err != nil { 212 | return nil, err 213 | } 214 | 215 | body = bytes.NewReader(jsonData) 216 | } 217 | 218 | request, err := http.NewRequestWithContext( 219 | ctx, method, c.baseUrl.JoinPath(path).String(), body, 220 | ) 221 | 222 | request.Header.Set("Content-Type", "application/json") 223 | request.Header.Set("Accept", "application/json") 224 | request.Header.Set("User-Agent", fmt.Sprintf("GoAI (%s %s) Go/%s", runtime.GOARCH, runtime.GOOS, runtime.Version())) 225 | 226 | return c.httpClient.Do(request) 227 | } 228 | 229 | const maxBufferSize = 512 * 1000 230 | 231 | func (c *Client) stream(ctx context.Context, method, path string, data any, fn func([]byte) error) error { 232 | httpResp, err := c.do(ctx, method, path, data) 233 | if err != nil { 234 | return err 235 | } 236 | 237 | defer httpResp.Body.Close() 238 | 239 | scanner := bufio.NewScanner(httpResp.Body) 240 | // increase the buffer size to avoid running out of space 241 | scanBuf := make([]byte, 0, maxBufferSize) 242 | scanner.Buffer(scanBuf, maxBufferSize) 243 | for scanner.Scan() { 244 | var errorResponse struct { 245 | Error string `json:"error,omitempty"` 246 | } 247 | 248 | bts := scanner.Bytes() 249 | if err := json.Unmarshal(bts, &errorResponse); err != nil { 250 | return fmt.Errorf("unmarshal: %w", err) 251 | } 252 | 253 | if errorResponse.Error != "" { 254 | return errors.New(errorResponse.Error) 255 | } 256 | 257 | if httpResp.StatusCode >= http.StatusBadRequest { 258 | return fmt.Errorf("http status code: %s", httpResp.Status) 259 | } 260 | 261 | if err := fn(bts); err != nil { 262 | return err 263 | } 264 | } 265 | 266 | return nil 267 | } 268 | 269 | type errorResponse struct { 270 | Error string `json:"error"` 271 | } 272 | 273 | func (c *Client) processResponse(httpResp *http.Response, response any) error { 274 | defer httpResp.Body.Close() 275 | 276 | respBody, err := io.ReadAll(httpResp.Body) 277 | if err != nil { 278 | return err 279 | } 280 | 281 | if httpResp.StatusCode < http.StatusBadRequest { 282 | return c.unMarshalJSON(respBody, response) 283 | } 284 | 285 | var e errorResponse 286 | err = c.unMarshalJSON(respBody, &e) 287 | if err != nil { 288 | return fmt.Errorf("http status code: %s, %w", httpResp.Status, err) 289 | } 290 | 291 | return fmt.Errorf("http status code: %s, %s", httpResp.Status, e.Error) 292 | } 293 | 294 | func (c *Client) marshalJSON(data any) ([]byte, error) { 295 | return json.Marshal(data) 296 | } 297 | 298 | func (c *Client) unMarshalJSON(data []byte, v any) error { 299 | return json.Unmarshal(data, v) 300 | } 301 | -------------------------------------------------------------------------------- /provider/ollama/client_test.go: -------------------------------------------------------------------------------- 1 | package ollama 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "errors" 7 | "fmt" 8 | "net/http" 9 | "net/http/httptest" 10 | "net/url" 11 | "reflect" 12 | "strings" 13 | "testing" 14 | ) 15 | 16 | func _handlerFunc(t *testing.T, wantCode int, wantResp any) func(http.ResponseWriter, *http.Request) { 17 | return func(w http.ResponseWriter, r *http.Request) { 18 | w.WriteHeader(wantCode) 19 | w.Header().Set("Content-Type", "application/json") 20 | var err error 21 | 22 | switch wantResp.(type) { 23 | case string: 24 | _, err = w.Write([]byte(wantResp.(string))) 25 | case []byte: 26 | _, err = w.Write(wantResp.([]byte)) 27 | default: 28 | err = json.NewEncoder(w).Encode(wantResp) 29 | } 30 | 31 | if err != nil { 32 | t.Fatal("failed to encode error response:", err) 33 | } 34 | } 35 | } 36 | 37 | func TestClient_Chat(t *testing.T) { 38 | tests := []struct { 39 | name string 40 | request *ChatRequest 41 | response *ChatResponse 42 | wantCode int 43 | wantErr error 44 | }{ 45 | { 46 | name: "test chat ok", 47 | request: &ChatRequest{ 48 | Model: "test-model", Messages: []Message{{Role: "user", Content: "request 1"}}, 49 | }, 50 | response: &ChatResponse{ 51 | Model: "test-model", Message: Message{Role: "user", Content: "response 1"}, 52 | }, 53 | wantCode: http.StatusOK, 54 | wantErr: nil, 55 | }, 56 | { 57 | name: "test chat model not exist", 58 | request: &ChatRequest{}, 59 | response: &ChatResponse{}, 60 | wantCode: http.StatusNotFound, 61 | wantErr: errors.New("model not found"), 62 | }, 63 | } 64 | var handlerFunc func(writer http.ResponseWriter, request *http.Request) 65 | ts := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { 66 | handlerFunc(writer, request) 67 | })) 68 | defer ts.Close() 69 | c := &Client{ 70 | baseUrl: &url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, 71 | httpClient: http.DefaultClient, 72 | } 73 | 74 | for _, tt := range tests { 75 | t.Run(tt.name, func(t *testing.T) { 76 | handlerFunc = _handlerFunc(t, tt.wantCode, fmt.Sprintf(`{"model": "%s", "message": {"role": "%s", "content": "%s"}}`, 77 | tt.response.Model, tt.response.Message.Role, tt.response.Message.Content)) 78 | if tt.wantCode > http.StatusBadRequest { 79 | handlerFunc = _handlerFunc(t, tt.wantCode, fmt.Sprintf(`{"error": "%s"}`, tt.wantErr.Error())) 80 | } 81 | 82 | got, err := c.Chat(context.Background(), tt.request) 83 | if !strings.Contains(fmt.Sprintf("%v", err), fmt.Sprintf("%v", tt.wantErr)) { 84 | t.Errorf("Chat() error = %v, wantErr %v", err, tt.wantErr) 85 | return 86 | } 87 | 88 | if err == nil && !reflect.DeepEqual(got, tt.response) { 89 | t.Errorf("Chat() got = %v, want %v", got, tt.response) 90 | } 91 | 92 | t.Logf("Chat() got = %v", got) 93 | }) 94 | } 95 | } 96 | 97 | func TestClient_Embed(t *testing.T) { 98 | type args struct { 99 | ctx context.Context 100 | req *EmbedRequest 101 | } 102 | tests := []struct { 103 | name string 104 | request *EmbedRequest 105 | response *EmbedResponse 106 | wantCode int 107 | wantErr error 108 | }{ 109 | { 110 | name: "test embed ok", 111 | request: &EmbedRequest{Model: "nomic-embed-text", Input: []string{"hello ai"}}, 112 | response: &EmbedResponse{Model: "nomic-embed-text", Embeddings: [][]float32{{0.25, 0.36}}}, 113 | wantCode: http.StatusOK, 114 | wantErr: nil, 115 | }, 116 | { 117 | name: "test embed model not exist", 118 | request: &EmbedRequest{}, 119 | response: &EmbedResponse{}, 120 | wantCode: http.StatusNotFound, 121 | wantErr: errors.New("model not found"), 122 | }, 123 | } 124 | 125 | var handlerFunc func(writer http.ResponseWriter, request *http.Request) 126 | ts := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { 127 | handlerFunc(writer, request) 128 | })) 129 | defer ts.Close() 130 | c := &Client{ 131 | baseUrl: &url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, 132 | httpClient: http.DefaultClient, 133 | } 134 | 135 | for _, tt := range tests { 136 | t.Run(tt.name, func(t *testing.T) { 137 | handlerFunc = _handlerFunc(t, tt.wantCode, fmt.Sprintf(`{"model": "%s", "embeddings": %s}`, 138 | tt.response.Model, strings.ReplaceAll(fmt.Sprintf("%v", tt.response.Embeddings), " ", ","))) 139 | if tt.wantCode > http.StatusBadRequest { 140 | handlerFunc = _handlerFunc(t, tt.wantCode, fmt.Sprintf(`{"error": "%s"}`, tt.wantErr.Error())) 141 | } 142 | 143 | got, err := c.Embed(context.Background(), tt.request) 144 | if !strings.Contains(fmt.Sprintf("%v", err), fmt.Sprintf("%v", tt.wantErr)) { 145 | t.Errorf("Embed() error = %v, wantErr %v", err, tt.wantErr) 146 | return 147 | } 148 | 149 | if err == nil && !reflect.DeepEqual(got, tt.response) { 150 | t.Errorf("Embed() got = %v, want %v", got, tt.response) 151 | } 152 | 153 | t.Logf("Embed() got = %v", got) 154 | }) 155 | } 156 | } 157 | -------------------------------------------------------------------------------- /provider/ollama/embedding_model.go: -------------------------------------------------------------------------------- 1 | package ollama 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/tech1024/goai/embedding" 7 | ) 8 | 9 | func NewEmbeddingModel(client *Client, model string) *EmbeddingModel { 10 | return &EmbeddingModel{ 11 | client: client, 12 | model: model, 13 | } 14 | } 15 | 16 | type EmbeddingModel struct { 17 | client *Client 18 | model string 19 | } 20 | 21 | func (embeddingModel *EmbeddingModel) Call(ctx context.Context, request embedding.Request) (embedding.Response, error) { 22 | er := EmbedRequest{ 23 | Model: embeddingModel.model, 24 | Input: request.Inputs, 25 | } 26 | 27 | var embeddingResponse embedding.Response 28 | resp, err := embeddingModel.client.Embed(ctx, &er) 29 | 30 | if err != nil { 31 | return embeddingResponse, err 32 | } 33 | 34 | embeddingResponse.Embeddings = make([]embedding.Embedding, len(resp.Embeddings)) 35 | for i, es := range resp.Embeddings { 36 | embeddingResponse.Embeddings[i] = embedding.Embedding{ 37 | Embedding: es, 38 | Index: i, 39 | } 40 | } 41 | 42 | return embeddingResponse, nil 43 | } 44 | -------------------------------------------------------------------------------- /provider/ollama/types.go: -------------------------------------------------------------------------------- 1 | package ollama 2 | 3 | import ( 4 | "encoding/json" 5 | "time" 6 | ) 7 | 8 | type ToolCall struct { 9 | Function ToolCallFunction `json:"function"` 10 | } 11 | 12 | type ToolCallFunction struct { 13 | Index int `json:"index,omitempty"` 14 | Name string `json:"name"` 15 | Arguments ToolCallFunctionArguments `json:"arguments"` 16 | } 17 | 18 | type ToolCallFunctionArguments map[string]any 19 | 20 | func (t *ToolCallFunctionArguments) String() string { 21 | bts, _ := json.Marshal(t) 22 | return string(bts) 23 | } 24 | 25 | type Tools []Tool 26 | 27 | func (t Tools) String() string { 28 | bts, _ := json.Marshal(t) 29 | return string(bts) 30 | } 31 | 32 | func (t Tool) String() string { 33 | bts, _ := json.Marshal(t) 34 | return string(bts) 35 | } 36 | 37 | type Tool struct { 38 | Type string `json:"type"` 39 | Function ToolFunction `json:"function"` 40 | } 41 | 42 | type ToolFunction struct { 43 | Name string `json:"name"` 44 | Description string `json:"description"` 45 | Parameters struct { 46 | Type string `json:"type"` 47 | Required []string `json:"required"` 48 | Properties map[string]struct { 49 | Type string `json:"type"` 50 | Description string `json:"description"` 51 | Enum []string `json:"enum,omitempty"` 52 | } `json:"properties"` 53 | } `json:"parameters"` 54 | } 55 | 56 | func (t *ToolFunction) String() string { 57 | bts, _ := json.Marshal(t) 58 | return string(bts) 59 | } 60 | 61 | type Duration struct { 62 | time.Duration 63 | } 64 | 65 | func (d Duration) MarshalJSON() ([]byte, error) { 66 | if d.Duration < 0 { 67 | return []byte("-1"), nil 68 | } 69 | return []byte("\"" + d.Duration.String() + "\""), nil 70 | } 71 | -------------------------------------------------------------------------------- /provider/openai/chat_model.go: -------------------------------------------------------------------------------- 1 | package openai 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "io" 7 | 8 | "github.com/sashabaranov/go-openai" 9 | "github.com/tech1024/goai/prompt" 10 | ) 11 | 12 | func NewChatModel(client *openai.Client, model string) *ChatModel { 13 | return &ChatModel{ 14 | client: client, 15 | model: model, 16 | } 17 | } 18 | 19 | type ChatModel struct { 20 | client *openai.Client 21 | model string 22 | } 23 | 24 | func (chatModel *ChatModel) Call(ctx context.Context, prompt prompt.Prompt) (string, error) { 25 | req, err := chatModel.buildChatRequest(prompt) 26 | if err != nil { 27 | return "", err 28 | } 29 | 30 | resp, err := chatModel.client.CreateChatCompletion(ctx, req) 31 | 32 | if err != nil { 33 | return "", err 34 | } 35 | 36 | return resp.Choices[0].Message.Content, nil 37 | } 38 | 39 | func (chatModel *ChatModel) Stream(ctx context.Context, prompt prompt.Prompt, fn func([]byte) error) error { 40 | req, err := chatModel.buildChatRequest(prompt) 41 | if err != nil { 42 | return err 43 | } 44 | 45 | stream, err := chatModel.client.CreateChatCompletionStream(ctx, req) 46 | if err != nil { 47 | return err 48 | } 49 | 50 | defer stream.Close() 51 | 52 | var resp openai.ChatCompletionStreamResponse 53 | for { 54 | resp, err = stream.Recv() 55 | if errors.Is(err, io.EOF) { 56 | break 57 | } 58 | if err != nil { 59 | return err 60 | } 61 | 62 | err = fn([]byte(resp.Choices[0].Delta.Content)) 63 | if err != nil { 64 | return err 65 | } 66 | } 67 | 68 | return nil 69 | } 70 | 71 | func (chatModel *ChatModel) buildChatRequest(prompt prompt.Prompt) (openai.ChatCompletionRequest, error) { 72 | request := openai.ChatCompletionRequest{ 73 | Model: chatModel.model, 74 | Messages: make([]openai.ChatCompletionMessage, len(prompt.Messages)), 75 | } 76 | for i, message := range prompt.Messages { 77 | request.Messages[i] = openai.ChatCompletionMessage{ 78 | Role: message.Type().String(), 79 | Content: message.Text(), 80 | } 81 | } 82 | if prompt.ChatOption.Model != "" { 83 | request.Model = prompt.ChatOption.Model 84 | } 85 | 86 | return request, nil 87 | } 88 | --------------------------------------------------------------------------------