first commit
This commit is contained in:
@@ -0,0 +1,165 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
"github.com/smallnest/pigo/internal/cli/run"
|
||||
"github.com/smallnest/pigo/internal/provider"
|
||||
"github.com/smallnest/pigo/internal/runtime"
|
||||
)
|
||||
|
||||
// Session is a single, stateful pigo agent conversation. Create one with New,
|
||||
// drive it with Prompt or Stream, and release its resources with Close. The
|
||||
// conversation history accumulates across calls, so follow-up prompts see the
|
||||
// earlier exchange; call Reset to start over on the same session.
|
||||
//
|
||||
// A Session is not safe for concurrent use. Drive it from a single goroutine, or
|
||||
// give each goroutine its own Session.
|
||||
type Session struct {
|
||||
env run.Env
|
||||
runCfg runtime.RunConfig
|
||||
agentCtx *agentcore.AgentContext
|
||||
model string
|
||||
}
|
||||
|
||||
// New builds a Session from the given options. It resolves the provider and
|
||||
// credentials, assembles the tool set, and validates the tool policy and
|
||||
// thinking level up front, so a configuration mistake (an unknown tool name, an
|
||||
// invalid thinking level, an unresolvable provider) is returned here rather than
|
||||
// surfacing on the first Prompt.
|
||||
//
|
||||
// No network call is made by New: the provider is only contacted when you call
|
||||
// Prompt or Stream. This makes New cheap and safe to use in tests.
|
||||
//
|
||||
// See the package documentation for the default tool, skill, and memory
|
||||
// behavior — in particular, that tools are enabled and auto-executed by default.
|
||||
func New(opts ...Option) (*Session, error) {
|
||||
c := config{model: "openrouter/free"}
|
||||
for _, o := range opts {
|
||||
o(&c)
|
||||
}
|
||||
|
||||
// Validate the reasoning-effort level through the same layered config chain
|
||||
// the CLI uses, so an invalid WithThinkingLevel value fails fast here.
|
||||
thinking, err := run.ResolveThinkingLevel(c.thinking)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// One ToolPolicy value carries both lists so they cannot be swapped; deny
|
||||
// always wins over allow inside run.ApplyToolPolicy.
|
||||
policy := run.NewToolPolicy(c.allowedTools, c.disallowedTools)
|
||||
|
||||
// SetupEnv resolves the provider, assembles the (policy-filtered) tool set,
|
||||
// builds the system prompt, and — because skills/memory are opt-in here —
|
||||
// leaves the machine's shared state untouched unless WithSkills/WithMemory
|
||||
// were passed. It also validates the tool policy against the real tool set,
|
||||
// so an unknown tool name is reported as an error.
|
||||
env, err := run.SetupEnv(
|
||||
c.model, c.baseURL, c.protocol, c.provider, c.apiKey,
|
||||
c.noTools, !c.skills, c.systemPrompt, c.appendSystemPrompt, c.memory, policy,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Resolve the API key by provider name: an explicit WithAPIKey overrides the
|
||||
// provider's environment variable. The key is held only in the credential
|
||||
// store and never logged.
|
||||
creds := provider.NewCredentialStore(nil)
|
||||
if c.apiKey != "" {
|
||||
creds.SetOverride(env.ProviderName, c.apiKey)
|
||||
}
|
||||
|
||||
runCfg := run.NewConfig(
|
||||
c.model, env.ProviderName, thinking, env.Provider, creds,
|
||||
run.ToolRegistry(env.Tools), run.TodoReminders(env.Tools),
|
||||
)
|
||||
|
||||
return &Session{
|
||||
env: env,
|
||||
runCfg: runCfg,
|
||||
agentCtx: &agentcore.AgentContext{
|
||||
SystemPrompt: env.SysPrompt,
|
||||
Tools: env.Tools,
|
||||
},
|
||||
model: c.model,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Prompt sends one user message, runs the agent loop to completion (executing
|
||||
// any tool calls the model makes along the way), and returns the assistant's
|
||||
// final text. The exchange is appended to the session history so later prompts
|
||||
// have this context.
|
||||
func (s *Session) Prompt(ctx context.Context, prompt string) (string, error) {
|
||||
return s.Stream(ctx, prompt, nil)
|
||||
}
|
||||
|
||||
// Stream is Prompt with incremental output: onText, if non-nil, is called with
|
||||
// each chunk of assistant text as it arrives, and the complete final text is
|
||||
// also returned. Tool calls still run automatically between text chunks. A nil
|
||||
// onText makes Stream behave exactly like Prompt.
|
||||
func (s *Session) Stream(ctx context.Context, prompt string, onText func(string)) (string, error) {
|
||||
// The loop expects the initiating user message already appended; it then
|
||||
// mutates agentCtx.Messages in place (assistant + tool results), which is
|
||||
// what carries the conversation forward across calls.
|
||||
s.agentCtx.Messages = append(s.agentCtx.Messages, agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent(prompt)},
|
||||
})
|
||||
|
||||
stream := runtime.StartRun(ctx, s.agentCtx, s.runCfg)
|
||||
final, err := runtime.DrainStream(ctx, stream, runtime.StreamHandler{OnText: onText})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if final == nil {
|
||||
return "", nil
|
||||
}
|
||||
return agentcore.ContentToText(final.Content), nil
|
||||
}
|
||||
|
||||
// Reset clears the conversation history, so the next Prompt starts a fresh
|
||||
// exchange. The provider, tool set, and system prompt are unchanged.
|
||||
func (s *Session) Reset() {
|
||||
s.agentCtx.Messages = nil
|
||||
}
|
||||
|
||||
// ToolNames returns the names of the tools available to this session, in the
|
||||
// order they are advertised to the model. It reflects the applied tool policy,
|
||||
// so it is a convenient way to confirm WithTools/WithDisallowedTools did what
|
||||
// you intended. The result is empty for a WithoutTools session.
|
||||
func (s *Session) ToolNames() []string {
|
||||
names := make([]string, len(s.env.Tools))
|
||||
for i, t := range s.env.Tools {
|
||||
names[i] = t.Name()
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// Model returns the model id the session was created with.
|
||||
func (s *Session) Model() string { return s.model }
|
||||
|
||||
// Provider returns the resolved provider name (e.g. "anthropic", "openrouter"),
|
||||
// which is inferred from the model id unless WithProvider was set.
|
||||
func (s *Session) Provider() string { return s.env.ProviderName }
|
||||
|
||||
// Close releases resources held by the session: any loaded plugin manager and
|
||||
// the persistent memory store (when WithMemory was used). It is safe to call
|
||||
// once, and safe to call on a session that holds neither. After Close the
|
||||
// session must not be used again.
|
||||
func (s *Session) Close() error {
|
||||
var firstErr error
|
||||
if s.env.Plugins != nil {
|
||||
if err := s.env.Plugins.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
if s.env.Memory != nil {
|
||||
if err := s.env.Memory.Close(); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package agent_test
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/agent"
|
||||
)
|
||||
|
||||
// hermetic points provider/skill/plugin discovery at throwaway dirs and supplies
|
||||
// a dummy key so New resolves fully without ever contacting a network. None of
|
||||
// these tests call Prompt/Stream, so no real request is made.
|
||||
func hermetic(t *testing.T) {
|
||||
t.Helper()
|
||||
t.Setenv("OPENROUTER_API_KEY", "test-key")
|
||||
t.Setenv("PIGO_HOME", t.TempDir())
|
||||
}
|
||||
|
||||
func contains(set []string, name string) bool {
|
||||
for _, n := range set {
|
||||
if n == name {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// TestNewDefaults is the zero-config path: the full built-in tool set is
|
||||
// advertised and the model id resolves to the openrouter provider.
|
||||
func TestNewDefaults(t *testing.T) {
|
||||
hermetic(t)
|
||||
sess, err := agent.New(agent.WithModel("openrouter/free"))
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
defer sess.Close()
|
||||
|
||||
if got := sess.Model(); got != "openrouter/free" {
|
||||
t.Errorf("Model() = %q, want %q", got, "openrouter/free")
|
||||
}
|
||||
if got := sess.Provider(); got != "openrouter" {
|
||||
t.Errorf("Provider() = %q, want %q", got, "openrouter")
|
||||
}
|
||||
for _, want := range []string{"read", "write", "edit", "grep", "find", "bash", "task"} {
|
||||
if !contains(sess.ToolNames(), want) {
|
||||
t.Errorf("default tool set missing %q: %q", want, sess.ToolNames())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithToolsAllowlist confirms an allowlist narrows the set to exactly the
|
||||
// named tools, in order.
|
||||
func TestWithToolsAllowlist(t *testing.T) {
|
||||
hermetic(t)
|
||||
sess, err := agent.New(
|
||||
agent.WithModel("openrouter/free"),
|
||||
agent.WithTools("read", "grep"),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
defer sess.Close()
|
||||
|
||||
if got := strings.Join(sess.ToolNames(), ","); got != "read,grep" {
|
||||
t.Errorf("ToolNames() = %q, want %q", got, "read,grep")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDenyWinsOverAllow is the fail-closed guarantee at the SDK layer: a tool on
|
||||
// both lists is removed.
|
||||
func TestDenyWinsOverAllow(t *testing.T) {
|
||||
hermetic(t)
|
||||
sess, err := agent.New(
|
||||
agent.WithModel("openrouter/free"),
|
||||
agent.WithTools("read", "bash"),
|
||||
agent.WithDisallowedTools("bash"),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
defer sess.Close()
|
||||
|
||||
if contains(sess.ToolNames(), "bash") {
|
||||
t.Errorf("bash was on both lists and must be removed: %q", sess.ToolNames())
|
||||
}
|
||||
if !contains(sess.ToolNames(), "read") {
|
||||
t.Errorf("read was allowed and not denied, so must survive: %q", sess.ToolNames())
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithoutTools yields an empty set — a pure text completion.
|
||||
func TestWithoutTools(t *testing.T) {
|
||||
hermetic(t)
|
||||
sess, err := agent.New(
|
||||
agent.WithModel("openrouter/free"),
|
||||
agent.WithoutTools(),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
defer sess.Close()
|
||||
|
||||
if len(sess.ToolNames()) != 0 {
|
||||
t.Errorf("WithoutTools must leave no tools, got %q", sess.ToolNames())
|
||||
}
|
||||
}
|
||||
|
||||
// TestUnknownToolIsError confirms a misspelled tool name fails construction
|
||||
// rather than silently dropping the boundary.
|
||||
func TestUnknownToolIsError(t *testing.T) {
|
||||
hermetic(t)
|
||||
_, err := agent.New(
|
||||
agent.WithModel("openrouter/free"),
|
||||
agent.WithTools("raed"),
|
||||
)
|
||||
if err == nil {
|
||||
t.Fatal("New = nil error, want a failure for the misspelled tool name")
|
||||
}
|
||||
}
|
||||
|
||||
// TestInvalidThinkingLevelIsError confirms the level is validated up front.
|
||||
func TestInvalidThinkingLevelIsError(t *testing.T) {
|
||||
hermetic(t)
|
||||
_, err := agent.New(
|
||||
agent.WithModel("openrouter/free"),
|
||||
agent.WithThinkingLevel("supersonic"),
|
||||
)
|
||||
if err == nil {
|
||||
t.Fatal("New = nil error, want a failure for the invalid thinking level")
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidThinkingLevels accepts every documented level.
|
||||
func TestValidThinkingLevels(t *testing.T) {
|
||||
hermetic(t)
|
||||
for _, level := range []string{"off", "minimal", "low", "medium", "high", "xhigh", "max"} {
|
||||
sess, err := agent.New(
|
||||
agent.WithModel("openrouter/free"),
|
||||
agent.WithThinkingLevel(level),
|
||||
)
|
||||
if err != nil {
|
||||
t.Errorf("WithThinkingLevel(%q): %v", level, err)
|
||||
continue
|
||||
}
|
||||
sess.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// TestCloseHermetic confirms Close is a no-op (nil) when the session holds no
|
||||
// plugin manager or memory store, and is safe to call.
|
||||
func TestCloseHermetic(t *testing.T) {
|
||||
hermetic(t)
|
||||
sess, err := agent.New(agent.WithModel("openrouter/free"), agent.WithoutTools())
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
if err := sess.Close(); err != nil {
|
||||
t.Errorf("Close() = %v, want nil", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
// Package agent is the public, embeddable SDK for driving a pigo agent from
|
||||
// your own Go program. It wraps pigo's internal run-assembly, provider, and
|
||||
// agent-loop packages behind a small surface whose every exported type is a Go
|
||||
// primitive (string, []string, bool, func) — so importing this package never
|
||||
// pulls an internal type into your code, and pigo can evolve its internals
|
||||
// without breaking you.
|
||||
//
|
||||
// # Quick start
|
||||
//
|
||||
// sess, err := agent.New(
|
||||
// agent.WithModel("claude-opus-4-8"),
|
||||
// agent.WithAPIKey(os.Getenv("ANTHROPIC_API_KEY")),
|
||||
// )
|
||||
// if err != nil {
|
||||
// log.Fatal(err)
|
||||
// }
|
||||
// defer sess.Close()
|
||||
//
|
||||
// reply, err := sess.Prompt(context.Background(), "Say hello in one word.")
|
||||
// fmt.Println(reply)
|
||||
//
|
||||
// # Model, provider, credentials
|
||||
//
|
||||
// The model id selects the provider the same way the pigo CLI does:
|
||||
// "claude-opus-4-8" resolves to Anthropic, "openrouter/free" to OpenRouter,
|
||||
// and so on. Point at any OpenAI- or Anthropic-compatible endpoint with
|
||||
// [WithBaseURL] + [WithProtocol], or a named provider from your config with
|
||||
// [WithProvider]. The API key comes from [WithAPIKey] or, if unset, the
|
||||
// provider's usual environment variable (e.g. ANTHROPIC_API_KEY). Keys are
|
||||
// never logged.
|
||||
//
|
||||
// # Tools run automatically — read this
|
||||
//
|
||||
// By default a session is created with pigo's full built-in tool set (read,
|
||||
// write, edit, bash, find, grep, and more) and those tools are executed WITHOUT
|
||||
// any per-call confirmation prompt — equivalent to running the CLI with
|
||||
// --approve. An agent can therefore read, modify, and delete files under its
|
||||
// working directory and run shell commands on the host. This is the right
|
||||
// default for an automated SDK, but it means you should only send prompts you
|
||||
// trust, and run in a directory (and, ideally, a sandbox) you are willing to let
|
||||
// the agent modify. To constrain or remove that capability use [WithTools] (an
|
||||
// allowlist), [WithDisallowedTools] (a denylist, which always wins), or
|
||||
// [WithoutTools] (a pure text completion with no tools at all).
|
||||
//
|
||||
// # Conversation state
|
||||
//
|
||||
// A [Session] keeps the running conversation: each [Session.Prompt] or
|
||||
// [Session.Stream] call appends to the same history, so follow-up prompts see
|
||||
// what came before. Call [Session.Reset] to start a fresh conversation on the
|
||||
// same session, or [Session.Close] when you are done. A Session is NOT safe for
|
||||
// concurrent use — drive it from one goroutine, or create one Session per
|
||||
// goroutine.
|
||||
//
|
||||
// # Defaults
|
||||
//
|
||||
// - Tools: on (full built-in set, auto-executed; see the safety note above).
|
||||
// - Skills: off — enable discovery of on-disk skills with [WithSkills].
|
||||
// - Memory: off — enable the persistent memory store with [WithMemory].
|
||||
// - Thinking: "medium" — override with [WithThinkingLevel].
|
||||
//
|
||||
// Skills and memory are off by default so an embedded session is hermetic: it
|
||||
// does not read or write the machine's shared pigo state unless you ask it to.
|
||||
package agent
|
||||
@@ -0,0 +1,123 @@
|
||||
package agent
|
||||
|
||||
// config is the resolved, unexported construction state for a Session. It is
|
||||
// populated only through Option values, so callers never name or mutate it
|
||||
// directly — the exported surface stays limited to With* constructors and the
|
||||
// Session methods.
|
||||
type config struct {
|
||||
model string
|
||||
baseURL string
|
||||
protocol string
|
||||
provider string
|
||||
apiKey string
|
||||
systemPrompt string
|
||||
appendSystemPrompt []string
|
||||
thinking string
|
||||
noTools bool
|
||||
allowedTools []string
|
||||
disallowedTools []string
|
||||
skills bool
|
||||
memory bool
|
||||
}
|
||||
|
||||
// Option configures a Session at construction time. Options are applied in the
|
||||
// order passed to New, so a later option overrides an earlier one that sets the
|
||||
// same field. Because config is unexported, the only way to produce an Option is
|
||||
// through the With* constructors below — which keeps the public surface free of
|
||||
// internal types.
|
||||
type Option func(*config)
|
||||
|
||||
// WithModel sets the model id, which also selects the provider the way the pigo
|
||||
// CLI does (e.g. "claude-opus-4-8" → Anthropic, "openrouter/free" → OpenRouter).
|
||||
// The default is "openrouter/free".
|
||||
func WithModel(model string) Option {
|
||||
return func(c *config) { c.model = model }
|
||||
}
|
||||
|
||||
// WithBaseURL points the session at a custom endpoint. Pair it with
|
||||
// [WithProtocol] to say whether that endpoint speaks the OpenAI or Anthropic
|
||||
// wire format.
|
||||
func WithBaseURL(baseURL string) Option {
|
||||
return func(c *config) { c.baseURL = baseURL }
|
||||
}
|
||||
|
||||
// WithProtocol selects the wire protocol for a custom endpoint: "openai" or
|
||||
// "anthropic". It is only consulted when [WithBaseURL] is set.
|
||||
func WithProtocol(protocol string) Option {
|
||||
return func(c *config) { c.protocol = protocol }
|
||||
}
|
||||
|
||||
// WithProvider selects a named provider from your pigo configuration instead of
|
||||
// inferring one from the model id.
|
||||
func WithProvider(name string) Option {
|
||||
return func(c *config) { c.provider = name }
|
||||
}
|
||||
|
||||
// WithAPIKey sets the API key for the resolved provider, overriding the
|
||||
// provider's environment variable. When unset, the provider's usual environment
|
||||
// variable is used (e.g. ANTHROPIC_API_KEY, OPENROUTER_API_KEY).
|
||||
func WithAPIKey(key string) Option {
|
||||
return func(c *config) { c.apiKey = key }
|
||||
}
|
||||
|
||||
// WithSystemPrompt replaces pigo's built-in base instruction with prompt. Use
|
||||
// this for full control over the agent's persona and rules; use
|
||||
// [WithAppendSystemPrompt] instead to keep the built-in instruction and add to
|
||||
// it.
|
||||
func WithSystemPrompt(prompt string) Option {
|
||||
return func(c *config) { c.systemPrompt = prompt }
|
||||
}
|
||||
|
||||
// WithAppendSystemPrompt appends one or more blocks to the system prompt,
|
||||
// leaving pigo's built-in instruction in place. Repeated calls accumulate.
|
||||
func WithAppendSystemPrompt(blocks ...string) Option {
|
||||
return func(c *config) {
|
||||
c.appendSystemPrompt = append(c.appendSystemPrompt, blocks...)
|
||||
}
|
||||
}
|
||||
|
||||
// WithThinkingLevel sets the reasoning-effort level. Valid values are "off",
|
||||
// "minimal", "low", "medium", "high", "xhigh", and "max". The default is
|
||||
// "medium". An invalid value makes New return an error.
|
||||
func WithThinkingLevel(level string) Option {
|
||||
return func(c *config) { c.thinking = level }
|
||||
}
|
||||
|
||||
// WithTools restricts the session to the named built-in tools (an allowlist,
|
||||
// e.g. WithTools("read", "grep")). Names are matched case-insensitively, so
|
||||
// "Read" and "read" are equivalent. A name that matches no tool makes New
|
||||
// return an error rather than silently ignoring it. Combine with
|
||||
// [WithDisallowedTools]; deny always wins over allow.
|
||||
func WithTools(names ...string) Option {
|
||||
return func(c *config) { c.allowedTools = append(c.allowedTools, names...) }
|
||||
}
|
||||
|
||||
// WithDisallowedTools removes the named built-in tools (a denylist, e.g.
|
||||
// WithDisallowedTools("bash")). Deny always wins: a tool named here is removed
|
||||
// even if it also appears in [WithTools]. As with WithTools, an unknown name
|
||||
// makes New return an error.
|
||||
func WithDisallowedTools(names ...string) Option {
|
||||
return func(c *config) { c.disallowedTools = append(c.disallowedTools, names...) }
|
||||
}
|
||||
|
||||
// WithoutTools removes every tool, producing a pure text-completion session that
|
||||
// cannot touch the filesystem or run commands. It overrides [WithTools] and
|
||||
// [WithDisallowedTools], which become inert once the tool set is empty.
|
||||
func WithoutTools() Option {
|
||||
return func(c *config) { c.noTools = true }
|
||||
}
|
||||
|
||||
// WithSkills enables discovery of on-disk skills, which are advertised to the
|
||||
// model and loadable during a run. Skills are off by default so an embedded
|
||||
// session stays independent of the machine's shared skills directory.
|
||||
func WithSkills() Option {
|
||||
return func(c *config) { c.skills = true }
|
||||
}
|
||||
|
||||
// WithMemory enables pigo's persistent memory store, letting the agent recall
|
||||
// context saved by earlier runs and record new memories. Memory is off by
|
||||
// default so an embedded session does not read or write shared state unless
|
||||
// asked.
|
||||
func WithMemory() Option {
|
||||
return func(c *config) { c.memory = true }
|
||||
}
|
||||
Reference in New Issue
Block a user