diff --git a/README.md b/README.md index 7946724..0f00c80 100644 --- a/README.md +++ b/README.md @@ -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 `` responses into OpenAI `tool_calls` payloads with `finish_reason: "tool_calls"`. -- **Reasoning Content**: Automatically parses `` and `` 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 `` outputs into standard OpenAI `tool_calls` payloads with `finish_reason: "tool_calls"`. +- **Reasoning Content**: Extracts `` and `` 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 diff --git a/bin/groqqer b/bin/groqqer index dd94681..6a23745 100755 Binary files a/bin/groqqer and b/bin/groqqer differ diff --git a/groqqer.go b/groqqer.go index dff8791..87eee60 100644 --- a/groqqer.go +++ b/groqqer.go @@ -1,30 +1,32 @@ -// groqqer: OpenAI-compatible LLM gateway for the Groq Streamlit space -// Reverse engineers https://dromerosm-groq-chatbot.hf.space into an OpenAI proxy +// groqqer: Pure Go OpenAI-compatible LLM gateway for the Groq Streamlit space +// Reverse engineers https://dromerosm-groq-chatbot.hf.space via direct WebSocket Protobuf wire format // Created by Luxferre in 2026, released into the public domain package main import ( "bufio" + "bytes" "context" "crypto/rand" + "crypto/tls" "encoding/base64" "encoding/binary" "encoding/json" + "errors" "flag" "fmt" "io" "log" "net" "net/http" + "net/url" "os" - "os/exec" "os/signal" "regexp" "strconv" "strings" "sync" - "sync/atomic" "syscall" "time" ) @@ -32,8 +34,7 @@ import ( var ( DefaultUserAgent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36" DefaultTargetURL = "https://dromerosm-groq-chatbot.hf.space" - DefaultModel = "llama-3.3-70b-versatile" - cdpCmdCounter int64 + DefaultModel = "qwen/qwen3.6-27b" ) // --------------------------------------------------------------------------- @@ -170,28 +171,6 @@ func GenerateUUID() string { return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]) } -func getFreePort() (string, error) { - listener, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - return "", err - } - defer listener.Close() - return strconv.Itoa(listener.Addr().(*net.TCPAddr).Port), nil -} - -func findFreeXDisplay() string { - for d := 100; d < 200; d++ { - lockFile := fmt.Sprintf("/tmp/.X%d-lock", d) - sockFile := fmt.Sprintf("/tmp/.X11-unix/X%d", d) - if _, err := os.Stat(lockFile); os.IsNotExist(err) { - if _, err2 := os.Stat(sockFile); os.IsNotExist(err2) { - return fmt.Sprintf(":%d", d) - } - } - } - return ":100" -} - // --------------------------------------------------------------------------- // SOCKS5 Proxy Client (RFC 1928 / RFC 1929) // --------------------------------------------------------------------------- @@ -369,6 +348,716 @@ func DialSOCKS5(ctx context.Context, proxyURL, targetAddr string) (net.Conn, err return conn, nil } +// --------------------------------------------------------------------------- +// Pure Go Protobuf Wire Formatter & Decoder +// --------------------------------------------------------------------------- + +func encodeVarint(val uint64) []byte { + var buf []byte + for val >= 0x80 { + buf = append(buf, byte(val|0x80)) + val >>= 7 + } + buf = append(buf, byte(val)) + return buf +} + +func encodeTag(fieldNum int, wireType int) []byte { + return encodeVarint(uint64((fieldNum << 3) | wireType)) +} + +func encodeLengthDelimited(fieldNum int, data []byte) []byte { + tag := encodeTag(fieldNum, 2) + length := encodeVarint(uint64(len(data))) + res := append(tag, length...) + return append(res, data...) +} + +func encodeString(fieldNum int, str string) []byte { + return encodeLengthDelimited(fieldNum, []byte(str)) +} + +type ProtoField struct { + Tag int + WireType int + Varint uint64 + Data []byte +} + +func decodeProtoFields(data []byte) []ProtoField { + var fields []ProtoField + r := bytes.NewReader(data) + for r.Len() > 0 { + rawTag, err := binary.ReadUvarint(r) + if err != nil { + break + } + fieldNum := int(rawTag >> 3) + wireType := int(rawTag & 0x07) + pf := ProtoField{Tag: fieldNum, WireType: wireType} + switch wireType { + case 0: + v, err := binary.ReadUvarint(r) + if err != nil { + return fields + } + pf.Varint = v + case 1: + buf := make([]byte, 8) + if _, err := io.ReadFull(r, buf); err != nil { + return fields + } + pf.Data = buf + case 2: + length, err := binary.ReadUvarint(r) + if err != nil { + return fields + } + buf := make([]byte, length) + if _, err := io.ReadFull(r, buf); err != nil { + return fields + } + pf.Data = buf + case 5: + buf := make([]byte, 4) + if _, err := io.ReadFull(r, buf); err != nil { + return fields + } + pf.Data = buf + default: + return fields + } + fields = append(fields, pf) + } + return fields +} + +type ParsedForwardMsg struct { + IsFinished bool + DeltaPath []uint64 + Markdown string + ChatInputID string + SelectboxID string + Options []string + AlertText string +} + +func parseForwardMsg(data []byte) ParsedForwardMsg { + var res ParsedForwardMsg + fields := decodeProtoFields(data) + for _, f := range fields { + if f.Tag == 6 && (f.Varint == 0 || f.Varint == 1 || f.Varint == 3 || f.Varint == 4 || f.Varint == 5) { + res.IsFinished = true + } + if f.Tag == 2 && f.WireType == 2 { // ForwardMsgMetadata + for _, mf := range decodeProtoFields(f.Data) { + if mf.Tag == 2 { // delta_path + if mf.WireType == 2 { // packed uint32 + r := bytes.NewReader(mf.Data) + for r.Len() > 0 { + v, _ := binary.ReadUvarint(r) + res.DeltaPath = append(res.DeltaPath, v) + } + } else if mf.WireType == 0 { + res.DeltaPath = append(res.DeltaPath, mf.Varint) + } + } + } + } + if f.Tag == 5 && f.WireType == 2 { // Delta + dFields := decodeProtoFields(f.Data) + for _, df := range dFields { + if df.Tag == 3 && df.WireType == 2 { // Element new_element + eFields := decodeProtoFields(df.Data) + for _, ef := range eFields { + if ef.Tag == 29 && ef.WireType == 2 { // Markdown + mFields := decodeProtoFields(ef.Data) + for _, mf := range mFields { + if mf.Tag == 1 && mf.WireType == 2 { + res.Markdown = string(mf.Data) + } + } + } else if ef.Tag == 49 && ef.WireType == 2 { // ChatInput + cFields := decodeProtoFields(ef.Data) + for _, cf := range cFields { + if cf.Tag == 1 && cf.WireType == 2 { + res.ChatInputID = string(cf.Data) + } + } + } else if ef.Tag == 25 && ef.WireType == 2 { // Selectbox + sFields := decodeProtoFields(ef.Data) + for _, sf := range sFields { + if sf.Tag == 1 && sf.WireType == 2 { + res.SelectboxID = string(sf.Data) + } else if sf.Tag == 4 && sf.WireType == 2 { + res.Options = append(res.Options, string(sf.Data)) + } + } + } else if ef.Tag == 30 && ef.WireType == 2 { // Alert + aFields := decodeProtoFields(ef.Data) + for _, af := range aFields { + if af.Tag == 1 && af.WireType == 2 { + res.AlertText = string(af.Data) + } + } + } + } + } + } + } + } + return res +} + +// --------------------------------------------------------------------------- +// RFC 6455 Pure Go WebSocket Client +// --------------------------------------------------------------------------- + +func sendWSBinaryFrame(conn io.Writer, payload []byte) error { + var header []byte + header = append(header, 0x82) // Binary frame (0x02) | FIN (0x80) + + length := len(payload) + maskKey := make([]byte, 4) + _, _ = rand.Read(maskKey) + + if length < 126 { + header = append(header, byte(length|0x80)) + } else if length < 65536 { + header = append(header, 126|0x80) + var b [2]byte + binary.BigEndian.PutUint16(b[:], uint16(length)) + header = append(header, b[:]...) + } else { + header = append(header, 127|0x80) + var b [8]byte + binary.BigEndian.PutUint64(b[:], uint64(length)) + header = append(header, b[:]...) + } + + header = append(header, maskKey...) + masked := make([]byte, length) + for i := 0; i < length; i++ { + masked[i] = payload[i] ^ maskKey[i%4] + } + + _, err := conn.Write(append(header, masked...)) + return err +} + +func sendWSPongFrame(conn io.Writer, payload []byte) error { + var header []byte + header = append(header, 0x8a) // Pong (0x0A) | FIN (0x80) + + length := len(payload) + maskKey := make([]byte, 4) + _, _ = rand.Read(maskKey) + + header = append(header, byte(length|0x80)) + header = append(header, maskKey...) + masked := make([]byte, length) + for i := 0; i < length; i++ { + masked[i] = payload[i] ^ maskKey[i%4] + } + + _, err := conn.Write(append(header, masked...)) + return err +} + +func readWSFrame(conn net.Conn, reader *bufio.Reader, timeout time.Duration) ([]byte, byte, error) { + if timeout > 0 { + _ = conn.SetReadDeadline(time.Now().Add(timeout)) + } else { + _ = conn.SetReadDeadline(time.Time{}) + } + + b1, err := reader.ReadByte() + if err != nil { + return nil, 0, err + } + b2, err := reader.ReadByte() + if err != nil { + return nil, 0, err + } + + opcode := b1 & 0x0f + isMasked := (b2 & 0x80) != 0 + length := int(b2 & 0x7f) + + if length == 126 { + var ext uint16 + if err := binary.Read(reader, binary.BigEndian, &ext); err != nil { + return nil, 0, err + } + length = int(ext) + } else if length == 127 { + var ext uint64 + if err := binary.Read(reader, binary.BigEndian, &ext); err != nil { + return nil, 0, err + } + length = int(ext) + } + + var maskKey [4]byte + if isMasked { + if _, err := io.ReadFull(reader, maskKey[:]); err != nil { + return nil, 0, err + } + } + + payload := make([]byte, length) + if _, err := io.ReadFull(reader, payload); err != nil { + return nil, 0, err + } + + if isMasked { + for i := 0; i < length; i++ { + payload[i] ^= maskKey[i%4] + } + } + + return payload, opcode, nil +} + +// --------------------------------------------------------------------------- +// Streamlit Session & Client +// --------------------------------------------------------------------------- + +type StreamlitSession struct { + conn net.Conn + reader *bufio.Reader + chatInputID string + selectboxID string + options []string + activeModel string + closed bool + mu sync.Mutex +} + +func (s *StreamlitSession) Close() { + s.mu.Lock() + defer s.mu.Unlock() + if !s.closed { + s.closed = true + if s.conn != nil { + _ = s.conn.Close() + } + } +} + +func dialStreamlitWebSocket(ctx context.Context, targetURL, proxyURL, userAgent string) (net.Conn, *bufio.Reader, error) { + u, err := url.Parse(targetURL) + if err != nil { + return nil, nil, fmt.Errorf("invalid target url: %w", err) + } + + host := u.Hostname() + port := u.Port() + isTLS := u.Scheme == "https" || u.Scheme == "wss" + if port == "" { + if isTLS { + port = "443" + } else { + port = "80" + } + } + + targetAddr := net.JoinHostPort(host, port) + var rawConn net.Conn + + if proxyURL != "" { + rawConn, err = DialSOCKS5(ctx, proxyURL, targetAddr) + } else { + var d net.Dialer + rawConn, err = d.DialContext(ctx, "tcp", targetAddr) + } + if err != nil { + return nil, nil, fmt.Errorf("failed to connect to %s: %w", targetAddr, err) + } + + var conn net.Conn = rawConn + if isTLS { + tlsConfig := &tls.Config{ + ServerName: host, + } + tlsConn := tls.Client(rawConn, tlsConfig) + if err := tlsConn.HandshakeContext(ctx); err != nil { + rawConn.Close() + return nil, nil, fmt.Errorf("tls handshake failed: %w", err) + } + conn = tlsConn + } + + keyBytes := make([]byte, 16) + _, _ = rand.Read(keyBytes) + secKey := base64.StdEncoding.EncodeToString(keyBytes) + + originScheme := "https" + if !isTLS { + originScheme = "http" + } + origin := fmt.Sprintf("%s://%s", originScheme, host) + + req := fmt.Sprintf( + "GET /_stcore/stream HTTP/1.1\r\n"+ + "Host: %s\r\n"+ + "Upgrade: websocket\r\n"+ + "Connection: Upgrade\r\n"+ + "Sec-WebSocket-Key: %s\r\n"+ + "Sec-WebSocket-Version: 13\r\n"+ + "Origin: %s\r\n"+ + "User-Agent: %s\r\n\r\n", + host, secKey, origin, userAgent, + ) + + if _, err := conn.Write([]byte(req)); err != nil { + conn.Close() + return nil, nil, fmt.Errorf("failed to write websocket handshake request: %w", err) + } + + reader := bufio.NewReader(conn) + statusLine, err := reader.ReadString('\n') + if err != nil || !strings.Contains(statusLine, "101") { + conn.Close() + return nil, nil, fmt.Errorf("websocket handshake rejected: %s", statusLine) + } + + for { + line, err := reader.ReadString('\n') + if err != nil || strings.TrimSpace(line) == "" { + break + } + } + + return conn, reader, nil +} + +func OpenStreamlitSession(ctx context.Context, targetURL, proxyURL, userAgent string) (*StreamlitSession, error) { + conn, reader, err := dialStreamlitWebSocket(ctx, targetURL, proxyURL, userAgent) + if err != nil { + return nil, err + } + + sess := &StreamlitSession{ + conn: conn, + reader: reader, + } + + // Initial rerun_script BackMsg (ClientState: query_string = "") + clientState := encodeString(1, "") + backMsg := encodeLengthDelimited(11, clientState) + + if err := sendWSBinaryFrame(conn, backMsg); err != nil { + sess.Close() + return nil, fmt.Errorf("failed to send initial rerun BackMsg: %w", err) + } + + // Read frames until script_finished + for { + payload, opcode, err := readWSFrame(conn, reader, 25*time.Second) + if err != nil { + sess.Close() + return nil, fmt.Errorf("error reading initial session frames: %w", err) + } + + if opcode == 0x09 { // Ping + _ = sendWSPongFrame(conn, payload) + continue + } + + parsed := parseForwardMsg(payload) + if parsed.ChatInputID != "" { + sess.chatInputID = parsed.ChatInputID + } + if parsed.SelectboxID != "" { + sess.selectboxID = parsed.SelectboxID + sess.options = parsed.Options + } + if parsed.IsFinished { + break + } + } + + if sess.chatInputID == "" { + // Default fallback chatInput widget ID if not matched in stream + sess.chatInputID = "$$ID-ae617304c8297cf4ddb1a23ee392ce5b-None" + } + if sess.selectboxID == "" { + sess.selectboxID = "$$ID-625687b40312ff8d73403e7858a2f552-chat_model_selector" + } + + return sess, nil +} + +func (s *StreamlitSession) SwitchModel(modelID string) error { + if modelID == "" || s.selectboxID == "" { + return nil + } + + cleanModel := strings.TrimSpace(modelID) + if idx := strings.Index(cleanModel, " ("); idx != -1 { + cleanModel = strings.TrimSpace(cleanModel[:idx]) + } + + // Match model against options + matchedModel := cleanModel + for _, opt := range s.options { + rawOpt := opt + if idx := strings.Index(rawOpt, " ("); idx != -1 { + rawOpt = strings.TrimSpace(rawOpt[:idx]) + } + if strings.EqualFold(rawOpt, cleanModel) || strings.Contains(strings.ToLower(rawOpt), strings.ToLower(cleanModel)) { + matchedModel = rawOpt + break + } + } + + if s.activeModel == matchedModel { + return nil + } + + // Streamlit Selectbox sends string_value (field 6) in WidgetState + sbWidget := append(encodeString(1, s.selectboxID), encodeString(6, matchedModel)...) + wsData := encodeLengthDelimited(1, sbWidget) + wStates := encodeLengthDelimited(2, wsData) + csData := append(encodeString(1, ""), wStates...) + switchBackMsg := encodeLengthDelimited(11, csData) + + if err := sendWSBinaryFrame(s.conn, switchBackMsg); err != nil { + return fmt.Errorf("failed to send model switch BackMsg: %w", err) + } + + for { + payload, opcode, err := readWSFrame(s.conn, s.reader, 25*time.Second) + if err != nil { + return fmt.Errorf("error reading model switch frames: %w", err) + } + if opcode == 0x09 { + _ = sendWSPongFrame(s.conn, payload) + continue + } + parsed := parseForwardMsg(payload) + if parsed.ChatInputID != "" { + s.chatInputID = parsed.ChatInputID + } + if parsed.SelectboxID != "" { + s.selectboxID = parsed.SelectboxID + } + if parsed.IsFinished { + break + } + } + + s.activeModel = matchedModel + return nil +} + +type GatewayAlertError struct { + Alert string + IsRateLimit bool + RetryAfter int +} + +func (e *GatewayAlertError) Error() string { + return e.Alert +} + +func parseRetryAfter(alert string) int { + re := regexp.MustCompile(`(?i)(?:retry\s+after|retry\s+in)\s+(\d+)`) + match := re.FindStringSubmatch(alert) + if len(match) >= 2 { + sec, _ := strconv.Atoi(match[1]) + if sec > 0 { + return sec + } + } + return 6 +} + +func isRateLimit(alert string) bool { + lower := strings.ToLower(alert) + return strings.Contains(lower, "rate limit") || strings.Contains(lower, "too many requests") || strings.Contains(lower, "429") +} + +func resolveAlert(alerts []string) (alertText string, isRL bool, retryAfter int) { + for _, a := range alerts { + if isRateLimit(a) { + return a, true, parseRetryAfter(a) + } + } + + for _, a := range alerts { + lower := strings.ToLower(a) + if !strings.Contains(lower, "the last response did not finish") { + return a, false, 0 + } + } + + if len(alerts) > 0 { + return alerts[len(alerts)-1], false, 0 + } + return "", false, 0 +} + +func isStaticUIMarkdown(text, promptText string) bool { + trimmed := strings.TrimSpace(text) + if trimmed == "" || trimmed == strings.TrimSpace(promptText) { + return true + } + staticPrefixes := []string{ + "Model:", + "Model :", + "- **Model ID:**", + "[Powered by Groq]", + "---", + "-----", + "Text models available", + "This model accepts and returns text", + "This model uses non-streaming responses", + "Select a vision-capable model", + "Attached document:", + "### Upload an Image", + "### Usage Summary", + "### Chat Interface", + "### Upload Context", + "### Comparison Context", + "A document is already attached", + "**Token Usage", + "**Latest API Rate Snapshot:**", + "**Request Max Tokens:**", + "**Conversation Total:**", + "**Important:**", + "Important:", + "Partial response: generation did not finish.", + } + for _, p := range staticPrefixes { + if strings.HasPrefix(trimmed, p) { + return true + } + } + if strings.Contains(trimmed, "**Important:**") || + strings.Contains(trimmed, "Important:") || + strings.Contains(trimmed, "older messages were left out of this request") || + strings.Contains(trimmed, "Response limit adjusted to") { + return true + } + return false +} + +func (s *StreamlitSession) SubmitPrompt(promptText string, onChunk func(delta string), onAlert func(alert string)) error { + // Construct BackMsg with chat_input_value and selected model in selectbox + chatInputVal := encodeString(1, promptText) + ciWidget := append(encodeString(1, s.chatInputID), encodeLengthDelimited(15, chatInputVal)...) + + var wsData []byte + if s.selectboxID != "" && s.activeModel != "" { + sbWidget := append(encodeString(1, s.selectboxID), encodeString(6, s.activeModel)...) + wsData = append(encodeLengthDelimited(1, sbWidget), encodeLengthDelimited(1, ciWidget)...) + } else { + wsData = encodeLengthDelimited(1, ciWidget) + } + + wStates := encodeLengthDelimited(2, wsData) + csData := append(encodeString(1, ""), wStates...) + promptBackMsg := encodeLengthDelimited(11, csData) + + if err := sendWSBinaryFrame(s.conn, promptBackMsg); err != nil { + return fmt.Errorf("failed to send prompt BackMsg: %w", err) + } + + lastLen := 0 + assistantPathKey := "" + userEchoSeen := false + var alerts []string + + for { + payload, opcode, err := readWSFrame(s.conn, s.reader, 90*time.Second) + if err != nil { + return fmt.Errorf("connection closed during response: %w", err) + } + + if opcode == 0x09 { // Ping + _ = sendWSPongFrame(s.conn, payload) + continue + } + + parsed := parseForwardMsg(payload) + + if parsed.AlertText != "" { + alerts = append(alerts, parsed.AlertText) + if onAlert != nil { + onAlert(parsed.AlertText) + } + } + + if parsed.Markdown != "" { + // Ignore sidebar elements (deltaPath[0] == 1) + if len(parsed.DeltaPath) > 0 && parsed.DeltaPath[0] == 1 { + continue + } + + // In this Streamlit space: + // DeltaPath[0] == 0 is the main page. + // DeltaPath[1] <= 5 are the header widgets (0-4) and the user chat echo container (5). + // The assistant response is always DeltaPath[1] >= 6. + if len(parsed.DeltaPath) >= 2 && parsed.DeltaPath[0] == 0 && parsed.DeltaPath[1] <= 5 { + continue + } + + // Metadata or captions inside message containers have index >= 1 (e.g. [0, 6, 1]) + if len(parsed.DeltaPath) >= 3 && parsed.DeltaPath[len(parsed.DeltaPath)-1] != 0 { + continue + } + + // Check if user echo arrived (fallback for environments without hierarchical delta paths) + trimmedMD := strings.TrimSpace(parsed.Markdown) + trimmedPrompt := strings.TrimSpace(promptText) + if !userEchoSeen { + if trimmedMD == trimmedPrompt || + strings.HasPrefix(trimmedPrompt, trimmedMD) || + strings.HasPrefix(trimmedMD, trimmedPrompt) { + userEchoSeen = true + continue + } + } + + // If text is not static UI, stream assistant delta + if !isStaticUIMarkdown(parsed.Markdown, promptText) { + pathKey := fmt.Sprintf("%v", parsed.DeltaPath) + if assistantPathKey == "" { + assistantPathKey = pathKey + } + if pathKey == assistantPathKey { + if len(parsed.Markdown) > lastLen { + delta := parsed.Markdown[lastLen:] + lastLen = len(parsed.Markdown) + if onChunk != nil { + onChunk(delta) + } + } + } + } + } + + if parsed.IsFinished { + break + } + } + + if lastLen == 0 { + if alertText, isRL, retrySec := resolveAlert(alerts); alertText != "" { + return &GatewayAlertError{ + Alert: alertText, + IsRateLimit: isRL, + RetryAfter: retrySec, + } + } + return fmt.Errorf("empty assistant response received from space") + } + + return nil +} + // --------------------------------------------------------------------------- // Tool and Message Processing // --------------------------------------------------------------------------- @@ -381,7 +1070,17 @@ func BuildToolInstruction(tools []Tool, toolChoice interface{}) string { var sb strings.Builder sb.WriteString("\n\n# Tool Calling Instructions\nYou have access to the following functions:\n\n") sb.WriteString(string(toolsBytes)) - sb.WriteString("\n\n\nWhen you need to call a function, respond ONLY with a block formatted exactly as follows:\n\n{\"name\": \"\", \"arguments\": {}}\n\n\nDo not include conversational filler before or after the tool call.") + sb.WriteString("\n\n\n") + sb.WriteString("When you need to call a function, respond ONLY with a block formatted exactly as follows:\n") + sb.WriteString("\n{\"name\": \"\", \"arguments\": {}}\n\n\n") + sb.WriteString("CRITICAL EXECUTION RULES:\n") + sb.WriteString("1. If you invoke a tool, output ONLY the block. Do not write conversational greetings, explanations, or filler text outside the tags.\n") + sb.WriteString("2. Put any reasoning or thought process inside ... tags.\n") + sb.WriteString("3. If multiple tools need to be called, output each inside its own ... block.\n") + sb.WriteString("4. When tool execution results are provided in blocks, inspect the output:\n") + sb.WriteString(" - If further steps or additional tools are needed to fulfill the user's request, emit the next block.\n") + sb.WriteString(" - If all needed information has been retrieved, synthesize the answers and deliver the final response to the user.\n") + sb.WriteString(" - Never stop or terminate the conversation prematurely while intermediate steps remain.\n") if toolChoice != nil { if choiceStr, ok := toolChoice.(string); ok { @@ -400,14 +1099,110 @@ func BuildToolInstruction(tools []Tool, toolChoice interface{}) string { return sb.String() } +func ExtractToolName(fn interface{}) string { + if fn == nil { + return "" + } + if fnMap, ok := fn.(map[string]interface{}); ok { + if name, ok := fnMap["name"].(string); ok { + return strings.TrimSpace(name) + } + } + if tcf, ok := fn.(ToolCallFunction); ok { + return strings.TrimSpace(tcf.Name) + } + b, err := json.Marshal(fn) + if err == nil { + var m map[string]interface{} + if json.Unmarshal(b, &m) == nil { + if name, ok := m["name"].(string); ok { + return strings.TrimSpace(name) + } + } + } + return "" +} + +func GetAllowedToolNames(tools []Tool, toolChoice interface{}) map[string]bool { + if len(tools) == 0 { + return nil + } + + if choiceStr, ok := toolChoice.(string); ok { + if choiceStr == "none" { + return nil + } + } + + if choiceMap, ok := toolChoice.(map[string]interface{}); ok { + if fnMap, ok := choiceMap["function"].(map[string]interface{}); ok { + if name, ok := fnMap["name"].(string); ok && strings.TrimSpace(name) != "" { + return map[string]bool{strings.TrimSpace(name): true} + } + } + } + + allowed := make(map[string]bool) + for _, t := range tools { + name := ExtractToolName(t.Function) + if name != "" { + allowed[name] = true + } + } + + if len(allowed) == 0 { + return nil + } + return allowed +} + +var toolCallTagRegex = regexp.MustCompile(`(?s).*?`) + +func stripToolCallTags(s string) string { + return strings.TrimSpace(toolCallTagRegex.ReplaceAllString(s, "")) +} + func FormatPrompt(req ChatCompletionRequest) string { toolInstruction := BuildToolInstruction(req.Tools, req.ToolChoice) + // Build mapping from tool_call_id to function name + toolCallIDToName := make(map[string]string) + for _, msg := range req.Messages { + if msg.Role == "assistant" { + for _, tc := range msg.ToolCalls { + if tc.ID != "" && tc.Function.Name != "" { + toolCallIDToName[tc.ID] = tc.Function.Name + } + } + } + } + + // Find where trailing tool responses start (if any) + firstTrailingToolIdx := len(req.Messages) + for i := len(req.Messages) - 1; i >= 0; i-- { + role := req.Messages[i].Role + if role == "tool" || role == "function" { + firstTrailingToolIdx = i + } else { + break + } + } + var systemInstructions []string var historyTurns []string var currentTurn string - for i, msg := range req.Messages { + // Process messages before the trailing tool responses + endHistoryIdx := firstTrailingToolIdx + if firstTrailingToolIdx == len(req.Messages) && len(req.Messages) > 0 { + if req.Messages[len(req.Messages)-1].Role == "user" { + endHistoryIdx = len(req.Messages) - 1 + currentTurn = req.Messages[len(req.Messages)-1].GetContentString() + } + } + + for i := 0; i < endHistoryIdx; i++ { + msg := req.Messages[i] contentStr := msg.GetContentString() switch msg.Role { case "system": @@ -415,9 +1210,16 @@ func FormatPrompt(req ChatCompletionRequest) string { systemInstructions = append(systemInstructions, contentStr) } case "assistant": + cleanAssistant := strings.TrimSpace(contentStr) + if strings.HasPrefix(cleanAssistant, "Model:") || strings.HasPrefix(cleanAssistant, "Model :") { + cleanAssistant = "" + } + if len(msg.ToolCalls) > 0 { + cleanAssistant = stripToolCallTags(cleanAssistant) + } var sb strings.Builder - if contentStr != "" { - sb.WriteString(contentStr) + if cleanAssistant != "" { + sb.WriteString(cleanAssistant) } for _, tc := range msg.ToolCalls { if sb.Len() > 0 { @@ -429,10 +1231,14 @@ func FormatPrompt(req ChatCompletionRequest) string { } sb.WriteString(fmt.Sprintf("\n{\"name\": %q, \"arguments\": %s}\n", tc.Function.Name, args)) } - historyTurns = append(historyTurns, "Assistant: "+sb.String()) + if sb.Len() > 0 { + historyTurns = append(historyTurns, "Assistant: "+sb.String()) + } case "tool", "function": toolName := msg.Name - if toolName == "" { + if resolved, ok := toolCallIDToName[msg.ToolCallID]; ok && resolved != "" { + toolName = resolved + } else if toolName == "" { toolName = msg.ToolCallID } var contentJSON []byte @@ -441,21 +1247,88 @@ func FormatPrompt(req ChatCompletionRequest) string { } else { contentJSON, _ = json.Marshal(contentStr) } - turnText := fmt.Sprintf("\n{\"name\": %q, \"content\": %s}\n", toolName, string(contentJSON)) - if i == len(req.Messages)-1 { - currentTurn = turnText - } else { - historyTurns = append(historyTurns, "Tool Result: "+turnText) + jsonStr := string(contentJSON) + if len(jsonStr) > 2000 { + jsonStr = jsonStr[:1800] + `"... [truncated]"` } + var turnText string + if msg.ToolCallID != "" { + turnText = fmt.Sprintf("\n{\"name\": %q, \"tool_call_id\": %q, \"content\": %s}\n", toolName, msg.ToolCallID, jsonStr) + } else { + turnText = fmt.Sprintf("\n{\"name\": %q, \"content\": %s}\n", toolName, jsonStr) + } + historyTurns = append(historyTurns, fmt.Sprintf("Tool Result (%s): %s", toolName, turnText)) case "user": - if i == len(req.Messages)-1 { - currentTurn = contentStr - } else { - historyTurns = append(historyTurns, "User: "+contentStr) - } + historyTurns = append(historyTurns, "User: "+contentStr) } } + // If we have trailing tool responses, group them together with a continuation directive + if firstTrailingToolIdx < len(req.Messages) { + var toolResSb strings.Builder + toolResSb.WriteString("[Tool Execution Results]\n") + for i := firstTrailingToolIdx; i < len(req.Messages); i++ { + msg := req.Messages[i] + toolName := msg.Name + if resolved, ok := toolCallIDToName[msg.ToolCallID]; ok && resolved != "" { + toolName = resolved + } else if toolName == "" { + toolName = msg.ToolCallID + } + contentStr := msg.GetContentString() + var contentJSON []byte + if json.Valid([]byte(contentStr)) { + contentJSON = []byte(contentStr) + } else { + contentJSON, _ = json.Marshal(contentStr) + } + jsonStr := string(contentJSON) + if len(jsonStr) > 3500 { + jsonStr = jsonStr[:3200] + `"... [truncated]"` + } + + toolResSb.WriteString("\n") + if msg.ToolCallID != "" { + toolResSb.WriteString(fmt.Sprintf("{\"name\": %q, \"tool_call_id\": %q, \"content\": %s}\n", toolName, msg.ToolCallID, jsonStr)) + } else { + toolResSb.WriteString(fmt.Sprintf("{\"name\": %q, \"content\": %s}\n", toolName, jsonStr)) + } + toolResSb.WriteString("\n") + } + toolResSb.WriteString("\n[Next Steps Directive]\n") + toolResSb.WriteString("You have received the results of the tool execution(s) above.\n") + toolResSb.WriteString("- Analyze these results in the context of the user request and conversation history.\n") + toolResSb.WriteString("- If additional tool calls are needed to complete the task, emit the next block immediately.\n") + toolResSb.WriteString("- If all information needed to fulfill the request is now available, provide the final comprehensive response to the user.\n") + toolResSb.WriteString("- Do not terminate or stop without answering the user or taking the next step.") + currentTurn = toolResSb.String() + } + + // Prune history turns if history exceeds budget (~12000 chars) + totalHistLen := 0 + for _, t := range historyTurns { + totalHistLen += len(t) + } + if totalHistLen > 12000 && len(historyTurns) > 3 { + firstTurn := historyTurns[0] + var pruned []string + pruned = append(pruned, firstTurn) + pruned = append(pruned, "[... earlier conversation turns omitted for brevity ...]") + budget := 10000 + var recent []string + currentBudget := 0 + for i := len(historyTurns) - 1; i >= 1; i-- { + tLen := len(historyTurns[i]) + if currentBudget+tLen > budget { + break + } + currentBudget += tLen + recent = append([]string{historyTurns[i]}, recent...) + } + pruned = append(pruned, recent...) + historyTurns = pruned + } + var promptBuilder strings.Builder if len(systemInstructions) > 0 || toolInstruction != "" { @@ -481,10 +1354,14 @@ func FormatPrompt(req ChatCompletionRequest) string { } if currentTurn != "" { - if len(historyTurns) > 0 || len(systemInstructions) > 0 { - promptBuilder.WriteString("[User]\n") + if strings.HasPrefix(currentTurn, "[Tool Execution Results]") { + promptBuilder.WriteString(currentTurn) + } else { + if len(historyTurns) > 0 || len(systemInstructions) > 0 { + promptBuilder.WriteString("[User]\n") + } + promptBuilder.WriteString(currentTurn) } - promptBuilder.WriteString(currentTurn) } return strings.TrimSpace(promptBuilder.String()) @@ -697,7 +1574,40 @@ func parseSingleToolCall(jsonStr string) (ToolCall, bool) { return repairToolCallJSON(cleaned) } -func parseXMLToolCall(block string) (ToolCall, bool) { +func parseToolCallsFromText(input string) ([]ToolCall, bool) { + cleaned := cleanJSONBlock(strings.TrimSpace(input)) + if cleaned == "" { + return nil, false + } + + // Check if it is a JSON array of tool calls: [...] + if strings.HasPrefix(cleaned, "[") && strings.HasSuffix(cleaned, "]") { + var rawSlice []interface{} + if err := json.Unmarshal([]byte(cleaned), &rawSlice); err == nil && len(rawSlice) > 0 { + var calls []ToolCall + for _, item := range rawSlice { + b, _ := json.Marshal(item) + if tc, ok := parseSingleToolCall(string(b)); ok { + calls = append(calls, tc) + } else if tc, ok := repairToolCallJSON(string(b)); ok { + calls = append(calls, tc) + } + } + if len(calls) > 0 { + return calls, true + } + } + } + + // Check single tool call + if tc, ok := parseSingleToolCall(cleaned); ok { + return []ToolCall{tc}, true + } + + return nil, false +} + +func parseXMLToolCall(block string) ([]ToolCall, bool) { inner := strings.TrimSpace(block) if strings.HasPrefix(inner, "") { inner = strings.TrimPrefix(inner, "") @@ -707,8 +1617,8 @@ func parseXMLToolCall(block string) (ToolCall, bool) { } inner = cleanJSONBlock(inner) - if tc, ok := parseSingleToolCall(inner); ok { - return tc, true + if calls, ok := parseToolCallsFromText(inner); ok { + return calls, true } var fnName string @@ -733,110 +1643,169 @@ func parseXMLToolCall(block string) (ToolCall, bool) { if argsStr == "" { argsStr = "{}" } - return ToolCall{ - ID: "call_" + GenerateUUID()[:8], - Type: "function", - Function: ToolCallFunction{ - Name: fnName, - Arguments: argsStr, + return []ToolCall{ + { + ID: "call_" + GenerateUUID()[:8], + Type: "function", + Function: ToolCallFunction{ + Name: fnName, + Arguments: argsStr, + }, }, }, true } - return ToolCall{}, false + return nil, false } func ExtractToolCallBlocks(content string) (blocks []string, remaining string) { + var b []string s := content - remaining = content + searchOffset := 0 + var remainingBuilder strings.Builder + lastPos := 0 - for strings.Contains(s, "") { - sIdx := strings.Index(s, "") - rest := s[sIdx+len(""):] + for { + relStart := strings.Index(s[searchOffset:], "") + if relStart == -1 { + break + } + startIdx := searchOffset + relStart + afterStart := startIdx + len("") - relNextSIdx := strings.Index(rest, "") - var nextSIdx int - if relNextSIdx != -1 { - nextSIdx = sIdx + len("") + relNextSIdx - } else { - nextSIdx = -1 + relEnd := strings.Index(s[afterStart:], "") + if relEnd == -1 { + // Unclosed , do not treat as block + searchOffset = afterStart + continue } - relEIdx := strings.Index(rest, "") - var eIdx int - if relEIdx != -1 { - eIdx = sIdx + len("") + relEIdx - } else { - eIdx = -1 + if innerNext := strings.Index(s[afterStart:afterStart+relEnd], ""); innerNext != -1 { + searchOffset = afterStart + innerNext + continue } - var blockText string - var blockEndPos int + endIdx := afterStart + relEnd + len("") + blockText := s[startIdx:endIdx] + b = append(b, blockText) - if eIdx != -1 && (nextSIdx == -1 || eIdx < nextSIdx) { - blockEndPos = eIdx + len("") - blockText = s[sIdx:blockEndPos] - s = s[blockEndPos:] - } else if nextSIdx != -1 { - blockEndPos = nextSIdx - blockText = s[sIdx:blockEndPos] - s = s[blockEndPos:] - } else { - blockText = s[sIdx:] - s = "" + if startIdx > lastPos { + remainingBuilder.WriteString(s[lastPos:startIdx]) } - - blocks = append(blocks, blockText) + lastPos = endIdx + searchOffset = endIdx } - for strings.Contains(remaining, "") { - st := strings.Index(remaining, "") - rest := remaining[st+len(""):] - - relNext := strings.Index(rest, "") - var nextSt int - if relNext != -1 { - nextSt = st + len("") + relNext - } else { - nextSt = -1 - } - - relEn := strings.Index(rest, "") - var en int - if relEn != -1 { - en = st + len("") + relEn - } else { - en = -1 - } - - if en != -1 && (nextSt == -1 || en < nextSt) { - remaining = strings.TrimSpace(remaining[:st] + remaining[en+len(""):]) - } else if nextSt != -1 { - remaining = strings.TrimSpace(remaining[:st] + remaining[nextSt:]) - } else { - remaining = strings.TrimSpace(remaining[:st]) - } + if lastPos < len(s) { + remainingBuilder.WriteString(s[lastPos:]) } - return blocks, remaining + return b, strings.TrimSpace(remainingBuilder.String()) } -func DetectToolCalls(content string) ([]ToolCall, string, bool) { - blocks, remaining := ExtractToolCallBlocks(content) - var calls []ToolCall +func DetectToolCalls(content string, allowedTools map[string]bool) ([]ToolCall, string, bool) { + if len(allowedTools) == 0 { + return nil, content, false + } - for _, block := range blocks { - if toolCall, ok := parseXMLToolCall(block); ok { - calls = append(calls, toolCall) + type validBlock struct { + startIdx int + endIdx int + calls []ToolCall + } + + var validBlocks []validBlock + searchOffset := 0 + + for { + relStart := strings.Index(content[searchOffset:], "") + if relStart == -1 { + break } + startIdx := searchOffset + relStart + afterStart := startIdx + len("") + + relEnd := strings.Index(content[afterStart:], "") + if relEnd == -1 { + // Unclosed tag - leave as plain content + searchOffset = afterStart + continue + } + + if innerNext := strings.Index(content[afterStart:afterStart+relEnd], ""); innerNext != -1 { + searchOffset = afterStart + innerNext + continue + } + + endIdx := afterStart + relEnd + len("") + blockText := content[startIdx:endIdx] + + if parsedCalls, ok := parseXMLToolCall(blockText); ok { + var callsForBlock []ToolCall + for _, c := range parsedCalls { + if allowedTools[c.Function.Name] { + callsForBlock = append(callsForBlock, c) + } + } + if len(callsForBlock) > 0 { + validBlocks = append(validBlocks, validBlock{ + startIdx: startIdx, + endIdx: endIdx, + calls: callsForBlock, + }) + searchOffset = endIdx + continue + } + } + + // Not a valid tool call or function name not allowed - keep in content + searchOffset = endIdx } - if len(calls) > 0 { - return calls, remaining, true + if len(validBlocks) > 0 { + var allCalls []ToolCall + var remainingBuilder strings.Builder + lastPos := 0 + for _, vb := range validBlocks { + if vb.startIdx > lastPos { + remainingBuilder.WriteString(content[lastPos:vb.startIdx]) + } + allCalls = append(allCalls, vb.calls...) + lastPos = vb.endIdx + } + if lastPos < len(content) { + remainingBuilder.WriteString(content[lastPos:]) + } + + for i := range allCalls { + idx := i + allCalls[i].Index = &idx + } + + cleanRemaining := strings.TrimSpace(remainingBuilder.String()) + return allCalls, cleanRemaining, true } - if tc, ok := parseSingleToolCall(strings.TrimSpace(content)); ok { - return []ToolCall{tc}, "", true + // Direct JSON fallback (only if every parsed call is present in allowedTools) + trimmed := strings.TrimSpace(content) + if (strings.HasPrefix(trimmed, "{") && strings.HasSuffix(trimmed, "}")) || + (strings.HasPrefix(trimmed, "[") && strings.HasSuffix(trimmed, "]")) || + (strings.HasPrefix(trimmed, "```json") && strings.HasSuffix(trimmed, "```")) { + if directCalls, ok := parseToolCallsFromText(trimmed); ok { + var validCalls []ToolCall + for _, c := range directCalls { + if allowedTools[c.Function.Name] { + validCalls = append(validCalls, c) + } + } + if len(validCalls) > 0 && len(validCalls) == len(directCalls) { + for i := range validCalls { + idx := i + validCalls[i].Index = &idx + } + return validCalls, "", true + } + } } return nil, content, false @@ -946,18 +1915,48 @@ func (f *StreamThinkingFilter) Flush(onContent func(string), onReasoning func(st } type StreamToolCallFilter struct { - inToolCall bool - buf string - toolCallBuf string - toolIndex int - emittedCall bool + allowedTools map[string]bool + inToolCall bool + buf string + toolCallBuf string + toolIndex int + emittedCall bool + preambleBuf string } -func NewStreamToolCallFilter() *StreamToolCallFilter { - return &StreamToolCallFilter{} +func NewStreamToolCallFilter(allowedTools map[string]bool) *StreamToolCallFilter { + return &StreamToolCallFilter{ + allowedTools: allowedTools, + } } -func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onToolCall func(ToolCall)) { +func (f *StreamToolCallFilter) appendPreambleOrContent(text string, onContent func(string)) { + if text == "" { + return + } + if f.emittedCall { + onContent(text) + return + } + if len(f.preambleBuf)+len(text) < 128 && !strings.Contains(f.preambleBuf+text, "\n\n") { + f.preambleBuf += text + } else { + if f.preambleBuf != "" { + onContent(f.preambleBuf) + f.preambleBuf = "" + } + onContent(text) + } +} + +func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onReasoning func(string), onToolCall func(ToolCall)) { + if len(f.allowedTools) == 0 { + if chunk != "" { + onContent(chunk) + } + return + } + f.buf += chunk toolStartTag := "" toolEndTag := "" @@ -969,21 +1968,18 @@ func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onTool if !f.inToolCall { if idx := strings.Index(f.buf, toolStartTag); idx != -1 { before := f.buf[:idx] - if before != "" { - onContent(before) - } f.inToolCall = true f.buf = f.buf[idx+len(toolStartTag):] + f.preambleBuf += before } else if matchLen := hasPrefixOf(f.buf, startPrefixes); matchLen > 0 { safe := f.buf[:len(f.buf)-matchLen] - if safe != "" { - onContent(safe) - } f.buf = f.buf[len(f.buf)-matchLen:] + f.appendPreambleOrContent(safe, onContent) break } else { - onContent(f.buf) + safe := f.buf f.buf = "" + f.appendPreambleOrContent(safe, onContent) break } } else { @@ -992,20 +1988,36 @@ func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onTool f.buf = f.buf[idx+len(toolEndTag):] f.inToolCall = false - if tc, ok := parseSingleToolCall(f.toolCallBuf); ok { - idxCopy := f.toolIndex - tc.Index = &idxCopy - f.toolIndex++ - f.emittedCall = true - onToolCall(tc) - } else if tc2, ok2 := parseXMLToolCall("" + f.toolCallBuf + ""); ok2 { - idxCopy := f.toolIndex - tc2.Index = &idxCopy - f.toolIndex++ - f.emittedCall = true - onToolCall(tc2) + rawBlock := "" + f.toolCallBuf + "" + tcs, ok := parseXMLToolCall(rawBlock) + var validCalls []ToolCall + if ok { + for _, tc := range tcs { + if f.allowedTools[tc.Function.Name] { + validCalls = append(validCalls, tc) + } + } + } + + if len(validCalls) > 0 { + if strings.TrimSpace(f.preambleBuf) != "" { + onReasoning(f.preambleBuf) + } + f.preambleBuf = "" + for _, tc := range validCalls { + idxCopy := f.toolIndex + tc.Index = &idxCopy + f.toolIndex++ + f.emittedCall = true + onToolCall(tc) + } } else { - onContent("" + f.toolCallBuf + "") + // Tool call invalid or function name not allowed - flush as content + if f.preambleBuf != "" { + onContent(f.preambleBuf) + f.preambleBuf = "" + } + onContent(rawBlock) } f.toolCallBuf = "" } else if matchLen := hasPrefixOf(f.buf, endPrefixes); matchLen > 0 { @@ -1022,22 +2034,60 @@ func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onTool } } -func (f *StreamToolCallFilter) Flush(onContent func(string), onToolCall func(ToolCall)) { +func (f *StreamToolCallFilter) Flush(onContent func(string), onReasoning func(string), onToolCall func(ToolCall)) { + if len(f.allowedTools) == 0 { + if f.buf != "" { + onContent(f.buf) + f.buf = "" + } + return + } + if f.inToolCall && len(f.toolCallBuf) > 0 { - if tc, ok := parseSingleToolCall(f.toolCallBuf); ok { - idxCopy := f.toolIndex - tc.Index = &idxCopy - f.emittedCall = true - onToolCall(tc) - } else if tc2, ok2 := parseXMLToolCall("" + f.toolCallBuf + ""); ok2 { - idxCopy := f.toolIndex - tc2.Index = &idxCopy - f.emittedCall = true - onToolCall(tc2) + rawBlock := "" + f.toolCallBuf + "" + tcs, ok := parseXMLToolCall(rawBlock) + var validCalls []ToolCall + if ok { + for _, tc := range tcs { + if f.allowedTools[tc.Function.Name] { + validCalls = append(validCalls, tc) + } + } + } + + if len(validCalls) > 0 { + if strings.TrimSpace(f.preambleBuf) != "" { + onReasoning(f.preambleBuf) + } + f.preambleBuf = "" + for _, tc := range validCalls { + idxCopy := f.toolIndex + tc.Index = &idxCopy + f.toolIndex++ + f.emittedCall = true + onToolCall(tc) + } } else { + if f.preambleBuf != "" { + onContent(f.preambleBuf) + f.preambleBuf = "" + } onContent("" + f.toolCallBuf) } f.toolCallBuf = "" + f.inToolCall = false + } + + if f.emittedCall { + if strings.TrimSpace(f.preambleBuf) != "" { + onReasoning(f.preambleBuf) + } + f.preambleBuf = "" + } else { + if f.preambleBuf != "" { + onContent(f.preambleBuf) + f.preambleBuf = "" + } } if len(f.buf) > 0 { onContent(f.buf) @@ -1114,7 +2164,11 @@ func WriteCompletionResponse(w http.ResponseWriter, completionID string, created finishReason = "stop" } var contentVal interface{} = content - if len(toolCalls) > 0 && strings.TrimSpace(content) == "" { + if len(toolCalls) > 0 { + finishReason = "tool_calls" + if reasoning == "" && strings.TrimSpace(content) != "" { + reasoning = strings.TrimSpace(content) + } contentVal = nil } @@ -1145,664 +2199,28 @@ func WriteCompletionResponse(w http.ResponseWriter, completionID string, created json.NewEncoder(w).Encode(resp) } -// --------------------------------------------------------------------------- -// Pure Go Chromium CDP Engine & BrowserBridge -// --------------------------------------------------------------------------- - -func findChromiumBinary(customPath string) string { - if customPath != "" { - if _, err := exec.LookPath(customPath); err == nil { - return customPath - } - } - candidates := []string{"chromium", "google-chrome", "chromium-browser", "chrome"} - for _, c := range candidates { - if p, err := exec.LookPath(c); err == nil { - return p - } - } - return "" +type APIErrorDetail struct { + Message string `json:"message"` + Type string `json:"type"` + Param *string `json:"param"` + Code int `json:"code"` } -func dialCDPWebSocket(wsURL string) (net.Conn, *bufio.Reader, error) { - parts := strings.TrimPrefix(wsURL, "ws://") - slashIdx := strings.Index(parts, "/") - if slashIdx == -1 { - return nil, nil, fmt.Errorf("invalid websocket url: %s", wsURL) - } - host := parts[:slashIdx] - path := parts[slashIdx:] - - conn, err := net.DialTimeout("tcp", host, 5*time.Second) - if err != nil { - return nil, nil, err - } - - nonce := make([]byte, 16) - _, _ = rand.Read(nonce) - secKey := base64.StdEncoding.EncodeToString(nonce) - - req := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n", path, host, secKey) - if _, err := conn.Write([]byte(req)); err != nil { - conn.Close() - return nil, nil, err - } - - reader := bufio.NewReader(conn) - statusLine, err := reader.ReadString('\n') - if err != nil || !strings.Contains(statusLine, "101") { - conn.Close() - return nil, nil, fmt.Errorf("websocket upgrade failed: %s", statusLine) - } - - for { - line, err := reader.ReadString('\n') - if err != nil || strings.TrimSpace(line) == "" { - break - } - } - return conn, reader, nil +type APIErrorResponse struct { + Error APIErrorDetail `json:"error"` } -func sendWSFrame(conn net.Conn, payload []byte) error { - var header []byte - header = append(header, 0x81) - - length := len(payload) - var maskKey [4]byte - _, _ = rand.Read(maskKey[:]) - - if length < 126 { - header = append(header, byte(length)|0x80) - } else if length < 65536 { - header = append(header, 126|0x80) - header = append(header, byte(length>>8), byte(length&0xff)) - } else { - header = append(header, 127|0x80) - for i := 7; i >= 0; i-- { - header = append(header, byte((length>>(i*8))&0xff)) - } - } - - header = append(header, maskKey[:]...) - masked := make([]byte, length) - for i := 0; i < length; i++ { - masked[i] = payload[i] ^ maskKey[i%4] - } - - _, err := conn.Write(append(header, masked...)) - return err -} - -func readWSFrame(conn net.Conn, reader *bufio.Reader) ([]byte, error) { - conn.SetReadDeadline(time.Now().Add(35 * time.Second)) - b1, err := reader.ReadByte() - if err != nil { - return nil, err - } - opcode := b1 & 0x0f - if opcode == 0x08 { - return nil, fmt.Errorf("websocket closed by server") - } - - b2, err := reader.ReadByte() - if err != nil { - return nil, err - } - - isMasked := (b2 & 0x80) != 0 - length := int(b2 & 0x7f) - - if length == 126 { - var extLen uint16 - if err := binary.Read(reader, binary.BigEndian, &extLen); err != nil { - return nil, err - } - length = int(extLen) - } else if length == 127 { - var extLen uint64 - if err := binary.Read(reader, binary.BigEndian, &extLen); err != nil { - return nil, err - } - length = int(extLen) - } - - var maskKey [4]byte - if isMasked { - if _, err := io.ReadFull(reader, maskKey[:]); err != nil { - return nil, err - } - } - - payload := make([]byte, length) - if _, err := io.ReadFull(reader, payload); err != nil { - return nil, err - } - - if isMasked { - for i := 0; i < length; i++ { - payload[i] ^= maskKey[i%4] - } - } - - return payload, nil -} - -func sendCDPCommand(conn net.Conn, reader *bufio.Reader, method string, params map[string]interface{}) (map[string]interface{}, error) { - cmdID := int(atomic.AddInt64(&cdpCmdCounter, 1)) - msg := map[string]interface{}{ - "id": cmdID, - "method": method, - "params": params, - } - b, _ := json.Marshal(msg) - if err := sendWSFrame(conn, b); err != nil { - return nil, err - } - - for { - frame, err := readWSFrame(conn, reader) - if err != nil { - return nil, err - } - var res map[string]interface{} - if err := json.Unmarshal(frame, &res); err == nil { - if idVal, ok := res["id"].(float64); ok && int(idVal) == cmdID { - return res, nil - } - } - } -} - -type BrowserBridge struct { - mu sync.Mutex - cmd *exec.Cmd - xvfbCmd *exec.Cmd - tmpDir string - port string - conn net.Conn - reader *bufio.Reader - browserBin string - userAgent string - headless bool - useXvfb bool - socksProxy string - targetURL string - availableModels []string - selectedModel string -} - -func NewBrowserBridge(browserBin, userAgent string, headless, useXvfb bool, socksProxy, targetURL string) *BrowserBridge { - if userAgent == "" { - userAgent = DefaultUserAgent - } - if targetURL == "" { - targetURL = DefaultTargetURL - } - return &BrowserBridge{ - browserBin: browserBin, - userAgent: userAgent, - headless: headless, - useXvfb: useXvfb, - socksProxy: socksProxy, - targetURL: targetURL, - } -} - -func (bb *BrowserBridge) Start() error { - bb.mu.Lock() - defer bb.mu.Unlock() - - if bb.conn != nil { - return nil - } - - binPath := findChromiumBinary(bb.browserBin) - if binPath == "" { - return fmt.Errorf("no chromium or google-chrome binary found on system") - } - - var effectiveDisplay string - if bb.useXvfb { - xvfbPath, err := exec.LookPath("Xvfb") - if err != nil { - xvfbPath, err = exec.LookPath("Xfbdev") - } - if err != nil { - log.Printf("warning: Xvfb not found on system, falling back to offscreen flags") - } else { - vDisplay := findFreeXDisplay() - xCmd := exec.Command(xvfbPath, vDisplay, "-screen", "0", "1280x800x24", "-ac", "-nolisten", "tcp") - if err := xCmd.Start(); err == nil { - bb.xvfbCmd = xCmd - effectiveDisplay = vDisplay - time.Sleep(300 * time.Millisecond) - log.Printf("Virtual X server started on display %s", vDisplay) - } else { - log.Printf("warning: failed to start virtual X server: %v", err) - } - } - } - - if effectiveDisplay == "" { - effectiveDisplay = os.Getenv("DISPLAY") - } - - tmpDir, err := os.MkdirTemp("", "groqqer_bridge_*") - if err != nil { - return fmt.Errorf("failed to create temp profile dir: %w", err) - } - bb.tmpDir = tmpDir - - port, err := getFreePort() - if err != nil { - port = "9559" - } - bb.port = port - - args := []string{ - "--remote-debugging-port=" + port, - "--user-data-dir=" + tmpDir, - "--disable-gpu", - "--no-sandbox", - "--no-first-run", - "--no-default-browser-check", - "--disable-extensions", - "--disable-default-apps", - "--disable-background-networking", - "--disable-sync", - "--disable-translate", - "--mute-audio", - "--hide-scrollbars", - "--window-size=1280,800", - "--user-agent=" + bb.userAgent, - } - - if bb.xvfbCmd != nil { - args = append(args, "--window-position=0,0") - } else { - args = append(args, "--window-position=-3000,-3000") - } - - if bb.headless && bb.xvfbCmd == nil { - args = append(args, "--headless=new") - } - - if bb.socksProxy != "" { - proxyArg := bb.socksProxy - if !strings.HasPrefix(proxyArg, "socks5://") && !strings.HasPrefix(proxyArg, "socks5h://") { - proxyArg = "socks5://" + proxyArg - } - args = append(args, "--proxy-server="+proxyArg) - } - - args = append(args, bb.targetURL) - - cmd := exec.Command(binPath, args...) - if effectiveDisplay != "" { - cmd.Env = append(os.Environ(), "DISPLAY="+effectiveDisplay) - } else { - cmd.Env = os.Environ() - } - - if err := cmd.Start(); err != nil { - return fmt.Errorf("failed to launch chromium: %w", err) - } - bb.cmd = cmd - - // Connect to CDP target - var pageWS string - for attempt := 0; attempt < 25; attempt++ { - time.Sleep(500 * time.Millisecond) - resp, err := http.Get("http://127.0.0.1:" + port + "/json/list") - if err != nil { - continue - } - var targets []map[string]interface{} - if err := json.NewDecoder(resp.Body).Decode(&targets); err == nil { - resp.Body.Close() - for _, t := range targets { - if t["type"] == "page" { - if ws, ok := t["webSocketDebuggerUrl"].(string); ok && ws != "" { - pageWS = ws - break - } - } - } - if pageWS != "" { - break - } - } else { - resp.Body.Close() - } - } - - if pageWS == "" { - return fmt.Errorf("timed out waiting for chromium CDP page target") - } - - conn, reader, err := dialCDPWebSocket(pageWS) - if err != nil { - return fmt.Errorf("failed to dial CDP websocket: %w", err) - } - bb.conn = conn - bb.reader = reader - - log.Printf("Browser bridge connected to %s", bb.targetURL) - return nil -} - -func (bb *BrowserBridge) Evaluate(js string) (interface{}, error) { - res, err := sendCDPCommand(bb.conn, bb.reader, "Runtime.evaluate", map[string]interface{}{ - "expression": js, - "returnByValue": true, - "awaitPromise": true, - }) - if err != nil { - return nil, err - } - if r, ok := res["result"].(map[string]interface{}); ok { - if r2, ok := r["result"].(map[string]interface{}); ok { - return r2["value"], nil - } - } - return nil, nil -} - -func (bb *BrowserBridge) WaitForReady(timeout time.Duration) error { - deadline := time.Now().Add(timeout) - for time.Now().Before(deadline) { - val, err := bb.Evaluate(`(() => { - const sb = document.querySelector('[data-testid="stSelectbox"]'); - const ta = document.querySelector('textarea'); - return sb && ta && !ta.disabled ? true : false; - })()`) - if err == nil && val == true { - return nil - } - time.Sleep(300 * time.Millisecond) - } - return fmt.Errorf("timed out waiting for Streamlit chat interface to render") -} - -func (bb *BrowserBridge) DiscoverModels() ([]string, error) { - if err := bb.WaitForReady(20 * time.Second); err != nil { - return nil, err - } - - res, err := bb.Evaluate(`(() => { - const listbox = document.querySelector('[role="listbox"]'); - if (!listbox) { - const btn = document.querySelector('[data-testid="stSelectbox"] button[aria-label="Open"]'); - if (btn) btn.click(); - } - return new Promise((resolve) => { - const start = Date.now(); - const interval = setInterval(() => { - const options = Array.from(document.querySelectorAll('[role="option"]')).map(el => el.innerText.trim()).filter(Boolean); - if (options.length > 0 || Date.now() - start > 3000) { - clearInterval(interval); - const btn = document.querySelector('[data-testid="stSelectbox"] button[aria-label="Open"]') || document.body; - btn.click(); - resolve(options); - } - }, 50); - }); - })()`) - if err != nil { - return nil, err - } - - var rawList []string - if arr, ok := res.([]interface{}); ok { - for _, item := range arr { - if str, ok := item.(string); ok && str != "" { - rawList = append(rawList, str) - } - } - } - - var modelIDs []string - for _, item := range rawList { - cleanID := item - if idx := strings.Index(item, " ("); idx != -1 { - cleanID = strings.TrimSpace(item[:idx]) - } - if cleanID != "" { - modelIDs = append(modelIDs, cleanID) - } - } - - curr, _ := bb.Evaluate(`(() => { - const input = document.querySelector('[data-testid="stSelectbox"] input'); - return input ? input.value : ''; - })()`) - if currStr, ok := curr.(string); ok && currStr != "" { - cleanCurr := currStr - if idx := strings.Index(currStr, " ("); idx != -1 { - cleanCurr = strings.TrimSpace(currStr[:idx]) - } - bb.selectedModel = cleanCurr - } - - if len(modelIDs) > 0 { - bb.availableModels = modelIDs - } - - return modelIDs, nil -} - -func (bb *BrowserBridge) SwitchModel(modelID string) error { - if modelID == "" { - return nil - } - - targetSubstr := strings.ToLower(strings.TrimSpace(modelID)) - if strings.Contains(targetSubstr, "/") { - parts := strings.Split(targetSubstr, "/") - targetSubstr = parts[len(parts)-1] - } - - targetJSON, _ := json.Marshal(targetSubstr) - js := fmt.Sprintf(`(() => { - const target = %s; - const input = document.querySelector('[data-testid="stSelectbox"] input'); - if (input && input.value.toLowerCase().includes(target)) { - return Promise.resolve({ success: true, model: input.value }); - } - - const listbox = document.querySelector('[role="listbox"]'); - if (!listbox) { - const btn = document.querySelector('[data-testid="stSelectbox"] button[aria-label="Open"]') || input; - if (btn) btn.click(); - } - - return new Promise((resolve) => { - const start = Date.now(); - const interval = setInterval(() => { - const options = Array.from(document.querySelectorAll('[role="option"]')); - if (options.length > 0) { - const matched = options.find(o => o.innerText.toLowerCase().includes(target)); - if (matched) { - clearInterval(interval); - const chosenText = matched.innerText; - matched.click(); - resolve({ success: true, model: chosenText }); - return; - } - } - - if (Date.now() - start > 3000) { - clearInterval(interval); - const btn = document.querySelector('[data-testid="stSelectbox"] button[aria-label="Open"]') || document.body; - btn.click(); - resolve({ success: false, model: '', error: 'Option not found' }); - } - }, 50); - }); - })()`, string(targetJSON)) - - res, err := bb.Evaluate(js) - if err != nil { - return err - } - - if rMap, ok := res.(map[string]interface{}); ok { - if success, _ := rMap["success"].(bool); success { - chosen, _ := rMap["model"].(string) - cleanCurr := chosen - if idx := strings.Index(chosen, " ("); idx != -1 { - cleanCurr = strings.TrimSpace(chosen[:idx]) - } - bb.selectedModel = cleanCurr - log.Printf("Successfully switched model to: %s", cleanCurr) - time.Sleep(1 * time.Second) - return nil - } - } - - log.Printf("Model %s not found in selectbox, keeping current model %s", modelID, bb.selectedModel) - return nil -} - -func (bb *BrowserBridge) ResetChat() error { - res, err := bb.Evaluate(`(() => { - const msgs = document.querySelectorAll('[data-testid="stChatMessage"]'); - const retryBtn = Array.from(document.querySelectorAll('button')).find(b => b.innerText.includes('Continue without retrying')); - const clearBtn = Array.from(document.querySelectorAll('button')).find(b => b.innerText.includes('Clear Chat')); - const ta = document.querySelector('textarea'); - - if (msgs.length === 0 && !retryBtn && ta && !ta.disabled) { - return 'clean'; - } - - if (clearBtn) { - clearBtn.click(); - return 'cleared'; - } - if (retryBtn) { - retryBtn.click(); - return 'retried'; - } - return 'none'; - })()`) - if err != nil { - return err - } - - if res == "clean" { - return nil - } - - deadline := time.Now().Add(6 * time.Second) - for time.Now().Before(deadline) { - time.Sleep(100 * time.Millisecond) - ready, _ := bb.Evaluate(`(() => { - const msgs = document.querySelectorAll('[data-testid="stChatMessage"]'); - const ta = document.querySelector('textarea[data-testid="stChatInputTextArea"]') || document.querySelector('textarea'); - const retryBtn = Array.from(document.querySelectorAll('button')).find(b => b.innerText.includes('Continue without retrying')); - if (retryBtn) { - retryBtn.click(); - return false; - } - return msgs.length === 0 && ta && !ta.disabled; - })()`) - if ready == true { - return nil - } - } - - return nil -} - -func (bb *BrowserBridge) SubmitPrompt(promptText string) error { - promptJSON, err := json.Marshal(promptText) - if err != nil { - return err - } - - // Ensure textarea is present and enabled - deadline := time.Now().Add(10 * time.Second) - for time.Now().Before(deadline) { - val, _ := bb.Evaluate(`(() => { - const ta = document.querySelector('textarea'); - return ta && !ta.disabled; - })()`) - if val == true { - break - } - time.Sleep(100 * time.Millisecond) - } - - js := fmt.Sprintf(`(() => { - const ta = document.querySelector('textarea[data-testid="stChatInputTextArea"]') || document.querySelector('textarea'); - if (!ta) return 'no textarea'; - ta.focus(); - const nativeSetter = Object.getOwnPropertyDescriptor(window.HTMLTextAreaElement.prototype, 'value').set; - nativeSetter.call(ta, %s); - ta.dispatchEvent(new Event('input', { bubbles: true })); - ta.dispatchEvent(new Event('change', { bubbles: true })); - - const btn = document.querySelector('button[data-testid="stChatInputSubmitButton"]'); - if (btn) { - btn.removeAttribute('disabled'); - btn.disabled = false; - btn.click(); - return 'clicked'; - } - - ta.dispatchEvent(new KeyboardEvent('keydown', { key: 'Enter', code: 'Enter', keyCode: 13, which: 13, bubbles: true })); - return 'enter'; - })()`, string(promptJSON)) - - _, err = bb.Evaluate(js) - if err != nil { - return err - } - - // Verify that new message is accepted by checking if count > 0 - confirmDeadline := time.Now().Add(4 * time.Second) - for time.Now().Before(confirmDeadline) { - currCountVal, _ := bb.Evaluate(`document.querySelectorAll('[data-testid="stChatMessage"]').length`) - if c, ok := currCountVal.(float64); ok && int(c) > 0 { - return nil - } - time.Sleep(100 * time.Millisecond) - } - - // Retry clicking if not registered yet - bb.Evaluate(`(() => { - const btn = document.querySelector('button[data-testid="stChatInputSubmitButton"]'); - if (btn) btn.click(); - const ta = document.querySelector('textarea'); - if (ta) ta.dispatchEvent(new KeyboardEvent('keydown', { key: 'Enter', code: 'Enter', keyCode: 13, which: 13, bubbles: true })); - })()`) - - return nil -} - -func (bb *BrowserBridge) Close() { - bb.mu.Lock() - defer bb.mu.Unlock() - - if bb.conn != nil { - bb.conn.Close() - bb.conn = nil - } - if bb.cmd != nil && bb.cmd.Process != nil { - _ = bb.cmd.Process.Kill() - _ = bb.cmd.Wait() - bb.cmd = nil - } - if bb.xvfbCmd != nil && bb.xvfbCmd.Process != nil { - _ = bb.xvfbCmd.Process.Kill() - _ = bb.xvfbCmd.Wait() - bb.xvfbCmd = nil - } - if bb.tmpDir != "" { - _ = os.RemoveAll(bb.tmpDir) - bb.tmpDir = "" +func WriteAPIError(w http.ResponseWriter, status int, message string, errType string, code int) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + resp := APIErrorResponse{ + Error: APIErrorDetail{ + Message: message, + Type: errType, + Code: code, + }, } + _ = json.NewEncoder(w).Encode(resp) } // --------------------------------------------------------------------------- @@ -1810,19 +2228,21 @@ func (bb *BrowserBridge) Close() { // --------------------------------------------------------------------------- type GroqqerGateway struct { - bridge *BrowserBridge - models []ModelItem - modelsMu sync.RWMutex - defaultModel string - port int + targetURL string + proxyURL string + userAgent string + defaultModel string + port int + models []ModelItem + modelsMu sync.RWMutex } func DefaultFallbackModels() []string { return []string{ - "openai/gpt-oss-20b", - "openai/gpt-oss-120b", "qwen/qwen3.6-27b", "qwen/qwen3.8-27b", + "openai/gpt-oss-120b", + "openai/gpt-oss-20b", "groq/compound", "groq/compound-mini", "allam-2-7b", @@ -1833,17 +2253,52 @@ func DefaultFallbackModels() []string { } func (g *GroqqerGateway) RefreshModels() { - discovered, err := g.bridge.DiscoverModels() - if err != nil || len(discovered) == 0 { - discovered = DefaultFallbackModels() - } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() - g.modelsMu.Lock() - defer g.modelsMu.Unlock() + sess, err := OpenStreamlitSession(ctx, g.targetURL, g.proxyURL, g.userAgent) + if err != nil { + log.Printf("warning: model discovery session failed: %v", err) + if len(g.GetModels()) == 0 { + g.setFallbackModels() + } + return + } + defer sess.Close() now := time.Now().Unix() var items []ModelItem - for _, id := range discovered { + for _, opt := range sess.options { + rawOpt := opt + if idx := strings.Index(rawOpt, " ("); idx != -1 { + rawOpt = strings.TrimSpace(rawOpt[:idx]) + } + if rawOpt != "" { + items = append(items, ModelItem{ + ID: rawOpt, + Object: "model", + Created: now, + OwnedBy: "groq", + }) + } + } + + if len(items) == 0 { + g.setFallbackModels() + return + } + + g.modelsMu.Lock() + g.models = items + g.modelsMu.Unlock() +} + +func (g *GroqqerGateway) setFallbackModels() { + g.modelsMu.Lock() + defer g.modelsMu.Unlock() + now := time.Now().Unix() + var items []ModelItem + for _, id := range DefaultFallbackModels() { items = append(items, ModelItem{ ID: id, Object: "model", @@ -1875,6 +2330,30 @@ func (g *GroqqerGateway) GetModels() []ModelItem { return out } +func (g *GroqqerGateway) MatchModel(requestedModel string) string { + if requestedModel == "" { + return g.defaultModel + } + models := g.GetModels() + for _, m := range models { + if strings.EqualFold(m.ID, requestedModel) { + return m.ID + } + } + // Try suffix match e.g. "qwen3.6-27b" for "qwen/qwen3.6-27b" + cleanReq := strings.ToLower(strings.TrimSpace(requestedModel)) + if strings.Contains(cleanReq, "/") { + parts := strings.Split(cleanReq, "/") + cleanReq = parts[len(parts)-1] + } + for _, m := range models { + if strings.Contains(strings.ToLower(m.ID), cleanReq) { + return m.ID + } + } + return requestedModel +} + func (g *GroqqerGateway) HandleModels(w http.ResponseWriter, r *http.Request) { w.Header().Set("Access-Control-Allow-Origin", "*") w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS") @@ -1903,131 +2382,92 @@ func (g *GroqqerGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re } if r.Method != "POST" { - http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + WriteAPIError(w, http.StatusMethodNotAllowed, "Method not allowed", "invalid_request_error", 405) return } var req ChatCompletionRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - http.Error(w, "Invalid JSON payload", http.StatusBadRequest) + WriteAPIError(w, http.StatusBadRequest, "Invalid JSON payload: "+err.Error(), "invalid_request_error", 400) return } if len(req.Messages) == 0 { - http.Error(w, "messages array must not be empty", http.StatusBadRequest) + WriteAPIError(w, http.StatusBadRequest, "messages array must not be empty", "invalid_request_error", 400) return } - requestedModel := req.Model - if requestedModel == "" { - requestedModel = g.defaultModel - } - + requestedModel := g.MatchModel(req.Model) promptText := FormatPrompt(req) completionID := "chatcmpl-" + GenerateUUID() created := time.Now().Unix() - - // Lock bridge for exclusive session execution - g.bridge.mu.Lock() - defer g.bridge.mu.Unlock() + allowedTools := GetAllowedToolNames(req.Tools, req.ToolChoice) log.Printf("Handling completion: model=%s, stream=%v, messages=%d", requestedModel, req.Stream, len(req.Messages)) - // 1. Switch model if needed - if err := g.bridge.SwitchModel(requestedModel); err != nil { - log.Printf("model switch warning: %v", err) - } + ctx, cancel := context.WithTimeout(r.Context(), 180*time.Second) + defer cancel() - // 2. Ensure chat session is clean before submitting - if err := g.bridge.ResetChat(); err != nil { - log.Printf("reset chat warning: %v", err) - } + start := time.Now() + maxAttempts := 3 - // 3. Submit the prompt - if err := g.bridge.SubmitPrompt(promptText); err != nil { - log.Printf("submit prompt error: %v", err) - http.Error(w, "Failed to submit prompt to target space: "+err.Error(), http.StatusInternalServerError) - return - } - - effectiveModel := g.bridge.selectedModel - if effectiveModel == "" { - effectiveModel = requestedModel - } - - targetMinCount := 2 - - // 4. Handle response: streaming or non-streaming - if req.Stream { - flusher, ok := w.(http.Flusher) - if !ok { - http.Error(w, "Streaming unsupported", http.StatusInternalServerError) + for attempt := 0; attempt < maxAttempts; attempt++ { + sess, err := OpenStreamlitSession(ctx, g.targetURL, g.proxyURL, g.userAgent) + if err != nil { + log.Printf("session connection error (attempt %d): %v", attempt+1, err) + if attempt < maxAttempts-1 { + time.Sleep(1 * time.Second) + continue + } + WriteAPIError(w, http.StatusBadGateway, "Failed to connect to target space: "+err.Error(), "api_error", 502) return } - streamer := NewStreamer(w, flusher, completionID, created, effectiveModel) - streamer.Role() + if err := sess.SwitchModel(requestedModel); err != nil { + log.Printf("model switch warning: %v", err) + } - thinkingFilter := NewStreamThinkingFilter() - toolFilter := NewStreamToolCallFilter() + effectiveModel := sess.activeModel + if effectiveModel == "" { + effectiveModel = requestedModel + } - lastLen := 0 - start := time.Now() - finishReason := "stop" - var fullCaptured strings.Builder - - evalJS := fmt.Sprintf(`(() => { - const msgs = Array.from(document.querySelectorAll('[data-testid="stChatMessage"]')); - if (msgs.length < %d) { - const alertEl = document.querySelector('[data-testid="stAlert"], [data-testid="stNotification"]'); - if (alertEl) { - return { text: '', done: true, error: alertEl.innerText }; - } - return { text: '', done: false, error: '' }; - } - - const lastMsg = msgs[msgs.length - 1]; - const md = lastMsg.querySelector('[data-testid="stMarkdownContainer"]'); - const text = md ? md.innerText : ''; - - const ta = document.querySelector('textarea'); - const hasModelCaption = lastMsg.innerText.includes('Model:') || !!lastMsg.querySelector('[data-testid="stCaptionContainer"]'); - const alertEl = lastMsg.querySelector('[data-testid="stAlert"], [data-testid="stNotification"]'); - const errText = alertEl ? alertEl.innerText : ''; - - const isDone = ((!ta || !ta.disabled) && hasModelCaption) || (errText !== ''); - - return { text, done: isDone, error: errText }; - })()`, targetMinCount) - - for time.Since(start) < 120*time.Second { - time.Sleep(35 * time.Millisecond) - - res, err := g.bridge.Evaluate(evalJS) - if err != nil { - continue + if req.Stream { + flusher, ok := w.(http.Flusher) + if !ok { + sess.Close() + WriteAPIError(w, http.StatusInternalServerError, "Streaming unsupported", "api_error", 500) + return } - if rMap, ok := res.(map[string]interface{}); ok { - txt, _ := rMap["text"].(string) - done, _ := rMap["done"].(bool) - errTxt, _ := rMap["error"].(string) + var streamer *Streamer + var thinkingFilter *StreamThinkingFilter + var toolFilter *StreamToolCallFilter + var streamStarted bool + var alertMsg string - if errTxt != "" && strings.TrimSpace(txt) == "" { - txt = errTxt + initStreamer := func() { + if !streamStarted { + streamer = NewStreamer(w, flusher, completionID, created, effectiveModel) + streamer.Role() + thinkingFilter = NewStreamThinkingFilter() + toolFilter = NewStreamToolCallFilter(allowedTools) + streamStarted = true } + } - if len(txt) > lastLen { - delta := txt[lastLen:] - lastLen = len(txt) - fullCaptured.WriteString(delta) - + err = sess.SubmitPrompt(promptText, + func(delta string) { + initStreamer() thinkingFilter.Feed(delta, func(cleanContent string) { toolFilter.Feed(cleanContent, func(userText string) { streamer.Content(userText) }, + func(reasoning string) { + streamer.Reasoning(reasoning) + }, func(tc ToolCall) { streamer.ToolCall(tc) }, @@ -2037,111 +2477,176 @@ func (g *GroqqerGateway) HandleChatCompletions(w http.ResponseWriter, r *http.Re streamer.Reasoning(reasoning) }, ) + }, + func(alert string) { + alertMsg = alert + }, + ) + sess.Close() + + if err != nil { + var alertErr *GatewayAlertError + if errors.As(err, &alertErr) && !streamStarted && alertErr.IsRateLimit { + if attempt < maxAttempts-1 { + waitSec := alertErr.RetryAfter + if waitSec <= 0 { + waitSec = 6 + } + if waitSec > 20 { + waitSec = 20 + } + log.Printf("Rate limit hit on streaming %s. Waiting %d seconds before retry (attempt %d)...", effectiveModel, waitSec, attempt+1) + select { + case <-time.After(time.Duration(waitSec)*time.Second + 500*time.Millisecond): + continue + case <-ctx.Done(): + WriteAPIError(w, http.StatusTooManyRequests, "Rate limit retry cancelled: context deadline exceeded", "rate_limit_error", 429) + return + } + } + w.Header().Set("Retry-After", strconv.Itoa(alertErr.RetryAfter)) + WriteAPIError(w, http.StatusTooManyRequests, alertErr.Alert, "rate_limit_error", 429) + return } - if done && (len(txt) > 0 || errTxt != "") { - break + if !streamStarted { + if errors.As(err, &alertErr) { + lower := strings.ToLower(alertErr.Alert) + if strings.Contains(lower, "too large") || strings.Contains(lower, "cannot fit") { + WriteAPIError(w, http.StatusBadRequest, alertErr.Alert, "invalid_request_error", 400) + return + } + WriteAPIError(w, http.StatusBadGateway, alertErr.Alert, "api_error", 502) + return + } + WriteAPIError(w, http.StatusInternalServerError, "Streaming error: "+err.Error(), "api_error", 500) + return + } + + if alertMsg != "" { + streamer.Content("\n\n" + alertMsg) } } - } - thinkingFilter.Flush( - func(cleanContent string) { - toolFilter.Feed(cleanContent, + if streamStarted { + finishReason := "stop" + thinkingFilter.Flush( + func(cleanContent string) { + toolFilter.Feed(cleanContent, + func(userText string) { + streamer.Content(userText) + }, + func(reasoning string) { + streamer.Reasoning(reasoning) + }, + func(tc ToolCall) { + streamer.ToolCall(tc) + }, + ) + }, + func(reasoning string) { + streamer.Reasoning(reasoning) + }, + ) + + toolFilter.Flush( func(userText string) { streamer.Content(userText) }, + func(reasoning string) { + streamer.Reasoning(reasoning) + }, func(tc ToolCall) { streamer.ToolCall(tc) }, ) - }, - func(reasoning string) { - streamer.Reasoning(reasoning) - }, - ) - toolFilter.Flush( - func(userText string) { - streamer.Content(userText) - }, - func(tc ToolCall) { - streamer.ToolCall(tc) - }, - ) + if toolFilter.emittedCall { + finishReason = "tool_calls" + } - if toolFilter.emittedCall { - finishReason = "tool_calls" + streamer.Finish(finishReason) + log.Printf("Streaming completion finished for %s in %v", effectiveModel, time.Since(start)) + return + } + + WriteAPIError(w, http.StatusInternalServerError, "Empty response from space", "api_error", 500) + return } - streamer.Finish(finishReason) - log.Printf("Streaming completion finished for %s in %v", effectiveModel, time.Since(start)) + // Non-streaming completion + var fullBuilder strings.Builder + var alertMsg string + + err = sess.SubmitPrompt(promptText, + func(delta string) { + fullBuilder.WriteString(delta) + }, + func(alert string) { + alertMsg = alert + }, + ) + sess.Close() + _ = alertMsg + + if err != nil { + var alertErr *GatewayAlertError + if errors.As(err, &alertErr) && alertErr.IsRateLimit { + if attempt < maxAttempts-1 { + waitSec := alertErr.RetryAfter + if waitSec <= 0 { + waitSec = 6 + } + if waitSec > 20 { + waitSec = 20 + } + log.Printf("Rate limit hit on %s. Waiting %d seconds before retry (attempt %d)...", effectiveModel, waitSec, attempt+1) + select { + case <-time.After(time.Duration(waitSec)*time.Second + 500*time.Millisecond): + continue + case <-ctx.Done(): + WriteAPIError(w, http.StatusTooManyRequests, "Rate limit retry cancelled: context deadline exceeded", "rate_limit_error", 429) + return + } + } + w.Header().Set("Retry-After", strconv.Itoa(alertErr.RetryAfter)) + WriteAPIError(w, http.StatusTooManyRequests, alertErr.Alert, "rate_limit_error", 429) + return + } + + if errors.As(err, &alertErr) { + lower := strings.ToLower(alertErr.Alert) + if strings.Contains(lower, "too large") || strings.Contains(lower, "cannot fit") { + WriteAPIError(w, http.StatusBadRequest, alertErr.Alert, "invalid_request_error", 400) + return + } + WriteAPIError(w, http.StatusBadGateway, alertErr.Alert, "api_error", 502) + return + } + + WriteAPIError(w, http.StatusInternalServerError, "Prompt execution error: "+err.Error(), "api_error", 500) + return + } + + fullText := fullBuilder.String() + if strings.TrimSpace(fullText) == "" { + WriteAPIError(w, http.StatusInternalServerError, "The model returned an empty response", "api_error", 500) + return + } + + cleanContent, reasoningContent := ExtractThinking(fullText) + toolCalls, remainingContent, hasTools := DetectToolCalls(cleanContent, allowedTools) + + finishReason := "stop" + if hasTools { + finishReason = "tool_calls" + cleanContent = remainingContent + } + + WriteCompletionResponse(w, completionID, created, effectiveModel, cleanContent, reasoningContent, toolCalls, finishReason) + log.Printf("Completion finished for %s in %v", effectiveModel, time.Since(start)) return } - - // Non-streaming response - start := time.Now() - var fullText string - - evalJS := fmt.Sprintf(`(() => { - const msgs = Array.from(document.querySelectorAll('[data-testid="stChatMessage"]')); - if (msgs.length < %d) { - const alertEl = document.querySelector('[data-testid="stAlert"], [data-testid="stNotification"]'); - if (alertEl) { - return { text: '', done: true, error: alertEl.innerText }; - } - return { text: '', done: false, error: '' }; - } - - const lastMsg = msgs[msgs.length - 1]; - const md = lastMsg.querySelector('[data-testid="stMarkdownContainer"]'); - const text = md ? md.innerText : ''; - - const ta = document.querySelector('textarea'); - const hasModelCaption = lastMsg.innerText.includes('Model:') || !!lastMsg.querySelector('[data-testid="stCaptionContainer"]'); - const alertEl = lastMsg.querySelector('[data-testid="stAlert"], [data-testid="stNotification"]'); - const errText = alertEl ? alertEl.innerText : ''; - - const isDone = ((!ta || !ta.disabled) && hasModelCaption) || (errText !== ''); - - return { text, done: isDone, error: errText }; - })()`, targetMinCount) - - for time.Since(start) < 120*time.Second { - time.Sleep(50 * time.Millisecond) - - res, err := g.bridge.Evaluate(evalJS) - if err != nil { - continue - } - - if rMap, ok := res.(map[string]interface{}); ok { - txt, _ := rMap["text"].(string) - done, _ := rMap["done"].(bool) - errTxt, _ := rMap["error"].(string) - - if errTxt != "" && strings.TrimSpace(txt) == "" { - txt = errTxt - } - - fullText = txt - if done && (len(txt) > 0 || errTxt != "") { - break - } - } - } - - cleanContent, reasoningContent := ExtractThinking(fullText) - toolCalls, remainingContent, hasTools := DetectToolCalls(cleanContent) - - finishReason := "stop" - if hasTools { - finishReason = "tool_calls" - cleanContent = remainingContent - } - - WriteCompletionResponse(w, completionID, created, effectiveModel, cleanContent, reasoningContent, toolCalls, finishReason) - log.Printf("Completion finished for %s in %v", effectiveModel, time.Since(start)) } // --------------------------------------------------------------------------- @@ -2152,24 +2657,20 @@ func main() { port := flag.Int("port", 8080, "HTTP server listening port") targetURL := flag.String("target", DefaultTargetURL, "Groq Streamlit space URL") defaultModel := flag.String("default-model", DefaultModel, "Default model ID if unspecified") - browserBin := flag.String("browser", "", "Path to Chromium/Chrome binary") userAgent := flag.String("user-agent", DefaultUserAgent, "User-Agent string") flag.StringVar(userAgent, "ua", DefaultUserAgent, "User-Agent string (alias)") - useXvfb := flag.Bool("xvfb", true, "Manage virtual X server (Xvfb) for display isolation") - noXvfb := flag.Bool("no-xvfb", false, "Disable virtual X server") - headless := flag.Bool("headless", true, "Run Chromium in headless mode") - socksProxy := flag.String("socks", "", "SOCKS5 proxy (e.g. socks5://127.0.0.1:1080)") flag.StringVar(socksProxy, "proxy", "", "Proxy address alias") flag.StringVar(socksProxy, "socks5", "", "SOCKS5 proxy alias") - flag.Parse() + // Retain flags for backward compatibility (browser/Xvfb are no longer required) + _ = flag.String("browser", "", "Ignored: browser is no longer required (headless WebSocket mode)") + _ = flag.Bool("xvfb", false, "Ignored: Xvfb is no longer required") + _ = flag.Bool("no-xvfb", false, "Ignored: Xvfb is no longer required") + _ = flag.Bool("headless", true, "Ignored: always runs headless via direct WebSocket") - effectiveXvfb := *useXvfb - if *noXvfb { - effectiveXvfb = false - } + flag.Parse() effectiveProxy := *socksProxy if effectiveProxy == "" { @@ -2181,22 +2682,19 @@ func main() { } } - bridge := NewBrowserBridge(*browserBin, *userAgent, *headless, effectiveXvfb, effectiveProxy, *targetURL) - - log.Println("Starting groqqer browser bridge...") - if err := bridge.Start(); err != nil { - log.Fatalf("failed to start browser bridge: %v", err) - } + log.Printf("Initializing groqqer gateway (pure Go WebSocket client, zero browser dependency)...") gateway := &GroqqerGateway{ - bridge: bridge, + targetURL: *targetURL, + proxyURL: effectiveProxy, + userAgent: *userAgent, defaultModel: *defaultModel, port: *port, } - // Initial discovery of live models + // Initial discovery of live models in background go func() { - log.Println("Discovering available models from Streamlit UI...") + log.Println("Discovering available models from Streamlit space...") gateway.RefreshModels() models := gateway.GetModels() log.Printf("Ready with %d active models. Default: %s", len(models), gateway.defaultModel) @@ -2217,6 +2715,7 @@ func main() { w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]interface{}{ "service": "groqqer", + "mode": "headless-websocket-protobuf", "status": "running", "target": *targetURL, "default": *defaultModel, @@ -2240,7 +2739,6 @@ func main() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = server.Shutdown(ctx) - bridge.Close() os.Exit(0) }() diff --git a/groqqer_test.go b/groqqer_test.go new file mode 100644 index 0000000..37bf081 --- /dev/null +++ b/groqqer_test.go @@ -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") { + 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 := ` +{"name": "search", "arguments": {"query": "golang"}} +` + 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 := ` +{"name": "search", "arguments": {"query": "golang"}} + + +{"name": "weather", "arguments": {"city": "Tokyo"}} +` + 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 := ` +[ + {"name": "tool_a", "arguments": {"a": 1}}, + {"name": "tool_b", "arguments": {"b": 2}} +] +` + 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}, + {"thinkingResult: 42", "hello", false}, + {"{\"name\":\"calc\"}", "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: + +{"name": "non_existing_tool", "arguments": {"foo": "bar"}} + +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 := ` +{"name": "calculator", "arguments": {"expr": "1+1"}} +` + + 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 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 := ` +{"name": "calculator", "arguments": {"expr": "2+2"}} + +And here is a fake tool: + +{"name": "fake_tool", "arguments": {}} +` + + 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", + "\n", + "{\"name\": \"unknown_tool\", ", + "\"arguments\": {}}\n", + "\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{"", `{"name":"calc"}`, ""} + 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 != "{\"name\":\"calc\"}" { + 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("\n{\"name\": \"calculator\", \"arguments\": {\"expr\": \"5*5\"}}\n", 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: `{"name": "calc", "arguments": {"expr": "2+2"}}`, + 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 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, "") + if count != 1 { + t.Errorf("Expected exactly 1 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 tag in my script?"}, + }, + } + + prompt := FormatPrompt(req) + if !strings.Contains(prompt, "How do I format a tag in my script?") { + t.Errorf("Expected user prompt to preserve text, got:\n%s", prompt) + } +}