feat: add Ollama provider with local and cloud support (#3636)
Add a dedicated Ollama AI provider (ai/ollama/) that auto-detects local vs cloud mode based on the base URL: - Local Ollama: native /api/chat endpoint with NDJSON streaming - Ollama Cloud: OpenAI-compatible /v1/chat/completions with SSE streaming Both modes support tool calls with a multi-round execution loop. Add agent.BaseURL option so agents can point at non-default LLM endpoints (e.g. local Ollama, proxies). Wire it through micro.AgentBaseURL at the top level. Include a complete example (examples/agent-ollama/) demonstrating a knowledge-base service with auto-discovered tools, a custom time tool, streaming, and env-var configuration for local vs cloud. Closes #3632
This commit is contained in:
YongSoo Park
@@ -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/
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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.")
|
||||
}
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user