modelsocket

package module
v0.0.0-...-e5b6f49 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jan 8, 2026 License: Apache-2.0 Imports: 11 Imported by: 0

README

modelsocket-go

Go client for the ModelSocket protocol - WebSocket-based LLM integration with streaming, tool calling, and sequence forking.

Install

go get github.com/chrisboulton/modelsocket-go

Requires Go 1.23+.

Quick Start

ctx := context.Background()

client, err := modelsocket.Connect(ctx, "wss://models.mixlayer.ai/ws", os.Getenv("API_KEY"))
if err != nil {
    log.Fatal(err)
}
defer client.Close(ctx)

seq, err := client.Open(ctx, "meta/llama3.1-8b-instruct-free")
if err != nil {
    log.Fatal(err)
}
defer seq.Close(ctx)

seq.Append(ctx, "Hello!", modelsocket.AsUser())

stream, _ := seq.Generate(ctx, modelsocket.GenerateAsAssistant())
for chunk, err := range stream.Chunks(ctx) {
    if err != nil {
        log.Fatal(err)
    }
    fmt.Print(chunk.Text)
}

Client

The Client manages the WebSocket connection and routes events to sequences. It's safe for concurrent use.

Client Options

Pass options to Connect() to configure the client:

client, err := modelsocket.Connect(ctx, url, apiKey,
    modelsocket.WithLogger(slog.Default()),
    modelsocket.WithOnSend(func(req *modelsocket.MSRequest) {
        // Called before each request is sent
    }),
    modelsocket.WithOnReceive(func(evt *modelsocket.MSEvent) {
        // Called after each event is received
    }),
)
Option Description
WithLogger(*slog.Logger) Structured logger for debug output
WithOnSend(func(*MSRequest)) Hook called before sending requests
WithOnReceive(func(*MSEvent)) Hook called after receiving events
Open Options

Configure sequences when calling client.Open():

seq, err := client.Open(ctx, "meta/llama3.1-8b-instruct-free",
    modelsocket.WithSkipPrelude(),
    modelsocket.WithToolbox(toolbox),
)
Option Description
WithSkipPrelude() Skip the model's default system prompt
WithToolbox(*Toolbox) Enable tool calling with the provided toolbox
Custom Transport

Use NewWithTransport() to provide your own transport implementation:

client := modelsocket.NewWithTransport(ctx, myTransport,
    modelsocket.WithLogger(logger),
)

Custom transports must implement the Transport interface:

type Transport interface {
    Send(ctx context.Context, req *MSRequest) error
    Receive(ctx context.Context) (*MSEvent, error)
    Close() error
}

Use cases for custom transports:

  • Testing - Mock transport for unit tests without network calls
  • Proxying - Route through a custom proxy or middleware
  • Alternative protocols - Use HTTP/SSE or other transports instead of WebSocket

Examples

export MODELSOCKET_API_KEY="your-key"

# Simple chat
go run ./examples/chat

# Tool calling
go run ./examples/tools

Tool Calling

// Define a tool
weatherTool := modelsocket.NewFuncTool(
    modelsocket.ToolDefinition{
        Name:        "get_weather",
        Description: "Get weather for a city",
        Parameters: modelsocket.ToolParameters{
            Type: "object",
            Properties: map[string]modelsocket.ToolProperty{
                "location": {Type: "string", Description: "City name"},
            },
            Required: []string{"location"},
        },
    },
    func(ctx context.Context, args string) (string, error) {
        return `{"temperature": 72, "units": "F"}`, nil
    },
)

// Register tools
toolbox := modelsocket.NewToolbox()
toolbox.Add(weatherTool)

// Optionally, you can customize the tool call instructions - what informs the
// model how to format tool calls. You should pretty much never need to do this
// toolbox.SetToolInstructions("....")

// You can also bring your own presenter for decorating tool definitions when
// passed into the model. It's up to you to ensure all tools with their definitions
// are listed.
// toolbox.SetToolDefinitionPrompt("You have access to the following)

// Open sequence with tools
seq, _ := client.Open(ctx, model,
    modelsocket.WithToolbox(toolbox),
)

// Handle tool calls in stream
for chunk, err := range stream.Chunks(ctx) {
    if err != nil {
        break
    }
    fmt.Print(chunk.Text)
    if len(chunk.ToolCalls) > 0 {
        results, _ := toolbox.CallTools(ctx, chunk.ToolCalls)
        seq.ToolReturn(ctx, results)
    }
}

Documentation

Overview

Package modelsocket provides a Go client for the ModelSocket protocol.

ModelSocket is a WebSocket-based protocol for efficiently integrating with Large Language Models (LLMs). It provides streaming text generation, tool calling, and sequence forking capabilities.

Thread Safety

Client and Seq are safe for concurrent use by multiple goroutines. However, only one Seq.Generate call can be active per sequence at a time. GenStream should only be consumed by a single goroutine.

Basic Usage 

ctx := context.Background()

// Connect to server
client, err := modelsocket.Connect(ctx, "wss://example.com/ws", "api-key")
if err != nil {
    log.Fatal(err)
}
defer client.Close(ctx)

// Open a sequence
seq, err := client.Open(ctx, "model-name")
if err != nil {
    log.Fatal(err)
}
defer seq.Close(ctx)

// Append user message
err = seq.Append(ctx, "Hello!", modelsocket.AsUser())
if err != nil {
    log.Fatal(err)
}

// Generate response using iterator
stream, err := seq.Generate(ctx, modelsocket.GenerateAsAssistant())
if err != nil {
    log.Fatal(err)
}

for chunk, err := range stream.Chunks(ctx) {
    if err != nil {
        log.Fatal(err)
    }
    fmt.Print(chunk.Text)
}

Observability

Use WithLogger, WithOnSend, and WithOnReceive to add logging and monitoring to the client: 

client, err := modelsocket.Connect(ctx, url, apiKey,
    modelsocket.WithLogger(slog.Default()),
    modelsocket.WithOnSend(func(req *modelsocket.MSRequest) {
        metrics.RequestsSent.Inc()
    }),
)

Index

Constants

This section is empty.

Variables

View Source 
var (
	ErrClosed          = errors.New("modelsocket: connection closed")
	ErrSeqClosed       = errors.New("modelsocket: sequence closed")
	ErrTimeout         = errors.New("modelsocket: operation timed out")
	ErrInvalidState    = errors.New("modelsocket: invalid sequence state")
	ErrToolNotFound    = errors.New("modelsocket: tool not found")
	ErrUnexpectedEvent = errors.New("modelsocket: unexpected event")
	ErrBufferFull      = errors.New("modelsocket: buffer full")
)

Sentinel errors for common conditions.

Functions

This section is empty.

Types

type AppendOption 

type AppendOption func(*appendConfig)

AppendOption configures text appending.

func AsAssistant 

func AsAssistant() AppendOption

AsAssistant marks the message as from the assistant.

func AsSystem 

func AsSystem() AppendOption

AsSystem marks the message as a system message.

func AsUser 

func AsUser() AppendOption

AsUser marks the message as from the user.

func WithEcho 

func WithEcho() AppendOption

WithEcho echoes the appended text back in events.

type Client 

type Client struct {
	// contains filtered or unexported fields
}

Client is the main client for connecting to a ModelSocket server. It is safe for concurrent use by multiple goroutines.

func Connect 

func Connect(ctx context.Context, url string, apiKey string, opts ...ClientOption) (*Client, error)

Connect establishes a connection to a ModelSocket server.

func NewWithTransport 

func NewWithTransport(ctx context.Context, transport Transport, opts ...ClientOption) *Client

NewWithTransport creates a Client with a custom transport. This is useful for testing or custom transport implementations.

func (*Client) Close 

func (c *Client) Close(ctx context.Context) error

Close closes the connection and all sequences.

func (*Client) Open 

func (c *Client) Open(ctx context.Context, model string, opts ...OpenOption) (*Seq, error)

Open creates a new sequence with the specified model.

type ClientOption 

type ClientOption func(*clientConfig)

ClientOption configures a ModelSocket client.

func WithLogger 

func WithLogger(logger *slog.Logger) ClientOption

WithLogger sets a structured logger for the client.

func WithOnReceive 

func WithOnReceive(fn func(*MSEvent)) ClientOption

WithOnReceive sets a callback invoked after each event is received.

func WithOnSend 

func WithOnSend(fn func(*MSRequest)) ClientOption

WithOnSend sets a callback invoked before each request is sent.

type ConnectionError 

type ConnectionError struct {
	Op  string
	URL string
	Err error
}

ConnectionError represents a connection-level error.

func (*ConnectionError) Error 

func (e *ConnectionError) Error() string

func (*ConnectionError) Unwrap 

func (e *ConnectionError) Unwrap() error

type DialOptions 

type DialOptions struct {
	// HTTPHeader specifies additional HTTP headers to send during handshake.
	HTTPHeader http.Header

	// HTTPClient is the HTTP client used for the handshake.
	// If nil, http.DefaultClient is used.
	HTTPClient *http.Client
}

DialOptions configures the WebSocket connection.

type FuncTool 

type FuncTool struct {
	// contains filtered or unexported fields
}

FuncTool wraps a function as a Tool.

func NewFuncTool 

func NewFuncTool(def ToolDefinition, fn func(ctx context.Context, args string) (string, error)) *FuncTool

NewFuncTool creates a tool from a function.

func (*FuncTool) Call 

func (f *FuncTool) Call(ctx context.Context, args string) (string, error)

Call invokes the tool function.

func (*FuncTool) Definition 

func (f *FuncTool) Definition() ToolDefinition

Definition returns the tool definition.

type GenChunk 

type GenChunk struct {
	Text      string
	Hidden    bool
	Tokens    []int
	ToolCalls []ToolCall
}

GenChunk represents a chunk of generated content.

type GenOption 

type GenOption func(*genConfig)

GenOption configures text generation.

func GenerateAsAssistant 

func GenerateAsAssistant() GenOption

GenerateAsAssistant generates text as the assistant role.

func GenerateAsSystem 

func GenerateAsSystem() GenOption

GenerateAsSystem generates text as the system role.

func GenerateAsUser 

func GenerateAsUser() GenOption

GenerateAsUser generates text as the user role.

func WithHidden 

func WithHidden() GenOption

WithHidden hides the generated text from the conversation history.

func WithMaxLength 

func WithMaxLength(n int) GenOption

WithMaxLength sets the maximum length in characters.

func WithMaxTokens 

func WithMaxTokens(n int) GenOption

WithMaxTokens sets the maximum number of tokens to generate.

func WithRegexMask 

func WithRegexMask(pattern string) GenOption

WithRegexMask constrains generation to match a regex pattern.

func WithRepeatPenalty 

func WithRepeatPenalty(p float64) GenOption

WithRepeatPenalty sets the repetition penalty.

func WithSeed 

func WithSeed(seed int64) GenOption

WithSeed sets the random seed for reproducible generation.

func WithStopStrings 

func WithStopStrings(stops ...string) GenOption

WithStopStrings sets strings that will stop generation.

func WithTemperature 

func WithTemperature(t float64) GenOption

WithTemperature sets the sampling temperature.

func WithTopK 

func WithTopK(k int) GenOption

WithTopK sets the top-k sampling parameter.

func WithTopP 

func WithTopP(p float64) GenOption

WithTopP sets the nucleus sampling parameter.

type GenStream 

type GenStream struct {
	// contains filtered or unexported fields
}

GenStream provides streaming access to generated content.

func (*GenStream) Chunks 

func (g *GenStream) Chunks(ctx context.Context) iter.Seq2[*GenChunk, error]

Chunks returns an iterator over all chunks in the stream.

func (*GenStream) InputTokens 

func (g *GenStream) InputTokens() int

InputTokens returns the input token count. Only valid after stream is exhausted.

func (*GenStream) Next 

func (g *GenStream) Next(ctx context.Context) (*GenChunk, error)

Next returns the next chunk, or nil if done. Returns an error if one occurred during generation. The context can be used to cancel waiting for the next chunk.

func (*GenStream) OutputTokens 

func (g *GenStream) OutputTokens() int

OutputTokens returns the output token count. Only valid after stream is exhausted.

func (*GenStream) Text 

func (g *GenStream) Text(ctx context.Context) (string, error)

Text collects all generated text and returns it.

func (*GenStream) TextAndTokens 

func (g *GenStream) TextAndTokens(ctx context.Context) (string, []int, error)

TextAndTokens collects all generated text and tokens.

type MSEvent 

type MSEvent struct {
	Event string `json:"event"`

	// Common fields
	SeqID string `json:"seq_id,omitempty"`
	CID   string `json:"cid,omitempty"`

	// SeqText fields
	Text            string `json:"text,omitempty"`
	Hidden          bool   `json:"hidden,omitempty"`
	NumInputTokens  int    `json:"num_input_tokens,omitempty"`
	NumOutputTokens int    `json:"num_output_tokens,omitempty"`
	Tokens          []int  `json:"tokens,omitempty"`

	// SeqToolCall fields
	ToolCalls []SeqToolCall `json:"tool_calls,omitempty"`

	// SeqForkFinish fields
	ChildSeqID string `json:"child_seq_id,omitempty"`

	// SeqState fields
	State SeqState `json:"state,omitempty"`

	// SeqClosed fields
	InputTokens  int    `json:"input_tokens,omitempty"`
	OutputTokens int    `json:"output_tokens,omitempty"`
	DurationMs   int64  `json:"duration_ms,omitempty"`
	ErrorMsg     string `json:"error,omitempty"`

	// Error fields
	Message string `json:"message,omitempty"`
}

MSEvent represents an event received from the server.

func (*MSEvent) IsError 

func (e *MSEvent) IsError() bool

IsError returns true if this is an error event.

func (*MSEvent) IsSeqAppendFinish 

func (e *MSEvent) IsSeqAppendFinish() bool

IsSeqAppendFinish returns true if this is a seq_append_finish event.

func (*MSEvent) IsSeqClosed 

func (e *MSEvent) IsSeqClosed() bool

IsSeqClosed returns true if this is a seq_closed event.

func (*MSEvent) IsSeqForkFinish 

func (e *MSEvent) IsSeqForkFinish() bool

IsSeqForkFinish returns true if this is a seq_fork_finish event.

func (*MSEvent) IsSeqGenFinish 

func (e *MSEvent) IsSeqGenFinish() bool

IsSeqGenFinish returns true if this is a seq_gen_finish event.

func (*MSEvent) IsSeqOpened 

func (e *MSEvent) IsSeqOpened() bool

IsSeqOpened returns true if this is a seq_opened event.

func (*MSEvent) IsSeqState 

func (e *MSEvent) IsSeqState() bool

IsSeqState returns true if this is a seq_state event.

func (*MSEvent) IsSeqText 

func (e *MSEvent) IsSeqText() bool

IsSeqText returns true if this is a seq_text event.

func (*MSEvent) IsSeqToolCall 

func (e *MSEvent) IsSeqToolCall() bool

IsSeqToolCall returns true if this is a seq_tool_call event.

func (*MSEvent) Type 

func (e *MSEvent) Type() string

Type returns the event type.

type MSRequest 

type MSRequest struct {
	Request string      `json:"request"`
	CID     string      `json:"cid"`
	SeqID   string      `json:"seq_id,omitempty"`
	Data    interface{} `json:"data"`
}

MSRequest represents a request sent to the server.

func NewAppendRequest 

func NewAppendRequest(cid, seqID string, data SeqAppendData) *MSRequest

NewAppendRequest creates a new append command request.

func NewCloseRequest 

func NewCloseRequest(cid, seqID string) *MSRequest

NewCloseRequest creates a new close command request.

func NewForkRequest 

func NewForkRequest(cid, seqID string) *MSRequest

NewForkRequest creates a new fork command request.

func NewGenRequest 

func NewGenRequest(cid, seqID string, data SeqGenData) *MSRequest

NewGenRequest creates a new gen command request.

func NewSeqOpenRequest 

func NewSeqOpenRequest(cid string, data SeqOpenData) *MSRequest

NewSeqOpenRequest creates a new seq_open request.

func NewToolReturnRequest 

func NewToolReturnRequest(cid, seqID string, results []ToolResult, genOpts SeqGenData) *MSRequest

NewToolReturnRequest creates a new tool_return command request.

type OpenOption 

type OpenOption func(*openConfig)

OpenOption configures sequence opening.

func WithSkipPrelude 

func WithSkipPrelude() OpenOption

WithSkipPrelude skips the model's default prelude/system prompt.

func WithToolbox 

func WithToolbox(tb *Toolbox) OpenOption

WithToolbox registers a toolbox for tool calling.

type ProtocolError 

type ProtocolError struct {
	Code    string
	Message string
	SeqID   string
	CID     string
}

ProtocolError represents a protocol-level error from the server.

func (*ProtocolError) Error 

func (e *ProtocolError) Error() string

type Role 

type Role string

Role represents the role of a message in a conversation. 

const (
	RoleUser      Role = "user"
	RoleAssistant Role = "assistant"
	RoleSystem    Role = "system"
)

type SendError 

type SendError struct {
	Op  string
	Err error
}

SendError represents an error during request sending.

func (*SendError) Error 

func (e *SendError) Error() string

func (*SendError) Unwrap 

func (e *SendError) Unwrap() error

type Seq 

type Seq struct {
	// contains filtered or unexported fields
}

Seq represents an active conversation sequence. It is safe for concurrent use by multiple goroutines. However, only one Generate call can be active at a time.

func (*Seq) Append 

func (s *Seq) Append(ctx context.Context, text string, opts ...AppendOption) error

Append adds text to the sequence.

func (*Seq) Close 

func (s *Seq) Close(ctx context.Context) error

Close closes the sequence.

func (*Seq) Fork 

func (s *Seq) Fork(ctx context.Context) (*Seq, error)

Fork creates a new sequence with the same conversation history.

func (*Seq) Generate 

func (s *Seq) Generate(ctx context.Context, opts ...GenOption) (*GenStream, error)

Generate starts text generation and returns a stream.

func (*Seq) ID 

func (s *Seq) ID() string

ID returns the sequence ID.

func (*Seq) State 

func (s *Seq) State() SeqState

State returns the current sequence state.

func (*Seq) ToolReturn 

func (s *Seq) ToolReturn(ctx context.Context, results []ToolResult) error

ToolReturn sends tool call results back to the model.

type SeqAppendData 

type SeqAppendData struct {
	Text   string `json:"text"`
	Role   string `json:"role,omitempty"`
	Echo   bool   `json:"echo,omitempty"`
	Hidden bool   `json:"hidden,omitempty"`
}

SeqAppendData is the data for an append command.

type SeqError 

type SeqError struct {
	SeqID   string
	Message string
}

SeqError represents a sequence-level error.

func (*SeqError) Error 

func (e *SeqError) Error() string

type SeqGenData 

type SeqGenData struct {
	Role          string   `json:"role,omitempty"`
	MaxTokens     *int     `json:"max_tokens,omitempty"`
	MaxLength     *int     `json:"max_length,omitempty"`
	Temperature   *float64 `json:"temperature,omitempty"`
	TopP          *float64 `json:"top_p,omitempty"`
	TopK          *int     `json:"top_k,omitempty"`
	RepeatPenalty *float64 `json:"repeat_penalty,omitempty"`
	Seed          *int64   `json:"seed,omitempty"`
	StopStrings   []string `json:"stop_strings,omitempty"`
	RegexMask     *string  `json:"regex_mask,omitempty"`
	Hidden        bool     `json:"hidden,omitempty"`
	PrefillText   *string  `json:"prefill_text,omitempty"`
	ReturnTokens  *bool    `json:"return_tokens,omitempty"`
}

SeqGenData is the data for a gen command.

type SeqOpenData 

type SeqOpenData struct {
	Model        string `json:"model"`
	ToolsEnabled bool   `json:"tools_enabled,omitempty"`
	ToolPrompt   string `json:"tool_prompt,omitempty"`
	SkipPrelude  bool   `json:"skip_prelude,omitempty"`
}

SeqOpenData is the data for a seq_open request.

type SeqState 

type SeqState string

SeqState represents the state of a sequence. 

const (
	StateReady      SeqState = "ready"
	StateAppending  SeqState = "appending"
	StateGenerating SeqState = "generating"
	StateToolCall   SeqState = "tool_call"
	StateForking    SeqState = "forking"
	StateClosed     SeqState = "closed"
)

type SeqToolCall 

type SeqToolCall struct {
	Name string `json:"name"`
	Args string `json:"args"`
}

SeqToolCall represents a tool call from the model.

type Tool 

type Tool interface {
	Definition() ToolDefinition
	Call(ctx context.Context, args string) (string, error)
}

Tool defines the interface for a callable tool.

type ToolCall 

type ToolCall struct {
	Name string
	Args string
}

ToolCall represents a tool call from the model.

type ToolDefinition 

type ToolDefinition struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	Parameters  ToolParameters `json:"parameters"`
}

ToolDefinition describes a tool for the model.

type ToolParameters 

type ToolParameters struct {
	Type       string                  `json:"type"`
	Properties map[string]ToolProperty `json:"properties,omitempty"`
	Required   []string                `json:"required,omitempty"`
}

ToolParameters defines the JSON Schema for tool parameters.

type ToolProperty 

type ToolProperty struct {
	Type        string   `json:"type"`
	Description string   `json:"description,omitempty"`
	Enum        []string `json:"enum,omitempty"`
}

ToolProperty defines a single parameter property.

type ToolResult 

type ToolResult struct {
	Name   string `json:"name"`
	Result string `json:"result"`
}

ToolResult represents the result of a tool call.

type Toolbox 

type Toolbox struct {
	// contains filtered or unexported fields
}

Toolbox manages a collection of tools.

func NewToolbox 

func NewToolbox() *Toolbox

NewToolbox creates an empty toolbox.

func (*Toolbox) Add 

func (t *Toolbox) Add(tool Tool)

Add registers a tool.

func (*Toolbox) Call 

func (t *Toolbox) Call(ctx context.Context, name string, args string) (string, error)

Call executes a tool by name with the given arguments.

func (*Toolbox) CallTools 

func (t *Toolbox) CallTools(ctx context.Context, calls []ToolCall) ([]ToolResult, error)

CallTools executes multiple tool calls and returns results.

func (*Toolbox) Definitions 

func (t *Toolbox) Definitions() []ToolDefinition

Definitions returns all tool definitions.

func (*Toolbox) Get 

func (t *Toolbox) Get(name string) (Tool, bool)

Get retrieves a tool by name.

func (*Toolbox) SetToolDefinitionPrompt 

func (t *Toolbox) SetToolDefinitionPrompt(prompt string)

func (*Toolbox) SetToolInstructions 

func (t *Toolbox) SetToolInstructions(instructions string)

func (*Toolbox) ToolDefinitionPrompt 

func (t *Toolbox) ToolDefinitionPrompt() string

ToolPrompt returns the tool prompt. If a custom prompt was set via SetToolPrompt, it returns that; otherwise it returns an auto-generated prompt describing all tools.

func (*Toolbox) ToolInstructions 

func (t *Toolbox) ToolInstructions() string

ToolInstructions returns the tool instructions.

type Transport 

type Transport interface {
	Send(ctx context.Context, req *MSRequest) error
	Receive(ctx context.Context) (*MSEvent, error)
	Close() error
}

Transport provides the interface for sending and receiving messages. Implementations must be safe for concurrent use.

func Dial 

func Dial(ctx context.Context, url string, apiKey string, opts *DialOptions) (Transport, error)

Dial connects to a ModelSocket server and returns a Transport.

Source Files

View all Source files

Directories

Path Synopsis
examples
chat command
Example chat application using ModelSocket.
 Example chat application using ModelSocket.
tools command
Example demonstrating tool calling with ModelSocket.
 Example demonstrating tool calling with ModelSocket.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.