feat: config, OpenAI-compatible streaming provider and safe tools
This commit is contained in:
OpenClaw Dev
commit
c1c7b0e005
13 changed files with 1031 additions and 0 deletions
53
internal/llm/llm.go
Normal file
53
internal/llm/llm.go
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
// Package llm defines the provider abstraction and message types.
|
||||
package llm
|
||||
|
||||
import "context"
|
||||
|
||||
type Role string
|
||||
|
||||
const (
|
||||
RoleSystem Role = "system"
|
||||
RoleUser Role = "user"
|
||||
RoleAssistant Role = "assistant"
|
||||
RoleTool Role = "tool"
|
||||
)
|
||||
|
||||
type ToolCall struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"` // raw JSON
|
||||
}
|
||||
|
||||
type Message struct {
|
||||
Role Role `json:"role"`
|
||||
Content string `json:"content"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
}
|
||||
|
||||
// ToolSpec describes a callable tool using JSON Schema parameters.
|
||||
type ToolSpec struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Parameters map[string]any `json:"parameters"`
|
||||
}
|
||||
|
||||
type Request struct {
|
||||
Messages []Message
|
||||
Tools []ToolSpec
|
||||
}
|
||||
|
||||
// Response is the final assembled assistant turn.
|
||||
type Response struct {
|
||||
Content string
|
||||
ToolCalls []ToolCall
|
||||
FinishReason string
|
||||
}
|
||||
|
||||
// DeltaFunc receives streamed text tokens. May be nil.
|
||||
type DeltaFunc func(text string)
|
||||
|
||||
type Provider interface {
|
||||
Chat(ctx context.Context, req Request, onDelta DeltaFunc) (Response, error)
|
||||
}
|
||||
205
internal/llm/openai.go
Normal file
205
internal/llm/openai.go
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
package llm
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// OpenAI is a provider for any OpenAI-compatible /chat/completions endpoint
|
||||
// (OpenAI, OpenRouter, Ollama, LM Studio, vLLM, ...). It always streams.
|
||||
type OpenAI struct {
|
||||
BaseURL string
|
||||
APIKey string
|
||||
Model string
|
||||
Temperature float64
|
||||
HTTP *http.Client
|
||||
}
|
||||
|
||||
func NewOpenAI(baseURL, apiKey, model string, temperature float64, timeout time.Duration) *OpenAI {
|
||||
// Default transport: TLS certificate verification stays ON.
|
||||
return &OpenAI{
|
||||
BaseURL: strings.TrimRight(baseURL, "/"), APIKey: apiKey, Model: model,
|
||||
Temperature: temperature, HTTP: &http.Client{Timeout: timeout},
|
||||
}
|
||||
}
|
||||
|
||||
type oaTool struct {
|
||||
Type string `json:"type"`
|
||||
Function ToolSpec `json:"function"`
|
||||
}
|
||||
|
||||
type oaFunc struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Arguments string `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
type oaToolCall struct {
|
||||
Index *int `json:"index,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function oaFunc `json:"function"`
|
||||
}
|
||||
|
||||
type oaMsg struct {
|
||||
Role string `json:"role"`
|
||||
Content *string `json:"content"`
|
||||
ToolCalls []oaToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
}
|
||||
|
||||
type oaReq struct {
|
||||
Model string `json:"model"`
|
||||
Messages []oaMsg `json:"messages"`
|
||||
Tools []oaTool `json:"tools,omitempty"`
|
||||
Stream bool `json:"stream"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
}
|
||||
|
||||
type oaChunk struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
ToolCalls []oaToolCall `json:"tool_calls"`
|
||||
} `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
func toOA(msgs []Message) []oaMsg {
|
||||
out := make([]oaMsg, 0, len(msgs))
|
||||
for _, m := range msgs {
|
||||
c := m.Content
|
||||
om := oaMsg{Role: string(m.Role), Content: &c, ToolCallID: m.ToolCallID, Name: m.Name}
|
||||
if m.Role == RoleAssistant && len(m.ToolCalls) > 0 && c == "" {
|
||||
om.Content = nil
|
||||
}
|
||||
for _, tc := range m.ToolCalls {
|
||||
om.ToolCalls = append(om.ToolCalls, oaToolCall{ID: tc.ID, Type: "function", Function: oaFunc{Name: tc.Name, Arguments: tc.Arguments}})
|
||||
}
|
||||
out = append(out, om)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (p *OpenAI) Chat(ctx context.Context, req Request, onDelta DeltaFunc) (Response, error) {
|
||||
body := oaReq{Model: p.Model, Messages: toOA(req.Messages), Stream: true, Temperature: p.Temperature}
|
||||
for _, t := range req.Tools {
|
||||
body.Tools = append(body.Tools, oaTool{Type: "function", Function: t})
|
||||
}
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
hr, err := http.NewRequestWithContext(ctx, http.MethodPost, p.BaseURL+"/chat/completions", bytes.NewReader(b))
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
hr.Header.Set("Content-Type", "application/json")
|
||||
hr.Header.Set("Accept", "text/event-stream")
|
||||
if p.APIKey != "" {
|
||||
hr.Header.Set("Authorization", "Bearer "+p.APIKey)
|
||||
}
|
||||
resp, err := p.HTTP.Do(hr)
|
||||
if err != nil {
|
||||
return Response{}, fmt.Errorf("llm request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode/100 != 2 {
|
||||
msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return Response{}, fmt.Errorf("llm http %d: %s", resp.StatusCode, strings.TrimSpace(string(msg)))
|
||||
}
|
||||
return parseStream(resp.Body, onDelta)
|
||||
}
|
||||
|
||||
type partialCall struct {
|
||||
id, name string
|
||||
args strings.Builder
|
||||
}
|
||||
|
||||
func parseStream(r io.Reader, onDelta DeltaFunc) (Response, error) {
|
||||
var out Response
|
||||
var text strings.Builder
|
||||
calls := map[int]*partialCall{}
|
||||
sc := bufio.NewScanner(r)
|
||||
sc.Buffer(make([]byte, 64*1024), 4<<20)
|
||||
for sc.Scan() {
|
||||
line := strings.TrimSpace(sc.Text())
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if data == "[DONE]" {
|
||||
break
|
||||
}
|
||||
var ch oaChunk
|
||||
if err := json.Unmarshal([]byte(data), &ch); err != nil {
|
||||
return out, fmt.Errorf("decode stream chunk: %w", err)
|
||||
}
|
||||
if ch.Error != nil {
|
||||
return out, fmt.Errorf("llm error: %s", ch.Error.Message)
|
||||
}
|
||||
for _, c := range ch.Choices {
|
||||
if c.Delta.Content != "" {
|
||||
text.WriteString(c.Delta.Content)
|
||||
if onDelta != nil {
|
||||
onDelta(c.Delta.Content)
|
||||
}
|
||||
}
|
||||
for i, tc := range c.Delta.ToolCalls {
|
||||
idx := i
|
||||
if tc.Index != nil {
|
||||
idx = *tc.Index
|
||||
}
|
||||
pc := calls[idx]
|
||||
if pc == nil {
|
||||
pc = &partialCall{}
|
||||
calls[idx] = pc
|
||||
}
|
||||
if tc.ID != "" {
|
||||
pc.id = tc.ID
|
||||
}
|
||||
if tc.Function.Name != "" {
|
||||
pc.name += tc.Function.Name
|
||||
}
|
||||
pc.args.WriteString(tc.Function.Arguments)
|
||||
}
|
||||
if c.FinishReason != nil {
|
||||
out.FinishReason = *c.FinishReason
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := sc.Err(); err != nil {
|
||||
return out, fmt.Errorf("read stream: %w", err)
|
||||
}
|
||||
out.Content = text.String()
|
||||
idxs := make([]int, 0, len(calls))
|
||||
for i := range calls {
|
||||
idxs = append(idxs, i)
|
||||
}
|
||||
sort.Ints(idxs)
|
||||
for _, i := range idxs {
|
||||
pc := calls[i]
|
||||
id := pc.id
|
||||
if id == "" {
|
||||
id = fmt.Sprintf("call_%d", i)
|
||||
}
|
||||
args := pc.args.String()
|
||||
if args == "" {
|
||||
args = "{}"
|
||||
}
|
||||
out.ToolCalls = append(out.ToolCalls, ToolCall{ID: id, Name: pc.name, Arguments: args})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
67
internal/llm/openai_test.go
Normal file
67
internal/llm/openai_test.go
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestOpenAIStreamTextAndToolCalls(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("Authorization") != "Bearer k" || r.URL.Path != "/v1/chat/completions" {
|
||||
http.Error(w, "bad", 400)
|
||||
return
|
||||
}
|
||||
var body map[string]any
|
||||
json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["stream"] != true || len(body["tools"].([]any)) != 1 {
|
||||
http.Error(w, "bad body", 400)
|
||||
return
|
||||
}
|
||||
chunks := []string{
|
||||
`{"choices":[{"delta":{"content":"Hel"}}]}`,
|
||||
`{"choices":[{"delta":{"content":"lo"}}]}`,
|
||||
`{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"time_now","arguments":"{\"tz\":"}}]}}]}`,
|
||||
`{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"UTC\"}"}}]},"finish_reason":"tool_calls"}]}`,
|
||||
}
|
||||
for _, c := range chunks {
|
||||
fmt.Fprintf(w, "data: %s\n\n", c)
|
||||
}
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
p := NewOpenAI(srv.URL+"/v1/", "k", "m", 0, 5*time.Second)
|
||||
var streamed strings.Builder
|
||||
resp, err := p.Chat(context.Background(), Request{
|
||||
Messages: []Message{{Role: RoleUser, Content: "hi"}},
|
||||
Tools: []ToolSpec{{Name: "time_now", Parameters: map[string]any{"type": "object"}}},
|
||||
}, func(s string) { streamed.WriteString(s) })
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp.Content != "Hello" || streamed.String() != "Hello" {
|
||||
t.Fatalf("content %q streamed %q", resp.Content, streamed.String())
|
||||
}
|
||||
if len(resp.ToolCalls) != 1 || resp.ToolCalls[0].Arguments != `{"tz":"UTC"}` || resp.ToolCalls[0].ID != "c1" {
|
||||
t.Fatalf("tool calls %+v", resp.ToolCalls)
|
||||
}
|
||||
if resp.FinishReason != "tool_calls" {
|
||||
t.Fatal(resp.FinishReason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIHTTPError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "nope", 401)
|
||||
}))
|
||||
defer srv.Close()
|
||||
_, err := NewOpenAI(srv.URL, "", "m", 0, time.Second).Chat(context.Background(), Request{}, nil)
|
||||
if err == nil || !strings.Contains(err.Error(), "401") {
|
||||
t.Fatalf("want 401 err, got %v", err)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue