diff --git a/.gitignore b/.gitignore index 872f28431..abb71641e 100644 --- a/.gitignore +++ b/.gitignore @@ -62,6 +62,7 @@ examples/mcp/hello/hello /plan-delegate /agent-plan-delegate /micro-mcp-gateway +/agent-ollama # Local Jekyll / Bundler artifacts internal/website/.bundle/ diff --git a/agent/agent.go b/agent/agent.go index edf6bb276..a5a7556d3 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -34,6 +34,7 @@ import ( _ "go-micro.dev/v6/ai/gemini" _ "go-micro.dev/v6/ai/groq" _ "go-micro.dev/v6/ai/mistral" + _ "go-micro.dev/v6/ai/ollama" _ "go-micro.dev/v6/ai/openai" _ "go-micro.dev/v6/ai/together" ) @@ -149,6 +150,9 @@ func (a *agentImpl) setupWithToolHandler(handler ai.ToolHandler) { if a.opts.Model != "" { modelOpts = append(modelOpts, ai.WithModel(a.opts.Model)) } + if a.opts.BaseURL != "" { + modelOpts = append(modelOpts, ai.WithBaseURL(a.opts.BaseURL)) + } // Reuse the existing tools instance: its name map is populated by // discoverTools, and rebuilding it here would orphan a base handler that diff --git a/agent/options.go b/agent/options.go index cff7265a1..f31c0fd0f 100644 --- a/agent/options.go +++ b/agent/options.go @@ -40,6 +40,7 @@ type Options struct { Provider string Model string APIKey string + BaseURL string Address string Registry registry.Registry Client client.Client @@ -168,6 +169,12 @@ func APIKey(k string) Option { return func(o *Options) { o.APIKey = k } } +// BaseURL sets the base URL for the LLM provider. Use this to point +// the provider at a non-default endpoint (e.g., local Ollama, a proxy). +func BaseURL(url string) Option { + return func(o *Options) { o.BaseURL = url } +} + // Address sets the network address for the agent's service endpoint. // Use "127.0.0.1:0" in local harnesses/tests to bind an ephemeral loopback // port and avoid advertising the default service address. diff --git a/ai/ollama/ollama.go b/ai/ollama/ollama.go new file mode 100644 index 000000000..146b05eff --- /dev/null +++ b/ai/ollama/ollama.go @@ -0,0 +1,729 @@ +// Package ollama implements the Ollama model provider. +// +// Ollama runs open-weight models locally (or via Ollama Cloud). This +// provider supports two API styles: +// +// - Native (/api/chat): local Ollama servers (default, http://localhost:11434) +// - OpenAI-compatible (/v1/chat/completions): Ollama Cloud (https://ollama.com/v1) +// +// The provider auto-detects which style to use based on the base URL. +// Set OLLAMA_BASE_URL to point at your server (local or cloud). +// +// Usage (local): +// +// import _ "go-micro.dev/v6/ai/ollama" +// +// m := ai.New("ollama", +// ai.WithBaseURL("http://localhost:11434"), +// ai.WithModel("llama3.2"), +// ) +// +// Usage (Ollama Cloud): +// +// m := ai.New("ollama", +// ai.WithBaseURL("https://ollama.com/v1"), +// ai.WithAPIKey("your-key"), +// ai.WithModel("gemma4:31b-cloud"), +// ) +package ollama + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + + "go-micro.dev/v6/ai" +) + +func init() { + ai.Register("ollama", func(opts ...ai.Option) ai.Model { + return NewProvider(opts...) + }) + ai.RegisterStream("ollama") +} + +// Provider implements the ai.Model interface for Ollama. +type Provider struct { + opts ai.Options + + // cloudOverride forces cloud mode for testing. When true, the provider + // uses the OpenAI-compatible endpoint regardless of the base URL. + cloudOverride bool +} + +// NewProvider creates a new Ollama provider. +func NewProvider(opts ...ai.Option) *Provider { + options := ai.NewOptions(opts...) + if options.Model == "" { + options.Model = "llama3.2" + } + if options.BaseURL == "" { + options.BaseURL = "http://localhost:11434" + } + return &Provider{opts: options} +} + +// Init initializes the provider with options. +func (p *Provider) Init(opts ...ai.Option) error { + for _, o := range opts { + o(&p.opts) + } + return nil +} + +// Options returns the provider options. +func (p *Provider) Options() ai.Options { return p.opts } + +// String returns the provider name. +func (p *Provider) String() string { return "ollama" } + +// isCloud returns true when the base URL points at Ollama Cloud (ollama.com), +// which uses the OpenAI-compatible /v1/chat/completions endpoint instead of +// the native /api/chat. +func (p *Provider) isCloud() bool { + if p.cloudOverride { + return true + } + return strings.Contains(p.opts.BaseURL, "ollama.com") +} + +// chatPath returns the API endpoint path for chat completions. +func (p *Provider) chatPath() string { + if p.isCloud() { + return "/v1/chat/completions" + } + return "/api/chat" +} + +// streamPath returns the API endpoint path for streaming chat. +// Ollama Cloud uses the same /v1/chat/completions with stream:true. +// Local Ollama uses /api/chat with stream:true. +func (p *Provider) streamPath() string { + return p.chatPath() +} + +// Generate generates a response from the Ollama model. +func (p *Provider) Generate(ctx context.Context, req *ai.Request, opts ...ai.GenerateOption) (*ai.Response, error) { + if p.isCloud() { + return p.generateOpenAI(ctx, req) + } + return p.generateNative(ctx, req) +} + +// Stream generates a streaming response. +func (p *Provider) Stream(ctx context.Context, req *ai.Request, opts ...ai.GenerateOption) (ai.Stream, error) { + if p.isCloud() { + return p.streamOpenAI(ctx, req) + } + return p.streamNative(ctx, req) +} + +// --------------------------------------------------------------------------- +// OpenAI-compatible mode (Ollama Cloud: ollama.com/v1) +// --------------------------------------------------------------------------- + +func (p *Provider) generateOpenAI(ctx context.Context, req *ai.Request) (*ai.Response, error) { + var tools []map[string]any + for _, t := range req.Tools { + tools = append(tools, map[string]any{ + "type": "function", + "function": map[string]any{ + "name": t.Name, + "description": t.Description, + "parameters": map[string]any{ + "type": "object", + "properties": t.Properties, + }, + }, + }) + } + + messages := buildOpenAIMessages(req) + apiReq := map[string]any{ + "model": p.opts.Model, + "messages": messages, + "stream": false, + } + if len(tools) > 0 { + apiReq["tools"] = tools + } + if p.opts.MaxTokens > 0 { + apiReq["max_tokens"] = p.opts.MaxTokens + } + + resp, rawMsg, err := p.callOpenAI(ctx, apiReq) + if err != nil { + return nil, err + } + + // No tool calls or no handler — return as-is. + if len(resp.ToolCalls) == 0 || p.opts.ToolHandler == nil { + return resp, nil + } + + // Tool execution loop. + convMessages := append(messages, map[string]any{ + "role": "assistant", + "content": rawMsg.content, + "tool_calls": rawMsg.toolCalls, + }) + + pendingCalls := resp.ToolCalls + for round := 0; round < 10; round++ { + for i := range pendingCalls { + result := p.opts.ToolHandler(ctx, pendingCalls[i]) + pendingCalls[i].Result = result.Content + convMessages = append(convMessages, map[string]any{ + "role": "tool", + "tool_call_id": pendingCalls[i].ID, + "content": result.Content, + }) + } + + followUpReq := map[string]any{ + "model": p.opts.Model, + "messages": convMessages, + "stream": false, + } + if len(tools) > 0 { + followUpReq["tools"] = tools + } + if p.opts.MaxTokens > 0 { + followUpReq["max_tokens"] = p.opts.MaxTokens + } + + followUpResp, followUpRaw, err := p.callOpenAI(ctx, followUpReq) + if err != nil { + break + } + + if len(followUpResp.ToolCalls) > 0 { + resp.ToolCalls = append(resp.ToolCalls, followUpResp.ToolCalls...) + pendingCalls = followUpResp.ToolCalls + convMessages = append(convMessages, map[string]any{ + "role": "assistant", + "content": followUpRaw.content, + "tool_calls": followUpRaw.toolCalls, + }) + continue + } + + if followUpResp.Reply != "" { + resp.Answer = followUpResp.Reply + } + break + } + + return resp, nil +} + +func (p *Provider) callOpenAI(ctx context.Context, req map[string]any) (*ai.Response, *rawChatMessage, error) { + reqBody, err := json.Marshal(req) + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal request: %w", err) + } + + apiURL := strings.TrimRight(p.opts.BaseURL, "/") + p.chatPath() + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(reqBody)) + if err != nil { + return nil, nil, fmt.Errorf("failed to create request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + if p.opts.APIKey != "" { + httpReq.Header.Set("Authorization", "Bearer "+p.opts.APIKey) + } + + httpResp, err := http.DefaultClient.Do(httpReq) + if err != nil { + return nil, nil, fmt.Errorf("API request failed: %w", err) + } + defer httpResp.Body.Close() + + respBody, _ := io.ReadAll(httpResp.Body) + if httpResp.StatusCode != http.StatusOK { + return nil, nil, fmt.Errorf("API error (%s): %s", httpResp.Status, string(respBody)) + } + + var chatResp struct { + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + Choices []struct { + Message struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []struct { + ID string `json:"id"` + Function struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"message"` + } `json:"choices"` + } + + if err := json.Unmarshal(respBody, &chatResp); err != nil { + return nil, nil, fmt.Errorf("failed to parse response: %w", err) + } + if len(chatResp.Choices) == 0 { + return nil, nil, fmt.Errorf("no response from API") + } + + choice := chatResp.Choices[0] + response := &ai.Response{ + Reply: choice.Message.Content, + Usage: ai.Usage{ + InputTokens: chatResp.Usage.PromptTokens, + OutputTokens: chatResp.Usage.CompletionTokens, + TotalTokens: chatResp.Usage.TotalTokens, + }, + } + + var rawToolCalls []map[string]any + for _, tc := range choice.Message.ToolCalls { + var input map[string]any + if err := json.Unmarshal([]byte(tc.Function.Arguments), &input); err != nil { + input = map[string]any{} + } + response.ToolCalls = append(response.ToolCalls, ai.ToolCall{ + ID: tc.ID, + Name: tc.Function.Name, + Input: input, + }) + rawToolCalls = append(rawToolCalls, map[string]any{ + "id": tc.ID, + "type": "function", + "function": map[string]any{ + "name": tc.Function.Name, + "arguments": tc.Function.Arguments, + }, + }) + } + + raw := &rawChatMessage{ + content: choice.Message.Content, + toolCalls: rawToolCalls, + } + return response, raw, nil +} + +func (p *Provider) streamOpenAI(ctx context.Context, req *ai.Request) (ai.Stream, error) { + messages := buildOpenAIMessages(req) + apiReq := map[string]any{ + "model": p.opts.Model, + "messages": messages, + "stream": true, + "stream_options": map[string]any{"include_usage": true}, + } + if p.opts.MaxTokens > 0 { + apiReq["max_tokens"] = p.opts.MaxTokens + } + + reqBody, err := json.Marshal(apiReq) + if err != nil { + return nil, fmt.Errorf("failed to marshal stream request: %w", err) + } + + apiURL := strings.TrimRight(p.opts.BaseURL, "/") + p.streamPath() + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(reqBody)) + if err != nil { + return nil, fmt.Errorf("failed to create stream request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "text/event-stream") + if p.opts.APIKey != "" { + httpReq.Header.Set("Authorization", "Bearer "+p.opts.APIKey) + } + + httpResp, err := http.DefaultClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("stream API request failed: %w", err) + } + if httpResp.StatusCode != http.StatusOK { + defer httpResp.Body.Close() + respBody, _ := io.ReadAll(httpResp.Body) + return nil, fmt.Errorf("stream API error (%s): %s", httpResp.Status, string(respBody)) + } + + return &sseStream{body: httpResp.Body, scanner: bufio.NewScanner(httpResp.Body)}, nil +} + +// buildOpenAIMessages converts an ai.Request into the OpenAI chat message format. +func buildOpenAIMessages(req *ai.Request) []map[string]any { + messages := []map[string]any{} + if req.SystemPrompt != "" { + messages = append(messages, map[string]any{"role": "system", "content": req.SystemPrompt}) + } + for _, m := range req.Messages { + messages = append(messages, map[string]any{"role": m.Role, "content": m.Content}) + } + if req.Prompt != "" { + messages = append(messages, map[string]any{"role": "user", "content": req.Prompt}) + } + return messages +} + +// sseStream reads OpenAI-style server-sent events (used by Ollama Cloud). +type sseStream struct { + body io.ReadCloser + scanner *bufio.Scanner + closed bool +} + +func (s *sseStream) Recv() (*ai.Response, error) { + for s.scanner.Scan() { + line := strings.TrimSpace(s.scanner.Text()) + if line == "" || strings.HasPrefix(line, ":") { + continue + } + if !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "[DONE]" { + return nil, io.EOF + } + var chunk struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + } `json:"choices"` + Usage *struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + } + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + return nil, fmt.Errorf("failed to parse stream chunk: %w", err) + } + if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" { + return &ai.Response{Reply: chunk.Choices[0].Delta.Content}, nil + } + if chunk.Usage != nil { + return &ai.Response{Usage: ai.Usage{ + InputTokens: chunk.Usage.PromptTokens, + OutputTokens: chunk.Usage.CompletionTokens, + TotalTokens: chunk.Usage.TotalTokens, + }}, nil + } + } + if err := s.scanner.Err(); err != nil { + return nil, err + } + return nil, io.EOF +} + +func (s *sseStream) Close() error { + if s.closed { + return nil + } + s.closed = true + return s.body.Close() +} + +// --------------------------------------------------------------------------- +// Native mode (local Ollama: localhost:11434/api/chat) +// --------------------------------------------------------------------------- + +func (p *Provider) generateNative(ctx context.Context, req *ai.Request) (*ai.Response, error) { + var tools []map[string]any + for _, t := range req.Tools { + tools = append(tools, map[string]any{ + "type": "function", + "function": map[string]any{ + "name": t.Name, + "description": t.Description, + "parameters": map[string]any{ + "type": "object", + "properties": t.Properties, + }, + }, + }) + } + + messages := []map[string]any{} + if req.SystemPrompt != "" { + messages = append(messages, map[string]any{"role": "system", "content": req.SystemPrompt}) + } + for _, m := range req.Messages { + messages = append(messages, map[string]any{"role": m.Role, "content": m.Content}) + } + if req.Prompt != "" { + messages = append(messages, map[string]any{"role": "user", "content": req.Prompt}) + } + + apiReq := map[string]any{ + "model": p.opts.Model, + "messages": messages, + "stream": false, + } + if len(tools) > 0 { + apiReq["tools"] = tools + } + if p.opts.MaxTokens > 0 { + apiReq["options"] = map[string]any{"num_predict": p.opts.MaxTokens} + } + + resp, rawMsg, err := p.callNative(ctx, apiReq) + if err != nil { + return nil, err + } + + if len(resp.ToolCalls) == 0 || p.opts.ToolHandler == nil { + return resp, nil + } + + convMessages := append(messages, map[string]any{ + "role": "assistant", + "content": rawMsg.content, + }) + if len(rawMsg.toolCalls) > 0 { + convMessages[len(convMessages)-1]["tool_calls"] = rawMsg.toolCalls + } + + pendingCalls := resp.ToolCalls + for round := 0; round < 10; round++ { + for i := range pendingCalls { + result := p.opts.ToolHandler(ctx, pendingCalls[i]) + pendingCalls[i].Result = result.Content + convMessages = append(convMessages, map[string]any{ + "role": "tool", + "content": result.Content, + }) + } + + followUpReq := map[string]any{ + "model": p.opts.Model, + "messages": convMessages, + "stream": false, + } + if len(tools) > 0 { + followUpReq["tools"] = tools + } + if p.opts.MaxTokens > 0 { + followUpReq["options"] = map[string]any{"num_predict": p.opts.MaxTokens} + } + + followUpResp, followUpRaw, err := p.callNative(ctx, followUpReq) + if err != nil { + break + } + + if len(followUpResp.ToolCalls) > 0 { + resp.ToolCalls = append(resp.ToolCalls, followUpResp.ToolCalls...) + pendingCalls = followUpResp.ToolCalls + convMessages = append(convMessages, map[string]any{ + "role": "assistant", + "content": followUpRaw.content, + }) + if len(followUpRaw.toolCalls) > 0 { + convMessages[len(convMessages)-1]["tool_calls"] = followUpRaw.toolCalls + } + continue + } + + if followUpResp.Reply != "" { + resp.Answer = followUpResp.Reply + } + break + } + + return resp, nil +} + +func (p *Provider) callNative(ctx context.Context, req map[string]any) (*ai.Response, *rawChatMessage, error) { + reqBody, err := json.Marshal(req) + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal request: %w", err) + } + + apiURL := strings.TrimRight(p.opts.BaseURL, "/") + p.chatPath() + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(reqBody)) + if err != nil { + return nil, nil, fmt.Errorf("failed to create request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + if p.opts.APIKey != "" { + httpReq.Header.Set("Authorization", "Bearer "+p.opts.APIKey) + } + + httpResp, err := http.DefaultClient.Do(httpReq) + if err != nil { + return nil, nil, fmt.Errorf("API request failed: %w", err) + } + defer httpResp.Body.Close() + + respBody, _ := io.ReadAll(httpResp.Body) + if httpResp.StatusCode != http.StatusOK { + return nil, nil, fmt.Errorf("API error (%s): %s", httpResp.Status, string(respBody)) + } + + var chatResp struct { + Message struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []struct { + Function struct { + Name string `json:"name"` + Arguments any `json:"arguments"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"message"` + Done bool `json:"done"` + PromptEvalCount int `json:"prompt_eval_count"` + EvalCount int `json:"eval_count"` + } + + if err := json.Unmarshal(respBody, &chatResp); err != nil { + return nil, nil, fmt.Errorf("failed to parse response: %w", err) + } + + response := &ai.Response{ + Reply: chatResp.Message.Content, + Usage: ai.Usage{ + InputTokens: chatResp.PromptEvalCount, + OutputTokens: chatResp.EvalCount, + TotalTokens: chatResp.PromptEvalCount + chatResp.EvalCount, + }, + } + + var rawToolCalls []map[string]any + for _, tc := range chatResp.Message.ToolCalls { + var input map[string]any + switch v := tc.Function.Arguments.(type) { + case string: + if err := json.Unmarshal([]byte(v), &input); err != nil { + input = map[string]any{} + } + case map[string]any: + input = v + default: + input = map[string]any{} + } + response.ToolCalls = append(response.ToolCalls, ai.ToolCall{ + Name: tc.Function.Name, + Input: input, + }) + rawToolCalls = append(rawToolCalls, map[string]any{ + "function": map[string]any{ + "name": tc.Function.Name, + "arguments": tc.Function.Arguments, + }, + }) + } + + raw := &rawChatMessage{ + content: chatResp.Message.Content, + toolCalls: rawToolCalls, + } + return response, raw, nil +} + +func (p *Provider) streamNative(ctx context.Context, req *ai.Request) (ai.Stream, error) { + messages := []map[string]any{} + if req.SystemPrompt != "" { + messages = append(messages, map[string]any{"role": "system", "content": req.SystemPrompt}) + } + for _, m := range req.Messages { + messages = append(messages, map[string]any{"role": m.Role, "content": m.Content}) + } + if req.Prompt != "" { + messages = append(messages, map[string]any{"role": "user", "content": req.Prompt}) + } + + apiReq := map[string]any{ + "model": p.opts.Model, + "messages": messages, + "stream": true, + } + if p.opts.MaxTokens > 0 { + apiReq["options"] = map[string]any{"num_predict": p.opts.MaxTokens} + } + + reqBody, err := json.Marshal(apiReq) + if err != nil { + return nil, fmt.Errorf("failed to marshal stream request: %w", err) + } + + apiURL := strings.TrimRight(p.opts.BaseURL, "/") + p.streamPath() + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(reqBody)) + if err != nil { + return nil, fmt.Errorf("failed to create stream request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + if p.opts.APIKey != "" { + httpReq.Header.Set("Authorization", "Bearer "+p.opts.APIKey) + } + + httpResp, err := http.DefaultClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("stream API request failed: %w", err) + } + if httpResp.StatusCode != http.StatusOK { + defer httpResp.Body.Close() + respBody, _ := io.ReadAll(httpResp.Body) + return nil, fmt.Errorf("stream API error (%s): %s", httpResp.Status, string(respBody)) + } + + return &ndjsonStream{body: httpResp.Body, scanner: bufio.NewScanner(httpResp.Body)}, nil +} + +// ndjsonStream reads newline-delimited JSON (used by local Ollama). +type ndjsonStream struct { + body io.ReadCloser + scanner *bufio.Scanner + closed bool +} + +func (s *ndjsonStream) Recv() (*ai.Response, error) { + for s.scanner.Scan() { + line := strings.TrimSpace(s.scanner.Text()) + if line == "" { + continue + } + var chunk struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + Done bool `json:"done"` + } + if err := json.Unmarshal([]byte(line), &chunk); err != nil { + return nil, fmt.Errorf("failed to parse stream chunk: %w", err) + } + if chunk.Done { + return nil, io.EOF + } + if chunk.Message.Content != "" { + return &ai.Response{Reply: chunk.Message.Content}, nil + } + } + if err := s.scanner.Err(); err != nil { + return nil, err + } + return nil, io.EOF +} + +func (s *ndjsonStream) Close() error { + if s.closed { + return nil + } + s.closed = true + return s.body.Close() +} + +// rawChatMessage holds the raw assistant content and tool calls for +// follow-up messages. +type rawChatMessage struct { + content string + toolCalls []map[string]any +} diff --git a/ai/ollama/ollama_test.go b/ai/ollama/ollama_test.go new file mode 100644 index 000000000..645f13e06 --- /dev/null +++ b/ai/ollama/ollama_test.go @@ -0,0 +1,333 @@ +package ollama + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "go-micro.dev/v6/ai" +) + +// --------------------------------------------------------------------------- +// Provider basics +// --------------------------------------------------------------------------- + +func TestProvider_String(t *testing.T) { + p := NewProvider() + if p.String() != "ollama" { + t.Errorf("Expected 'ollama', got '%s'", p.String()) + } +} + +func TestProvider_Init(t *testing.T) { + p := NewProvider() + err := p.Init( + ai.WithModel("test-model"), + ai.WithAPIKey("test-key"), + ai.WithBaseURL("https://test.com"), + ) + if err != nil { + t.Fatalf("Init failed: %v", err) + } + opts := p.Options() + if opts.Model != "test-model" { + t.Errorf("Expected model 'test-model', got '%s'", opts.Model) + } + if opts.APIKey != "test-key" { + t.Errorf("Expected API key 'test-key', got '%s'", opts.APIKey) + } + if opts.BaseURL != "https://test.com" { + t.Errorf("Expected base URL 'https://test.com', got '%s'", opts.BaseURL) + } +} + +func TestProvider_Defaults(t *testing.T) { + p := NewProvider() + opts := p.Options() + if opts.Model != "llama3.2" { + t.Errorf("Expected default model 'llama3.2', got '%s'", opts.Model) + } + if opts.BaseURL != "http://localhost:11434" { + t.Errorf("Expected default base URL 'http://localhost:11434', got '%s'", opts.BaseURL) + } +} + +func TestProvider_IsCloud(t *testing.T) { + local := NewProvider(ai.WithBaseURL("http://localhost:11434")) + if local.isCloud() { + t.Error("localhost should not be cloud") + } + cloud := NewProvider(ai.WithBaseURL("https://ollama.com/v1")) + if !cloud.isCloud() { + t.Error("ollama.com should be cloud") + } +} + +// --------------------------------------------------------------------------- +// Native mode (local Ollama: /api/chat) +// --------------------------------------------------------------------------- + +func TestNative_Generate(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/chat" { + t.Errorf("Expected /api/chat, got %s", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "model": "llama3.2", + "message": {"role": "assistant", "content": "Hello from local Ollama!"}, + "done": true, + "prompt_eval_count": 10, + "eval_count": 5 + }`)) + })) + defer srv.Close() + + p := NewProvider(ai.WithBaseURL(srv.URL), ai.WithModel("llama3.2")) + resp, err := p.Generate(context.Background(), &ai.Request{ + Prompt: "Hi", + SystemPrompt: "You are helpful", + }) + if err != nil { + t.Fatalf("Generate failed: %v", err) + } + if resp.Reply != "Hello from local Ollama!" { + t.Errorf("Expected 'Hello from local Ollama!', got '%s'", resp.Reply) + } + if resp.Usage.TotalTokens != 15 { + t.Errorf("Expected total tokens 15, got %d", resp.Usage.TotalTokens) + } +} + +func TestNative_GenerateWithToolCall(t *testing.T) { + callCount := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + callCount++ + w.Header().Set("Content-Type", "application/json") + if callCount == 1 { + w.Write([]byte(`{ + "model": "llama3.2", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "get_weather", "arguments": "{\"city\":\"Seoul\"}"}}] + }, + "done": true + }`)) + } else { + w.Write([]byte(`{ + "model": "llama3.2", + "message": {"role": "assistant", "content": "The weather in Seoul is sunny."}, + "done": true + }`)) + } + })) + defer srv.Close() + + handler := func(ctx context.Context, call ai.ToolCall) ai.ToolResult { + if call.Name != "get_weather" { + t.Errorf("Expected tool 'get_weather', got '%s'", call.Name) + } + return ai.ToolResult{ID: call.ID, Content: `{"temp": 22, "condition": "sunny"}`} + } + + p := NewProvider( + ai.WithBaseURL(srv.URL), + ai.WithModel("llama3.2"), + ai.WithToolHandler(handler), + ) + resp, err := p.Generate(context.Background(), &ai.Request{ + Prompt: "What's the weather?", + Tools: []ai.Tool{{ + Name: "get_weather", + Description: "Get weather", + Properties: map[string]any{"city": map[string]any{"type": "string"}}, + }}, + }) + if err != nil { + t.Fatalf("Generate failed: %v", err) + } + if len(resp.ToolCalls) == 0 { + t.Error("Expected tool calls") + } + if resp.Answer != "The weather in Seoul is sunny." { + t.Errorf("Expected final answer, got '%s'", resp.Answer) + } +} + +func TestNative_Stream(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"message":{"role":"assistant","content":"Hello"},"done":false}` + "\n")) + w.Write([]byte(`{"message":{"role":"assistant","content":" world"},"done":false}` + "\n")) + w.Write([]byte(`{"message":{"role":"assistant","content":""},"done":true}` + "\n")) + })) + defer srv.Close() + + p := NewProvider(ai.WithBaseURL(srv.URL), ai.WithModel("llama3.2")) + stream, err := p.Stream(context.Background(), &ai.Request{Prompt: "Hi"}) + if err != nil { + t.Fatalf("Stream failed: %v", err) + } + defer stream.Close() + + var chunks []string + for { + resp, err := stream.Recv() + if err != nil { + break + } + if resp.Reply != "" { + chunks = append(chunks, resp.Reply) + } + } + result := strings.Join(chunks, "") + if result != "Hello world" { + t.Errorf("Expected 'Hello world', got '%s'", result) + } +} + +// --------------------------------------------------------------------------- +// Cloud mode (Ollama Cloud: /v1/chat/completions) +// --------------------------------------------------------------------------- + +func TestCloud_Generate(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/chat/completions" { + t.Errorf("Expected /v1/chat/completions, got %s", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + "choices": [{"message": {"role": "assistant", "content": "Hello from Ollama Cloud!"}}] + }`)) + })) + defer srv.Close() + + p := NewProvider(ai.WithBaseURL(srv.URL), ai.WithModel("gemma4:31b-cloud"), ai.WithAPIKey("test-key")) + p.cloudOverride = true + resp, err := p.Generate(context.Background(), &ai.Request{ + Prompt: "Hi", + SystemPrompt: "You are helpful", + }) + if err != nil { + t.Fatalf("Generate failed: %v", err) + } + if resp.Reply != "Hello from Ollama Cloud!" { + t.Errorf("Expected 'Hello from Ollama Cloud!', got '%s'", resp.Reply) + } + if resp.Usage.TotalTokens != 15 { + t.Errorf("Expected total tokens 15, got %d", resp.Usage.TotalTokens) + } +} + +func TestCloud_GenerateWithToolCall(t *testing.T) { + callCount := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + callCount++ + w.Header().Set("Content-Type", "application/json") + if callCount == 1 { + w.Write([]byte(`{ + "choices": [{"message": { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "call_1", "function": {"name": "search", "arguments": "{\"query\":\"go interfaces\"}"}}] + }}] + }`)) + } else { + w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "Go interfaces are implicit."}}] + }`)) + } + })) + defer srv.Close() + + handler := func(ctx context.Context, call ai.ToolCall) ai.ToolResult { + return ai.ToolResult{ID: call.ID, Content: `{"results": ["Go interfaces are implicit"]}`} + } + + p := NewProvider( + ai.WithBaseURL(srv.URL), + ai.WithModel("gemma4:31b-cloud"), + ai.WithAPIKey("test-key"), + ai.WithToolHandler(handler), + ) + p.cloudOverride = true + resp, err := p.Generate(context.Background(), &ai.Request{ + Prompt: "Search for Go interfaces", + Tools: []ai.Tool{{ + Name: "search", + Description: "Search the knowledge base", + Properties: map[string]any{"query": map[string]any{"type": "string"}}, + }}, + }) + if err != nil { + t.Fatalf("Generate failed: %v", err) + } + if len(resp.ToolCalls) == 0 { + t.Error("Expected tool calls") + } + if resp.Answer != "Go interfaces are implicit." { + t.Errorf("Expected final answer, got '%s'", resp.Answer) + } +} + +func TestCloud_Stream(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" cloud\"}}]}\n\n")) + w.Write([]byte("data: [DONE]\n\n")) + })) + defer srv.Close() + + p := NewProvider( + ai.WithBaseURL(srv.URL), + ai.WithModel("gemma4:31b-cloud"), + ai.WithAPIKey("test-key"), + ) + p.cloudOverride = true + stream, err := p.Stream(context.Background(), &ai.Request{Prompt: "Hi"}) + if err != nil { + t.Fatalf("Stream failed: %v", err) + } + defer stream.Close() + + var chunks []string + for { + resp, err := stream.Recv() + if err != nil { + break + } + if resp.Reply != "" { + chunks = append(chunks, resp.Reply) + } + } + result := strings.Join(chunks, "") + if result != "Hello cloud" { + t.Errorf("Expected 'Hello cloud', got '%s'", result) + } +} + +// --------------------------------------------------------------------------- +// Error handling +// --------------------------------------------------------------------------- + +func TestProvider_APIError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error": "model not found"}`)) + })) + defer srv.Close() + + p := NewProvider(ai.WithBaseURL(srv.URL), ai.WithModel("nonexistent")) + _, err := p.Generate(context.Background(), &ai.Request{Prompt: "Hi"}) + if err == nil { + t.Error("Expected error on API failure") + } + if !strings.Contains(err.Error(), "API error") { + t.Errorf("Expected 'API error' in message, got '%s'", err.Error()) + } +} diff --git a/examples/agent-ollama/main.go b/examples/agent-ollama/main.go new file mode 100644 index 000000000..17ecc9004 --- /dev/null +++ b/examples/agent-ollama/main.go @@ -0,0 +1,246 @@ +// Agent Ollama — a self-contained agent powered by Ollama Cloud. +// +// This example demonstrates the full harness loop — service tools, custom +// tools, agent memory, guardrails, and streaming — using the Ollama +// provider with gemma4:31b-cloud on Ollama Cloud. +// +// It creates a "knowledge" service with two endpoints (Add, Search) that +// the agent discovers as tools, plus a custom "current_time" tool. The +// agent answers natural-language questions by calling those tools. +// +// Run (Ollama Cloud — default): +// +// OLLAMA_API_KEY=your-key go run main.go +// +// Run (local Ollama): +// +// OLLAMA_BASE_URL=http://localhost:11434 \ +// OLLAMA_MODEL=llama3.2 \ +// go run main.go +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "time" + + "go-micro.dev/v6" + "go-micro.dev/v6/agent" +) + +// --------------------------------------------------------------------------- +// knowledge service — a tiny in-memory knowledge base +// --------------------------------------------------------------------------- + +type KnowledgeEntry struct { + ID string `json:"id" description:"Unique entry identifier"` + Topic string `json:"topic" description:"Topic or category"` + Content string `json:"content" description:"The knowledge content"` +} + +type AddKnowledgeRequest struct { + Topic string `json:"topic" description:"Topic or category (required)"` + Content string `json:"content" description:"The knowledge content (required)"` +} + +type AddKnowledgeResponse struct { + Entry *KnowledgeEntry `json:"entry" description:"The added entry"` +} + +type SearchKnowledgeRequest struct { + Topic string `json:"topic,omitempty" description:"Filter by topic (optional)"` + Keyword string `json:"keyword,omitempty" description:"Search keyword in content (optional)"` +} + +type SearchKnowledgeResponse struct { + Entries []*KnowledgeEntry `json:"entries" description:"Matching entries"` +} + +type KnowledgeService struct { + mu sync.RWMutex + entries []*KnowledgeEntry + nextID int +} + +// Add stores a new knowledge entry. +// +// @example {"topic": "go", "content": "Go interfaces are implicit."} +func (s *KnowledgeService) Add(ctx context.Context, req *AddKnowledgeRequest, rsp *AddKnowledgeResponse) error { + s.mu.Lock() + defer s.mu.Unlock() + s.nextID++ + e := &KnowledgeEntry{ + ID: fmt.Sprintf("kb-%d", s.nextID), + Topic: req.Topic, + Content: req.Content, + } + s.entries = append(s.entries, e) + rsp.Entry = e + return nil +} + +// Search finds knowledge entries by topic or keyword. +// +// @example {"topic": "go"} +// @example {"keyword": "interface"} +func (s *KnowledgeService) Search(ctx context.Context, req *SearchKnowledgeRequest, rsp *SearchKnowledgeResponse) error { + s.mu.RLock() + defer s.mu.RUnlock() + for _, e := range s.entries { + if req.Topic != "" && !strings.EqualFold(e.Topic, req.Topic) { + continue + } + if req.Keyword != "" && !strings.Contains(strings.ToLower(e.Content), strings.ToLower(req.Keyword)) { + continue + } + rsp.Entries = append(rsp.Entries, e) + } + return nil +} + +// --------------------------------------------------------------------------- +// main +// --------------------------------------------------------------------------- + +func main() { + // Ollama Cloud is the default. Override with env vars for local Ollama. + baseURL := os.Getenv("OLLAMA_BASE_URL") + if baseURL == "" { + baseURL = "https://ollama.com/v1" + } + model := os.Getenv("OLLAMA_MODEL") + if model == "" { + model = "gemma4:31b-cloud" + } + apiKey := os.Getenv("OLLAMA_API_KEY") + + fmt.Println("╔══════════════════════════════════════════╗") + fmt.Println("║ Ollama-Powered Go Micro Agent ║") + fmt.Println("╚══════════════════════════════════════════╝") + fmt.Println() + fmt.Printf(" Ollama URL: %s\n", baseURL) + fmt.Printf(" Model: %s\n", model) + if apiKey != "" { + fmt.Printf(" API Key: (set)\n") + } else { + fmt.Printf(" API Key: (none — set OLLAMA_API_KEY)\n") + } + fmt.Println() + + // 1. Start the knowledge service. Its handlers become agent tools. + svc := micro.NewService("knowledge") + svc.Handle(new(KnowledgeService)) + go svc.Run() + + // Give the service a moment to register. + time.Sleep(2 * time.Second) + + // 2. Create the agent. It discovers the knowledge service endpoints + // as tools automatically, plus gets a custom "current_time" tool. + ag := micro.NewAgent("ollama-assistant", + micro.AgentServices("knowledge"), + micro.AgentPrompt( + "You are a helpful knowledge assistant. You can search and add to "+ + "a knowledge base using the knowledge service tools. "+ + "When asked about the current time, use the current_time tool. "+ + "Be concise and factual.", + ), + micro.AgentProvider("ollama"), + micro.AgentModel(model), + micro.AgentAPIKey(apiKey), + micro.AgentBaseURL(baseURL), + micro.AgentMaxSteps(10), + micro.AgentLoopLimit(3), + // Custom tool — any function, not tied to a service. + agent.WithTool( + "current_time", + "Get the current date and time in a human-readable format", + map[string]any{ + "timezone": map[string]any{ + "type": "string", + "description": "Optional timezone (defaults to local)", + }, + }, + func(ctx context.Context, input map[string]any) (string, error) { + tz, _ := input["timezone"].(string) + if tz == "" { + return time.Now().Format("2006-01-02 15:04:05 MST"), nil + } + loc, err := time.LoadLocation(tz) + if err != nil { + return "", fmt.Errorf("unknown timezone: %s", tz) + } + return time.Now().In(loc).Format("2006-01-02 15:04:05 MST"), nil + }, + ), + ) + + // 3. Seed initial knowledge via the agent's first question. + questions := []string{ + "What time is it now?", + "Add a new knowledge entry: topic 'go', content 'Go interfaces are implicit — a type implements an interface by having the required methods.'", + "Add another entry: topic 'go', content 'Go is a statically typed, compiled language designed at Google.'", + "Add another entry: topic 'ai', content 'Large language models generate text by predicting the next token in a sequence.'", + "Search the knowledge base for entries about Go.", + "Search for everything in the knowledge base.", + } + + fmt.Println("─── Agent Demo ───") + fmt.Println() + + for i, q := range questions { + fmt.Printf("Q%d: %s\n", i+1, q) + fmt.Print("A: ") + + resp, err := ag.Ask(context.Background(), q) + if err != nil { + fmt.Printf("error: %v\n", err) + fmt.Println() + continue + } + + // Show tool calls the agent made. + if len(resp.ToolCalls) > 0 { + for _, tc := range resp.ToolCalls { + args, _ := json.Marshal(tc.Input) + fmt.Printf(" [tool] %s(%s)\n", tc.Name, string(args)) + } + } + + fmt.Println(resp.Reply) + if resp.Reply == "" && len(resp.ToolCalls) == 0 { + fmt.Println("(no response)") + } + fmt.Println() + } + + // 4. Streaming demonstration. + fmt.Println("─── Streaming Demo ───") + fmt.Println() + streamQ := "Explain what Go Micro is in two sentences." + fmt.Printf("Q: %s\n", streamQ) + fmt.Print("A: ") + + stream, err := ag.Stream(context.Background(), streamQ) + if err != nil { + fmt.Printf("stream error: %v\n", err) + } else { + for { + chunk, err := stream.Recv() + if err != nil { + break + } + if chunk.Reply != "" { + fmt.Print(chunk.Reply) + } + } + fmt.Println() + } + + fmt.Println() + fmt.Println("Done.") +} diff --git a/micro.go b/micro.go index 12a32048d..3fda8df03 100644 --- a/micro.go +++ b/micro.go @@ -100,6 +100,10 @@ func AgentModel(m string) AgentOption { return agent.Model(m) } // AgentAPIKey sets the API key for the LLM provider. func AgentAPIKey(k string) AgentOption { return agent.APIKey(k) } +// AgentBaseURL sets the base URL for the LLM provider. Use this to point +// the provider at a non-default endpoint (e.g., local Ollama, a proxy). +func AgentBaseURL(url string) AgentOption { return agent.BaseURL(url) } + // ApproveFunc gates an agent's tool calls before they run. type ApproveFunc = agent.ApproveFunc