Add go-llm v2: redesigned API for simpler LLM abstraction
v2 is a new Go module (v2/) with a dramatically simpler API: - Unified Message type (no more Input marker interface) - Define[T] for ergonomic tool creation with standard context.Context - Chat session with automatic tool-call loop (agent loop) - Streaming via pull-based StreamReader - MCP one-call connect (MCPStdioServer, MCPHTTPServer, MCPSSEServer) - Middleware support (logging, retry, timeout, usage tracking) - Decoupled JSON Schema (map[string]any, no provider coupling) - Sample tools: WebSearch, Browser, Exec, ReadFile, WriteFile, HTTP - Providers: OpenAI, Anthropic, Google (all with streaming) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
264
v2/mcp.go
Normal file
264
v2/mcp.go
Normal file
@@ -0,0 +1,264 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"sync"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
// MCPTransport specifies how to connect to an MCP server.
|
||||
type MCPTransport string
|
||||
|
||||
const (
|
||||
MCPStdio MCPTransport = "stdio"
|
||||
MCPSSE MCPTransport = "sse"
|
||||
MCPHTTP MCPTransport = "http"
|
||||
)
|
||||
|
||||
// MCPServer represents a connection to an MCP server.
|
||||
type MCPServer struct {
|
||||
name string
|
||||
transport MCPTransport
|
||||
|
||||
// stdio fields
|
||||
command string
|
||||
args []string
|
||||
env []string
|
||||
|
||||
// network fields
|
||||
url string
|
||||
|
||||
// internal
|
||||
client *mcp.Client
|
||||
session *mcp.ClientSession
|
||||
tools map[string]*mcp.Tool
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// MCPOption configures an MCP server.
|
||||
type MCPOption func(*MCPServer)
|
||||
|
||||
// WithMCPEnv adds environment variables for the subprocess.
|
||||
func WithMCPEnv(env ...string) MCPOption {
|
||||
return func(s *MCPServer) { s.env = env }
|
||||
}
|
||||
|
||||
// WithMCPName sets a friendly name for logging.
|
||||
func WithMCPName(name string) MCPOption {
|
||||
return func(s *MCPServer) { s.name = name }
|
||||
}
|
||||
|
||||
// MCPStdioServer creates and connects to an MCP server via stdio transport.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// server, err := llm.MCPStdioServer(ctx, "npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp")
|
||||
func MCPStdioServer(ctx context.Context, command string, args ...string) (*MCPServer, error) {
|
||||
s := &MCPServer{
|
||||
name: command,
|
||||
transport: MCPStdio,
|
||||
command: command,
|
||||
args: args,
|
||||
}
|
||||
if err := s.connect(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// MCPHTTPServer creates and connects to an MCP server via streamable HTTP transport.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// server, err := llm.MCPHTTPServer(ctx, "https://mcp.example.com")
|
||||
func MCPHTTPServer(ctx context.Context, url string, opts ...MCPOption) (*MCPServer, error) {
|
||||
s := &MCPServer{
|
||||
name: url,
|
||||
transport: MCPHTTP,
|
||||
url: url,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(s)
|
||||
}
|
||||
if err := s.connect(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// MCPSSEServer creates and connects to an MCP server via SSE transport.
|
||||
func MCPSSEServer(ctx context.Context, url string, opts ...MCPOption) (*MCPServer, error) {
|
||||
s := &MCPServer{
|
||||
name: url,
|
||||
transport: MCPSSE,
|
||||
url: url,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(s)
|
||||
}
|
||||
if err := s.connect(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *MCPServer) connect(ctx context.Context) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.session != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
s.client = mcp.NewClient(&mcp.Implementation{
|
||||
Name: "go-llm-v2",
|
||||
Version: "2.0.0",
|
||||
}, nil)
|
||||
|
||||
var transport mcp.Transport
|
||||
|
||||
switch s.transport {
|
||||
case MCPSSE:
|
||||
transport = &mcp.SSEClientTransport{
|
||||
Endpoint: s.url,
|
||||
}
|
||||
case MCPHTTP:
|
||||
transport = &mcp.StreamableClientTransport{
|
||||
Endpoint: s.url,
|
||||
}
|
||||
default: // stdio
|
||||
cmd := exec.Command(s.command, s.args...)
|
||||
cmd.Env = append(os.Environ(), s.env...)
|
||||
transport = &mcp.CommandTransport{
|
||||
Command: cmd,
|
||||
}
|
||||
}
|
||||
|
||||
session, err := s.client.Connect(ctx, transport, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect to MCP server %s: %w", s.name, err)
|
||||
}
|
||||
|
||||
s.session = session
|
||||
|
||||
// Load tools
|
||||
s.tools = make(map[string]*mcp.Tool)
|
||||
for tool, err := range session.Tools(ctx, nil) {
|
||||
if err != nil {
|
||||
s.session.Close()
|
||||
s.session = nil
|
||||
return fmt.Errorf("failed to list tools from %s: %w", s.name, err)
|
||||
}
|
||||
s.tools[tool.Name] = tool
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close closes the connection to the MCP server.
|
||||
func (s *MCPServer) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.session == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
err := s.session.Close()
|
||||
s.session = nil
|
||||
s.tools = nil
|
||||
return err
|
||||
}
|
||||
|
||||
// IsConnected returns true if the server is connected.
|
||||
func (s *MCPServer) IsConnected() bool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.session != nil
|
||||
}
|
||||
|
||||
// ListTools returns Tool definitions for all tools this server provides.
|
||||
func (s *MCPServer) ListTools() []Tool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
var tools []Tool
|
||||
for _, t := range s.tools {
|
||||
tools = append(tools, s.toTool(t))
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// CallTool invokes a tool on the server.
|
||||
func (s *MCPServer) CallTool(ctx context.Context, name string, arguments map[string]any) (string, error) {
|
||||
s.mu.RLock()
|
||||
session := s.session
|
||||
s.mu.RUnlock()
|
||||
|
||||
if session == nil {
|
||||
return "", fmt.Errorf("%w: %s", ErrNotConnected, s.name)
|
||||
}
|
||||
|
||||
result, err := session.CallTool(ctx, &mcp.CallToolParams{
|
||||
Name: name,
|
||||
Arguments: arguments,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if len(result.Content) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
return contentToString(result.Content), nil
|
||||
}
|
||||
|
||||
func (s *MCPServer) toTool(t *mcp.Tool) Tool {
|
||||
var inputSchema map[string]any
|
||||
if t.InputSchema != nil {
|
||||
data, err := json.Marshal(t.InputSchema)
|
||||
if err == nil {
|
||||
_ = json.Unmarshal(data, &inputSchema)
|
||||
}
|
||||
}
|
||||
|
||||
if inputSchema == nil {
|
||||
inputSchema = map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{},
|
||||
}
|
||||
}
|
||||
|
||||
return Tool{
|
||||
Name: t.Name,
|
||||
Description: t.Description,
|
||||
Schema: inputSchema,
|
||||
isMCP: true,
|
||||
mcpServer: s,
|
||||
}
|
||||
}
|
||||
|
||||
func contentToString(content []mcp.Content) string {
|
||||
var parts []string
|
||||
for _, c := range content {
|
||||
switch tc := c.(type) {
|
||||
case *mcp.TextContent:
|
||||
parts = append(parts, tc.Text)
|
||||
default:
|
||||
if data, err := json.Marshal(c); err == nil {
|
||||
parts = append(parts, string(data))
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(parts) == 1 {
|
||||
return parts[0]
|
||||
}
|
||||
data, _ := json.Marshal(parts)
|
||||
return string(data)
|
||||
}
|
||||
Reference in New Issue
Block a user