Initial commit
This commit is contained in:
@@ -0,0 +1,15 @@
|
||||
# Go OpenAI-Compatible Client for llama.cpp
|
||||
|
||||
|
||||
This gem provides a Go implementation of an OpenAI-compatible API client that interfaces with llama.cpp models.
|
||||
It enables seamless integration with LLMs deployed via llama.cpp, supporting common APIs like chat, embeddings, and image generation.
|
||||
Features include:
|
||||
- REST API compatibility with OpenAI's model
|
||||
- Support for various model formats (GGUF, GGML)
|
||||
- Configuration via environment variables
|
||||
- Request/response streaming support
|
||||
- Token counting and usage metrics
|
||||
- Batch processing capabilities
|
||||
- Flexible prompt templating
|
||||
- Error handling with retry logic
|
||||
```
|
||||
@@ -0,0 +1,191 @@
|
||||
# Go OpenAPI Compatibility Layer
|
||||
|
||||
A Go library and CLI tool that provides compatibility with OpenAI API, enabling interactive sessions with local LLMs via llama.cpp.
|
||||
|
||||
## Overview
|
||||
|
||||
This project bridges the gap between OpenAI's API and local LLM inference engines (llama.cpp). It provides:
|
||||
|
||||
- **OpenAPI-compatible client** - Full support for ChatCompletion API
|
||||
- **Streaming responses** - Real-time token streaming like OpenAI's API
|
||||
- **Interactive CLI** - Session-based interaction with system prompts
|
||||
- **Tool support** - Custom tools for reasoning (abstraction analysis, instruction updates)
|
||||
- **Standard tools** - Built-in shell and file system tools
|
||||
- **Extensible registry** - Easy addition of custom tools
|
||||
|
||||
## Features
|
||||
|
||||
### Core Functionality
|
||||
|
||||
- **Local LLM Integration**: Works with GGUF model files loaded via llama.cpp
|
||||
- **System Prompts**: Configurable context for consistent responses
|
||||
- **Streaming Output**: Chunked responses for real-time interaction
|
||||
- **Reasoning Support**: Built-in reasoning content field for CoT models
|
||||
|
||||
### CLI Tools
|
||||
|
||||
| Tool | Description |
|
||||
|------|-------------|
|
||||
| `analyze_abstraction` | Analyzes code patterns (Decorator/Strategy) |
|
||||
| `update_instructions` | Updates the system prompt for the session |
|
||||
| `run_command` | Execute a shell command |
|
||||
| `read_file` | Read the contents of a file |
|
||||
| `write_file` | Write content to a file |
|
||||
| `list_dir` | List contents of a directory |
|
||||
| `replace_content` | Replace a block of text in a file |
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ CLI Interface │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Session Manager │
|
||||
│ - System prompts │
|
||||
│ - Session state │
|
||||
│ - Hook system (BeforeToolCall, OnReasoning, etc.) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Registry System │
|
||||
│ - MapRegistry (custom tools) │
|
||||
│ - ShellRegistry (command execution) │
|
||||
│ - FileRegistry (filesystem operations) │
|
||||
│ - CompositeRegistry (tool composition) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ OpenAPI Client │
|
||||
│ - ChatCompletion API implementation │
|
||||
│ - Request/Response handling │
|
||||
│ - Streaming support │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ llama.cpp (llama-cpp-go) │
|
||||
│ - GGUF model loading │
|
||||
│ - LLM inference │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
Environment variables:
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `OPENAI_BASE_URL` | API endpoint | `http://10.10.11.33:8082/v1/` |
|
||||
| `OPENAI_API_KEY` | API key for authentication | (empty) |
|
||||
|
||||
Model configuration (default):
|
||||
```
|
||||
Jackrong/Qwen3.5-4B-Claude-4.6-Opus-Reasoning-Distilled-GGUF
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Quick Start
|
||||
|
||||
```bash
|
||||
# Run the CLI (uses default settings)
|
||||
./main.go
|
||||
|
||||
# Run with custom model
|
||||
OPENAI_BASE_URL=http://localhost:8080/v1/ \
|
||||
OPENAI_API_KEY=your-api-key \
|
||||
./main.go
|
||||
```
|
||||
|
||||
### CLI Interaction
|
||||
|
||||
```bash
|
||||
$ ./main.go
|
||||
Interactive Architect Session (BaseURL: http://10.10.11.33:8082/v1/)
|
||||
Type your message and press Enter. Type 'exit' to quit.
|
||||
|
||||
User > What is the purpose of this project?
|
||||
Assistant > This is an OpenAPI compatibility layer that enables...
|
||||
|
||||
User > Analyze my current package structure
|
||||
Assistant > [Tool] Executing: analyze_abstraction({"name":"...", ...})...
|
||||
[Tool] Returning results...
|
||||
|
||||
User > exit
|
||||
```
|
||||
|
||||
### Custom Tool Registration
|
||||
|
||||
```go
|
||||
registry := openapi.NewMapRegistry()
|
||||
|
||||
// Register a custom tool
|
||||
registry.Register("my_custom_tool", "Describe the given function",
|
||||
map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"code": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
func(ctx context.Context, s *openapi.Session, args json.RawMessage) (any, error) {
|
||||
var params struct{ Code string }
|
||||
json.Unmarshal(args, ¶ms)
|
||||
// Your tool logic here
|
||||
return "Analysis complete.", nil
|
||||
})
|
||||
```
|
||||
|
||||
## Architecture Details
|
||||
|
||||
### Session Hooks
|
||||
|
||||
The session supports a powerful hook system for intercepting and modifying various events:
|
||||
|
||||
```go
|
||||
session.SetHooks(openapi.SessionHooks{
|
||||
BeforeToolCall: func(ctx context.Context, s *openapi.Session, tc *openapi.ToolCall) error {
|
||||
// Called before each tool is invoked
|
||||
return nil
|
||||
},
|
||||
OnReasoning: func(ctx context.Context, s *openapi.Session, reasoning string) error {
|
||||
// Called for each reasoning step
|
||||
return nil
|
||||
},
|
||||
OnDone: func(ctx context.Context, s *openapi.Session) error {
|
||||
// Called when the session completes
|
||||
return nil
|
||||
},
|
||||
OnError: func(ctx context.Context, s *openapi.Session, err error) {
|
||||
// Called on errors
|
||||
return nil
|
||||
},
|
||||
})
|
||||
```
|
||||
|
||||
### Registry Pattern
|
||||
|
||||
The composite registry allows combining multiple tool sources:
|
||||
|
||||
```go
|
||||
composite := &openapi.CompositeRegistry{}
|
||||
|
||||
// Base registry for custom tools
|
||||
baseRegistry := openapi.NewMapRegistry()
|
||||
composite.Add(baseRegistry)
|
||||
|
||||
// Add standard tools
|
||||
composite.Add(openapi.NewShellRegistry())
|
||||
composite.Add(openapi.NewFileRegistry("."))
|
||||
|
||||
// Apply to session
|
||||
session.SetRegistry(composite)
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,134 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/matst80/go-openapi-compat/pkg/openapi"
|
||||
)
|
||||
|
||||
func main() {
|
||||
baseURL := os.Getenv("OPENAI_BASE_URL")
|
||||
if baseURL == "" {
|
||||
baseURL = "http://10.10.11.33:8082/v1/"
|
||||
}
|
||||
apiKey := os.Getenv("OPENAI_API_KEY")
|
||||
|
||||
client := openapi.NewClient(baseURL, apiKey)
|
||||
model := "Jackrong/Qwen3.5-4B-Claude-4.6-Opus-Reasoning-Distilled-GGUF"
|
||||
|
||||
session := openapi.NewSession(client, model)
|
||||
session.DefaultRequestOptions.Seed = 42
|
||||
session.SetSystemPrompt("You are an expert software architect. You have access to shell and file tools to explore, analyze, and modify the codebase. Use them to provide deep insights and implement changes. blocks with git diffs will be applied")
|
||||
|
||||
// Setup composite registry with standard tools
|
||||
composite := &openapi.CompositeRegistry{}
|
||||
|
||||
// Base registry for custom tools
|
||||
baseRegistry := openapi.NewMapRegistry()
|
||||
composite.Add(baseRegistry)
|
||||
|
||||
// Add standard shell and file tools
|
||||
composite.Add(openapi.NewShellRegistry())
|
||||
composite.Add(openapi.NewFileRegistry(".")) // Use current directory as base
|
||||
|
||||
session.SetRegistry(composite)
|
||||
|
||||
// Tools for reasoning about abstractions
|
||||
baseRegistry.Register("analyze_abstraction", "Analyze a specific code abstraction",
|
||||
map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"name": map[string]any{"type": "string"},
|
||||
"description": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
func(ctx context.Context, s *openapi.Session, args json.RawMessage) (any, error) {
|
||||
var params struct{ Name, Description string }
|
||||
json.Unmarshal(args, ¶ms)
|
||||
return map[string]string{
|
||||
"analysis": fmt.Sprintf("The abstraction '%s' (%s) follows the 'Decorator' and 'Strategy' patterns by using hooks for extensibility.", params.Name, params.Description),
|
||||
"score": "9/10",
|
||||
}, nil
|
||||
})
|
||||
|
||||
baseRegistry.Register("update_instructions", "Update the architect's system instructions",
|
||||
map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"new_instructions": map[string]any{"type": "string"},
|
||||
},
|
||||
},
|
||||
func(ctx context.Context, s *openapi.Session, args json.RawMessage) (any, error) {
|
||||
var params struct{ NewInstructions string }
|
||||
json.Unmarshal(args, ¶ms)
|
||||
s.SetSystemPrompt(params.NewInstructions)
|
||||
return "Instructions updated successfully.", nil
|
||||
})
|
||||
|
||||
// Setup hooks for UI updates
|
||||
session.SetHooks(openapi.SessionHooks{
|
||||
BeforeToolCall: func(ctx context.Context, s *openapi.Session, tc *openapi.ToolCall) error {
|
||||
fmt.Printf("\n[Tool] Executing: %s(%s)...\n", tc.Function.Name, tc.Function.Arguments)
|
||||
return nil
|
||||
},
|
||||
OnReasoning: func(ctx context.Context, s *openapi.Session, reasoning string) error {
|
||||
fmt.Printf("\n[Thinking] %s", reasoning)
|
||||
return nil
|
||||
},
|
||||
OnDone: func(ctx context.Context, s *openapi.Session) error {
|
||||
fmt.Printf("\nDone.\n")
|
||||
return nil
|
||||
},
|
||||
OnError: func(ctx context.Context, s *openapi.Session, err error) {
|
||||
fmt.Printf("\n[Error] %v\n", err)
|
||||
},
|
||||
})
|
||||
|
||||
fmt.Printf("Interactive Architect Session (BaseURL: %s)\n", baseURL)
|
||||
fmt.Println("Type your message and press Enter. Type 'exit' to quit.")
|
||||
|
||||
scanner := bufio.NewScanner(os.Stdin)
|
||||
for {
|
||||
fmt.Print("\nUser > ")
|
||||
if !scanner.Scan() {
|
||||
break
|
||||
}
|
||||
input := scanner.Text()
|
||||
if input == "exit" || input == "quit" {
|
||||
break
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
|
||||
chunks, errs := session.Chat(ctx, input)
|
||||
|
||||
fmt.Print("Assistant > ")
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
for chunk := range chunks {
|
||||
if len(chunk.Choices) > 0 {
|
||||
fmt.Print(chunk.Choices[0].Delta.Content)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
fmt.Print("\nDone.\n")
|
||||
// Stream completed
|
||||
case err := <-errs:
|
||||
if err != nil {
|
||||
fmt.Printf("\nError: %v\n", err)
|
||||
}
|
||||
case <-ctx.Done():
|
||||
fmt.Println("\nTimeout reached.")
|
||||
}
|
||||
cancel()
|
||||
fmt.Println()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
BaseURL string
|
||||
APIKey string
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
func NewClient(baseURL, apiKey string) *Client {
|
||||
if !strings.HasSuffix(baseURL, "/") {
|
||||
baseURL += "/"
|
||||
}
|
||||
return &Client{
|
||||
BaseURL: baseURL,
|
||||
APIKey: apiKey,
|
||||
HTTPClient: &http.Client{},
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) CreateChatCompletion(ctx context.Context, req ChatCompletionRequest) (*ChatCompletionResponse, error) {
|
||||
req.Stream = false
|
||||
data, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal request: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"chat/completions", bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
|
||||
c.setHeaders(httpReq)
|
||||
|
||||
resp, err := c.HTTPClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("do request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var result ChatCompletionResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return nil, fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
func (c *Client) CreateChatCompletionStream(ctx context.Context, req ChatCompletionRequest) (<-chan ChatCompletionChunk, <-chan error) {
|
||||
req.Stream = true
|
||||
chunks := make(chan ChatCompletionChunk)
|
||||
errs := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
defer close(chunks)
|
||||
defer close(errs)
|
||||
|
||||
data, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
errs <- fmt.Errorf("marshal request: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"chat/completions", bytes.NewReader(data))
|
||||
if err != nil {
|
||||
errs <- fmt.Errorf("create request: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
c.setHeaders(httpReq)
|
||||
|
||||
resp, err := c.HTTPClient.Do(httpReq)
|
||||
if err != nil {
|
||||
errs <- fmt.Errorf("do request: %w", err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
errs <- fmt.Errorf("unexpected status %d: %s", resp.StatusCode, string(body))
|
||||
return
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(line, "data: ") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimPrefix(line, "data: ")
|
||||
if data == "[DONE]" {
|
||||
break
|
||||
}
|
||||
|
||||
var chunk ChatCompletionChunk
|
||||
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
||||
errs <- fmt.Errorf("unmarshal chunk: %w", err)
|
||||
return
|
||||
}
|
||||
select {
|
||||
case chunks <- chunk:
|
||||
case <-ctx.Done():
|
||||
errs <- ctx.Err()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
errs <- fmt.Errorf("scanner error: %w", err)
|
||||
}
|
||||
}()
|
||||
|
||||
return chunks, errs
|
||||
}
|
||||
|
||||
func (c *Client) setHeaders(req *http.Request) {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if c.APIKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCreateChatCompletion(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/chat/completions" {
|
||||
t.Errorf("expected path /chat/completions, got %s", r.URL.Path)
|
||||
}
|
||||
if r.Header.Get("Authorization") != "Bearer test-key" {
|
||||
t.Errorf("expected bearer token, got %s", r.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
var req ChatCompletionRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if !req.CachePrompt {
|
||||
t.Error("expected cache_prompt to be true")
|
||||
}
|
||||
|
||||
resp := ChatCompletionResponse{
|
||||
ID: "test-id",
|
||||
Model: "test-model",
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
Content: "Hello!",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
resp, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "Hi"},
|
||||
},
|
||||
CachePrompt: true,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.Content != "Hello!" {
|
||||
t.Errorf("expected Hello!, got %s", resp.Choices[0].Message.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionWithReasoning(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
resp := ChatCompletionResponse{
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
Content: "Hello!",
|
||||
ReasoningContent: "Thinking...",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
resp, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "Hi"},
|
||||
},
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if resp.Choices[0].Message.ReasoningContent != "Thinking..." {
|
||||
t.Errorf("expected Thinking..., got %s", resp.Choices[0].Message.ReasoningContent)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionWithToolCalls(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
resp := ChatCompletionResponse{
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
ToolCalls: []ToolCall{
|
||||
{
|
||||
ID: "call_123",
|
||||
Type: "function",
|
||||
Function: ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location":"Stockholm"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
resp, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "What is the weather in Stockholm?"},
|
||||
},
|
||||
Tools: []Tool{
|
||||
{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: "get_weather",
|
||||
Description: "Get the current weather",
|
||||
Parameters: json.RawMessage(`{"type":"object","properties":{"location":{"type":"string"}}}`),
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if len(resp.Choices[0].Message.ToolCalls) != 1 {
|
||||
t.Fatalf("expected 1 tool call, got %d", len(resp.Choices[0].Message.ToolCalls))
|
||||
}
|
||||
|
||||
tc := resp.Choices[0].Message.ToolCalls[0]
|
||||
if tc.Function.Name != "get_weather" {
|
||||
t.Errorf("expected get_weather, got %s", tc.Function.Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionWithToolResponse(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var req ChatCompletionRequest
|
||||
json.NewDecoder(r.Body).Decode(&req)
|
||||
|
||||
// Check if the tool response is present
|
||||
foundToolResponse := false
|
||||
for _, msg := range req.Messages {
|
||||
if msg.Role == RoleTool && msg.ToolCallID == "call_123" {
|
||||
foundToolResponse = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !foundToolResponse {
|
||||
t.Error("expected tool response message in request")
|
||||
}
|
||||
|
||||
resp := ChatCompletionResponse{
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
Content: "The weather in Stockholm is 15°C.",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
_, err := client.CreateChatCompletion(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "What is the weather in Stockholm?"},
|
||||
{
|
||||
Role: RoleAssistant,
|
||||
ToolCalls: []ToolCall{
|
||||
{
|
||||
ID: "call_123",
|
||||
Type: "function",
|
||||
Function: ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location":"Stockholm"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Role: RoleTool,
|
||||
ToolCallID: "call_123",
|
||||
Content: `{"temp":15}`,
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChatCompletionStreamWithReasoning(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"reasoning_content":"Thinking"}}]} `)
|
||||
fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"Hello"}}]} `)
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
chunks, errs := client.CreateChatCompletionStream(context.Background(), ChatCompletionRequest{
|
||||
Model: "test-model",
|
||||
Messages: []Message{
|
||||
{Role: RoleUser, Content: "Hi"},
|
||||
},
|
||||
})
|
||||
|
||||
var reasoning, content string
|
||||
for chunk := range chunks {
|
||||
if len(chunk.Choices) > 0 {
|
||||
reasoning += chunk.Choices[0].Delta.ReasoningContent
|
||||
content += chunk.Choices[0].Delta.Content
|
||||
}
|
||||
}
|
||||
|
||||
if err := <-errs; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if reasoning != "Thinking" {
|
||||
t.Errorf("expected Thinking, got %s", reasoning)
|
||||
}
|
||||
if content != "Hello" {
|
||||
t.Errorf("expected Hello, got %s", content)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// DiffBlock represents a single file change in a git diff.
|
||||
type DiffBlock struct {
|
||||
File string
|
||||
Content string // The full diff content for this file
|
||||
}
|
||||
|
||||
// DiffParser parses a stream of git diff output.
|
||||
type DiffParser struct {
|
||||
onBlock func(DiffBlock)
|
||||
current *DiffBlock
|
||||
sb strings.Builder
|
||||
}
|
||||
|
||||
// NewDiffParser creates a new DiffParser.
|
||||
func NewDiffParser(onBlock func(DiffBlock)) *DiffParser {
|
||||
return &DiffParser{onBlock: onBlock}
|
||||
}
|
||||
|
||||
// Parse reads from r and calls onBlock for each completed diff block.
|
||||
func (p *DiffParser) Parse(r io.Reader) error {
|
||||
scanner := bufio.NewScanner(r)
|
||||
for scanner.Scan() {
|
||||
p.Feed(scanner.Text())
|
||||
}
|
||||
p.Flush()
|
||||
return scanner.Err()
|
||||
}
|
||||
|
||||
// Feed processes a single line of diff output.
|
||||
func (p *DiffParser) Feed(line string) {
|
||||
if strings.HasPrefix(line, "diff --git ") {
|
||||
p.Flush()
|
||||
p.current = &DiffBlock{}
|
||||
// Extract file path from "diff --git a/path/to/file b/path/to/file"
|
||||
parts := strings.Split(line, " ")
|
||||
if len(parts) >= 4 {
|
||||
p.current.File = strings.TrimPrefix(parts[3], "b/")
|
||||
}
|
||||
}
|
||||
|
||||
if p.current != nil {
|
||||
p.sb.WriteString(line)
|
||||
p.sb.WriteString("\n")
|
||||
}
|
||||
}
|
||||
|
||||
// Flush finalizes the current block and emits it.
|
||||
func (p *DiffParser) Flush() {
|
||||
if p.current != nil {
|
||||
p.current.Content = p.sb.String()
|
||||
p.onBlock(*p.current)
|
||||
p.current = nil
|
||||
p.sb.Reset()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDiffParser(t *testing.T) {
|
||||
input := `diff --git a/pkg/openapi/session.go b/pkg/openapi/session.go
|
||||
index 1234567..89abcdef 100644
|
||||
--- a/pkg/openapi/session.go
|
||||
+++ b/pkg/openapi/session.go
|
||||
@@ -1,3 +1,4 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
+ "os"
|
||||
diff --git a/main.go b/main.go
|
||||
index 0000000..1111111
|
||||
--- a/main.go
|
||||
+++ b/main.go
|
||||
@@ -10,1 +10,1 @@
|
||||
-old
|
||||
+new
|
||||
`
|
||||
|
||||
var blocks []DiffBlock
|
||||
parser := NewDiffParser(func(b DiffBlock) {
|
||||
blocks = append(blocks, b)
|
||||
})
|
||||
|
||||
err := parser.Parse(strings.NewReader(input))
|
||||
if err != nil {
|
||||
t.Fatalf("Parse failed: %v", err)
|
||||
}
|
||||
|
||||
if len(blocks) != 2 {
|
||||
t.Errorf("Expected 2 blocks, got %d", len(blocks))
|
||||
}
|
||||
|
||||
if blocks[0].File != "pkg/openapi/session.go" {
|
||||
t.Errorf("Expected first file pkg/openapi/session.go, got %s", blocks[0].File)
|
||||
}
|
||||
|
||||
if blocks[1].File != "main.go" {
|
||||
t.Errorf("Expected second file main.go, got %s", blocks[1].File)
|
||||
}
|
||||
|
||||
if !strings.Contains(blocks[1].Content, "+new") {
|
||||
t.Errorf("Expected block 1 to contain '+new', got %s", blocks[1].Content)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SessionHooks provides hooks into the session lifecycle.
|
||||
type SessionHooks struct {
|
||||
BeforeMessage func(ctx context.Context, s *Session, msg *Message) error
|
||||
AfterMessage func(ctx context.Context, s *Session, msg *Message) error
|
||||
BeforeToolCall func(ctx context.Context, s *Session, tc *ToolCall) error
|
||||
AfterToolCall func(ctx context.Context, s *Session, tc *ToolCall, result any, err error) error
|
||||
OnReasoning func(ctx context.Context, s *Session, reasoning string) error
|
||||
OnDone func(ctx context.Context, s *Session) error
|
||||
OnError func(ctx context.Context, s *Session, err error)
|
||||
}
|
||||
|
||||
// Session manages a conversation history and tool execution logic.
|
||||
type Session struct {
|
||||
client *Client
|
||||
model string
|
||||
messages []Message
|
||||
registry ToolRegistry
|
||||
hooks SessionHooks
|
||||
|
||||
// TurnTimeout is the timeout for a single LLM request.
|
||||
TurnTimeout time.Duration
|
||||
|
||||
// DefaultRequestOptions allows setting default parameters for all requests in the session.
|
||||
DefaultRequestOptions ChatCompletionRequest
|
||||
}
|
||||
|
||||
// NewSession creates a new session with the given client and model.
|
||||
func NewSession(client *Client, model string) *Session {
|
||||
return &Session{
|
||||
client: client,
|
||||
model: model,
|
||||
registry: NewMapRegistry(),
|
||||
TurnTimeout: 2 * time.Minute,
|
||||
DefaultRequestOptions: ChatCompletionRequest{
|
||||
CachePrompt: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// SetRegistry replaces the tool registry.
|
||||
func (s *Session) SetRegistry(r ToolRegistry) {
|
||||
s.registry = r
|
||||
}
|
||||
|
||||
// RegisterTool is a helper for the default MapRegistry.
|
||||
func (s *Session) RegisterTool(name, description string, parameters any, handler ToolHandler) error {
|
||||
if mr, ok := s.registry.(*MapRegistry); ok {
|
||||
return mr.Register(name, description, parameters, handler)
|
||||
}
|
||||
return fmt.Errorf("current registry does not support direct registration")
|
||||
}
|
||||
|
||||
// AddMessage adds a message to the conversation history, triggering hooks.
|
||||
func (s *Session) AddMessage(ctx context.Context, msg Message) error {
|
||||
if s.hooks.BeforeMessage != nil {
|
||||
if err := s.hooks.BeforeMessage(ctx, s, &msg); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
s.messages = append(s.messages, msg)
|
||||
if s.hooks.AfterMessage != nil {
|
||||
return s.hooks.AfterMessage(ctx, s, &msg)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetSystemPrompt sets the initial system prompt.
|
||||
func (s *Session) SetSystemPrompt(prompt string) {
|
||||
if len(s.messages) > 0 && s.messages[0].Role == RoleSystem {
|
||||
s.messages[0].Content = prompt
|
||||
} else {
|
||||
s.messages = append([]Message{{Role: RoleSystem, Content: prompt}}, s.messages...)
|
||||
}
|
||||
}
|
||||
|
||||
// Chat sends a user prompt and handles the response loop asynchronously.
|
||||
func (s *Session) Chat(ctx context.Context, userPrompt string) (<-chan ChatCompletionChunk, <-chan error) {
|
||||
chunks := make(chan ChatCompletionChunk)
|
||||
errs := make(chan error, 1)
|
||||
go s.runLoop(ctx, userPrompt, chunks, errs)
|
||||
return chunks, errs
|
||||
}
|
||||
|
||||
func (s *Session) runLoop(ctx context.Context, userPrompt string, chunks chan<- ChatCompletionChunk, errs chan<- error) {
|
||||
defer close(chunks)
|
||||
defer close(errs)
|
||||
|
||||
if userPrompt != "" {
|
||||
if err := s.AddMessage(ctx, Message{Role: RoleUser, Content: userPrompt}); err != nil {
|
||||
s.handleError(ctx, err, errs)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
for {
|
||||
resp, err := s.nextCompletion(ctx)
|
||||
if err != nil {
|
||||
s.handleError(ctx, err, errs)
|
||||
return
|
||||
}
|
||||
|
||||
choice := resp.Choices[0]
|
||||
if err := s.processResponse(ctx, choice.Message); err != nil {
|
||||
s.handleError(ctx, err, errs)
|
||||
return
|
||||
}
|
||||
|
||||
if len(choice.Message.ToolCalls) > 0 {
|
||||
if err := s.handleToolCalls(ctx, choice.Message.ToolCalls); err != nil {
|
||||
s.handleError(ctx, err, errs)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if s.hooks.OnDone != nil {
|
||||
s.hooks.OnDone(ctx, s)
|
||||
}
|
||||
chunks <- s.asChunk(resp, choice)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) nextCompletion(ctx context.Context) (*ChatCompletionResponse, error) {
|
||||
turnCtx, cancel := s.extendContext(ctx)
|
||||
defer cancel()
|
||||
|
||||
req := s.DefaultRequestOptions
|
||||
req.Model = s.model
|
||||
req.Messages = s.messages
|
||||
req.Tools = s.registry.Tools(turnCtx)
|
||||
return s.client.CreateChatCompletion(turnCtx, req)
|
||||
}
|
||||
|
||||
func (s *Session) processResponse(ctx context.Context, msg Message) error {
|
||||
if msg.ReasoningContent != "" && s.hooks.OnReasoning != nil {
|
||||
if err := s.hooks.OnReasoning(ctx, s, msg.ReasoningContent); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return s.AddMessage(ctx, msg)
|
||||
}
|
||||
|
||||
func (s *Session) handleToolCalls(ctx context.Context, toolCalls []ToolCall) error {
|
||||
for i := range toolCalls {
|
||||
tc := &toolCalls[i]
|
||||
if s.hooks.BeforeToolCall != nil {
|
||||
if err := s.hooks.BeforeToolCall(ctx, s, tc); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Use extended context for tool execution to prevent timeouts on long tasks
|
||||
toolCtx, cancel := s.extendContext(ctx)
|
||||
result, err := s.registry.Execute(toolCtx, s, tc.Function.Name, json.RawMessage(tc.Function.Arguments))
|
||||
cancel()
|
||||
|
||||
if s.hooks.AfterToolCall != nil {
|
||||
if hErr := s.hooks.AfterToolCall(ctx, s, tc, result, err); hErr != nil {
|
||||
return hErr
|
||||
}
|
||||
}
|
||||
|
||||
content := ""
|
||||
if err != nil {
|
||||
content = fmt.Sprintf("Error: %v", err)
|
||||
} else {
|
||||
resJSON, _ := json.Marshal(result)
|
||||
content = string(resJSON)
|
||||
}
|
||||
|
||||
if err := s.AddMessage(ctx, Message{
|
||||
Role: RoleTool,
|
||||
Name: tc.Function.Name,
|
||||
ToolCallID: tc.ID,
|
||||
Content: content,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// extendContext creates a context that resets the timeout for a new "turn",
|
||||
// ignoring the parent's deadline but respecting its cancellation.
|
||||
func (s *Session) extendContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
||||
extCtx, cancel := context.WithTimeout(context.Background(), s.TurnTimeout)
|
||||
|
||||
stop := make(chan struct{})
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
cancel()
|
||||
case <-stop:
|
||||
case <-extCtx.Done():
|
||||
}
|
||||
}()
|
||||
|
||||
return extCtx, func() {
|
||||
close(stop)
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) asChunk(resp *ChatCompletionResponse, choice Choice) ChatCompletionChunk {
|
||||
return ChatCompletionChunk{
|
||||
ID: resp.ID,
|
||||
Object: "chat.completion.chunk",
|
||||
Created: resp.Created,
|
||||
Model: resp.Model,
|
||||
Choices: []struct {
|
||||
Index int `json:"index"`
|
||||
Delta Message `json:"delta"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}{
|
||||
{
|
||||
Index: choice.Index,
|
||||
Delta: Message{
|
||||
Role: choice.Message.Role,
|
||||
Content: choice.Message.Content,
|
||||
ReasoningContent: choice.Message.ReasoningContent,
|
||||
},
|
||||
FinishReason: choice.FinishReason,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) handleError(ctx context.Context, err error, errs chan<- error) {
|
||||
if s.hooks.OnError != nil {
|
||||
s.hooks.OnError(ctx, s, err)
|
||||
}
|
||||
errs <- err
|
||||
}
|
||||
|
||||
func (s *Session) Messages() []Message { return s.messages }
|
||||
func (s *Session) SetHooks(h SessionHooks) { s.hooks = h }
|
||||
@@ -0,0 +1,135 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSessionToolTurns(t *testing.T) {
|
||||
callCount := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
var resp ChatCompletionResponse
|
||||
|
||||
if callCount == 1 {
|
||||
// First call: Model decides to call a tool
|
||||
resp = ChatCompletionResponse{
|
||||
ID: "resp1",
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
ToolCalls: []ToolCall{
|
||||
{
|
||||
ID: "call_1",
|
||||
Type: "function",
|
||||
Function: ToolCallFunction{
|
||||
Name: "get_weather",
|
||||
Arguments: `{"location":"Stockholm"}`,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
FinishReason: "tool_calls",
|
||||
},
|
||||
},
|
||||
}
|
||||
} else {
|
||||
// Second call: Model gives final response after tool result
|
||||
resp = ChatCompletionResponse{
|
||||
ID: "resp2",
|
||||
Choices: []Choice{
|
||||
{
|
||||
Message: Message{
|
||||
Role: RoleAssistant,
|
||||
Content: "The weather in Stockholm is sunny.",
|
||||
},
|
||||
FinishReason: "stop",
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewClient(server.URL, "test-key")
|
||||
session := NewSession(client, "test-model")
|
||||
|
||||
type weatherArgs struct {
|
||||
Location string `json:"location"`
|
||||
}
|
||||
|
||||
toolExecuted := false
|
||||
session.RegisterTool("get_weather", "Get weather",
|
||||
map[string]any{"type": "object", "properties": map[string]any{"location": map[string]any{"type": "string"}}},
|
||||
func(ctx context.Context, s *Session, args json.RawMessage) (any, error) {
|
||||
var w weatherArgs
|
||||
json.Unmarshal(args, &w)
|
||||
if w.Location != "Stockholm" {
|
||||
t.Errorf("expected Stockholm, got %s", w.Location)
|
||||
}
|
||||
toolExecuted = true
|
||||
return map[string]any{"temp": 20, "condition": "sunny"}, nil
|
||||
})
|
||||
|
||||
hooksCalled := make(map[string]bool)
|
||||
session.SetHooks(SessionHooks{
|
||||
BeforeMessage: func(ctx context.Context, s *Session, msg *Message) error {
|
||||
hooksCalled["BeforeMessage"] = true
|
||||
return nil
|
||||
},
|
||||
AfterMessage: func(ctx context.Context, s *Session, msg *Message) error {
|
||||
hooksCalled["AfterMessage"] = true
|
||||
return nil
|
||||
},
|
||||
BeforeToolCall: func(ctx context.Context, s *Session, tc *ToolCall) error {
|
||||
hooksCalled["BeforeToolCall"] = true
|
||||
return nil
|
||||
},
|
||||
AfterToolCall: func(ctx context.Context, s *Session, tc *ToolCall, result any, err error) error {
|
||||
hooksCalled["AfterToolCall"] = true
|
||||
return nil
|
||||
},
|
||||
OnDone: func(ctx context.Context, s *Session) error {
|
||||
hooksCalled["OnDone"] = true
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
chunks, errs := session.Chat(ctx, "How is the weather in Stockholm?")
|
||||
|
||||
var finalContent string
|
||||
for chunk := range chunks {
|
||||
if len(chunk.Choices) > 0 {
|
||||
finalContent += chunk.Choices[0].Delta.Content
|
||||
}
|
||||
}
|
||||
|
||||
if err := <-errs; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if !toolExecuted {
|
||||
t.Error("expected tool to be executed")
|
||||
}
|
||||
|
||||
hookNames := []string{"BeforeMessage", "AfterMessage", "BeforeToolCall", "AfterToolCall", "OnDone"}
|
||||
for _, name := range hookNames {
|
||||
if !hooksCalled[name] {
|
||||
t.Errorf("expected hook %s to be called", name)
|
||||
}
|
||||
}
|
||||
|
||||
if finalContent != "The weather in Stockholm is sunny." {
|
||||
t.Errorf("expected final response, got %s", finalContent)
|
||||
}
|
||||
|
||||
if len(session.Messages()) != 4 {
|
||||
t.Errorf("expected 4 messages in history, got %d", len(session.Messages()))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// RunCommandArgs defines the arguments for the run_command tool.
|
||||
type RunCommandArgs struct {
|
||||
Command string `json:"command" desc:"The command to run"`
|
||||
}
|
||||
|
||||
// NewShellRegistry returns a registry with shell-related tools.
|
||||
func NewShellRegistry() ToolRegistry {
|
||||
r := NewMapRegistry()
|
||||
|
||||
RegisterTool(r, "run_command", "Execute a shell command",
|
||||
func(ctx context.Context, s *Session, params RunCommandArgs) (any, error) {
|
||||
cmd := exec.CommandContext(ctx, "sh", "-c", params.Command)
|
||||
output, err := cmd.CombinedOutput()
|
||||
return map[string]any{
|
||||
"output": string(output),
|
||||
"error": fmt.Sprintf("%v", err),
|
||||
}, nil
|
||||
})
|
||||
|
||||
return r
|
||||
}
|
||||
|
||||
// ReadFileArgs defines the arguments for the read_file tool.
|
||||
type ReadFileArgs struct {
|
||||
Path string `json:"path" desc:"Path to the file"`
|
||||
Start int `json:"start,omitempty" desc:"Start line number"`
|
||||
Lines int `json:"lines,omitempty" desc:"number of lines to read, default 200"`
|
||||
}
|
||||
|
||||
// WriteFileArgs defines the arguments for the write_file tool.
|
||||
type WriteFileArgs struct {
|
||||
Path string `json:"path" desc:"Path to the file"`
|
||||
Content string `json:"content" desc:"Content to write"`
|
||||
}
|
||||
|
||||
// ListDirArgs defines the arguments for the list_dir tool.
|
||||
type ListDirArgs struct {
|
||||
Path string `json:"path" desc:"Path to the directory"`
|
||||
}
|
||||
|
||||
// ReplaceContentArgs defines the arguments for the replace_content tool.
|
||||
type ReplaceContentArgs struct {
|
||||
Path string `json:"path" desc:"Path to the file"`
|
||||
TargetContent string `json:"target_content" desc:"The exact content to replace"`
|
||||
NewContent string `json:"new_content" desc:"The new content"`
|
||||
}
|
||||
|
||||
// NewFileRegistry returns a registry with file-related tools.
|
||||
func NewFileRegistry(baseDir string) ToolRegistry {
|
||||
r := NewMapRegistry()
|
||||
|
||||
RegisterTool(r, "read_file", "Read the contents of a file",
|
||||
func(ctx context.Context, s *Session, params ReadFileArgs) (any, error) {
|
||||
fullPath := filepath.Join(baseDir, params.Path)
|
||||
content, err := os.ReadFile(fullPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
lines := strings.Split(string(content), "\n")
|
||||
if params.Start < 1 {
|
||||
params.Start = 1
|
||||
}
|
||||
if params.Lines == 0 {
|
||||
params.Lines = 200
|
||||
}
|
||||
|
||||
start := params.Start - 1
|
||||
end := start + params.Lines
|
||||
if end > len(lines) {
|
||||
end = len(lines)
|
||||
}
|
||||
|
||||
if start >= end {
|
||||
return nil, fmt.Errorf("invalid start or lines parameters")
|
||||
}
|
||||
|
||||
return strings.Join(lines[start:end], "\n"), nil
|
||||
})
|
||||
|
||||
RegisterTool(r, "write_file", "Write content to a file",
|
||||
func(ctx context.Context, s *Session, params WriteFileArgs) (any, error) {
|
||||
fullPath := filepath.Join(baseDir, params.Path)
|
||||
if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := os.WriteFile(fullPath, []byte(params.Content), 0644); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return "File written successfully", nil
|
||||
})
|
||||
|
||||
RegisterTool(r, "list_dir", "List contents of a directory",
|
||||
func(ctx context.Context, s *Session, params ListDirArgs) (any, error) {
|
||||
fullPath := filepath.Join(baseDir, params.Path)
|
||||
entries, err := os.ReadDir(fullPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result []string
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.IsDir() {
|
||||
name += "/"
|
||||
}
|
||||
result = append(result, name)
|
||||
}
|
||||
return result, nil
|
||||
})
|
||||
|
||||
RegisterTool(r, "replace_content", "Replace a block of text in a file",
|
||||
func(ctx context.Context, s *Session, params ReplaceContentArgs) (any, error) {
|
||||
fullPath := filepath.Join(baseDir, params.Path)
|
||||
content, err := os.ReadFile(fullPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
oldContent := string(content)
|
||||
if !strings.Contains(oldContent, params.TargetContent) {
|
||||
return nil, fmt.Errorf("target content not found in file")
|
||||
}
|
||||
|
||||
newContent := strings.Replace(oldContent, params.TargetContent, params.NewContent, 1)
|
||||
if err := os.WriteFile(fullPath, []byte(newContent), 0644); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return "Content replaced successfully", nil
|
||||
})
|
||||
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestShellRegistry(t *testing.T) {
|
||||
registry := NewShellRegistry()
|
||||
ctx := context.Background()
|
||||
session := &Session{}
|
||||
|
||||
args := json.RawMessage(`{"command": "echo 'hello world'"}`)
|
||||
result, err := registry.Execute(ctx, session, "run_command", args)
|
||||
if err != nil {
|
||||
t.Fatalf("Execute failed: %v", err)
|
||||
}
|
||||
|
||||
resMap := result.(map[string]any)
|
||||
if resMap["output"] != "hello world\n" {
|
||||
t.Errorf("Expected 'hello world\n', got %q", resMap["output"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestFileRegistry(t *testing.T) {
|
||||
tempDir, err := os.MkdirTemp("", "file-registry-test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tempDir)
|
||||
|
||||
registry := NewFileRegistry(tempDir)
|
||||
ctx := context.Background()
|
||||
session := &Session{}
|
||||
|
||||
// Test write_file
|
||||
writeArgs := json.RawMessage(`{"path": "test.txt", "content": "hello file"}`)
|
||||
_, err = registry.Execute(ctx, session, "write_file", writeArgs)
|
||||
if err != nil {
|
||||
t.Fatalf("write_file failed: %v", err)
|
||||
}
|
||||
|
||||
content, err := os.ReadFile(filepath.Join(tempDir, "test.txt"))
|
||||
if err != nil || string(content) != "hello file" {
|
||||
t.Errorf("File content mismatch: %q", string(content))
|
||||
}
|
||||
|
||||
// Test read_file
|
||||
readArgs := json.RawMessage(`{"path": "test.txt"}`)
|
||||
res, err := registry.Execute(ctx, session, "read_file", readArgs)
|
||||
if err != nil {
|
||||
t.Fatalf("read_file failed: %v", err)
|
||||
}
|
||||
if res.(string) != "hello file" {
|
||||
t.Errorf("read_file output mismatch: %q", res)
|
||||
}
|
||||
|
||||
// Test list_dir
|
||||
listArgs := json.RawMessage(`{"path": "."}`)
|
||||
res, err = registry.Execute(ctx, session, "list_dir", listArgs)
|
||||
if err != nil {
|
||||
t.Fatalf("list_dir failed: %v", err)
|
||||
}
|
||||
files := res.([]string)
|
||||
if len(files) != 1 || files[0] != "test.txt" {
|
||||
t.Errorf("list_dir output mismatch: %v", files)
|
||||
}
|
||||
|
||||
// Test replace_content
|
||||
replaceArgs := json.RawMessage(`{"path": "test.txt", "target_content": "file", "new_content": "replacement"}`)
|
||||
_, err = registry.Execute(ctx, session, "replace_content", replaceArgs)
|
||||
if err != nil {
|
||||
t.Fatalf("replace_content failed: %v", err)
|
||||
}
|
||||
|
||||
content, err = os.ReadFile(filepath.Join(tempDir, "test.txt"))
|
||||
if err != nil || string(content) != "hello replacement" {
|
||||
t.Errorf("File content mismatch after replacement: %q", string(content))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ReflectRegistry provides helper methods for registering tools using reflection.
|
||||
type ReflectRegistry struct {
|
||||
*MapRegistry
|
||||
}
|
||||
|
||||
// NewReflectRegistry creates a new ReflectRegistry wrapping a MapRegistry.
|
||||
func NewReflectRegistry(r *MapRegistry) *ReflectRegistry {
|
||||
return &ReflectRegistry{MapRegistry: r}
|
||||
}
|
||||
|
||||
// RegisterTool registers a tool with a typed handler.
|
||||
// The handler's third argument must be a struct that defines the tool's parameters.
|
||||
func RegisterTool[T any](r *MapRegistry, name, description string, handler func(context.Context, *Session, T) (any, error)) error {
|
||||
var arg T
|
||||
schema, err := generateSchema(reflect.TypeOf(arg))
|
||||
if err != nil {
|
||||
return fmt.Errorf("generate schema: %w", err)
|
||||
}
|
||||
|
||||
wrappedHandler := func(ctx context.Context, s *Session, args json.RawMessage) (any, error) {
|
||||
var params T
|
||||
if err := json.Unmarshal(args, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal arguments: %w", err)
|
||||
}
|
||||
return handler(ctx, s, params)
|
||||
}
|
||||
|
||||
return r.Register(name, description, schema, wrappedHandler)
|
||||
}
|
||||
|
||||
func generateSchema(t reflect.Type) (map[string]any, error) {
|
||||
if t.Kind() == reflect.Ptr {
|
||||
t = t.Elem()
|
||||
}
|
||||
|
||||
if t.Kind() != reflect.Struct {
|
||||
return nil, fmt.Errorf("expected struct, got %s", t.Kind())
|
||||
}
|
||||
|
||||
properties := make(map[string]any)
|
||||
var required []string
|
||||
|
||||
for i := 0; i < t.NumField(); i++ {
|
||||
field := t.Field(i)
|
||||
jsonTag := field.Tag.Get("json")
|
||||
if jsonTag == "" || jsonTag == "-" {
|
||||
continue
|
||||
}
|
||||
|
||||
// Split json tag to handle omitempty etc.
|
||||
parts := strings.Split(jsonTag, ",")
|
||||
name := parts[0]
|
||||
|
||||
prop := make(map[string]any)
|
||||
|
||||
switch field.Type.Kind() {
|
||||
case reflect.String:
|
||||
prop["type"] = "string"
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
|
||||
prop["type"] = "integer"
|
||||
case reflect.Float32, reflect.Float64:
|
||||
prop["type"] = "number"
|
||||
case reflect.Bool:
|
||||
prop["type"] = "boolean"
|
||||
case reflect.Slice:
|
||||
prop["type"] = "array"
|
||||
// Simple slice handling, could be improved
|
||||
itemType, _ := getJsonType(field.Type.Elem())
|
||||
prop["items"] = map[string]string{"type": itemType}
|
||||
case reflect.Struct:
|
||||
subSchema, err := generateSchema(field.Type)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prop = subSchema
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported field type: %s", field.Type.Kind())
|
||||
}
|
||||
|
||||
if desc := field.Tag.Get("desc"); desc != "" {
|
||||
prop["description"] = desc
|
||||
}
|
||||
|
||||
properties[name] = prop
|
||||
|
||||
// Determine if field is required.
|
||||
// Priority: 'required' tag > 'omitempty' in 'json' tag.
|
||||
isRequired := true
|
||||
for _, p := range parts[1:] {
|
||||
if p == "omitempty" {
|
||||
isRequired = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if reqTag := field.Tag.Get("required"); reqTag != "" {
|
||||
isRequired = (reqTag == "true")
|
||||
}
|
||||
|
||||
if isRequired {
|
||||
required = append(required, name)
|
||||
}
|
||||
}
|
||||
|
||||
schema := map[string]any{
|
||||
"type": "object",
|
||||
"properties": properties,
|
||||
}
|
||||
if len(required) > 0 {
|
||||
schema["required"] = required
|
||||
}
|
||||
|
||||
return schema, nil
|
||||
}
|
||||
|
||||
func getJsonType(t reflect.Type) (string, error) {
|
||||
switch t.Kind() {
|
||||
case reflect.String:
|
||||
return "string", nil
|
||||
case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
|
||||
return "integer", nil
|
||||
case reflect.Float32, reflect.Float64:
|
||||
return "number", nil
|
||||
case reflect.Bool:
|
||||
return "boolean", nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported item type: %s", t.Kind())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type TestArgs struct {
|
||||
Name string `json:"name" desc:"The name of the user"`
|
||||
Age int `json:"age" desc:"The age of the user"`
|
||||
Tags []string `json:"tags,omitempty" desc:"Optional tags"`
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func TestRegisterTool(t *testing.T) {
|
||||
registry := NewMapRegistry()
|
||||
|
||||
err := RegisterTool(registry, "test_tool", "A test tool", func(ctx context.Context, s *Session, args TestArgs) (any, error) {
|
||||
return args.Name, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("failed to register tool: %v", err)
|
||||
}
|
||||
|
||||
tools := registry.Tools(context.Background())
|
||||
if len(tools) != 1 {
|
||||
t.Fatalf("expected 1 tool, got %d", len(tools))
|
||||
}
|
||||
|
||||
tool := tools[0]
|
||||
if tool.Function.Name != "test_tool" {
|
||||
t.Errorf("expected tool name 'test_tool', got '%s'", tool.Function.Name)
|
||||
}
|
||||
|
||||
var schema map[string]any
|
||||
if err := json.Unmarshal(tool.Function.Parameters, &schema); err != nil {
|
||||
t.Fatalf("failed to unmarshal parameters: %v", err)
|
||||
}
|
||||
|
||||
// Check schema
|
||||
if schema["type"] != "object" {
|
||||
t.Errorf("expected type 'object', got '%v'", schema["type"])
|
||||
}
|
||||
|
||||
properties := schema["properties"].(map[string]any)
|
||||
if _, ok := properties["name"]; !ok {
|
||||
t.Error("missing 'name' property")
|
||||
}
|
||||
if properties["name"].(map[string]any)["description"] != "The name of the user" {
|
||||
t.Errorf("wrong description for 'name': %v", properties["name"].(map[string]any)["description"])
|
||||
}
|
||||
|
||||
required := schema["required"].([]any)
|
||||
foundName := false
|
||||
for _, r := range required {
|
||||
if r == "name" {
|
||||
foundName = true
|
||||
}
|
||||
}
|
||||
if !foundName {
|
||||
t.Error("'name' should be required")
|
||||
}
|
||||
|
||||
// Test execution
|
||||
args := json.RawMessage(`{"name": "Alice", "age": 30, "enabled": true}`)
|
||||
res, err := registry.Execute(context.Background(), nil, "test_tool", args)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to execute tool: %v", err)
|
||||
}
|
||||
if res != "Alice" {
|
||||
t.Errorf("expected result 'Alice', got '%v'", res)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequiredTag(t *testing.T) {
|
||||
type RequiredArgs struct {
|
||||
ExplicitTrue string `json:"explicit_true,omitempty" required:"true"`
|
||||
ExplicitFalse string `json:"explicit_false" required:"false"`
|
||||
DefaultTrue string `json:"default_true"`
|
||||
DefaultFalse string `json:"default_false,omitempty"`
|
||||
}
|
||||
|
||||
registry := NewMapRegistry()
|
||||
err := RegisterTool(registry, "req_tool", "A tool with required tags", func(ctx context.Context, s *Session, args RequiredArgs) (any, error) {
|
||||
return nil, nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("failed to register tool: %v", err)
|
||||
}
|
||||
|
||||
tools := registry.Tools(context.Background())
|
||||
tool := tools[0]
|
||||
var schema map[string]any
|
||||
if err := json.Unmarshal(tool.Function.Parameters, &schema); err != nil {
|
||||
t.Fatalf("failed to unmarshal parameters: %v", err)
|
||||
}
|
||||
|
||||
required := schema["required"].([]any)
|
||||
reqMap := make(map[string]bool)
|
||||
for _, r := range required {
|
||||
reqMap[r.(string)] = true
|
||||
}
|
||||
|
||||
if !reqMap["explicit_true"] {
|
||||
t.Error("expected 'explicit_true' to be required")
|
||||
}
|
||||
if reqMap["explicit_false"] {
|
||||
t.Error("expected 'explicit_false' to NOT be required")
|
||||
}
|
||||
if !reqMap["default_true"] {
|
||||
t.Error("expected 'default_true' to be required")
|
||||
}
|
||||
if reqMap["default_false"] {
|
||||
t.Error("expected 'default_false' to NOT be required")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package openapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// ToolHandler is a function that implements a tool's logic.
|
||||
type ToolHandler func(ctx context.Context, s *Session, args json.RawMessage) (any, error)
|
||||
|
||||
// ToolDefinition pairs the OpenAI-compatible tool schema with a Go handler.
|
||||
type ToolDefinition struct {
|
||||
Tool Tool
|
||||
Handler ToolHandler
|
||||
}
|
||||
|
||||
// ToolRegistry defines the interface for tool discovery and execution.
|
||||
type ToolRegistry interface {
|
||||
Tools(ctx context.Context) []Tool
|
||||
Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error)
|
||||
}
|
||||
|
||||
// MapRegistry is a simple map-based implementation of ToolRegistry.
|
||||
type MapRegistry struct {
|
||||
tools map[string]ToolDefinition
|
||||
}
|
||||
|
||||
// NewMapRegistry creates a new MapRegistry.
|
||||
func NewMapRegistry() *MapRegistry {
|
||||
return &MapRegistry{tools: make(map[string]ToolDefinition)}
|
||||
}
|
||||
|
||||
// Register adds a tool to the registry.
|
||||
func (r *MapRegistry) Register(name, description string, parameters any, handler ToolHandler) error {
|
||||
paramsJSON, err := json.Marshal(parameters)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal parameters: %w", err)
|
||||
}
|
||||
|
||||
r.tools[name] = ToolDefinition{
|
||||
Tool: Tool{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: name,
|
||||
Description: description,
|
||||
Parameters: paramsJSON,
|
||||
},
|
||||
},
|
||||
Handler: handler,
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Tools returns the list of tools in the registry.
|
||||
func (r *MapRegistry) Tools(ctx context.Context) []Tool {
|
||||
tools := make([]Tool, 0, len(r.tools))
|
||||
for _, td := range r.tools {
|
||||
tools = append(tools, td.Tool)
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// Execute looks up and runs a tool by name.
|
||||
func (r *MapRegistry) Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error) {
|
||||
td, ok := r.tools[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("tool %s not found", name)
|
||||
}
|
||||
return td.Handler(ctx, s, args)
|
||||
}
|
||||
|
||||
// CompositeRegistry allows combining multiple tool registries.
|
||||
type CompositeRegistry struct {
|
||||
registries []ToolRegistry
|
||||
}
|
||||
|
||||
// Add adds a registry to the composite.
|
||||
func (r *CompositeRegistry) Add(registry ToolRegistry) {
|
||||
r.registries = append(r.registries, registry)
|
||||
}
|
||||
|
||||
// Tools returns the union of all tools from all registries.
|
||||
func (r *CompositeRegistry) Tools(ctx context.Context) []Tool {
|
||||
var all []Tool
|
||||
for _, reg := range r.registries {
|
||||
all = append(all, reg.Tools(ctx)...)
|
||||
}
|
||||
return all
|
||||
}
|
||||
|
||||
// Execute attempts to execute the tool by trying each registry in order.
|
||||
func (r *CompositeRegistry) Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error) {
|
||||
for _, reg := range r.registries {
|
||||
res, err := reg.Execute(ctx, s, name, args)
|
||||
if err == nil {
|
||||
return res, nil
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("tool %s not found in any registry", name)
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package openapi
|
||||
|
||||
import "encoding/json"
|
||||
|
||||
type Role string
|
||||
|
||||
const (
|
||||
RoleSystem Role = "system"
|
||||
RoleUser Role = "user"
|
||||
RoleAssistant Role = "assistant"
|
||||
RoleTool Role = "tool"
|
||||
)
|
||||
|
||||
type Message struct {
|
||||
Role Role `json:"role"`
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
}
|
||||
|
||||
type ToolCall struct {
|
||||
Index *int `json:"index,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function ToolCallFunction `json:"function,omitempty"`
|
||||
}
|
||||
|
||||
type ToolCallFunction struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Arguments string `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
type Tool struct {
|
||||
Type string `json:"type"`
|
||||
Function FunctionDefinition `json:"function"`
|
||||
}
|
||||
|
||||
type FunctionDefinition struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters json.RawMessage `json:"parameters,omitempty"`
|
||||
}
|
||||
|
||||
type ChatCompletionRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []Message `json:"messages"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
N *int `json:"n,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Stop []string `json:"stop,omitempty"`
|
||||
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||
PresencePenalty *float64 `json:"presence_penalty,omitempty"`
|
||||
FrequencyPenalty *float64 `json:"frequency_penalty,omitempty"`
|
||||
LogitBias map[string]int `json:"logit_bias,omitempty"`
|
||||
User string `json:"user,omitempty"`
|
||||
|
||||
// llama.cpp specific
|
||||
CachePrompt bool `json:"cache_prompt,omitempty"`
|
||||
NKeep int `json:"n_keep,omitempty"`
|
||||
Seed int `json:"seed,omitempty"`
|
||||
}
|
||||
|
||||
type Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
}
|
||||
|
||||
type Choice struct {
|
||||
Index int `json:"index"`
|
||||
Message Message `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
type ChatCompletionResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []Choice `json:"choices"`
|
||||
Usage Usage `json:"usage"`
|
||||
}
|
||||
|
||||
type ChatCompletionChunk struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []struct {
|
||||
Index int `json:"index"`
|
||||
Delta Message `json:"delta"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
Reference in New Issue
Block a user