Unverified Commit 91935631 authored by Grace's avatar Grace Committed by GitHub
Browse files

Renderer for Cogito v2 (#13139)

parent 8de30b56
package renderers
import (
"encoding/json"
"strings"
"github.com/ollama/ollama/api"
)
type CogitoRenderer struct {
isThinking bool
}
func (r *CogitoRenderer) Render(messages []api.Message, tools []api.Tool, thinkValue *api.ThinkValue) (string, error) {
var sb strings.Builder
defaultPrompt := "You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco."
// thinking is enabled: model must support it AND user must request it (true)
enableThinking := r.isThinking && (thinkValue != nil && thinkValue.Bool())
var systemPrompt string
var conversationMessages []api.Message
if len(messages) > 0 && messages[0].Role == "system" {
systemPrompt = messages[0].Content
conversationMessages = messages[1:]
} else {
conversationMessages = messages
}
var finalSystemPrompt string
if enableThinking {
finalSystemPrompt = "Enable deep thinking subroutine.\n\n" + defaultPrompt
if systemPrompt != "" {
finalSystemPrompt += "\n\n" + systemPrompt + "\n\n"
}
} else {
finalSystemPrompt = defaultPrompt
if systemPrompt != "" {
finalSystemPrompt += "\n\n" + systemPrompt
}
}
if len(tools) > 0 {
if finalSystemPrompt != "" {
finalSystemPrompt += "\nYou have the following functions available:\n"
} else {
finalSystemPrompt = "You have the following functions available:\n"
}
for _, tool := range tools {
toolJSON, _ := json.MarshalIndent(tool, "", " ") // TODO(gguo): double check json format
finalSystemPrompt += "```json\n" + string(toolJSON) + "\n```\n"
}
}
sb.WriteString("<|begin▁of▁sentence|>" + finalSystemPrompt)
outputsOpen := false
isLastUser := false
for i, message := range conversationMessages {
switch message.Role {
case "user":
isLastUser = true
sb.WriteString("<|User|>" + message.Content + "<|Assistant|>")
case "assistant":
isLastUser = false
if len(message.ToolCalls) > 0 {
if message.Content != "" {
sb.WriteString(message.Content)
}
sb.WriteString("<|tool▁calls▁begin|>")
for j, toolCall := range message.ToolCalls {
sb.WriteString("<|tool▁call▁begin|>function<|tool▁sep|>" + toolCall.Function.Name)
argsJSON, _ := json.Marshal(toolCall.Function.Arguments)
sb.WriteString("\n```json\n" + string(argsJSON) + "\n```")
sb.WriteString("<|tool▁call▁end|>")
if j < len(message.ToolCalls)-1 {
sb.WriteString("\n")
}
}
sb.WriteString("<|tool▁calls▁end|><|end▁of▁sentence|>")
} else {
sb.WriteString(message.Content + "<|end▁of▁sentence|>")
}
case "tool":
isLastUser = false
if !outputsOpen {
sb.WriteString("<|tool▁outputs▁begin|>")
outputsOpen = true
}
sb.WriteString("<|tool▁output▁begin|>" + message.Content + "<|tool▁output▁end|>")
hasNextTool := i+1 < len(conversationMessages) && conversationMessages[i+1].Role == "tool"
if hasNextTool {
sb.WriteString("\n")
} else {
sb.WriteString("<|tool▁outputs▁end|>")
outputsOpen = false
}
}
}
if outputsOpen {
sb.WriteString("<|tool▁outputs▁end|>")
}
if !isLastUser {
sb.WriteString("<|Assistant|>")
}
if enableThinking {
sb.WriteString("<think>\n")
}
return sb.String(), nil
}
package renderers
import (
"testing"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/api"
)
func TestCogitoRenderer(t *testing.T) {
tests := []struct {
name string
messages []api.Message
tools []api.Tool
thinkValue *api.ThinkValue
expected string
}{
{
name: "basic user message",
messages: []api.Message{
{Role: "user", Content: "Hello, how are you?"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello, how are you?<|Assistant|>`,
},
{
name: "basic with system message",
messages: []api.Message{
{Role: "system", Content: "You are a helpful assistant."},
{Role: "user", Content: "Hello, how are you?"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You are a helpful assistant.<|User|>Hello, how are you?<|Assistant|>`,
},
{
name: "conversation with assistant response",
messages: []api.Message{
{Role: "user", Content: "What is the capital of France?"},
{Role: "assistant", Content: "The capital of France is Paris."},
{Role: "user", Content: "Fantastic!"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>What is the capital of France?<|Assistant|>The capital of France is Paris.<|end▁of▁sentence|><|User|>Fantastic!<|Assistant|>`,
},
{
name: "thinking enabled without system",
messages: []api.Message{
{Role: "user", Content: "Hello, how are you?"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello, how are you?<|Assistant|><think>
`,
},
{
name: "thinking enabled with system",
messages: []api.Message{
{Role: "system", Content: "You are a helpful assistant."},
{Role: "user", Content: "Hello, how are you?"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You are a helpful assistant.
<|User|>Hello, how are you?<|Assistant|><think>
`,
},
{
name: "thinking disabled",
messages: []api.Message{
{Role: "user", Content: "Hello, how are you?"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello, how are you?<|Assistant|>`,
},
{
name: "with tools",
messages: []api.Message{
{Role: "user", Content: "What's the weather like?"},
},
thinkValue: &api.ThinkValue{Value: false},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Description: "Get current weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: map[string]api.ToolProperty{
"location": {
Type: api.PropertyType{"string"},
Description: "City name",
},
},
Required: []string{"location"},
},
},
},
},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You have the following functions available:
` + "```json\n" + `{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get current weather",
"parameters": {
"type": "object",
"required": [
"location"
],
"properties": {
"location": {
"type": "string",
"description": "City name"
}
}
}
}
}
` + "```\n" + `<|User|>What's the weather like?<|Assistant|>`,
},
{
name: "assistant with tool calls",
messages: []api.Message{
{Role: "user", Content: "What's the weather in Paris?"},
{
Role: "assistant",
Content: "I'll check the weather in Paris for you.",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: api.ToolCallFunctionArguments{
"location": "Paris",
},
},
},
},
},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>What's the weather in Paris?<|Assistant|>I'll check the weather in Paris for you.<|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>get_weather
` + "```json\n" + `{"location":"Paris"}
` + "```" + `<|tool▁call▁end|><|tool▁calls▁end|><|end▁of▁sentence|><|Assistant|>`,
},
{
name: "tool response",
messages: []api.Message{
{Role: "user", Content: "What's the weather in Paris?"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: api.ToolCallFunctionArguments{
"location": "Paris",
},
},
},
},
},
{Role: "tool", Content: "Temperature: 22°C, Sunny"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>What's the weather in Paris?<|Assistant|><|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>get_weather
` + "```json\n" + `{"location":"Paris"}
` + "```" + `<|tool▁call▁end|><|tool▁calls▁end|><|end▁of▁sentence|><|tool▁outputs▁begin|><|tool▁output▁begin|>Temperature: 22°C, Sunny<|tool▁output▁end|><|tool▁outputs▁end|><|Assistant|>`,
},
{
name: "multiple tool responses",
messages: []api.Message{
{Role: "user", Content: "Get weather for Paris and London"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: api.ToolCallFunctionArguments{
"location": "Paris",
},
},
},
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: api.ToolCallFunctionArguments{
"location": "London",
},
},
},
},
},
{Role: "tool", Content: "Paris: 22°C, Sunny"},
{Role: "tool", Content: "London: 18°C, Cloudy"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Get weather for Paris and London<|Assistant|><|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>get_weather
` + "```json\n" + `{"location":"Paris"}
` + "```" + `<|tool▁call▁end|>
<|tool▁call▁begin|>function<|tool▁sep|>get_weather
` + "```json\n" + `{"location":"London"}
` + "```" + `<|tool▁call▁end|><|tool▁calls▁end|><|end▁of▁sentence|><|tool▁outputs▁begin|><|tool▁output▁begin|>Paris: 22°C, Sunny<|tool▁output▁end|>
<|tool▁output▁begin|>London: 18°C, Cloudy<|tool▁output▁end|><|tool▁outputs▁end|><|Assistant|>`,
},
{
name: "thinking with tools",
messages: []api.Message{
{Role: "user", Content: "What's the weather like?"},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Description: "Get current weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: map[string]api.ToolProperty{
"location": {
Type: api.PropertyType{"string"},
Description: "City name",
},
},
Required: []string{"location"},
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You have the following functions available:
` + "```json\n" + `{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get current weather",
"parameters": {
"type": "object",
"required": [
"location"
],
"properties": {
"location": {
"type": "string",
"description": "City name"
}
}
}
}
}
` + "```\n" + `<|User|>What's the weather like?<|Assistant|><think>
`,
},
// test cases based on cogito
{
name: "single_turn_thinking_false",
messages: []api.Message{
{Role: "user", Content: "Hello"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello<|Assistant|>`,
},
{
name: "single_turn_thinking_true",
messages: []api.Message{
{Role: "user", Content: "Hello"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello<|Assistant|><think>
`,
},
{
name: "multi_turn_thinking_false",
messages: []api.Message{
{Role: "user", Content: "Hello"},
{Role: "assistant", Content: "Hi there!"},
{Role: "user", Content: "How are you?"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello<|Assistant|>Hi there!<|end▁of▁sentence|><|User|>How are you?<|Assistant|>`,
},
{
name: "multi_turn_thinking_true",
messages: []api.Message{
{Role: "user", Content: "Hello"},
{Role: "assistant", Content: "Hi there!"},
{Role: "user", Content: "How are you?"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello<|Assistant|>Hi there!<|end▁of▁sentence|><|User|>How are you?<|Assistant|><think>
`,
},
{
name: "multi_with_system_thinking_false",
messages: []api.Message{
{Role: "system", Content: "You are a helpful assistant"},
{Role: "user", Content: "Start"},
{Role: "assistant", Content: "Okay"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You are a helpful assistant<|User|>Start<|Assistant|>Okay<|end▁of▁sentence|><|Assistant|>`,
},
{
name: "multi_with_system_thinking_true",
messages: []api.Message{
{Role: "system", Content: "You are a helpful assistant"},
{Role: "user", Content: "Start"},
{Role: "assistant", Content: "Okay"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You are a helpful assistant
<|User|>Start<|Assistant|>Okay<|end▁of▁sentence|><|Assistant|><think>
`,
},
{
name: "multi_with_system2_thinking_false",
messages: []api.Message{
{Role: "system", Content: "You are a pirate chatbot who always responds in pirate speak!"},
{Role: "user", Content: "Give me a short introduction to LLMs."},
{Role: "assistant", Content: "Arrr! I'm a pirate"},
{Role: "user", Content: "Tell me more about LLMs."},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You are a pirate chatbot who always responds in pirate speak!<|User|>Give me a short introduction to LLMs.<|Assistant|>Arrr! I'm a pirate<|end▁of▁sentence|><|User|>Tell me more about LLMs.<|Assistant|>`,
},
{
name: "multi_with_system2_thinking_true",
messages: []api.Message{
{Role: "system", Content: "You are a pirate chatbot who always responds in pirate speak!"},
{Role: "user", Content: "Give me a short introduction to LLMs."},
{Role: "assistant", Content: "Arrr! I'm a pirate"},
{Role: "user", Content: "Tell me more about LLMs."},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
You are a pirate chatbot who always responds in pirate speak!
<|User|>Give me a short introduction to LLMs.<|Assistant|>Arrr! I'm a pirate<|end▁of▁sentence|><|User|>Tell me more about LLMs.<|Assistant|><think>
`,
},
// tools
{
name: "tool_calls_only_no_content",
messages: []api.Message{
{Role: "user", Content: "Get weather for Paris"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: api.ToolCallFunctionArguments{
"location": "Paris",
},
},
},
},
},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Get weather for Paris<|Assistant|><|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>get_weather
` + "```json\n" + `{"location":"Paris"}
` + "```" + `<|tool▁call▁end|><|tool▁calls▁end|><|end▁of▁sentence|><|Assistant|>`,
},
{
name: "complex_tool_arguments",
messages: []api.Message{
{Role: "user", Content: "Process complex data"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "process_data",
Arguments: api.ToolCallFunctionArguments{
"items": []any{"item1", "item2", "item3"},
"config": map[string]any{
"enabled": true,
"threshold": 0.95,
"tags": []string{"important", "urgent"},
},
},
},
},
},
},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Process complex data<|Assistant|><|tool▁calls▁begin|><|tool▁call▁begin|>function<|tool▁sep|>process_data
` + "```json\n" + `{"config":{"enabled":true,"tags":["important","urgent"],"threshold":0.95},"items":["item1","item2","item3"]}
` + "```" + `<|tool▁call▁end|><|tool▁calls▁end|><|end▁of▁sentence|><|Assistant|>`,
},
{
name: "empty_messages",
messages: []api.Message{
{Role: "system", Content: ""},
{Role: "user", Content: "Hello"},
{Role: "assistant", Content: ""},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hello<|Assistant|><|end▁of▁sentence|><|Assistant|>`,
},
{
name: "thinking_with_empty_assistant_content",
messages: []api.Message{
{Role: "user", Content: "Think about this"},
{Role: "assistant", Content: ""},
},
thinkValue: &api.ThinkValue{Value: true},
expected: `<|begin▁of▁sentence|>Enable deep thinking subroutine.
You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Think about this<|Assistant|><|end▁of▁sentence|><|Assistant|><think>
`,
},
{
name: "multiple_system_messages",
messages: []api.Message{
{Role: "system", Content: "First instruction"},
{Role: "system", Content: "Second instruction"},
{Role: "user", Content: "Hello"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.
First instruction<|User|>Hello<|Assistant|>`,
},
{
name: "special_characters_in_content",
messages: []api.Message{
{Role: "user", Content: "What about <|special|> tokens and \"quotes\"?"},
{Role: "assistant", Content: "They're handled normally in content."},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>What about <|special|> tokens and "quotes"?<|Assistant|>They're handled normally in content.<|end▁of▁sentence|><|Assistant|>`,
},
{
name: "long_conversation_multiple_rounds",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "Hello!"},
{Role: "user", Content: "How are you?"},
{Role: "assistant", Content: "Good, thanks!"},
{Role: "user", Content: "What's the weather?"},
},
thinkValue: &api.ThinkValue{Value: false},
expected: `<|begin▁of▁sentence|>You are Cogito, an AI assistant created by Deep Cogito, which is an AI research lab based in San Francisco.<|User|>Hi<|Assistant|>Hello!<|end▁of▁sentence|><|User|>How are you?<|Assistant|>Good, thanks!<|end▁of▁sentence|><|User|>What's the weather?<|Assistant|>`,
},
}
renderer := &CogitoRenderer{isThinking: true}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
rendered, err := renderer.Render(tt.messages, tt.tools, tt.thinkValue)
if err != nil {
t.Fatalf("Render() error = %v", err)
}
if diff := cmp.Diff(tt.expected, rendered); diff != "" {
t.Errorf("Render() mismatch (-want +got):\n%s", diff)
}
})
}
}
...@@ -56,6 +56,9 @@ func rendererForName(name string) Renderer { ...@@ -56,6 +56,9 @@ func rendererForName(name string) Renderer {
case "qwen3-vl-thinking": case "qwen3-vl-thinking":
renderer := &Qwen3VLRenderer{isThinking: true, useImgTags: RenderImgTags} renderer := &Qwen3VLRenderer{isThinking: true, useImgTags: RenderImgTags}
return renderer return renderer
case "cogito":
renderer := &CogitoRenderer{isThinking: true}
return renderer
default: default:
return nil return nil
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment