From a20eba286ed39d5a1250dfc0da3358759c2f759f Mon Sep 17 00:00:00 2001 From: Ameer Deen Date: Sun, 24 May 2026 22:47:35 +0400 Subject: [PATCH] Refactor progressive package layout --- .github/workflows/ci.yml | 4 + README.md | 11 +++ _examples/chat-console/README.md | 31 ++++++++ _examples/chat-console/go.mod | 17 +++++ _examples/chat-console/go.sum | 23 ++++++ _examples/chat-console/main.go | 120 ++++++++++++++++++++++++++++++ agent_test.go | 16 ++++ docs/design/design.md | 36 ++++++--- executor.go | 49 +----------- executor/executor.go | 12 +++ executor/sequential/sequential.go | 48 ++++++++++++ memory.go | 58 ++------------- memory/inmem/inmem.go | 56 ++++++++++++++ memory/memory.go | 14 ++++ message/message.go | 27 +++++++ metadata.go | 22 ++---- metadata/metadata.go | 25 +++++++ middleware.go | 10 +-- middleware/middleware.go | 14 ++++ provider.go | 23 +----- provider/provider.go | 43 +++++++++++ provider/xai/xai.go | 10 +-- registry.go | 46 +----------- tool.go | 58 ++------------- tool/call.go | 23 ++++++ tool/registry/registry.go | 49 ++++++++++++ tool/tool.go | 63 ++++++++++++++++ types.go | 69 ++++++----------- 28 files changed, 681 insertions(+), 296 deletions(-) create mode 100644 _examples/chat-console/README.md create mode 100644 _examples/chat-console/go.mod create mode 100644 _examples/chat-console/go.sum create mode 100644 _examples/chat-console/main.go create mode 100644 executor/executor.go create mode 100644 executor/sequential/sequential.go create mode 100644 memory/inmem/inmem.go create mode 100644 memory/memory.go create mode 100644 message/message.go create mode 100644 metadata/metadata.go create mode 100644 middleware/middleware.go create mode 100644 provider/provider.go create mode 100644 tool/call.go create mode 100644 tool/registry/registry.go create mode 100644 tool/tool.go diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 71e10f0..d97d17f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -25,6 +25,8 @@ jobs: run: | cd _examples/hello-world go vet ./... + cd ../chat-console + go vet ./... cd ../calculator go vet ./... @@ -35,6 +37,8 @@ jobs: run: | cd _examples/hello-world go test ./... + cd ../chat-console + go test ./... cd ../calculator go test ./... diff --git a/README.md b/README.md index 6b9d599..9fefe68 100644 --- a/README.md +++ b/README.md @@ -213,6 +213,17 @@ type MemoryStore interface { } ``` +Or opt into a supplied memory implementation explicitly: + +```go +import "github.com/katasec/forge/memory/inmem" + +agent, _ := forge.NewAgent(forge.Config{ + Provider: myProvider, + Memory: inmem.New(), +}) +``` + ### Metadata Attach arbitrary key-value data to the context, accessible by tools and middleware: diff --git a/_examples/chat-console/README.md b/_examples/chat-console/README.md new file mode 100644 index 0000000..4de8d54 --- /dev/null +++ b/_examples/chat-console/README.md @@ -0,0 +1,31 @@ +# Chat Console + +An interactive console app that shows the intended forge developer experience: + +- configure an agent once with `forge.Config` +- talk to it with `Ask` +- read the latest answer with `LastText` +- keep conversation context in memory across turns + +## Run + +```bash +export ANTHROPIC_API_KEY=sk-ant-... +go run . +``` + +Use xAI's OpenAI-compatible endpoint: + +```bash +export XAI_API_KEY=xai-... +go run . -provider xai +``` + +Use xAI Responses API with web search: + +```bash +export XAI_API_KEY=xai-... +go run . -provider xai-search +``` + +Type `exit` or `quit` to stop. diff --git a/_examples/chat-console/go.mod b/_examples/chat-console/go.mod new file mode 100644 index 0000000..d325442 --- /dev/null +++ b/_examples/chat-console/go.mod @@ -0,0 +1,17 @@ +module github.com/katasec/forge/_examples/chat-console + +go 1.25.6 + +replace github.com/katasec/forge => ../../ + +require github.com/katasec/forge v0.0.0-00010101000000-000000000000 + +require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/invopop/jsonschema v0.13.0 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/_examples/chat-console/go.sum b/_examples/chat-console/go.sum new file mode 100644 index 0000000..a8f1621 --- /dev/null +++ b/_examples/chat-console/go.sum @@ -0,0 +1,23 @@ +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/_examples/chat-console/main.go b/_examples/chat-console/main.go new file mode 100644 index 0000000..df54f54 --- /dev/null +++ b/_examples/chat-console/main.go @@ -0,0 +1,120 @@ +// Chat Console demonstrates the progressive forge developer experience: +// configure an agent once, then talk to it with Ask. +package main + +import ( + "bufio" + "context" + "flag" + "fmt" + "log" + "os" + "strings" + + "github.com/katasec/forge" + "github.com/katasec/forge/memory/inmem" + "github.com/katasec/forge/provider/anthropic" + "github.com/katasec/forge/provider/openai" + "github.com/katasec/forge/provider/xai" +) + +func main() { + providerFlag := flag.String("provider", "anthropic", "Provider to use: anthropic, xai, or xai-search") + flag.Parse() + + provider, citations := buildProvider(*providerFlag) + agent, err := forge.NewAgent(forge.Config{ + Provider: provider, + Memory: inmem.New(), + SystemPrompt: "You are a concise assistant. Keep answers short and useful.", + }) + if err != nil { + log.Fatal(err) + } + + runConsole(context.Background(), agent, citations) +} + +func buildProvider(name string) (forge.Provider, func()) { + switch name { + case "anthropic": + key := requireEnv("ANTHROPIC_API_KEY") + return anthropic.New(key, "claude-sonnet-4-20250514"), func() {} + case "xai": + key := requireEnv("XAI_API_KEY") + return openai.New("https://api.x.ai/v1", key, "grok-3-mini"), func() {} + case "xai-search": + key := requireEnv("XAI_API_KEY") + provider := xai.New(key, "grok-4-1-fast-non-reasoning", xai.WithWebSearch()) + return provider, func() { printCitations(provider.LastCitations()) } + default: + log.Fatalf("unknown provider %q; use anthropic, xai, or xai-search", name) + return nil, nil + } +} + +func runConsole(ctx context.Context, agent *forge.Agent, afterAnswer func()) { + scanner := bufio.NewScanner(os.Stdin) + + fmt.Println("Forge chat console") + fmt.Println("Type exit or quit to stop.") + fmt.Println() + + for { + fmt.Print("You: ") + if !scanner.Scan() { + break + } + + input := strings.TrimSpace(scanner.Text()) + if shouldExit(input) { + break + } + if input == "" { + continue + } + + resp, err := agent.Ask(ctx, input) + if err != nil { + log.Printf("agent error: %v", err) + continue + } + + fmt.Printf("Assistant: %s\n", resp.LastText()) + afterAnswer() + fmt.Println() + } + + if err := scanner.Err(); err != nil { + log.Printf("read input: %v", err) + } +} + +func shouldExit(input string) bool { + switch strings.ToLower(input) { + case "exit", "quit": + return true + default: + return false + } +} + +func requireEnv(name string) string { + value := os.Getenv(name) + if value == "" { + log.Fatalf("set %s", name) + } + return value +} + +func printCitations(citations []xai.Citation) { + if len(citations) == 0 { + return + } + + fmt.Println() + fmt.Println("Sources:") + for i, citation := range citations { + fmt.Printf(" [%d] %s - %s\n", i+1, citation.Title, citation.URL) + } +} diff --git a/agent_test.go b/agent_test.go index 9250af6..513c195 100644 --- a/agent_test.go +++ b/agent_test.go @@ -5,6 +5,8 @@ import ( "encoding/json" "errors" "testing" + + "github.com/katasec/forge/memory/inmem" ) // mockProvider is a test double that returns pre-configured responses. @@ -97,6 +99,20 @@ func TestNewAgentDisableMemory(t *testing.T) { } } +func TestNewAgentAcceptsExplicitMemoryStore(t *testing.T) { + store := inmem.New() + agent, err := NewAgent(Config{ + Provider: &mockProvider{}, + Memory: store, + }) + if err != nil { + t.Fatalf("NewAgent error: %v", err) + } + if agent.memory != store { + t.Fatalf("memory = %T, want explicit store", agent.memory) + } +} + func TestAgentAskPreservesDefaultConversation(t *testing.T) { provider := &recordingProvider{ responses: []*ProviderResponse{ diff --git a/docs/design/design.md b/docs/design/design.md index 41ba9eb..7c427c1 100644 --- a/docs/design/design.md +++ b/docs/design/design.md @@ -10,15 +10,25 @@ module github.com/katasec/forge go 1.23 ``` -Root package: `package forge` — all interfaces, core types, and default implementations. +Root package: `package forge` — the primary user-facing facade: `Config`, `Agent`, `Ask`, `AskIn`, core interfaces, and type aliases for common concepts. -Sub-packages hold swappable backend implementations (following the `database/sql` pattern): +Sub-packages hold focused concepts and swappable implementations. The root package re-exports the common API so developers can start with `forge.Config` and only import subpackages when they opt into a specific implementation: ``` -forge/ root — interfaces, types, Agent, Config, defaults +forge/ root facade: Agent, Config, common aliases, default helpers +forge/message/ message and role types +forge/tool/ tool interface, typed Func helper, calls, results +forge/tool/registry/ tool registry implementation +forge/provider/ provider interface, requests, responses, usage, finish reasons forge/provider/anthropic/ Anthropic Messages API provider forge/provider/openai/ OpenAI-compatible provider (OpenAI, xAI, Together, Groq) forge/provider/xai/ xAI Responses API provider (web search, X search, citations) +forge/memory/ memory store interface +forge/memory/inmem/ in-memory memory store +forge/executor/ tool executor interface +forge/executor/sequential/ sequential tool executor +forge/middleware/ provider middleware types +forge/metadata/ context metadata helpers ``` Future sub-packages (created when implementations exist, not preemptively): @@ -27,6 +37,8 @@ Future sub-packages (created when implementations exist, not preemptively): forge/memory/sqlite/ SQLite-backed MemoryStore forge/memory/redis/ Redis-backed MemoryStore forge/executor/concurrent/ Parallel tool executor +forge/middleware/logging/ Logging middleware +forge/middleware/retry/ Retry middleware ``` ## Code Style & API Philosophy @@ -37,11 +49,11 @@ Top-level functions should read like intent. Complex behavior should be composed The public API should make the common path obvious before exposing the lower-level machinery. Developers should be able to start with `Ask(ctx, prompt)`, use `AskIn(ctx, conversationID, prompt)` when they need named conversations, and drop to `Run(ctx, AgentRequest{...})` only when they need full control over roles, message history, or advanced orchestration. -Package layout should follow the same rule as the code. The root package may remain a friendly facade, but implementation-heavy defaults should move toward focused packages as the library grows: agent orchestration, messages, tools, providers, memory, executors, and metadata. +Package layout follows the same rule as the code. The root package remains a friendly facade, while focused packages organize agent concepts: messages, tools, providers, memory, executors, middleware, and metadata. Developers configure an agent runtime once, then interact with it through `Ask`, `AskIn`, or `Run`. --- -## 1. Core Types (`types.go`) +## 1. Core Types (`message`, `tool`, `provider`; re-exported by root) ### Role @@ -174,7 +186,7 @@ func MetadataFromContext(ctx context.Context) (Metadata, bool) --- -## 3. Tool System (`tool.go`, `registry.go`) +## 3. Tool System (`tool`, `tool/registry`; re-exported by root) ### Tool Interface @@ -234,10 +246,11 @@ func (r *ToolRegistry) Definitions() []ToolDefinition - `Register` adds tools. Duplicate names overwrite silently (last-write-wins). - `Get` returns a tool by name. - `Definitions` returns all registered tools as `ToolDefinition` for passing to a provider. +- The concrete implementation lives in `tool/registry`; root re-exports `ToolRegistry` and `NewToolRegistry` for the common path. --- -## 4. Executor (`executor.go`) +## 4. Executor (`executor`, `executor/sequential`; re-exported by root) ```go type ToolExecutor interface { @@ -262,7 +275,7 @@ func (e *SequentialExecutor) Execute(ctx context.Context, calls []ToolCall) []To --- -## 5. Memory (`memory.go`) +## 5. Memory (`memory`, `memory/inmem`; re-exported by root) ```go type MemoryStore interface { @@ -280,10 +293,12 @@ func NewInMemoryStore() *InMemoryStore - `Save` replaces the entire message history for that conversation ID. - `Clear` deletes the conversation. - `InMemoryStore` is safe for concurrent use. +- The `MemoryStore` contract lives in `memory`; the default implementation lives in `memory/inmem`. +- Root re-exports `MemoryStore`, `InMemoryStore`, and `NewInMemoryStore` to preserve the simple `forge.Config` experience. --- -## 6. Provider & Middleware (`provider.go`, `middleware.go`) +## 6. Provider & Middleware (`provider`, `middleware`; re-exported by root) ### Provider @@ -308,6 +323,7 @@ type Provider interface { - `Generate` makes a single LLM call. It does **not** loop. - The provider is responsible for translating `ProviderRequest` into the LLM's native API format and back. - A provider error (non-nil error return) always terminates the agent loop. +- The provider contract lives in `provider`; root re-exports `Provider`, `ProviderRequest`, and `ProviderResponse`. ### Middleware @@ -336,6 +352,8 @@ for i := len(middlewares) - 1; i >= 0; i-- { Use cases: logging, token counting, rate limiting, retry, prompt injection. +The middleware contract lives in `middleware`; root re-exports `RunFunc` and `Middleware`. + --- ## 7. Config & Agent (`config.go`, `agent.go`) diff --git a/executor.go b/executor.go index dee0fc0..cf76dc5 100644 --- a/executor.go +++ b/executor.go @@ -1,50 +1,9 @@ package forge import ( - "context" - "fmt" + executorpkg "github.com/katasec/forge/executor" + "github.com/katasec/forge/executor/sequential" ) -// ToolExecutor executes a batch of tool calls and returns the results. -type ToolExecutor interface { - Execute(ctx context.Context, calls []ToolCall) []ToolResult -} - -// SequentialExecutor invokes tools one at a time via a ToolRegistry. -type SequentialExecutor struct { - Registry *ToolRegistry -} - -// Execute processes each tool call in order. Missing tools and invocation -// errors are returned as ToolResults with IsError set to true. -func (e *SequentialExecutor) Execute(ctx context.Context, calls []ToolCall) []ToolResult { - results := make([]ToolResult, 0, len(calls)) - for _, call := range calls { - tool, ok := e.Registry.Get(call.Name) - if !ok { - results = append(results, ToolResult{ - CallID: call.ID, - Content: fmt.Sprintf("tool not found: %s", call.Name), - IsError: true, - }) - continue - } - - content, err := tool.Invoke(ctx, call.Arguments) - if err != nil { - results = append(results, ToolResult{ - CallID: call.ID, - Content: err.Error(), - IsError: true, - }) - continue - } - - results = append(results, ToolResult{ - CallID: call.ID, - Content: content, - IsError: false, - }) - } - return results -} +type ToolExecutor = executorpkg.Executor +type SequentialExecutor = sequential.Executor diff --git a/executor/executor.go b/executor/executor.go new file mode 100644 index 0000000..74053e5 --- /dev/null +++ b/executor/executor.go @@ -0,0 +1,12 @@ +package executor + +import ( + "context" + + "github.com/katasec/forge/tool" +) + +// Executor executes a batch of tool calls and returns the results. +type Executor interface { + Execute(ctx context.Context, calls []tool.Call) []tool.Result +} diff --git a/executor/sequential/sequential.go b/executor/sequential/sequential.go new file mode 100644 index 0000000..f7fc12d --- /dev/null +++ b/executor/sequential/sequential.go @@ -0,0 +1,48 @@ +package sequential + +import ( + "context" + "fmt" + + "github.com/katasec/forge/tool" + "github.com/katasec/forge/tool/registry" +) + +// Executor invokes tools one at a time via a tool registry. +type Executor struct { + Registry *registry.Registry +} + +// Execute processes each tool call in order. Missing tools and invocation +// errors are returned as Results with IsError set to true. +func (e *Executor) Execute(ctx context.Context, calls []tool.Call) []tool.Result { + results := make([]tool.Result, 0, len(calls)) + for _, call := range calls { + t, ok := e.Registry.Get(call.Name) + if !ok { + results = append(results, tool.Result{ + CallID: call.ID, + Content: fmt.Sprintf("tool not found: %s", call.Name), + IsError: true, + }) + continue + } + + content, err := t.Invoke(ctx, call.Arguments) + if err != nil { + results = append(results, tool.Result{ + CallID: call.ID, + Content: err.Error(), + IsError: true, + }) + continue + } + + results = append(results, tool.Result{ + CallID: call.ID, + Content: content, + IsError: false, + }) + } + return results +} diff --git a/memory.go b/memory.go index fccaf86..b6ea000 100644 --- a/memory.go +++ b/memory.go @@ -1,61 +1,13 @@ package forge import ( - "context" - "sync" + memorypkg "github.com/katasec/forge/memory" + "github.com/katasec/forge/memory/inmem" ) -// MemoryStore persists conversation message history. -type MemoryStore interface { - Load(ctx context.Context, conversationID string) ([]Message, error) - Save(ctx context.Context, conversationID string, messages []Message) error - Clear(ctx context.Context, conversationID string) error -} - -// InMemoryStore is a thread-safe in-memory implementation of MemoryStore. -type InMemoryStore struct { - mu sync.RWMutex - data map[string][]Message -} +type MemoryStore = memorypkg.Store +type InMemoryStore = inmem.Store -// NewInMemoryStore creates an empty InMemoryStore. func NewInMemoryStore() *InMemoryStore { - return &InMemoryStore{ - data: make(map[string][]Message), - } -} - -// Load returns a copy of the stored messages for the given conversation. -func (s *InMemoryStore) Load(_ context.Context, conversationID string) ([]Message, error) { - s.mu.RLock() - defer s.mu.RUnlock() - - msgs, ok := s.data[conversationID] - if !ok { - return nil, nil - } - - cp := make([]Message, len(msgs)) - copy(cp, msgs) - return cp, nil -} - -// Save replaces the entire message history for the given conversation. -func (s *InMemoryStore) Save(_ context.Context, conversationID string, messages []Message) error { - s.mu.Lock() - defer s.mu.Unlock() - - cp := make([]Message, len(messages)) - copy(cp, messages) - s.data[conversationID] = cp - return nil -} - -// Clear deletes the conversation history. -func (s *InMemoryStore) Clear(_ context.Context, conversationID string) error { - s.mu.Lock() - defer s.mu.Unlock() - - delete(s.data, conversationID) - return nil + return inmem.New() } diff --git a/memory/inmem/inmem.go b/memory/inmem/inmem.go new file mode 100644 index 0000000..63adfe5 --- /dev/null +++ b/memory/inmem/inmem.go @@ -0,0 +1,56 @@ +package inmem + +import ( + "context" + "sync" + + "github.com/katasec/forge/message" +) + +// Store is a thread-safe in-memory memory store. +type Store struct { + mu sync.RWMutex + data map[string][]message.Message +} + +// New creates an empty in-memory store. +func New() *Store { + return &Store{ + data: make(map[string][]message.Message), + } +} + +// Load returns a copy of the stored messages for the given conversation. +func (s *Store) Load(_ context.Context, conversationID string) ([]message.Message, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + msgs, ok := s.data[conversationID] + if !ok { + return nil, nil + } + + cp := make([]message.Message, len(msgs)) + copy(cp, msgs) + return cp, nil +} + +// Save replaces the entire message history for the given conversation. +func (s *Store) Save(_ context.Context, conversationID string, messages []message.Message) error { + s.mu.Lock() + defer s.mu.Unlock() + + cp := make([]message.Message, len(messages)) + copy(cp, messages) + s.data[conversationID] = cp + return nil +} + +// Clear deletes the conversation history. +func (s *Store) Clear(_ context.Context, conversationID string) error { + s.mu.Lock() + defer s.mu.Unlock() + + delete(s.data, conversationID) + return nil +} diff --git a/memory/memory.go b/memory/memory.go new file mode 100644 index 0000000..db57c77 --- /dev/null +++ b/memory/memory.go @@ -0,0 +1,14 @@ +package memory + +import ( + "context" + + "github.com/katasec/forge/message" +) + +// Store persists conversation message history. +type Store interface { + Load(ctx context.Context, conversationID string) ([]message.Message, error) + Save(ctx context.Context, conversationID string, messages []message.Message) error + Clear(ctx context.Context, conversationID string) error +} diff --git a/message/message.go b/message/message.go new file mode 100644 index 0000000..a83e39c --- /dev/null +++ b/message/message.go @@ -0,0 +1,27 @@ +package message + +import "github.com/katasec/forge/tool" + +// Role identifies the sender of a message in a conversation. +type Role string + +const ( + RoleUser Role = "user" + RoleAssistant Role = "assistant" + RoleTool Role = "tool" + RoleSystem Role = "system" +) + +// Message represents a single message in a conversation. +type Message struct { + ID string `json:"id"` + Role Role `json:"role"` + Content string `json:"content"` + ToolCalls []tool.Call `json:"tool_calls,omitempty"` + ToolResults []tool.Result `json:"tool_results,omitempty"` +} + +// UserMessage creates a user-role message with the given content. +func UserMessage(content string) Message { + return Message{Role: RoleUser, Content: content} +} diff --git a/metadata.go b/metadata.go index ef6d0ac..0b4c5d9 100644 --- a/metadata.go +++ b/metadata.go @@ -1,25 +1,17 @@ package forge -import "context" +import ( + "context" -type metadataKey struct{} + metadatapkg "github.com/katasec/forge/metadata" +) -// Metadata holds arbitrary key-value pairs attached to a context. -type Metadata struct { - Values map[string]string -} +type Metadata = metadatapkg.Metadata -// WithMetadata returns a new context with the given Metadata stored in it. func WithMetadata(ctx context.Context, m Metadata) context.Context { - if m.Values == nil { - m.Values = make(map[string]string) - } - return context.WithValue(ctx, metadataKey{}, m) + return metadatapkg.WithMetadata(ctx, m) } -// MetadataFromContext retrieves Metadata from the context. -// Returns false if no Metadata is present. func MetadataFromContext(ctx context.Context) (Metadata, bool) { - m, ok := ctx.Value(metadataKey{}).(Metadata) - return m, ok + return metadatapkg.FromContext(ctx) } diff --git a/metadata/metadata.go b/metadata/metadata.go new file mode 100644 index 0000000..a61e579 --- /dev/null +++ b/metadata/metadata.go @@ -0,0 +1,25 @@ +package metadata + +import "context" + +type key struct{} + +// Metadata holds arbitrary key-value pairs attached to a context. +type Metadata struct { + Values map[string]string +} + +// WithMetadata returns a new context with the given Metadata stored in it. +func WithMetadata(ctx context.Context, m Metadata) context.Context { + if m.Values == nil { + m.Values = make(map[string]string) + } + return context.WithValue(ctx, key{}, m) +} + +// FromContext retrieves Metadata from the context. +// Returns false if no Metadata is present. +func FromContext(ctx context.Context) (Metadata, bool) { + m, ok := ctx.Value(key{}).(Metadata) + return m, ok +} diff --git a/middleware.go b/middleware.go index 85b4ac1..1049e9e 100644 --- a/middleware.go +++ b/middleware.go @@ -1,10 +1,6 @@ package forge -import "context" +import middlewarepkg "github.com/katasec/forge/middleware" -// RunFunc is the signature for a single provider call, used by middleware. -type RunFunc func(ctx context.Context, req ProviderRequest) (*ProviderResponse, error) - -// Middleware wraps a RunFunc to intercept provider calls. -// Composition order: given [A, B, C], request flows A → B → C → provider → C → B → A. -type Middleware func(next RunFunc) RunFunc +type RunFunc = middlewarepkg.RunFunc +type Middleware = middlewarepkg.Middleware diff --git a/middleware/middleware.go b/middleware/middleware.go new file mode 100644 index 0000000..60c6be0 --- /dev/null +++ b/middleware/middleware.go @@ -0,0 +1,14 @@ +package middleware + +import ( + "context" + + "github.com/katasec/forge/provider" +) + +// RunFunc is the signature for a single provider call, used by middleware. +type RunFunc func(ctx context.Context, req provider.Request) (*provider.Response, error) + +// Middleware wraps a RunFunc to intercept provider calls. +// Composition order: given [A, B, C], request flows A -> B -> C -> provider -> C -> B -> A. +type Middleware func(next RunFunc) RunFunc diff --git a/provider.go b/provider.go index 0f19d8c..59c5dd9 100644 --- a/provider.go +++ b/provider.go @@ -1,22 +1,7 @@ package forge -import "context" +import providerpkg "github.com/katasec/forge/provider" -// ProviderRequest is the input to a single LLM call. -type ProviderRequest struct { - Messages []Message `json:"messages"` - Tools []ToolDefinition `json:"tools,omitempty"` - SystemPrompt string `json:"system_prompt,omitempty"` -} - -// ProviderResponse is the output of a single LLM call. -type ProviderResponse struct { - Message Message `json:"message"` - FinishReason FinishReason `json:"finish_reason"` - Usage TokenUsage `json:"usage"` -} - -// Provider makes a single LLM call. It does not loop. -type Provider interface { - Generate(ctx context.Context, req ProviderRequest) (*ProviderResponse, error) -} +type ProviderRequest = providerpkg.Request +type ProviderResponse = providerpkg.Response +type Provider = providerpkg.Provider diff --git a/provider/provider.go b/provider/provider.go new file mode 100644 index 0000000..0410edc --- /dev/null +++ b/provider/provider.go @@ -0,0 +1,43 @@ +package provider + +import ( + "context" + + "github.com/katasec/forge/message" + "github.com/katasec/forge/tool" +) + +// FinishReason indicates why the agent loop terminated. +type FinishReason string + +const ( + FinishReasonStop FinishReason = "stop" + FinishReasonToolUse FinishReason = "tool_use" + FinishReasonIterLimit FinishReason = "iter_limit" + FinishReasonError FinishReason = "error" +) + +// TokenUsage tracks token consumption across provider calls. +type TokenUsage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` +} + +// Request is the input to a single LLM call. +type Request struct { + Messages []message.Message `json:"messages"` + Tools []tool.Definition `json:"tools,omitempty"` + SystemPrompt string `json:"system_prompt,omitempty"` +} + +// Response is the output of a single LLM call. +type Response struct { + Message message.Message `json:"message"` + FinishReason FinishReason `json:"finish_reason"` + Usage TokenUsage `json:"usage"` +} + +// Provider makes a single LLM call. It does not loop. +type Provider interface { + Generate(ctx context.Context, req Request) (*Response, error) +} diff --git a/provider/xai/xai.go b/provider/xai/xai.go index 4571599..50cf1ca 100644 --- a/provider/xai/xai.go +++ b/provider/xai/xai.go @@ -146,8 +146,8 @@ func (p *Provider) LastCitations() []Citation { // --- xAI Responses API wire types --- type request struct { - Model string `json:"model"` - Input []inputItem `json:"input"` + Model string `json:"model"` + Input []inputItem `json:"input"` Tools []requestTool `json:"tools,omitempty"` } @@ -156,7 +156,7 @@ type inputItem struct { Role string `json:"role,omitempty"` Content string `json:"content,omitempty"` // Tool result fields - Type string `json:"type,omitempty"` // "function_call_output" + Type string `json:"type,omitempty"` // "function_call_output" CallID string `json:"call_id,omitempty"` Output string `json:"output,omitempty"` } @@ -292,8 +292,8 @@ func parseResponse(resp *response) (*forge.ProviderResponse, []Citation) { } } } - // Server-side tool calls (web_search_call, x_search_call, etc.) - // are auto-executed by xAI — we don't surface them. + // Server-side tool calls (web_search_call, x_search_call, etc.) + // are auto-executed by xAI — we don't surface them. } } diff --git a/registry.go b/registry.go index de8a213..cfa3acc 100644 --- a/registry.go +++ b/registry.go @@ -1,47 +1,9 @@ package forge -// ToolRegistry stores tools and provides lookup by name. -type ToolRegistry struct { - tools map[string]Tool - order []string -} - -// NewToolRegistry creates an empty ToolRegistry. -func NewToolRegistry() *ToolRegistry { - return &ToolRegistry{ - tools: make(map[string]Tool), - } -} +import "github.com/katasec/forge/tool/registry" -// Register adds one or more tools to the registry. -// Duplicate names overwrite silently (last-write-wins). -func (r *ToolRegistry) Register(tools ...Tool) { - for _, t := range tools { - name := t.Name() - if _, exists := r.tools[name]; !exists { - r.order = append(r.order, name) - } - r.tools[name] = t - } -} - -// Get returns a tool by name. Returns false if not found. -func (r *ToolRegistry) Get(name string) (Tool, bool) { - t, ok := r.tools[name] - return t, ok -} +type ToolRegistry = registry.Registry -// Definitions returns all registered tools as ToolDefinitions, -// in the order they were first registered. -func (r *ToolRegistry) Definitions() []ToolDefinition { - defs := make([]ToolDefinition, 0, len(r.order)) - for _, name := range r.order { - t := r.tools[name] - defs = append(defs, ToolDefinition{ - Name: t.Name(), - Description: t.Description(), - Schema: t.Schema(), - }) - } - return defs +func NewToolRegistry() *ToolRegistry { + return registry.New() } diff --git a/tool.go b/tool.go index d19de56..daf4949 100644 --- a/tool.go +++ b/tool.go @@ -2,62 +2,14 @@ package forge import ( "context" - "encoding/json" - "github.com/invopop/jsonschema" + toolpkg "github.com/katasec/forge/tool" ) -// Tool defines a callable tool that can be invoked by an agent. -type Tool interface { - Name() string - Description() string - Schema() ToolSchema - Invoke(ctx context.Context, args json.RawMessage) (string, error) -} - -// ToolSchema describes the JSON Schema for a tool's parameters. -type ToolSchema struct { - Parameters json.RawMessage `json:"parameters"` -} +type Tool = toolpkg.Tool +type ToolSchema = toolpkg.Schema +type ToolDefinition = toolpkg.Definition -// ToolDefinition is the wire format sent to providers so they know which tools are available. -type ToolDefinition struct { - Name string `json:"name"` - Description string `json:"description"` - Schema ToolSchema `json:"schema"` -} - -// funcTool is the unexported implementation backing the Func helper. -type funcTool[T any] struct { - name string - description string - schema ToolSchema - fn func(ctx context.Context, input T) (string, error) -} - -// Func creates a Tool from a typed function. The JSON Schema for parameters -// is derived from T using invopop/jsonschema at construction time. func Func[T any](name, description string, fn func(ctx context.Context, input T) (string, error)) Tool { - r := new(jsonschema.Reflector) - s := r.Reflect(new(T)) - params, _ := json.Marshal(s) - - return &funcTool[T]{ - name: name, - description: description, - schema: ToolSchema{Parameters: params}, - fn: fn, - } -} - -func (t *funcTool[T]) Name() string { return t.name } -func (t *funcTool[T]) Description() string { return t.description } -func (t *funcTool[T]) Schema() ToolSchema { return t.schema } - -func (t *funcTool[T]) Invoke(ctx context.Context, args json.RawMessage) (string, error) { - var input T - if err := json.Unmarshal(args, &input); err != nil { - return "", err - } - return t.fn(ctx, input) + return toolpkg.Func(name, description, fn) } diff --git a/tool/call.go b/tool/call.go new file mode 100644 index 0000000..c304d72 --- /dev/null +++ b/tool/call.go @@ -0,0 +1,23 @@ +package tool + +import "encoding/json" + +// Call represents a request from the LLM to invoke a tool. +type Call struct { + ID string `json:"id"` + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` +} + +// Result represents the outcome of a tool invocation. +type Result struct { + CallID string `json:"call_id"` + Content string `json:"content"` + IsError bool `json:"is_error"` +} + +// Error wraps a tool invocation failure. +type Error struct { + CallID string `json:"call_id"` + Message string `json:"message"` +} diff --git a/tool/registry/registry.go b/tool/registry/registry.go new file mode 100644 index 0000000..4cb824e --- /dev/null +++ b/tool/registry/registry.go @@ -0,0 +1,49 @@ +package registry + +import "github.com/katasec/forge/tool" + +// Registry stores tools and provides lookup by name. +type Registry struct { + tools map[string]tool.Tool + order []string +} + +// New creates an empty Registry. +func New() *Registry { + return &Registry{ + tools: make(map[string]tool.Tool), + } +} + +// Register adds one or more tools to the registry. +// Duplicate names overwrite silently (last-write-wins). +func (r *Registry) Register(tools ...tool.Tool) { + for _, t := range tools { + name := t.Name() + if _, exists := r.tools[name]; !exists { + r.order = append(r.order, name) + } + r.tools[name] = t + } +} + +// Get returns a tool by name. Returns false if not found. +func (r *Registry) Get(name string) (tool.Tool, bool) { + t, ok := r.tools[name] + return t, ok +} + +// Definitions returns all registered tools as Definitions, +// in the order they were first registered. +func (r *Registry) Definitions() []tool.Definition { + defs := make([]tool.Definition, 0, len(r.order)) + for _, name := range r.order { + t := r.tools[name] + defs = append(defs, tool.Definition{ + Name: t.Name(), + Description: t.Description(), + Schema: t.Schema(), + }) + } + return defs +} diff --git a/tool/tool.go b/tool/tool.go new file mode 100644 index 0000000..9647bc2 --- /dev/null +++ b/tool/tool.go @@ -0,0 +1,63 @@ +package tool + +import ( + "context" + "encoding/json" + + "github.com/invopop/jsonschema" +) + +// Tool defines a callable tool that can be invoked by an agent. +type Tool interface { + Name() string + Description() string + Schema() Schema + Invoke(ctx context.Context, args json.RawMessage) (string, error) +} + +// Schema describes the JSON Schema for a tool's parameters. +type Schema struct { + Parameters json.RawMessage `json:"parameters"` +} + +// Definition is the wire format sent to providers so they know which tools are available. +type Definition struct { + Name string `json:"name"` + Description string `json:"description"` + Schema Schema `json:"schema"` +} + +// funcTool is the unexported implementation backing the Func helper. +type funcTool[T any] struct { + name string + description string + schema Schema + fn func(ctx context.Context, input T) (string, error) +} + +// Func creates a Tool from a typed function. The JSON Schema for parameters +// is derived from T using invopop/jsonschema at construction time. +func Func[T any](name, description string, fn func(ctx context.Context, input T) (string, error)) Tool { + r := new(jsonschema.Reflector) + s := r.Reflect(new(T)) + params, _ := json.Marshal(s) + + return &funcTool[T]{ + name: name, + description: description, + schema: Schema{Parameters: params}, + fn: fn, + } +} + +func (t *funcTool[T]) Name() string { return t.name } +func (t *funcTool[T]) Description() string { return t.description } +func (t *funcTool[T]) Schema() Schema { return t.schema } + +func (t *funcTool[T]) Invoke(ctx context.Context, args json.RawMessage) (string, error) { + var input T + if err := json.Unmarshal(args, &input); err != nil { + return "", err + } + return t.fn(ctx, input) +} diff --git a/types.go b/types.go index 080ad97..6bb671b 100644 --- a/types.go +++ b/types.go @@ -1,71 +1,44 @@ package forge -import "encoding/json" +import ( + "github.com/katasec/forge/message" + "github.com/katasec/forge/provider" + "github.com/katasec/forge/tool" +) -// Role identifies the sender of a message in a conversation. -type Role string +type Role = message.Role const ( - RoleUser Role = "user" - RoleAssistant Role = "assistant" - RoleTool Role = "tool" - RoleSystem Role = "system" + RoleUser = message.RoleUser + RoleAssistant = message.RoleAssistant + RoleTool = message.RoleTool + RoleSystem = message.RoleSystem ) -// Message represents a single message in a conversation. -type Message struct { - ID string `json:"id"` - Role Role `json:"role"` - Content string `json:"content"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` - ToolResults []ToolResult `json:"tool_results,omitempty"` -} +type Message = message.Message -// UserMessage creates a user-role message with the given content. func UserMessage(content string) Message { - return Message{Role: RoleUser, Content: content} -} - -// ToolCall represents a request from the LLM to invoke a tool. -type ToolCall struct { - ID string `json:"id"` - Name string `json:"name"` - Arguments json.RawMessage `json:"arguments"` + return message.UserMessage(content) } -// ToolResult represents the outcome of a tool invocation. -type ToolResult struct { - CallID string `json:"call_id"` - Content string `json:"content"` - IsError bool `json:"is_error"` -} +type ToolCall = tool.Call +type ToolResult = tool.Result +type ToolError = tool.Error -// FinishReason indicates why the agent loop terminated. -type FinishReason string +type FinishReason = provider.FinishReason const ( - FinishReasonStop FinishReason = "stop" - FinishReasonToolUse FinishReason = "tool_use" - FinishReasonIterLimit FinishReason = "iter_limit" - FinishReasonError FinishReason = "error" + FinishReasonStop = provider.FinishReasonStop + FinishReasonToolUse = provider.FinishReasonToolUse + FinishReasonIterLimit = provider.FinishReasonIterLimit + FinishReasonError = provider.FinishReasonError ) -// TokenUsage tracks token consumption across provider calls. -type TokenUsage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` -} +type TokenUsage = provider.TokenUsage -// ErrorPolicy controls agent behavior when a tool invocation fails. type ErrorPolicy string const ( ErrorPolicyStop ErrorPolicy = "stop" ErrorPolicyContinue ErrorPolicy = "continue" ) - -// ToolError wraps a tool invocation failure. -type ToolError struct { - CallID string `json:"call_id"` - Message string `json:"message"` -}