From 2eac753a4d03c6eb26f9196b204498e3cba7af8a Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 1 Mar 2026 16:35:17 +0000 Subject: [PATCH] Add OpenAI-compatible backend for local LLM support (GPT-OSS) Implement OpenAIClient backend that speaks the OpenAI Chat Completions API (/v1/chat/completions), enabling llm9p to work with any local model server: Ollama, vLLM, llama-server, LocalAI, or LM Studio. Primary target is GPT-OSS (OpenAI's open-weight MoE models), but any model served via these platforms works. The backend supports: - Blocking and streaming chat completions - Tool/function calling with STOP:/TOOL: formatting - Token counting from API usage (with estimation fallback) - Conversation history, system prompts, temperature control - Stateless AskWithRequest for session isolation Usage: ./llm9p -backend openai -openai-url http://localhost:11434/v1 -model gpt-oss:20b Also fixes pre-existing stale mock backends in test files (AskWithRequest signature was out of date with the Backend interface). https://claude.ai/code/session_017qVVZUUhfCCvNkMYmXDZAa --- cmd/llm9p/main.go | 16 +- docs/local-llm-feasibility.md | 53 +- go.mod | 5 +- go.sum | 2 + internal/llm/backend.go | 3 +- internal/llm/client.go | 6 +- internal/llm/client_test.go | 2 +- internal/llm/openai_client.go | 825 +++++++++++++++++++++++++++ internal/llm/session.go | 8 +- internal/llm/session_compact_test.go | 43 +- internal/llmfs/mock_backend_test.go | 12 +- internal/protocol/fs.go | 4 +- 12 files changed, 934 insertions(+), 45 deletions(-) create mode 100644 internal/llm/openai_client.go diff --git a/cmd/llm9p/main.go b/cmd/llm9p/main.go index d9f5718d6f16d63c95f766619453b8cdf9e91988..feb5e8a64f1042a637429d86f0491082308a276b 100644 --- a/cmd/llm9p/main.go +++ b/cmd/llm9p/main.go @@ -37,7 +37,9 @@ import ( func main() { addr := flag.String("addr", ":5640", "Address to listen on") debug := flag.Bool("debug", false, "Enable debug logging") - backend := flag.String("backend", "api", "Backend to use: 'api' (Anthropic API) or 'cli' (Claude Code CLI for Max subscription)") + backend := flag.String("backend", "api", "Backend to use: 'api' (Anthropic API), 'cli' (Claude Code CLI), or 'openai' (OpenAI-compatible local server)") + openaiURL := flag.String("openai-url", "http://localhost:11434/v1", "Base URL for OpenAI-compatible server (e.g. http://localhost:11434/v1 for Ollama)") + model := flag.String("model", "", "Model name for OpenAI-compatible backend (e.g. gpt-oss:20b)") flag.Parse() var client llm.Backend @@ -64,8 +66,18 @@ func main() { client = llm.NewClient(apiKey) log.Println("Using Anthropic API backend") + case "openai": + if *model == "" { + fmt.Fprintln(os.Stderr, "Error: -model flag is required for OpenAI-compatible backend") + fmt.Fprintln(os.Stderr, "Example: -backend openai -model gpt-oss:20b") + os.Exit(1) + } + apiKey := os.Getenv("OPENAI_API_KEY") + client = llm.NewOpenAIClient(*openaiURL, apiKey, *model) + log.Printf("Using OpenAI-compatible backend at %s (model: %s)", *openaiURL, *model) + default: - fmt.Fprintf(os.Stderr, "Error: unknown backend '%s' (use 'api' or 'cli')\n", *backend) + fmt.Fprintf(os.Stderr, "Error: unknown backend '%s' (use 'api', 'cli', or 'openai')\n", *backend) os.Exit(1) } diff --git a/docs/local-llm-feasibility.md b/docs/local-llm-feasibility.md index 33d73cf97d3bbb792cde4fa8d6005632d25502a2..b39822a0363c6836be3bfb94d94a98fe8381cba1 100644 --- a/docs/local-llm-feasibility.md +++ b/docs/local-llm-feasibility.md @@ -2,7 +2,50 @@ ## Summary -Adding local LLM support to llm9p is **highly feasible** and architecturally straightforward. The existing `Backend` interface already provides the right abstraction, and the ecosystem has converged on OpenAI-compatible APIs as the standard interface for local model servers. A new `OpenAIClient` backend (~400-500 lines of Go) would enable llm9p to work with Ollama, llama.cpp, vLLM, LocalAI, LM Studio, and any other server exposing `/v1/chat/completions`. +Adding local LLM support to llm9p is **highly feasible** and architecturally straightforward. The primary target is **GPT-OSS** (OpenAI's open-weight models), but the implementation covers any model served via the OpenAI-compatible `/v1/chat/completions` endpoint. The existing `Backend` interface already provides the right abstraction. A new `OpenAIClient` backend (~400-500 lines of Go) would enable llm9p to work with GPT-OSS (via Ollama, vLLM, or llama.cpp), as well as Llama, Mistral, Qwen, and any other model served by these platforms. + +## Primary Target: GPT-OSS + +GPT-OSS is OpenAI's first open-weight model release since GPT-2, released August 2025 under the Apache 2.0 license. It consists of two Mixture-of-Experts (MoE) models: + +| Model | Total Params | Active/Token | VRAM | Context | Target Hardware | +|---|---|---|---|---|---| +| **gpt-oss-20b** | 21B | 3.6B | ~14-16 GB | 128K | Consumer GPUs (RTX 4090), Apple Silicon | +| **gpt-oss-120b** | 117B | 5.1B | ~80 GB | 128K | H100/H200/B200 | + +### Why GPT-OSS is a good fit for llm9p + +- **Strong tool calling**: Outperforms o4-mini on TauBench -- important for llm9p's tool_use protocol +- **128K context**: Matches Claude's context window, so the existing compaction logic works well +- **Runs on consumer hardware**: gpt-oss-20b needs only ~14 GB (Apple M-series or a single RTX 4090) +- **Standard API**: All serving backends expose it via OpenAI `/v1/chat/completions` -- no special handling needed +- **Harmony format abstracted away**: GPT-OSS uses a new token format called "Harmony" internally, but Ollama/vLLM/llama.cpp handle the conversion. The client just uses the standard chat completions API + +### Serving GPT-OSS locally + +```bash +# Ollama (simplest -- auto-downloads the model) +ollama pull gpt-oss:20b +# API at http://localhost:11434/v1 + +# vLLM (production, GPU servers) +vllm serve openai/gpt-oss-20b --tool-call-parser openai +# API at http://localhost:8000/v1 + +# llama.cpp (GGUF quantization, partial GPU offload) +llama-server -hf ggml-org/gpt-oss-20b-GGUF --jinja +# API at http://localhost:8080/v1 +``` + +### Using GPT-OSS with llm9p (planned) + +```bash +# GPT-OSS via Ollama +./llm9p -backend openai -openai-url http://localhost:11434/v1 -model gpt-oss:20b + +# GPT-OSS via vLLM +./llm9p -backend openai -openai-url http://localhost:8000/v1 -model openai/gpt-oss-20b +``` ## Current Architecture @@ -205,14 +248,14 @@ Total: **~700 lines of new code**, mostly mechanical since it follows the existi | Model | Size | Context | Tool Calling | Notes | |---|---|---|---|---| +| **GPT-OSS 20B** | ~14-16 GB | 128K | Yes (strong) | Primary target; MoE, only 3.6B active | +| **GPT-OSS 120B** | ~80 GB | 128K | Yes (strong) | For GPU servers; beats o4-mini | | Llama 3.1 8B Instruct | 4-8 GB | 128K | Yes | Best balance of quality and speed | | Qwen 2.5 7B Instruct | 4-8 GB | 32K | Yes | Strong multilingual, good at tools | | Mistral 7B Instruct | 4-8 GB | 32K | Yes | Fast, good instruction following | -| Phi-3 Mini 3.8B | 2-4 GB | 128K | Limited | Smallest usable model | -| DeepSeek-R1 7B | 4-8 GB | 64K | No | Strong reasoning, no tool calling | ## Conclusion -Adding local LLM support is a well-scoped, low-risk enhancement. The `Backend` interface is already designed for exactly this kind of extension. The OpenAI-compatible API is the clear integration point since the entire ecosystem has standardized on it. The `sashabaranov/go-openai` Go library provides everything needed. Implementation would take roughly a day of development, following the established patterns in `cli_client.go`. +Adding local LLM support is a well-scoped, low-risk enhancement. The `Backend` interface is already designed for exactly this kind of extension. GPT-OSS is the primary target -- it offers strong tool calling, 128K context, and runs on consumer hardware. The OpenAI-compatible API is the clear integration point since the entire ecosystem (including GPT-OSS serving via Ollama/vLLM/llama.cpp) has standardized on it. The `sashabaranov/go-openai` Go library provides everything needed. -The result would make llm9p usable in fully offline/air-gapped environments, eliminate API costs for development and experimentation, and open the door to any model ecosystem (not just Claude). +The result would make llm9p usable with GPT-OSS and other open-weight models in fully offline/air-gapped environments, eliminate API costs for development and experimentation, and open the door to any model ecosystem. diff --git a/go.mod b/go.mod index fea183c5aaa44bc74c34037dd783b2ac2e258d99..5ce13a889edf97ef4455cd79ec38dd91881399d0 100644 --- a/go.mod +++ b/go.mod @@ -2,7 +2,10 @@ module github.com/NERVsystems/llm9p go 1.21 -require github.com/anthropics/anthropic-sdk-go v0.2.0-beta.3 +require ( + github.com/anthropics/anthropic-sdk-go v0.2.0-beta.3 + github.com/sashabaranov/go-openai v1.41.2 +) require ( github.com/tidwall/gjson v1.14.4 // indirect diff --git a/go.sum b/go.sum index 78e0cb45f057862bb959b86a069281921d61f8d3..ce03bd89355ad1962b366a1a36469c9af20d36eb 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ github.com/anthropics/anthropic-sdk-go v0.2.0-beta.3 h1:b5t1ZJMvV/l99y4jbz7kRFdUp3BSDkI8EhSlHczivtw= github.com/anthropics/anthropic-sdk-go v0.2.0-beta.3/go.mod h1:AapDW22irxK2PSumZiQXYUFvsdQgkwIWlpESweWZI/c= +github.com/sashabaranov/go-openai v1.41.2 h1:vfPRBZNMpnqu8ELsclWcAvF19lDNgh1t6TVfFFOPiSM= +github.com/sashabaranov/go-openai v1.41.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM= github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= diff --git a/internal/llm/backend.go b/internal/llm/backend.go index 791d1fe954e42944714060d52f10d058faaa61c0..0e1ff9bea11c33ec9a5d1021bc5e923f1430cf38 100644 --- a/internal/llm/backend.go +++ b/internal/llm/backend.go @@ -94,6 +94,7 @@ type Backend interface { WaitStream() } -// Verify that both clients implement Backend +// Verify that all clients implement Backend var _ Backend = (*Client)(nil) var _ Backend = (*CLIClient)(nil) +var _ Backend = (*OpenAIClient)(nil) diff --git a/internal/llm/client.go b/internal/llm/client.go index 3122c300e2877aa4258cb6d239c79d2677f2a166..dfbbf6e70a9b753946dd3fbb77bfc1480124bb89 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -17,9 +17,9 @@ import ( // StructuredContent, when non-empty, holds a JSON array of content blocks // for proper API replay of tool-use turns. Plain-text turns leave it empty. type Message struct { - Role string `json:"role"` // "user" or "assistant" - Content string `json:"content"` // text content (always set) - StructuredContent string `json:"sc,omitempty"` // JSON content blocks (tool turns only) + Role string `json:"role"` // "user" or "assistant" + Content string `json:"content"` // text content (always set) + StructuredContent string `json:"sc,omitempty"` // JSON content blocks (tool turns only) } // MetricsCallback is called after each LLM request with performance data diff --git a/internal/llm/client_test.go b/internal/llm/client_test.go index 7324d24c0943e1f319655e43bf842d99ccaee86a..27c47c27bc8554ca64e9226fdb1100e23cec5ecc 100644 --- a/internal/llm/client_test.go +++ b/internal/llm/client_test.go @@ -13,7 +13,7 @@ func TestContextLimitForModel(t *testing.T) { {"claude-3-sonnet-20240229", 200000}, {"claude-3-haiku-20240307", 200000}, {"claude-sonnet-4-20250514", 200000}, - {"CLAUDE-3-OPUS", 200000}, // case insensitive + {"CLAUDE-3-OPUS", 200000}, // case insensitive {"some-sonnet-model", 200000}, // substring match {"unknown-model", 200000}, // default } diff --git a/internal/llm/openai_client.go b/internal/llm/openai_client.go new file mode 100644 index 0000000000000000000000000000000000000000..1032b7d162a01e9eb08918ada8d64364df71a4ca --- /dev/null +++ b/internal/llm/openai_client.go @@ -0,0 +1,825 @@ +// OpenAI-compatible backend for local LLM servers (Ollama, llama.cpp, vLLM, etc.). +// +// This backend speaks the OpenAI Chat Completions API (/v1/chat/completions), +// which is the de facto standard for local model serving. Primary target is +// GPT-OSS, but any model served by Ollama, vLLM, llama-server, LocalAI, or +// LM Studio works. +package llm + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + "sync" + "time" + + openai "github.com/sashabaranov/go-openai" +) + +// OpenAIClient uses an OpenAI-compatible API for LLM requests. +// Works with any server that implements /v1/chat/completions: +// Ollama, vLLM, llama-server, LocalAI, LM Studio, etc. +type OpenAIClient struct { + client *openai.Client + mu sync.RWMutex + model string + temperature float64 + systemPrompt string + prefill string + messages []Message + lastTokens int + totalTokens int + thinkingTokens int + contextLimit int + streaming bool + streamChan chan string + streamDone chan struct{} +} + +// NewOpenAIClient creates a new OpenAI-compatible LLM client. +// baseURL should include /v1 (e.g. "http://localhost:11434/v1" for Ollama). +// apiKey can be empty or a dummy value for local servers that don't require auth. +// model is the model name as the server expects it (e.g. "gpt-oss:20b"). +func NewOpenAIClient(baseURL, apiKey, model string) *OpenAIClient { + if apiKey == "" { + apiKey = "not-needed" + } + config := openai.DefaultConfig(apiKey) + config.BaseURL = baseURL + return &OpenAIClient{ + client: openai.NewClientWithConfig(config), + model: model, + temperature: 0.7, + messages: make([]Message, 0), + contextLimit: 128000, // Default for GPT-OSS; overridable + } +} + +// Model returns the current model name. +func (c *OpenAIClient) Model() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.model +} + +// SetModel sets the model for subsequent requests. +func (c *OpenAIClient) SetModel(model string) { + c.mu.Lock() + defer c.mu.Unlock() + c.model = model +} + +// Temperature returns the current temperature. +func (c *OpenAIClient) Temperature() float64 { + c.mu.RLock() + defer c.mu.RUnlock() + return c.temperature +} + +// SetTemperature sets the temperature for subsequent requests. +func (c *OpenAIClient) SetTemperature(temp float64) error { + if temp < 0.0 || temp > 2.0 { + return fmt.Errorf("temperature must be between 0.0 and 2.0") + } + c.mu.Lock() + defer c.mu.Unlock() + c.temperature = temp + return nil +} + +// ThinkingTokens returns the thinking token budget. +// Not used by OpenAI-compatible backends; stored for interface compliance. +func (c *OpenAIClient) ThinkingTokens() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.thinkingTokens +} + +// SetThinkingTokens sets the thinking token budget. +// Not used by OpenAI-compatible backends. +func (c *OpenAIClient) SetThinkingTokens(tokens int) { + c.mu.Lock() + defer c.mu.Unlock() + c.thinkingTokens = tokens +} + +// Prefill returns the assistant response prefill string. +func (c *OpenAIClient) Prefill() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.prefill +} + +// SetPrefill sets a string to prefill the assistant response. +// OpenAI API doesn't support native prefill, so we prepend to the response. +func (c *OpenAIClient) SetPrefill(prefill string) { + c.mu.Lock() + defer c.mu.Unlock() + c.prefill = prefill +} + +// SystemPrompt returns the current system prompt. +func (c *OpenAIClient) SystemPrompt() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.systemPrompt +} + +// SetSystemPrompt sets the system prompt for subsequent requests. +func (c *OpenAIClient) SetSystemPrompt(prompt string) { + c.mu.Lock() + defer c.mu.Unlock() + c.systemPrompt = prompt +} + +// LastTokens returns the token count from the last response. +func (c *OpenAIClient) LastTokens() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.lastTokens +} + +// TotalTokens returns cumulative token count for this conversation. +func (c *OpenAIClient) TotalTokens() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.totalTokens +} + +// ContextLimit returns the model's context window limit. +func (c *OpenAIClient) ContextLimit() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.contextLimit +} + +// Messages returns a copy of the conversation history. +func (c *OpenAIClient) Messages() []Message { + c.mu.RLock() + defer c.mu.RUnlock() + result := make([]Message, len(c.messages)) + copy(result, c.messages) + return result +} + +// MessagesJSON returns the conversation history as JSON. +func (c *OpenAIClient) MessagesJSON() ([]byte, error) { + c.mu.RLock() + defer c.mu.RUnlock() + return json.MarshalIndent(c.messages, "", " ") +} + +// AddSystemMessage adds a system message to the context. +func (c *OpenAIClient) AddSystemMessage(content string) { + c.mu.Lock() + defer c.mu.Unlock() + c.messages = append([]Message{{Role: "system", Content: content}}, c.messages...) +} + +// Reset clears the conversation history. +func (c *OpenAIClient) Reset() { + c.mu.Lock() + defer c.mu.Unlock() + c.messages = make([]Message, 0) + c.lastTokens = 0 + c.totalTokens = 0 +} + +// Compact summarizes the conversation to reduce token usage. +func (c *OpenAIClient) Compact(ctx context.Context) error { + c.mu.Lock() + if len(c.messages) < 4 { + c.mu.Unlock() + return nil + } + + var conversationText string + for _, msg := range c.messages { + if msg.Role == "system" { + continue + } + conversationText += fmt.Sprintf("%s: %s\n\n", msg.Role, msg.Content) + } + + model := c.model + c.mu.Unlock() + + summaryPrompt := "Summarize this conversation concisely, preserving key facts, decisions, and context needed to continue:\n\n" + conversationText + + req := openai.ChatCompletionRequest{ + Model: model, + MaxTokens: 2048, + Temperature: 0.3, + Messages: []openai.ChatCompletionMessage{ + {Role: openai.ChatMessageRoleUser, Content: summaryPrompt}, + }, + } + + resp, err := c.client.CreateChatCompletion(ctx, req) + if err != nil { + return fmt.Errorf("compaction failed: %w", err) + } + + if len(resp.Choices) == 0 { + return fmt.Errorf("compaction returned no choices") + } + + summary := resp.Choices[0].Message.Content + tokens := resp.Usage.TotalTokens + + c.mu.Lock() + c.messages = []Message{{Role: "system", Content: "Previous conversation summary: " + summary}} + c.totalTokens = tokens + c.mu.Unlock() + + return nil +} + +// buildChatMessages converts internal Message history to OpenAI API format. +// Returns system messages separately (as a single system message) and the +// conversation messages. +func buildChatMessages(systemPrompt string, msgs []Message) []openai.ChatCompletionMessage { + apiMsgs := make([]openai.ChatCompletionMessage, 0, len(msgs)+2) + + // Collect all system content into a single system message + var systemParts []string + if systemPrompt != "" { + systemParts = append(systemParts, systemPrompt) + } + for _, msg := range msgs { + if msg.Role == "system" { + systemParts = append(systemParts, msg.Content) + } + } + if len(systemParts) > 0 { + apiMsgs = append(apiMsgs, openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleSystem, + Content: strings.Join(systemParts, "\n\n"), + }) + } + + // Add conversation messages + for _, msg := range msgs { + switch msg.Role { + case "user": + apiMsgs = append(apiMsgs, openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleUser, + Content: msg.Content, + }) + case "assistant": + // Check if this message has structured content with tool calls + if msg.StructuredContent != "" { + apiMsg := rebuildAssistantToolMessage(msg) + apiMsgs = append(apiMsgs, apiMsg) + } else { + apiMsgs = append(apiMsgs, openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleAssistant, + Content: msg.Content, + }) + } + } + } + + return apiMsgs +} + +// rebuildAssistantToolMessage reconstructs an assistant message with tool calls +// from the stored StructuredContent JSON. +func rebuildAssistantToolMessage(msg Message) openai.ChatCompletionMessage { + type rawBlock struct { + Type string `json:"type"` + Text string `json:"text"` + ID string `json:"id"` + Name string `json:"name"` + Input json.RawMessage `json:"input"` + ToolUseID string `json:"tool_use_id"` + Content string `json:"content"` + } + + var blocks []rawBlock + if err := json.Unmarshal([]byte(msg.StructuredContent), &blocks); err != nil { + // Fallback to plain text + return openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleAssistant, + Content: msg.Content, + } + } + + apiMsg := openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleAssistant, + } + + for _, b := range blocks { + switch b.Type { + case "text": + apiMsg.Content += b.Text + case "tool_use": + idx := len(apiMsg.ToolCalls) + apiMsg.ToolCalls = append(apiMsg.ToolCalls, openai.ToolCall{ + Index: &idx, + ID: b.ID, + Type: openai.ToolTypeFunction, + Function: openai.FunctionCall{ + Name: b.Name, + Arguments: string(b.Input), + }, + }) + } + } + + return apiMsg +} + +// Ask sends a prompt to the LLM and returns the response. +func (c *OpenAIClient) Ask(ctx context.Context, prompt string) (string, error) { + c.mu.Lock() + c.messages = append(c.messages, Message{Role: "user", Content: prompt}) + apiMsgs := buildChatMessages(c.systemPrompt, c.messages) + model := c.model + temp := c.temperature + c.mu.Unlock() + + req := openai.ChatCompletionRequest{ + Model: model, + MaxTokens: 4096, + Temperature: float32(temp), + Messages: apiMsgs, + } + + startTime := time.Now() + resp, err := c.client.CreateChatCompletion(ctx, req) + latencyMs := time.Since(startTime).Milliseconds() + + if err != nil { + c.mu.Lock() + if len(c.messages) > 0 { + c.messages = c.messages[:len(c.messages)-1] + } + c.mu.Unlock() + return "", fmt.Errorf("OpenAI API error: %w", err) + } + + if len(resp.Choices) == 0 { + c.mu.Lock() + if len(c.messages) > 0 { + c.messages = c.messages[:len(c.messages)-1] + } + c.mu.Unlock() + return "", fmt.Errorf("OpenAI API returned no choices") + } + + responseText := resp.Choices[0].Message.Content + tokens := resp.Usage.TotalTokens + if tokens == 0 { + tokens = estimateTokens(prompt) + estimateTokens(responseText) + } + + c.mu.Lock() + c.messages = append(c.messages, Message{Role: "assistant", Content: responseText}) + c.lastTokens = tokens + c.totalTokens += tokens + c.mu.Unlock() + + RecordMetrics(resp.Usage.PromptTokens, resp.Usage.CompletionTokens, latencyMs) + + return responseText, nil +} + +// AskWithHistory sends a prompt with explicit message history for per-fid isolation. +func (c *OpenAIClient) AskWithHistory(ctx context.Context, history []Message, prompt string) (string, int, error) { + c.mu.RLock() + model := c.model + temp := c.temperature + systemPrompt := c.systemPrompt + prefill := c.prefill + c.mu.RUnlock() + + // Build messages from history + new prompt + combined := make([]Message, len(history)) + copy(combined, history) + combined = append(combined, Message{Role: "user", Content: prompt}) + + apiMsgs := buildChatMessages(systemPrompt, combined) + + req := openai.ChatCompletionRequest{ + Model: model, + MaxTokens: 4096, + Temperature: float32(temp), + Messages: apiMsgs, + } + + startTime := time.Now() + resp, err := c.client.CreateChatCompletion(ctx, req) + latencyMs := time.Since(startTime).Milliseconds() + + if err != nil { + return "", 0, fmt.Errorf("OpenAI API error: %w", err) + } + + if len(resp.Choices) == 0 { + return "", 0, fmt.Errorf("OpenAI API returned no choices") + } + + responseText := resp.Choices[0].Message.Content + + if prefill != "" && !strings.HasPrefix(responseText, prefill) { + responseText = prefill + responseText + } + + tokens := resp.Usage.TotalTokens + if tokens == 0 { + tokens = estimateTokens(prompt) + estimateTokens(responseText) + } + + RecordMetrics(resp.Usage.PromptTokens, resp.Usage.CompletionTokens, latencyMs) + + return responseText, tokens, nil +} + +// AskWithRequest sends a prompt with all settings from the request (CSP - no client state). +// This is the primary method for the clone-based session architecture. +// When req.ToolDefs is non-nil, uses OpenAI function calling and returns a +// STOP:-prefixed response. Otherwise returns plain text (backward-compatible). +func (c *OpenAIClient) AskWithRequest(ctx context.Context, req AskRequest) (AskResponse, error) { + // Build messages from request history + combined := make([]Message, len(req.Messages)) + copy(combined, req.Messages) + + // Add tool results as tool-role messages + if len(req.ToolResults) > 0 { + for _, r := range req.ToolResults { + combined = append(combined, Message{ + Role: "user", + Content: fmt.Sprintf("Tool result for %s: %s", r.ToolUseID, r.Content), + }) + } + } else if req.Prompt != "" { + combined = append(combined, Message{Role: "user", Content: req.Prompt}) + } + + apiMsgs := buildChatMessages(req.SystemPrompt, combined) + + // Handle tool results: convert to OpenAI tool message format + // We need to rebuild the last messages if we have tool results + if len(req.ToolResults) > 0 { + // Remove the placeholder user messages we added above + apiMsgs = apiMsgs[:len(apiMsgs)-len(req.ToolResults)] + // Add proper tool result messages + for _, r := range req.ToolResults { + apiMsgs = append(apiMsgs, openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleTool, + Content: r.Content, + ToolCallID: r.ToolUseID, + }) + } + } + + model := req.Model + if model == "" { + c.mu.RLock() + model = c.model + c.mu.RUnlock() + } + + temp := req.Temperature + + chatReq := openai.ChatCompletionRequest{ + Model: model, + MaxTokens: 4096, + Temperature: float32(temp), + Messages: apiMsgs, + } + + // Attach tool definitions when present + if len(req.ToolDefs) > 0 { + chatReq.Tools = buildOpenAITools(req.ToolDefs) + chatReq.ToolChoice = "auto" + } + + var ( + responseText string + toolCalls []openai.ToolCall + finishReason openai.FinishReason + promptTokens int + completionToks int + totalTokens int + latencyMs int64 + ) + + startTime := time.Now() + + if req.StreamFunc != nil { + // Streaming path + chatReq.StreamOptions = &openai.StreamOptions{IncludeUsage: true} + stream, err := c.client.CreateChatCompletionStream(ctx, chatReq) + if err != nil { + return AskResponse{}, fmt.Errorf("OpenAI streaming error: %w", err) + } + defer stream.Close() + + var textParts []string + toolCallMap := make(map[int]*openai.ToolCall) + + for { + chunk, err := stream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return AskResponse{}, fmt.Errorf("OpenAI streaming error: %w", err) + } + + // Capture usage from the final chunk + if chunk.Usage != nil { + promptTokens = chunk.Usage.PromptTokens + completionToks = chunk.Usage.CompletionTokens + totalTokens = chunk.Usage.TotalTokens + } + + if len(chunk.Choices) == 0 { + continue + } + + choice := chunk.Choices[0] + finishReason = choice.FinishReason + + // Text delta + if choice.Delta.Content != "" { + textParts = append(textParts, choice.Delta.Content) + req.StreamFunc(choice.Delta.Content) + } + + // Tool call deltas + for _, tc := range choice.Delta.ToolCalls { + idx := 0 + if tc.Index != nil { + idx = *tc.Index + } + existing, ok := toolCallMap[idx] + if !ok { + toolCallMap[idx] = &openai.ToolCall{ + ID: tc.ID, + Type: tc.Type, + Function: openai.FunctionCall{ + Name: tc.Function.Name, + Arguments: tc.Function.Arguments, + }, + } + } else { + if tc.ID != "" { + existing.ID = tc.ID + } + if tc.Function.Name != "" { + existing.Function.Name += tc.Function.Name + } + existing.Function.Arguments += tc.Function.Arguments + } + } + } + + latencyMs = time.Since(startTime).Milliseconds() + responseText = strings.Join(textParts, "") + + // Collect tool calls in order + for i := 0; i < len(toolCallMap); i++ { + if tc, ok := toolCallMap[i]; ok { + toolCalls = append(toolCalls, *tc) + } + } + } else { + // Blocking path + resp, err := c.client.CreateChatCompletion(ctx, chatReq) + latencyMs = time.Since(startTime).Milliseconds() + if err != nil { + return AskResponse{}, fmt.Errorf("OpenAI API error: %w", err) + } + + if len(resp.Choices) == 0 { + return AskResponse{}, fmt.Errorf("OpenAI API returned no choices") + } + + responseText = resp.Choices[0].Message.Content + toolCalls = resp.Choices[0].Message.ToolCalls + finishReason = resp.Choices[0].FinishReason + promptTokens = resp.Usage.PromptTokens + completionToks = resp.Usage.CompletionTokens + totalTokens = resp.Usage.TotalTokens + } + + if totalTokens == 0 { + totalTokens = estimateTokens(responseText) + } + + RecordMetrics(promptTokens, completionToks, latencyMs) + + // Plain-text mode (no tools): return text as before + if len(req.ToolDefs) == 0 { + if req.Prefill != "" && !strings.HasPrefix(responseText, req.Prefill) { + responseText = req.Prefill + responseText + } + return AskResponse{Response: responseText, Tokens: totalTokens}, nil + } + + // Tool mode: format STOP: response and build structured JSON for history + var textParts []string + var toolCallEntries []struct{ id, name, args string } + var structBlocks []string + + if responseText != "" { + textParts = append(textParts, responseText) + escaped := jsonEscapeString(responseText) + structBlocks = append(structBlocks, fmt.Sprintf(`{"type":"text","text":"%s"}`, escaped)) + } + + for _, tc := range toolCalls { + toolCallEntries = append(toolCallEntries, struct{ id, name, args string }{ + tc.ID, tc.Function.Name, tc.Function.Arguments, + }) + inputJSON := tc.Function.Arguments + if inputJSON == "" { + inputJSON = "{}" + } + idEsc := jsonEscapeString(tc.ID) + nameEsc := jsonEscapeString(tc.Function.Name) + structBlocks = append(structBlocks, + fmt.Sprintf(`{"type":"tool_use","id":"%s","name":"%s","input":%s}`, idEsc, nameEsc, inputJSON)) + } + + structuredJSON := "" + if len(structBlocks) > 0 { + structuredJSON = "[" + strings.Join(structBlocks, ",") + "]" + } + + var sb strings.Builder + if finishReason == openai.FinishReasonToolCalls { + sb.WriteString("STOP:tool_use\n") + for _, tc := range toolCallEntries { + safeArgs := strings.ReplaceAll(tc.args, "\n", `\n`) + sb.WriteString(fmt.Sprintf("TOOL:%s:%s:%s\n", tc.id, tc.name, safeArgs)) + } + } else { + sb.WriteString("STOP:end_turn\n") + } + sb.WriteString(strings.Join(textParts, "")) + + return AskResponse{Response: sb.String(), StructuredJSON: structuredJSON, Tokens: totalTokens}, nil +} + +// buildOpenAITools converts ToolDef slice to OpenAI SDK tool params. +func buildOpenAITools(defs []ToolDef) []openai.Tool { + tools := make([]openai.Tool, 0, len(defs)) + for _, d := range defs { + tools = append(tools, openai.Tool{ + Type: openai.ToolTypeFunction, + Function: &openai.FunctionDefinition{ + Name: d.Name, + Description: d.Description, + Parameters: d.InputSchema, + }, + }) + } + return tools +} + +// StartStream begins streaming a response for the given prompt. +func (c *OpenAIClient) StartStream(ctx context.Context, prompt string) error { + c.mu.Lock() + if c.streaming { + c.mu.Unlock() + return fmt.Errorf("stream already in progress") + } + + c.messages = append(c.messages, Message{Role: "user", Content: prompt}) + apiMsgs := buildChatMessages(c.systemPrompt, c.messages) + model := c.model + temp := c.temperature + + c.streaming = true + c.streamChan = make(chan string, 100) + c.streamDone = make(chan struct{}) + c.mu.Unlock() + + go func() { + var fullResponse string + + defer func() { + c.mu.Lock() + if fullResponse != "" { + c.messages = append(c.messages, Message{Role: "assistant", Content: fullResponse}) + } + c.streaming = false + close(c.streamChan) + close(c.streamDone) + c.mu.Unlock() + }() + + chatReq := openai.ChatCompletionRequest{ + Model: model, + MaxTokens: 4096, + Temperature: float32(temp), + Messages: apiMsgs, + StreamOptions: &openai.StreamOptions{IncludeUsage: true}, + } + + stream, err := c.client.CreateChatCompletionStream(ctx, chatReq) + if err != nil { + select { + case c.streamChan <- fmt.Sprintf("[Error: %v]", err): + case <-ctx.Done(): + } + c.mu.Lock() + if len(c.messages) > 0 { + c.messages = c.messages[:len(c.messages)-1] + } + c.mu.Unlock() + return + } + defer stream.Close() + + var inputTokens, outputTokens int + + for { + chunk, err := stream.Recv() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + select { + case c.streamChan <- fmt.Sprintf("\n[Error: %v]", err): + case <-ctx.Done(): + } + if fullResponse == "" { + c.mu.Lock() + if len(c.messages) > 0 { + c.messages = c.messages[:len(c.messages)-1] + } + c.mu.Unlock() + } + return + } + + if chunk.Usage != nil { + inputTokens = chunk.Usage.PromptTokens + outputTokens = chunk.Usage.CompletionTokens + } + + if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" { + text := chunk.Choices[0].Delta.Content + fullResponse += text + select { + case c.streamChan <- text: + case <-ctx.Done(): + return + } + } + } + + tokens := inputTokens + outputTokens + if tokens == 0 { + tokens = estimateTokens(fullResponse) + } + + // Update token counts (fullResponse and messages handled in defer) + c.mu.Lock() + c.lastTokens = tokens + c.totalTokens += tokens + c.mu.Unlock() + }() + + return nil +} + +// ReadStreamChunk reads the next chunk from the stream. +func (c *OpenAIClient) ReadStreamChunk() (string, bool) { + c.mu.RLock() + streamChan := c.streamChan + c.mu.RUnlock() + + if streamChan == nil { + return "", false + } + + chunk, ok := <-streamChan + return chunk, ok +} + +// IsStreaming returns whether a stream is currently in progress. +func (c *OpenAIClient) IsStreaming() bool { + c.mu.RLock() + defer c.mu.RUnlock() + return c.streaming +} + +// WaitStream waits for the current stream to complete. +func (c *OpenAIClient) WaitStream() { + c.mu.RLock() + done := c.streamDone + c.mu.RUnlock() + + if done != nil { + <-done + } +} diff --git a/internal/llm/session.go b/internal/llm/session.go index f0459962fef446b890bbaadd00591b2df65ba3e4..8a05ce9740f974f78ca2385e0ab5e7e7030aa080 100644 --- a/internal/llm/session.go +++ b/internal/llm/session.go @@ -52,9 +52,9 @@ type Session struct { refs int32 // atomic reference count; session closed when it drops to 0 // Async generation + streaming support - streamCh chan string // raw text chunks during generation; nil when idle - doneCh chan struct{} // closed when generation completes; nil when idle - streamMu sync.Mutex // guards streamCh and doneCh + streamCh chan string // raw text chunks during generation; nil when idle + doneCh chan struct{} // closed when generation completes; nil when idle + streamMu sync.Mutex // guards streamCh and doneCh } // NewSession creates a new session with the given ID and defaults. @@ -662,7 +662,7 @@ func (sm *SessionManager) Compact(ctx context.Context, id int) error { // AskRequest contains all parameters for an API call. type AskRequest struct { Messages []Message - Prompt string // empty when ToolResults is set (tool_results IS the new user turn) + Prompt string // empty when ToolResults is set (tool_results IS the new user turn) Model string Temperature float64 SystemPrompt string diff --git a/internal/llm/session_compact_test.go b/internal/llm/session_compact_test.go index 073cb0df08a2719dc73efdf39d4d2d19f85b8ca9..7de15cf75b49c10449ee135fcc583d67d6d3e994 100644 --- a/internal/llm/session_compact_test.go +++ b/internal/llm/session_compact_test.go @@ -12,32 +12,35 @@ type mockAPIClient struct { askError error } -func (m *mockAPIClient) Model() string { return "claude-sonnet-4-20250514" } -func (m *mockAPIClient) SetModel(string) {} -func (m *mockAPIClient) Temperature() float64 { return 0.7 } -func (m *mockAPIClient) SetTemperature(float64) error { return nil } -func (m *mockAPIClient) SystemPrompt() string { return "" } -func (m *mockAPIClient) SetSystemPrompt(string) {} -func (m *mockAPIClient) ThinkingTokens() int { return 0 } -func (m *mockAPIClient) SetThinkingTokens(int) {} -func (m *mockAPIClient) Prefill() string { return "" } -func (m *mockAPIClient) SetPrefill(string) {} -func (m *mockAPIClient) LastTokens() int { return m.askTokens } -func (m *mockAPIClient) TotalTokens() int { return 0 } -func (m *mockAPIClient) ContextLimit() int { return 200000 } -func (m *mockAPIClient) Compact(context.Context) error { return nil } -func (m *mockAPIClient) Messages() []Message { return nil } -func (m *mockAPIClient) MessagesJSON() ([]byte, error) { return []byte("[]"), nil } -func (m *mockAPIClient) AddSystemMessage(string) {} -func (m *mockAPIClient) Reset() {} +func (m *mockAPIClient) Model() string { return "claude-sonnet-4-20250514" } +func (m *mockAPIClient) SetModel(string) {} +func (m *mockAPIClient) Temperature() float64 { return 0.7 } +func (m *mockAPIClient) SetTemperature(float64) error { return nil } +func (m *mockAPIClient) SystemPrompt() string { return "" } +func (m *mockAPIClient) SetSystemPrompt(string) {} +func (m *mockAPIClient) ThinkingTokens() int { return 0 } +func (m *mockAPIClient) SetThinkingTokens(int) {} +func (m *mockAPIClient) Prefill() string { return "" } +func (m *mockAPIClient) SetPrefill(string) {} +func (m *mockAPIClient) LastTokens() int { return m.askTokens } +func (m *mockAPIClient) TotalTokens() int { return 0 } +func (m *mockAPIClient) ContextLimit() int { return 200000 } +func (m *mockAPIClient) Compact(context.Context) error { return nil } +func (m *mockAPIClient) Messages() []Message { return nil } +func (m *mockAPIClient) MessagesJSON() ([]byte, error) { return []byte("[]"), nil } +func (m *mockAPIClient) AddSystemMessage(string) {} +func (m *mockAPIClient) Reset() {} func (m *mockAPIClient) Ask(_ context.Context, _ string) (string, error) { return m.askResponse, m.askError } func (m *mockAPIClient) AskWithHistory(_ context.Context, _ []Message, _ string) (string, int, error) { return m.askResponse, m.askTokens, m.askError } -func (m *mockAPIClient) AskWithRequest(_ context.Context, _ AskRequest) (string, int, error) { - return m.askResponse, m.askTokens, m.askError +func (m *mockAPIClient) AskWithRequest(_ context.Context, req AskRequest) (AskResponse, error) { + if m.askError != nil { + return AskResponse{}, m.askError + } + return AskResponse{Response: m.askResponse, Tokens: m.askTokens}, nil } func (m *mockAPIClient) StartStream(context.Context, string) error { return nil } func (m *mockAPIClient) ReadStreamChunk() (string, bool) { return "", false } diff --git a/internal/llmfs/mock_backend_test.go b/internal/llmfs/mock_backend_test.go index 567180aa3a35a7e90d0a873bcb5fcfbcb087dad0..8096973557b86ecb2bac2b61240a64eaceae6639 100644 --- a/internal/llmfs/mock_backend_test.go +++ b/internal/llmfs/mock_backend_test.go @@ -33,9 +33,9 @@ func NewMockBackend() *MockBackend { } } -func (m *MockBackend) Model() string { return m.model } -func (m *MockBackend) SetModel(model string) { m.model = model } -func (m *MockBackend) Temperature() float64 { return m.temperature } +func (m *MockBackend) Model() string { return m.model } +func (m *MockBackend) SetModel(model string) { m.model = model } +func (m *MockBackend) Temperature() float64 { return m.temperature } func (m *MockBackend) SetTemperature(temp float64) error { if temp < 0 || temp > 2 { return fmt.Errorf("invalid temperature") @@ -103,12 +103,12 @@ func (m *MockBackend) AskWithHistory(ctx context.Context, history []llm.Message, return m.askResponse, tokens, nil } -func (m *MockBackend) AskWithRequest(ctx context.Context, req llm.AskRequest) (string, int, error) { +func (m *MockBackend) AskWithRequest(ctx context.Context, req llm.AskRequest) (llm.AskResponse, error) { if m.askError != nil { - return "", 0, m.askError + return llm.AskResponse{}, m.askError } tokens := len(req.Prompt) + len(m.askResponse) - return m.askResponse, tokens, nil + return llm.AskResponse{Response: m.askResponse, Tokens: tokens}, nil } func (m *MockBackend) StartStream(ctx context.Context, prompt string) error { diff --git a/internal/protocol/fs.go b/internal/protocol/fs.go index e137af1e1b5f9feb873a491228890f42138b4494..9daa920ad8e372d894caacc7c4cf6cc690d2c229 100644 --- a/internal/protocol/fs.go +++ b/internal/protocol/fs.go @@ -91,8 +91,8 @@ func (f *BaseFile) Stat() Stat { } } -func (f *BaseFile) Open(mode uint8) error { return nil } -func (f *BaseFile) Close() error { return nil } +func (f *BaseFile) Open(mode uint8) error { return nil } +func (f *BaseFile) Close() error { return nil } func (f *BaseFile) Read(p []byte, offset int64) (int, error) { return 0, io.EOF } func (f *BaseFile) Write(p []byte, offset int64) (int, error) { return 0, ErrPermission }