fix: enforce tool whitelisting and isolate prompt tool call markers

This commit is contained in:
Luxferre
2026-09-08 16:59:17 +03:00
parent 354db0afcb
commit c99ffcac23
4 changed files with 2210 additions and 1092 deletions
+65 -63
View File
@@ -1,21 +1,21 @@
# groqqer
`groqqer` is a lightweight, zero-dependency Go OpenAI-compatible API gateway and proxy for the Groq Streamlit chatbot space (`https://dromerosm-groq-chatbot.hf.space`).
`groqqer` is a lightweight, zero-dependency Go OpenAI-compatible API gateway and proxy for the Groq Streamlit space (`https://dromerosm-groq-chatbot.hf.space`).
It converts standard OpenAI chat completions and model requests into automated headless browser interactions, providing fast inference with Groq models without requiring an API key.
It connects directly to Streamlit's binary WebSocket engine (`/_stcore/stream`) using a pure Go Protobuf wire-format encoder/decoder. It requires **no browser**, **no Chromium**, and **no Xvfb**, allowing it to run smoothly on minimal headless servers and low-resource containers.
## Features
## Key Features
- **OpenAI Compatible API**: Exposes standard `/v1/models` and `/v1/chat/completions` endpoints. Drop-in replacement for OpenAI SDKs, Open-WebUI, LiteLLM, LibreChat, and LangChain.
- **Zero Third-Party Dependencies**: Written entirely in pure Go using only the standard library (`net/http`, `encoding/json`, `os/exec`, `crypto/rand`, etc.). Includes a built-in RFC 6455 WebSocket CDP client and RFC 1928 SOCKS5 proxy client.
- **System Prompt Support**: Translates developer/system messages into natural instruction framing.
- **Streaming & Non-Streaming**: Supports Server-Sent Events (`stream: true`) with real-time token delivery, as well as synchronous JSON responses.
- **Tool / Function Calling**: Fully supports OpenAI `tools`, `tool_choice`, and multi-turn execution (`role: "tool"`). Automatically injects schemas and parses `<tool_call>` responses into OpenAI `tool_calls` payloads with `finish_reason: "tool_calls"`.
- **Reasoning Content**: Automatically parses `<think>` and `</think>` tags from reasoning models (such as Qwen 3.6/3.8) and streams or populates `reasoning_content` in accordance with OpenAI O-series conventions.
- **Dynamic Model Discovery**: Scrapes active models directly from the Streamlit UI on startup and allows seamless model switching between requests (`qwen/qwen3.6-27b`, `openai/gpt-oss-120b`, `openai/gpt-oss-20b`, `groq/compound`, etc.).
- **Display Isolation via Xvfb**: Automatically discovers and spawns a virtual X display (`:100` - `:199`) to keep browser automation completely isolated from host desktop displays. Gracefully falls back to offscreen coordinates if Xvfb is not present.
- **SOCKS5 Proxy Tunneling**: Tunnel browser traffic through SOCKS5 proxies using the `-socks` flag or `ALL_PROXY` / `SOCKS5_PROXY` environment variables.
- **Clean Session Isolation**: Automatically clears Streamlit conversation state between requests to prevent cumulative token window exhaustion or stale error states.
- **Direct WebSocket Protobuf Protocol**: Communicates directly with Streamlit's internal engine over RFC 6455 WebSockets and Protocol Buffers wire format.
- **Completely Headless & Zero Dependencies**: Written entirely in pure Go using only the standard library (`net/http`, `crypto/tls`, `encoding/binary`, `encoding/json`, etc.). No Chromium, Chrome, Xvfb, Puppeteer, or external Go modules required.
- **Sub-Second Latency**: Bypasses browser rendering and DOM parsing entirely, delivering responses with minimal overhead.
- **OpenAI Compatible API**: Exposes standard `/v1/models` and `/v1/chat/completions` endpoints. Drop-in replacement for OpenAI SDKs, LiteLLM, Open-WebUI, LibreChat, and LangChain.
- **System Prompt Support**: Formats developer/system instructions cleanly for the target model.
- **Streaming & Non-Streaming**: Supports Server-Sent Events (`stream: true`) with real-time token streaming and synchronous JSON responses.
- **Tool / Function Calling**: Fully supports OpenAI `tools`, `tool_choice`, and multi-turn execution (`role: "tool"`). Automatically injects schemas and parses `<tool_call>` outputs into standard OpenAI `tool_calls` payloads with `finish_reason: "tool_calls"`.
- **Reasoning Content**: Extracts `<think>` and `</think>` tags from reasoning models (e.g., Qwen 3.6/3.8) and streams or populates `reasoning_content` following OpenAI O-series conventions.
- **Dynamic Model Discovery & Switching**: Discovers active models directly from the Streamlit space on startup and allows seamless model switching between requests (`qwen/qwen3.6-27b`, `openai/gpt-oss-120b`, `openai/gpt-oss-20b`, `groq/compound`, etc.).
- **SOCKS5 Proxy Support**: Native pure Go SOCKS5 proxy client supporting authentication (`-socks` flag or `ALL_PROXY` / `SOCKS5_PROXY` environment variables).
## Architecture
@@ -28,14 +28,9 @@ It converts standard OpenAI chat completions and model requests into automated h
+---------------------------+
| groqqer |
| (Pure Go Standard Lib) |
| RFC 6455 WS + Protobuf |
+---------------------------+
| RFC 6455 WebSocket CDP
v
+---------------------------+
| Chromium / Chrome |
| (Display :100+ via Xvfb) |
+---------------------------+
| HTTPS
| WSS (TLS WebSocket)
v
+---------------------------+
| Groq Streamlit Space |
@@ -45,10 +40,9 @@ It converts standard OpenAI chat completions and model requests into automated h
## Requirements
- **Linux** (x86_64 or aarch64)
- **Go 1.20+**
- **Chromium** or **Google Chrome** installed (`/usr/bin/chromium`, `/usr/bin/google-chrome`, or in `$PATH`)
- *(Optional)* **Xvfb** (`xorg-server-xvfb`) for display isolation
- **Linux / macOS / Windows**
- **Go 1.20+** (for building from source)
- No browser or graphics packages needed.
## Installation & Build
@@ -58,7 +52,7 @@ Clone or navigate to the repository, then build the binary:
make
```
The optimized binary will be placed at `bin/groqqer`.
The optimized binary will be created at `bin/groqqer`.
## Running the Server
@@ -74,12 +68,11 @@ Start the gateway with default settings:
|------|---------|-------------|
| `-port` | `8080` | HTTP server listening port |
| `-target` | `https://dromerosm-groq-chatbot.hf.space` | Target Groq Streamlit space URL |
| `-default-model` | `llama-3.3-70b-versatile` | Fallback model if request does not specify one |
| `-browser` | `""` | Custom path to Chromium/Chrome binary |
| `-xvfb` | `true` | Enable virtual X server display isolation |
| `-no-xvfb` | `false` | Disable virtual X server (run offscreen) |
| `-default-model` | `qwen/qwen3.6-27b` | Fallback model if request does not specify one |
| `-socks` | `""` | SOCKS5 proxy URL (e.g. `socks5://127.0.0.1:1080`) |
| `-user-agent` | *(Chrome 133 UA)* | Custom browser User-Agent string |
| `-user-agent` | *(Chrome UA)* | Custom User-Agent string |
*(Note: legacy flags `-browser`, `-xvfb`, and `-headless` are retained for backward compatibility but are ignored, as groqqer operates completely headless via direct WebSocket).*
## API Usage Examples
@@ -94,10 +87,13 @@ curl -s http://127.0.0.1:8080/v1/models
{
"object": "list",
"data": [
{"id": "qwen/qwen3.6-27b", "object": "model", "created": 1788869470, "owned_by": "groq"},
{"id": "openai/gpt-oss-120b", "object": "model", "created": 1788869470, "owned_by": "groq"},
{"id": "openai/gpt-oss-20b", "object": "model", "created": 1788869470, "owned_by": "groq"},
{"id": "groq/compound", "object": "model", "created": 1788869470, "owned_by": "groq"}
{"id": "qwen/qwen3.6-27b", "object": "model", "created": 1788870683, "owned_by": "groq"},
{"id": "qwen/qwen3.8-27b", "object": "model", "created": 1788870683, "owned_by": "groq"},
{"id": "openai/gpt-oss-120b", "object": "model", "created": 1788870683, "owned_by": "groq"},
{"id": "openai/gpt-oss-20b", "object": "model", "created": 1788870683, "owned_by": "groq"},
{"id": "groq/compound", "object": "model", "created": 1788870683, "owned_by": "groq"},
{"id": "groq/compound-mini", "object": "model", "created": 1788870683, "owned_by": "groq"},
{"id": "allam-2-7b", "object": "model", "created": 1788870683, "owned_by": "groq"}
]
}
```
@@ -110,7 +106,7 @@ curl -s http://127.0.0.1:8080/v1/models
curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "qwen/qwen3.6-27b",
"model": "openai/gpt-oss-20b",
"messages": [
{"role": "system", "content": "You are a concise assistant. Reply in one sentence."},
{"role": "user", "content": "What is the capital of France?"}
@@ -121,17 +117,16 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
**Response:**
```json
{
"id": "chatcmpl-9691bbf8-8623-4155-a40f-a00e6b6903ac",
"id": "chatcmpl-893498c9-2808-4bec-8c6a-2a0352ed69e9",
"object": "chat.completion",
"created": 1788869478,
"model": "qwen/qwen3.6-27b",
"created": 1788870700,
"model": "openai/gpt-oss-20b",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "The capital of France is Paris.",
"reasoning_content": "Analyze User Input: Question: 'What is the capital of France?'..."
"content": "The capital of France is Paris."
},
"finish_reason": "stop"
}
@@ -154,7 +149,7 @@ curl -N -s -X POST http://127.0.0.1:8080/v1/chat/completions \
-d '{
"model": "qwen/qwen3.6-27b",
"messages": [
{"role": "user", "content": "Count from 1 to 4 with commas."}
{"role": "user", "content": "Count from 1 to 4 separated by commas."}
],
"stream": true
}'
@@ -162,15 +157,19 @@ curl -N -s -X POST http://127.0.0.1:8080/v1/chat/completions \
**Output:**
```
data: {"id":"chatcmpl-341d6d05","object":"chat.completion.chunk","created":1788869487,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"role":"assistant"}}]}
data: {"id":"chatcmpl-7a06d0fb","object":"chat.completion.chunk","created":1788870705,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"role":"assistant"}}]}
data: {"id":"chatcmpl-341d6d05","object":"chat.completion.chunk","created":1788869487,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"reasoning_content":"Thinking process..."}}]}
data: {"id":"chatcmpl-7a06d0fb","object":"chat.completion.chunk","created":1788870705,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"reasoning_content":"Thinking process..."}}]}
data: {"id":"chatcmpl-341d6d05","object":"chat.completion.chunk","created":1788869487,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"content":"1, 2, 3,"}}]}
data: {"id":"chatcmpl-7a06d0fb","object":"chat.completion.chunk","created":1788870705,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"content":"1,"}}]}
data: {"id":"chatcmpl-341d6d05","object":"chat.completion.chunk","created":1788869487,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"content":" 4"}}]}
data: {"id":"chatcmpl-7a06d0fb","object":"chat.completion.chunk","created":1788870705,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"content":" 2,"}}]}
data: {"id":"chatcmpl-341d6d05","object":"chat.completion.chunk","created":1788869487,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
data: {"id":"chatcmpl-7a06d0fb","object":"chat.completion.chunk","created":1788870705,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"content":" 3,"}}]}
data: {"id":"chatcmpl-7a06d0fb","object":"chat.completion.chunk","created":1788870705,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{"content":" 4"}}]}
data: {"id":"chatcmpl-7a06d0fb","object":"chat.completion.chunk","created":1788870705,"model":"qwen/qwen3.6-27b","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}
data: [DONE]
```
@@ -185,7 +184,7 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
-d '{
"model": "qwen/qwen3.6-27b",
"messages": [
{"role": "user", "content": "What is the current weather in Tokyo?"}
{"role": "user", "content": "What is the weather in Tokyo right now?"}
],
"tools": [
{
@@ -196,23 +195,23 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string", "description": "The city and country, e.g. Tokyo, Japan"}
"location": {"type": "string", "description": "City and country"}
},
"required": ["location"]
}
}
}
],
"tool_choice": "auto"
"tool_choice": "required"
}'
```
**Response:**
```json
{
"id": "chatcmpl-cb050ba0-2573-403f-a1db-e76cdfb8f99d",
"id": "chatcmpl-29f63b3b-7ab6-444c-8579-49f906007d43",
"object": "chat.completion",
"created": 1788869493,
"created": 1788870711,
"model": "qwen/qwen3.6-27b",
"choices": [
{
@@ -220,13 +219,14 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
"message": {
"role": "assistant",
"content": null,
"reasoning_content": "The user is asking for the current weather in Tokyo...",
"tool_calls": [
{
"id": "call_ff437271",
"id": "call_163e8f97",
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": "{\"location\":\"Tokyo, Japan\"}"
"arguments": "{\"location\":\"Tokyo\"}"
}
}
]
@@ -241,7 +241,7 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
### 5. Multi-Turn Tool Response Execution
Once you run your tool locally, send the tool result back using `role: "tool"`:
Send back the tool execution results using `role: "tool"`:
```bash
curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
@@ -249,25 +249,25 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
-d '{
"model": "qwen/qwen3.6-27b",
"messages": [
{"role": "user", "content": "What is the current weather in Tokyo?"},
{"role": "user", "content": "What is the weather in Tokyo right now?"},
{
"role": "assistant",
"tool_calls": [
{
"id": "call_ff437271",
"id": "call_163e8f97",
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": "{\"location\":\"Tokyo, Japan\"}"
"arguments": "{\"location\":\"Tokyo\"}"
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_ff437271",
"tool_call_id": "call_163e8f97",
"name": "get_current_weather",
"content": "{\"temperature\": \"18°C\", \"condition\": \"Sunny\", \"humidity\": \"45%\"}"
"content": "{\"temperature\": \"19°C\", \"weather\": \"Sunny with clear skies\"}"
}
]
}'
@@ -276,16 +276,16 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
**Response:**
```json
{
"id": "chatcmpl-aad3ed3e-b401-41b8-9b3b-a9455070def5",
"id": "chatcmpl-23354bba-33c7-45be-8807-9c84243037b8",
"object": "chat.completion",
"created": 1788869501,
"created": 1788870718,
"model": "qwen/qwen3.6-27b",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "It is currently sunny in Tokyo with a temperature of 18°C and 45% humidity."
"content": "It's currently 19°C and sunny with clear skies in Tokyo"
},
"finish_reason": "stop"
}
@@ -293,7 +293,9 @@ curl -s -X POST http://127.0.0.1:8080/v1/chat/completions \
}
```
## Python OpenAI SDK Example
---
### 6. Python OpenAI SDK Example
```python
from openai import OpenAI
BIN
View File
Binary file not shown.
+1450 -952
View File
File diff suppressed because it is too large Load Diff
+618
View File
@@ -0,0 +1,618 @@
package main
import (
"net/http/httptest"
"strings"
"testing"
)
func TestFormatPromptToolResolutionAndGrouping(t *testing.T) {
req := ChatCompletionRequest{
Model: "qwen/qwen3.6-27b",
Tools: []Tool{
{
Type: "function",
Function: map[string]interface{}{
"name": "calculator",
"description": "Calculate math expression",
},
},
},
Messages: []ChatMessage{
{
Role: "user",
Content: "Calculate 25 * 4 and then tell me what the square root of that is.",
},
{
Role: "assistant",
ToolCalls: []ToolCall{
{
ID: "call_abc123",
Type: "function",
Function: ToolCallFunction{
Name: "calculator",
Arguments: `{"expr":"25*4"}`,
},
},
},
},
{
Role: "tool",
ToolCallID: "call_abc123",
Content: "100",
},
},
}
prompt := FormatPrompt(req)
// Verify tool name was resolved from assistant tool call
if !strings.Contains(prompt, `"name": "calculator"`) {
t.Errorf("Expected prompt to contain resolved tool name calculator, got:\n%s", prompt)
}
// Verify tool_call_id is preserved
if !strings.Contains(prompt, `"tool_call_id": "call_abc123"`) {
t.Errorf("Expected prompt to contain tool_call_id, got:\n%s", prompt)
}
// Verify Tool Execution Results block exists
if !strings.Contains(prompt, "[Tool Execution Results]") {
t.Errorf("Expected prompt to have [Tool Execution Results], got:\n%s", prompt)
}
// Verify Next Steps Directive exists
if !strings.Contains(prompt, "[Next Steps Directive]") {
t.Errorf("Expected prompt to have [Next Steps Directive], got:\n%s", prompt)
}
// Verify tool result was not placed under [User]
if strings.Contains(prompt, "[User]\n<tool_response>") {
t.Errorf("Tool result should not be placed under [User], got:\n%s", prompt)
}
}
func TestFormatPromptParallelToolResults(t *testing.T) {
req := ChatCompletionRequest{
Model: "qwen/qwen3.6-27b",
Messages: []ChatMessage{
{
Role: "user",
Content: "What is the weather in Tokyo and Paris?",
},
{
Role: "assistant",
ToolCalls: []ToolCall{
{
ID: "call_tokyo",
Type: "function",
Function: ToolCallFunction{
Name: "get_weather",
Arguments: `{"city":"Tokyo"}`,
},
},
{
ID: "call_paris",
Type: "function",
Function: ToolCallFunction{
Name: "get_weather",
Arguments: `{"city":"Paris"}`,
},
},
},
},
{
Role: "tool",
ToolCallID: "call_tokyo",
Content: `{"temp": 18, "condition": "sunny"}`,
},
{
Role: "tool",
ToolCallID: "call_paris",
Content: `{"temp": 15, "condition": "rainy"}`,
},
},
}
prompt := FormatPrompt(req)
// Both tool results must be grouped under [Tool Execution Results]
execIdx := strings.Index(prompt, "[Tool Execution Results]")
if execIdx == -1 {
t.Fatalf("Expected [Tool Execution Results], got:\n%s", prompt)
}
execBlock := prompt[execIdx:]
if !strings.Contains(execBlock, "call_tokyo") {
t.Errorf("Expected call_tokyo in Tool Execution Results block, got:\n%s", execBlock)
}
if !strings.Contains(execBlock, "call_paris") {
t.Errorf("Expected call_paris in Tool Execution Results block, got:\n%s", execBlock)
}
// Neither should be in [Conversation History] as an orphan turn
historyIdx := strings.Index(prompt, "[Conversation History]")
if historyIdx != -1 {
historyBlock := prompt[historyIdx:execIdx]
if strings.Contains(historyBlock, "call_tokyo") || strings.Contains(historyBlock, "call_paris") {
t.Errorf("Trailing tool results should not be in history block, got:\n%s", historyBlock)
}
}
}
func TestDetectToolCallsFormats(t *testing.T) {
allowedTools := map[string]bool{
"search": true,
"weather": true,
"tool_a": true,
"tool_b": true,
"direct_tool": true,
}
// Test 1: XML single tool call
xmlSingle := `<tool_call>
{"name": "search", "arguments": {"query": "golang"}}
</tool_call>`
calls, rem, ok := DetectToolCalls(xmlSingle, allowedTools)
if !ok || len(calls) != 1 || calls[0].Function.Name != "search" {
t.Errorf("Failed to detect XML single tool call: ok=%v, calls=%v, rem=%q", ok, calls, rem)
}
if calls[0].Index == nil || *calls[0].Index != 0 {
t.Errorf("Expected index 0, got %v", calls[0].Index)
}
// Test 2: XML multiple tool calls
xmlMulti := `<tool_call>
{"name": "search", "arguments": {"query": "golang"}}
</tool_call>
<tool_call>
{"name": "weather", "arguments": {"city": "Tokyo"}}
</tool_call>`
calls2, _, ok2 := DetectToolCalls(xmlMulti, allowedTools)
if !ok2 || len(calls2) != 2 {
t.Fatalf("Failed to detect XML multiple tool calls: ok=%v, calls=%v", ok2, calls2)
}
if calls2[0].Function.Name != "search" || calls2[1].Function.Name != "weather" {
t.Errorf("Tool call names mismatch: %v", calls2)
}
if calls2[0].Index == nil || *calls2[0].Index != 0 || calls2[1].Index == nil || *calls2[1].Index != 1 {
t.Errorf("Indices mismatch: %v, %v", calls2[0].Index, calls2[1].Index)
}
// Test 3: JSON array inside XML
xmlArray := `<tool_call>
[
{"name": "tool_a", "arguments": {"a": 1}},
{"name": "tool_b", "arguments": {"b": 2}}
]
</tool_call>`
calls3, _, ok3 := DetectToolCalls(xmlArray, allowedTools)
if !ok3 || len(calls3) != 2 {
t.Fatalf("Failed to detect JSON array inside XML tool call: ok=%v, calls=%v", ok3, calls3)
}
if calls3[0].Function.Name != "tool_a" || calls3[1].Function.Name != "tool_b" {
t.Errorf("Tool names mismatch: %v", calls3)
}
// Test 4: Direct JSON array without tags
jsonArray := `[{"name": "direct_tool", "arguments": {"key": "val"}}]`
calls4, _, ok4 := DetectToolCalls(jsonArray, allowedTools)
if !ok4 || len(calls4) != 1 || calls4[0].Function.Name != "direct_tool" {
t.Errorf("Failed to detect direct JSON array: ok=%v, calls=%v", ok4, calls4)
}
}
func TestWriteCompletionResponseToolCallsNilContent(t *testing.T) {
rec := httptest.NewRecorder()
idx := 0
calls := []ToolCall{
{
Index: &idx,
ID: "call_test",
Type: "function",
Function: ToolCallFunction{
Name: "test_func",
Arguments: `{}`,
},
},
}
WriteCompletionResponse(rec, "cmpl-1", 123456, "qwen/qwen3.6-27b", "Some intermediate thought", "", calls, "tool_calls")
body := rec.Body.String()
if !strings.Contains(body, `"content":null`) {
t.Errorf("Expected content:null when tool_calls present, got:\n%s", body)
}
if !strings.Contains(body, `"reasoning_content":"Some intermediate thought"`) {
t.Errorf("Expected preamble text to be preserved in reasoning_content, got:\n%s", body)
}
if !strings.Contains(body, `"finish_reason":"tool_calls"`) {
t.Errorf("Expected finish_reason tool_calls, got:\n%s", body)
}
}
func TestIsStaticUIMarkdown(t *testing.T) {
cases := []struct {
text string
prompt string
expected bool
}{
{"Model: qwen/qwen3.6-27b", "hello", true},
{"Model: openai/gpt-oss-20b", "hello", true},
{"Model : deepseek-r1-distill-llama-70b", "hello", true},
{"- **Model ID:** `qwen/qwen3.6-27b`", "hello", true},
{"[Powered by Groq](https://groq.com)", "hello", true},
{"Partial response: generation did not finish.", "hello", true},
{"[Notice] **Important:** Do not rely on AI without checking", "hello", true},
{"Important: Do not rely on AI without checking", "hello", true},
{"3 older messages were left out of this request.", "hello", true},
{"Here is the solution to your question.", "hello", false},
{"<think>thinking</think>Result: 42", "hello", false},
{"<tool_call>{\"name\":\"calc\"}</tool_call>", "hello", false},
}
for _, c := range cases {
got := isStaticUIMarkdown(c.text, c.prompt)
if got != c.expected {
t.Errorf("isStaticUIMarkdown(%q) = %v; expected %v", c.text, got, c.expected)
}
}
}
func TestFormatPromptCorruptedCaptionFiltering(t *testing.T) {
req := ChatCompletionRequest{
Model: "qwen/qwen3.6-27b",
Messages: []ChatMessage{
{Role: "user", Content: "Hello"},
{Role: "assistant", Content: "Model: qwen/qwen3.6-27b"},
{Role: "user", Content: "What is 2+2?"},
},
}
prompt := FormatPrompt(req)
if strings.Contains(prompt, "Model: qwen/qwen3.6-27b") {
t.Errorf("Expected corrupted Model: caption to be filtered out of history, got:\n%s", prompt)
}
if !strings.Contains(prompt, "What is 2+2?") {
t.Errorf("Expected user question to remain, got:\n%s", prompt)
}
}
func TestFormatPromptHistoryPruning(t *testing.T) {
var messages []ChatMessage
messages = append(messages, ChatMessage{Role: "user", Content: "Initial user query to remember"})
// Add 30 turns of long conversation
for i := 1; i <= 30; i++ {
messages = append(messages, ChatMessage{Role: "assistant", Content: strings.Repeat("Long assistant response ", 30)})
messages = append(messages, ChatMessage{Role: "user", Content: strings.Repeat("Long user question ", 30)})
}
messages = append(messages, ChatMessage{Role: "user", Content: "Final question"})
req := ChatCompletionRequest{
Model: "qwen/qwen3.6-27b",
Messages: messages,
}
prompt := FormatPrompt(req)
// Verify initial user query is preserved
if !strings.Contains(prompt, "Initial user query to remember") {
t.Errorf("Expected prompt to preserve initial user query")
}
// Verify truncation marker is present
if !strings.Contains(prompt, "[... earlier conversation turns omitted for brevity ...]") {
t.Errorf("Expected prompt to contain omission notice")
}
// Verify prompt size is kept bounded
if len(prompt) > 20000 {
t.Errorf("Prompt size too large after pruning: %d chars", len(prompt))
}
}
func TestParseRetryAfterAndRateLimit(t *testing.T) {
alert1 := "Rate limit reached. Please retry in a few seconds. Retry after 10 seconds."
if !isRateLimit(alert1) {
t.Errorf("Expected alert1 to be identified as rate limit")
}
if sec := parseRetryAfter(alert1); sec != 10 {
t.Errorf("Expected retry seconds 10, got %d", sec)
}
alert2 := "429 Too Many Requests: retry in 5"
if !isRateLimit(alert2) {
t.Errorf("Expected alert2 to be identified as rate limit")
}
if sec := parseRetryAfter(alert2); sec != 5 {
t.Errorf("Expected retry seconds 5, got %d", sec)
}
alert3 := "Your message or image is too large for this model."
if isRateLimit(alert3) {
t.Errorf("Expected alert3 NOT to be rate limit")
}
}
func TestWriteAPIError(t *testing.T) {
rec := httptest.NewRecorder()
WriteAPIError(rec, 429, "Rate limit reached. Retry after 10 seconds.", "rate_limit_error", 429)
if rec.Code != 429 {
t.Errorf("Expected status 429, got %d", rec.Code)
}
body := rec.Body.String()
if !strings.Contains(body, `"rate_limit_error"`) || !strings.Contains(body, `"code":429`) {
t.Errorf("Unexpected error response body: %s", body)
}
}
func TestDetectToolCallsWithUnregisteredTool(t *testing.T) {
allowed := map[string]bool{"calculator": true}
content := `Here is an example:
<tool_call>
{"name": "non_existing_tool", "arguments": {"foo": "bar"}}
</tool_call>
Do not execute it.`
calls, rem, ok := DetectToolCalls(content, allowed)
if ok || len(calls) > 0 {
t.Errorf("Expected unregistered tool to not be detected as tool call, got ok=%v, calls=%v", ok, calls)
}
if !strings.Contains(rem, "non_existing_tool") {
t.Errorf("Expected remaining text to keep unregistered tool text intact, got:\n%s", rem)
}
}
func TestDetectToolCallsWithEmptyAllowedTools(t *testing.T) {
content := `<tool_call>
{"name": "calculator", "arguments": {"expr": "1+1"}}
</tool_call>`
calls, rem, ok := DetectToolCalls(content, nil)
if ok || len(calls) > 0 {
t.Errorf("Expected tool calls to be disabled when allowedTools is nil, got ok=%v, calls=%v", ok, calls)
}
if rem != content {
t.Errorf("Expected content to remain unchanged when allowedTools is nil")
}
}
func TestDetectToolCallsWithUnclosedTag(t *testing.T) {
allowed := map[string]bool{"calculator": true}
content := "Please explain what <tool_call> means in your prompt. My name is Alice."
calls, rem, ok := DetectToolCalls(content, allowed)
if ok || len(calls) > 0 {
t.Errorf("Expected unclosed tag to not be detected as tool call, got ok=%v, calls=%v", ok, calls)
}
if rem != content {
t.Errorf("Expected content to remain completely unchanged, got:\n%s", rem)
}
}
func TestDetectToolCallsPartiallyValid(t *testing.T) {
allowed := map[string]bool{"calculator": true}
content := `<tool_call>
{"name": "calculator", "arguments": {"expr": "2+2"}}
</tool_call>
And here is a fake tool:
<tool_call>
{"name": "fake_tool", "arguments": {}}
</tool_call>`
calls, rem, ok := DetectToolCalls(content, allowed)
if !ok || len(calls) != 1 {
t.Fatalf("Expected 1 valid tool call, got ok=%v, calls=%v", ok, calls)
}
if calls[0].Function.Name != "calculator" {
t.Errorf("Expected tool name calculator, got %s", calls[0].Function.Name)
}
if !strings.Contains(rem, "fake_tool") {
t.Errorf("Expected remaining content to keep fake_tool block, got:\n%s", rem)
}
if strings.Contains(rem, "calculator") {
t.Errorf("Expected valid calculator block to be removed from remaining content, got:\n%s", rem)
}
}
func TestStreamToolCallFilterNonExistingTool(t *testing.T) {
allowed := map[string]bool{"calculator": true}
filter := NewStreamToolCallFilter(allowed)
var contentParts []string
var reasoningParts []string
var toolCalls []ToolCall
onContent := func(s string) { contentParts = append(contentParts, s) }
onReasoning := func(s string) { reasoningParts = append(reasoningParts, s) }
onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }
chunks := []string{
"Here is ",
"an example:\n",
"<tool_call>\n",
"{\"name\": \"unknown_tool\", ",
"\"arguments\": {}}\n",
"</tool_call>\n",
"Done.",
}
for _, chunk := range chunks {
filter.Feed(chunk, onContent, onReasoning, onToolCall)
}
filter.Flush(onContent, onReasoning, onToolCall)
if len(toolCalls) > 0 {
t.Errorf("Expected zero tool calls emitted for unknown tool, got %d", len(toolCalls))
}
if filter.emittedCall {
t.Errorf("Expected emittedCall to be false")
}
fullContent := strings.Join(contentParts, "")
if !strings.Contains(fullContent, "unknown_tool") {
t.Errorf("Expected fullContent to contain unknown_tool block, got:\n%s", fullContent)
}
if !strings.Contains(fullContent, "Done.") {
t.Errorf("Expected fullContent to contain Done., got:\n%s", fullContent)
}
}
func TestStreamToolCallFilterEmptyAllowedTools(t *testing.T) {
filter := NewStreamToolCallFilter(nil)
var contentParts []string
var toolCalls []ToolCall
onContent := func(s string) { contentParts = append(contentParts, s) }
onReasoning := func(s string) {}
onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }
chunks := []string{"<tool_call>", `{"name":"calc"}`, "</tool_call>"}
for _, ch := range chunks {
filter.Feed(ch, onContent, onReasoning, onToolCall)
}
filter.Flush(onContent, onReasoning, onToolCall)
if len(toolCalls) > 0 {
t.Errorf("Expected zero tool calls when allowedTools is nil")
}
fullContent := strings.Join(contentParts, "")
if fullContent != "<tool_call>{\"name\":\"calc\"}</tool_call>" {
t.Errorf("Expected direct pass-through, got:\n%s", fullContent)
}
}
func TestStreamToolCallFilterValidTool(t *testing.T) {
allowed := map[string]bool{"calculator": true}
filter := NewStreamToolCallFilter(allowed)
var contentParts []string
var reasoningParts []string
var toolCalls []ToolCall
onContent := func(s string) { contentParts = append(contentParts, s) }
onReasoning := func(s string) { reasoningParts = append(reasoningParts, s) }
onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }
filter.Feed("Let me calculate that for you.\n", onContent, onReasoning, onToolCall)
filter.Feed("<tool_call>\n{\"name\": \"calculator\", \"arguments\": {\"expr\": \"5*5\"}}\n</tool_call>", onContent, onReasoning, onToolCall)
filter.Flush(onContent, onReasoning, onToolCall)
if len(toolCalls) != 1 {
t.Fatalf("Expected 1 tool call, got %d", len(toolCalls))
}
if toolCalls[0].Function.Name != "calculator" {
t.Errorf("Expected tool name calculator, got %s", toolCalls[0].Function.Name)
}
if !filter.emittedCall {
t.Errorf("Expected emittedCall to be true")
}
fullReasoning := strings.Join(reasoningParts, "")
if !strings.Contains(fullReasoning, "Let me calculate that") {
t.Errorf("Expected preamble text in reasoning, got:\n%s", fullReasoning)
}
}
func TestGetAllowedToolNames(t *testing.T) {
tools := []Tool{
{
Type: "function",
Function: map[string]interface{}{
"name": "search",
},
},
{
Type: "function",
Function: ToolCallFunction{
Name: "calculator",
},
},
}
// Auto choice
m := GetAllowedToolNames(tools, "auto")
if !m["search"] || !m["calculator"] || len(m) != 2 {
t.Errorf("Expected both tools allowed for auto, got %v", m)
}
// None choice
mNone := GetAllowedToolNames(tools, "none")
if mNone != nil {
t.Errorf("Expected nil for tool_choice none, got %v", mNone)
}
// Specific choice
mSpecific := GetAllowedToolNames(tools, map[string]interface{}{
"type": "function",
"function": map[string]interface{}{
"name": "calculator",
},
})
if !mSpecific["calculator"] || mSpecific["search"] || len(mSpecific) != 1 {
t.Errorf("Expected only calculator allowed, got %v", mSpecific)
}
// Empty tools
if GetAllowedToolNames(nil, "auto") != nil {
t.Errorf("Expected nil for nil tools")
}
}
func TestFormatPromptAssistantDeduplication(t *testing.T) {
req := ChatCompletionRequest{
Model: "qwen/qwen3.6-27b",
Tools: []Tool{
{Type: "function", Function: map[string]interface{}{"name": "calc"}},
},
Messages: []ChatMessage{
{Role: "user", Content: "Calculate 2+2"},
{
Role: "assistant",
Content: `<tool_call>{"name": "calc", "arguments": {"expr": "2+2"}}</tool_call>`,
ToolCalls: []ToolCall{
{
ID: "call_1",
Type: "function",
Function: ToolCallFunction{
Name: "calc",
Arguments: `{"expr": "2+2"}`,
},
},
},
},
{Role: "tool", ToolCallID: "call_1", Content: "4"},
},
}
prompt := FormatPrompt(req)
// Verify there is only one <tool_call> block in the Assistant turn
assistIdx := strings.Index(prompt, "Assistant:")
if assistIdx == -1 {
t.Fatalf("Expected Assistant turn in prompt, got:\n%s", prompt)
}
toolExecIdx := strings.Index(prompt, "[Tool Execution Results]")
assistSection := prompt[assistIdx:toolExecIdx]
count := strings.Count(assistSection, "<tool_call>")
if count != 1 {
t.Errorf("Expected exactly 1 <tool_call> block in Assistant turn, got %d:\n%s", count, assistSection)
}
}
func TestFormatPromptUserMessageWithToolCall(t *testing.T) {
req := ChatCompletionRequest{
Model: "qwen/qwen3.6-27b",
Messages: []ChatMessage{
{Role: "user", Content: "How do I format a <tool_call> tag in my script?"},
},
}
prompt := FormatPrompt(req)
if !strings.Contains(prompt, "How do I format a <tool_call> tag in my script?") {
t.Errorf("Expected user prompt to preserve <tool_call> text, got:\n%s", prompt)
}
}