Skip to content

Instantly share code, notes, and snippets.

@m0n0x41d

m0n0x41d/schema-guided-reasoning.go

Last active November 17, 2025 20:41
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save m0n0x41d/ae1d1b8003f0094e2f504fd8e53e8f5f to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save m0n0x41d/ae1d1b8003f0094e2f504fd8e53e8f5f to your computer and use it in GitHub Desktop.
Download ZIP
schema-guided-reasoning.go
Raw
schema-guided-reasoning.go
/*
This Go code demonstrates Schema-Guided Reasoning (SGR) with OpenAI.
This is a port of the original Python example by Rinat Abdullin, which you can find here:
https://gist.github.com/abdullin/46caec6cba361b9e8e8a00b2c48ee07c
This code snippet:
- implements a business agent capable of planning and reasoning
- implements tool calling using only SGR and simple dispatch
- uses with a simple (inexpensive) non-reasoning model for that
To give this agent something to work with, we ask it to help with running
a small business - selling courses to help to achieve AGI faster.
Once this script starts, it will emulate in-memory CRM with invoices,
emails, products and rules. Then it will execute sequentially a set of
tasks (see TASKS below). In order to carry them out, Agent will have to use
tools to issue invoices, create rules, send emails, and a few others.
Read more about SGR: http://abdullin.com/schema-guided-reasoning/
This demo is described in more detail here: https://abdullin.com/schema-guided-reasoning/demo
*/
package main
import (
"context"
"encoding/json"
"fmt"
"os"
"strings"
"github.com/invopop/jsonschema"
"github.com/openai/openai-go"
"github.com/openai/openai-go/option"
)
// Let's start by implementing our customer management system. For the sake of
// simplicity it will live in memory and have a very simple DB structure
type Product struct {
Name string `json:"name"`
Price float64 `json:"price"`
}
type Invoice struct {
ID string `json:"id"`
Email string `json:"email"`
File string `json:"file"`
SKUs []string `json:"skus"`
DiscountAmount float64 `json:"discount_amount"`
DiscountPercent int `json:"discount_percent"`
Total float64 `json:"total"`
Void bool `json:"void"`
}
type Email struct {
To string `json:"to"`
Subject string `json:"subject"`
Message string `json:"message"`
}
type Rule struct {
Email string `json:"email"`
Rule string `json:"rule"`
}
type Database struct {
Rules []Rule `json:"rules"`
Invoices map[string]*Invoice `json:"invoices"`
Emails []Email `json:"emails"`
Products map[string]Product `json:"products"`
}
var DB = &Database{
Rules: []Rule{},
Invoices: make(map[string]*Invoice),
Emails: []Email{},
Products: map[string]Product{
"SKU-205": {Name: "AGI 101 Course Personal", Price: 258},
"SKU-210": {Name: "AGI 101 Course Team (5 seats)", Price: 1290},
"SKU-220": {Name: "Building AGI - online exercises", Price: 315},
},
}
// Now, let's define a few tools which could be used by LLM to do something
// useful with this customer management system. We need tools to issue invoices,
// send emails, create rules and memorize new rules. Maybe a tool to cancel invoices.
// Tool: Sends an email with subject, message, attachments to a recipient
type SendEmail struct {
Tool string `json:"tool" jsonschema:"enum=send_email"`
Subject string `json:"subject" jsonschema_description:"Email subject"`
Message string `json:"message" jsonschema_description:"Email body"`
Files []string `json:"files" jsonschema_description:"Attachments"`
RecipientEmail string `json:"recipient_email" jsonschema_description:"Recipient email address"`
}
// Tool: Retrieves customer data such as rules, invoices, and emails from the database
type GetCustomerData struct {
Tool string `json:"tool" jsonschema:"enum=get_customer_data"`
Email string `json:"email" jsonschema_description:"Customer email address"`
}
// Tool: Issues an invoice to a customer, allowing up to a 50% discount
type IssueInvoice struct {
Tool string `json:"tool" jsonschema:"enum=issue_invoice"`
Email string `json:"email" jsonschema_description:"Customer email address"`
SKUs []string `json:"skus" jsonschema_description:"List of product SKUs"`
DiscountPercent int `json:"discount_percent" jsonschema:"minimum=0,maximum=50" jsonschema_description:"Discount percentage (max 50%)"` // never more than 50% discount
}
// Tool: Cancels (voids) an existing invoice and records the reason
type VoidInvoice struct {
Tool string `json:"tool" jsonschema:"enum=void_invoice"`
InvoiceID string `json:"invoice_id" jsonschema_description:"Invoice ID to void"`
Reason string `json:"reason" jsonschema_description:"Reason for voiding"`
}
// Tool: Saves a custom rule for interacting with a specific customer
type CreateRule struct {
Tool string `json:"tool" jsonschema:"enum=remember"`
Email string `json:"email" jsonschema_description:"Customer email address"`
Rule string `json:"rule" jsonschema_description:"Rule to remember"`
}
// let's define one more special command. LLM can use it whenever
// it thinks that its task is completed. It will report results with that.
type ReportTaskCompletion struct {
Tool string `json:"tool" jsonschema:"enum=report_completion"`
CompletedStepsLaconic []string `json:"completed_steps_laconic" jsonschema_description:"Summary of completed steps"`
Code string `json:"code" jsonschema:"enum=completed,enum=failed" jsonschema_description:"Completion status"`
}
// now we have all sub-schemas in place, let's define SGR schema for the agent
type NextStep struct {
CurrentState string `json:"current_state" jsonschema_description:"Current state of the task"` // we'll give some thinking space here
PlanRemainingStepsBrief []string `json:"plan_remaining_steps_brief" jsonschema:"minItems=1,maxItems=5" jsonschema_description:"List of remaining steps (1-5 items)"` // Cycle to think about what remains to be done. at least 1 at most 5 steps - we'll use only the first step, discarding all the rest.
TaskCompleted bool `json:"task_completed" jsonschema_description:"Whether the task is completed"` // now let's continue the cascade and check with LLM if the task is done
Function json.RawMessage `json:"function" jsonschema_description:"Function to execute for the first remaining step"` // Routing to one of the tools to execute the first remaining step - if task is completed, model will pick ReportTaskCompletion
}
// Generate JSON schema for a type
func GenerateSchema[T any]() any {
reflector := jsonschema.Reflector{
AllowAdditionalProperties: false,
DoNotReference: true,
}
var value T
schema := reflector.Reflect(value)
schemaBytes, _ := json.Marshal(schema)
var schemaMap map[string]any
json.Unmarshal(schemaBytes, &schemaMap)
sendEmailSchema := reflector.Reflect(SendEmail{})
getCustomerDataSchema := reflector.Reflect(GetCustomerData{})
issueInvoiceSchema := reflector.Reflect(IssueInvoice{})
voidInvoiceSchema := reflector.Reflect(VoidInvoice{})
createRuleSchema := reflector.Reflect(CreateRule{})
reportCompletionSchema := reflector.Reflect(ReportTaskCompletion{})
toMap := func(s *jsonschema.Schema) map[string]any {
b, _ := json.Marshal(s)
var m map[string]any
json.Unmarshal(b, &m)
return m
}
anyOf := []any{
toMap(reportCompletionSchema),
toMap(sendEmailSchema),
toMap(getCustomerDataSchema),
toMap(issueInvoiceSchema),
toMap(voidInvoiceSchema),
toMap(createRuleSchema),
}
if props, ok := schemaMap["properties"].(map[string]any); ok {
props["function"] = map[string]any{
"anyOf": anyOf,
}
}
return schemaMap
}
// This function handles executing commands issued by the agent. It simulates
// operations like sending emails, managing invoices, and updating customer
// rules within the in-memory database.
func dispatch(functionJSON json.RawMessage) (any, string, error) {
var toolType struct {
Tool string `json:"tool"`
}
if err := json.Unmarshal(functionJSON, &toolType); err != nil {
return nil, "", err
}
switch toolType.Tool {
case "send_email":
// here is how we can simulate email sending
// just append to the DB (for future reading), return composed email
// and pretend that we sent something
var cmd SendEmail
if err := json.Unmarshal(functionJSON, &cmd); err != nil {
return nil, "", err
}
email := Email{
To: cmd.RecipientEmail,
Subject: cmd.Subject,
Message: cmd.Message,
}
DB.Emails = append(DB.Emails, email)
return email, "send_email", nil
case "remember":
// likewize rule creation just stores rule associated with customer
var cmd CreateRule
if err := json.Unmarshal(functionJSON, &cmd); err != nil {
return nil, "", err
}
rule := Rule{
Email: cmd.Email,
Rule: cmd.Rule,
}
DB.Rules = append(DB.Rules, rule)
return rule, "remember", nil
case "get_customer_data":
// customer data reading - doesn't change anything. It queries DB for all
// records associated with the customer
var cmd GetCustomerData
if err := json.Unmarshal(functionJSON, &cmd); err != nil {
return nil, "", err
}
addr := cmd.Email
result := map[string]any{
"rules": []Rule{},
"invoices": [][]any{},
"emails": []Email{},
}
for _, r := range DB.Rules {
if r.Email == addr {
result["rules"] = append(result["rules"].([]Rule), r)
}
}
for id, inv := range DB.Invoices {
if inv.Email == addr {
result["invoices"] = append(result["invoices"].([][]any), []any{id, inv})
}
}
for _, e := range DB.Emails {
if e.To == addr {
result["emails"] = append(result["emails"].([]Email), e)
}
}
return result, "get_customer_data", nil
case "issue_invoice":
// invoice generation is going to be more tricky
// it will demonstrate discount calculation (we know that LLMs shouldn't be trusted
// with math. It also shows how to report problems back to LLM.
// ultimately, it computes a new invoice number and stores it in the DB
var cmd IssueInvoice
if err := json.Unmarshal(functionJSON, &cmd); err != nil {
return nil, "", err
}
total := 0.0
for _, sku := range cmd.SKUs {
product, ok := DB.Products[sku]
if !ok {
return fmt.Sprintf("Product %s not found", sku), "issue_invoice", nil
}
total += product.Price
}
discount := float64(cmd.DiscountPercent) * total / 100.0
invoiceID := fmt.Sprintf("INV-%d", len(DB.Invoices)+1)
invoice := &Invoice{
ID: invoiceID,
Email: cmd.Email,
File: "/invoices/" + invoiceID + ".pdf",
SKUs: cmd.SKUs,
DiscountAmount: discount,
DiscountPercent: cmd.DiscountPercent,
Total: total,
Void: false,
}
DB.Invoices[invoiceID] = invoice
return invoice, "issue_invoice", nil
case "void_invoice":
// invoice cancellation marks a specific invoice as void
var cmd VoidInvoice
if err := json.Unmarshal(functionJSON, &cmd); err != nil {
return nil, "", err
}
invoice, ok := DB.Invoices[cmd.InvoiceID]
if !ok {
return fmt.Sprintf("Invoice %s not found", cmd.InvoiceID), "void_invoice", nil
}
invoice.Void = true
return invoice, "void_invoice", nil
case "report_completion":
var cmd ReportTaskCompletion
if err := json.Unmarshal(functionJSON, &cmd); err != nil {
return nil, "", err
}
return &cmd, "report_completion", nil
default:
return nil, "", fmt.Errorf("unknown tool type: %s", toolType.Tool)
}
}
// Now, having such DB and tools, we could come up with a list of tasks
// that we can carry out sequentially
var TASKS = []string{
"Rule: address [email protected] as 'The SAMA', always give him 5% discount.",
// 1. this one should create a new rule for sama
"Rule for [email protected]: Email his invoices to [email protected]",
// 2. this should create a rule for elon
"[email protected] wants one of each product. Email him the invoice",
// 3. now, this task should create an invoice for sama that includes one of each product.
// But it should also remember to give discount and address him properly
"[email protected] wants 2x of what [email protected] got. Send invoice",
// 4. Even more tricky - we need to create the invoice for Musk based on the invoice of sama, but twice.
// Plus LLM needs to remeber to use the proper email address for invoices - [email protected]
"redo last [email protected] invoice: use 3x discount of [email protected]",
// 5. even more tricky. Need to cancel old invoice (we never told LLMs how) and issue the new invoice.
// BUT it should pull the discount from sama and triple it. Obviously the model should also remember
// to send invoice not to [email protected] but to [email protected]
"Add rule for [email protected] - politely reject all requests to buy SKU-220",
// let's demonstrate how the agent can change its plans after discovering new information - first we plant a new memory
"[email protected] and [email protected] wrote emails asking to buy 'Building AGI - online exercises', handle that",
// now let's give another task (agent will not have the memory above in the context UNTIL it is pulled from memory store)
}
// here is the prompt with some core context
// since the list of products is small, we can merge it with prompt
// In a bigger system, could add a tool to load things conditionally
// now we just need to implement the method to bring that all together
// we will use rich for pretty printing in console
const (
colorReset = "\033[0m"
colorCyan = "\033[36m"
colorYellow = "\033[33m"
colorBlue = "\033[34m"
colorGreen = "\033[32m"
colorGray = "\033[90m"
)
func printPanel(title, content string) {
fmt.Printf("\n%s=== %s ===%s\n%s\n", colorCyan, title, colorReset, content)
}
func printRule(title string) {
if title == "" {
fmt.Println(strings.Repeat("─", 60))
} else {
fmt.Printf("%s─── %s ───%s\n", colorGray, title, colorReset)
}
}
func prettyJSON(data any) string {
bytes, _ := json.MarshalIndent(data, "", " ")
return string(bytes)
}
// Runs each defined task sequentially. The AI agent uses reasoning to determine
// what steps are required to complete each task, executing tools as needed.
func main() {
apiKey := os.Getenv("OPENAI_API_KEY")
if apiKey == "" {
fmt.Println("Warning: OPENAI_API_KEY environment variable not set")
fmt.Println("Please set it with: export OPENAI_API_KEY=your-api-key")
os.Exit(1)
}
client := openai.NewClient(option.WithAPIKey(apiKey))
productsJSON, _ := json.Marshal(DB.Products)
systemPrompt := fmt.Sprintf(`You are a business assistant helping Rinat Abdullin with customer interactions.
- Clearly report when tasks are done.
- Always send customers emails after issuing invoices (with invoice attached).
- Be laconic. Especially in emails
- No need to wait for payment confirmation before proceeding.
- Always check customer data before issuing invoices or making changes.
Products: %s`, string(productsJSON))
var nextStepSchema = GenerateSchema[NextStep]()
schemaParam := openai.ResponseFormatJSONSchemaJSONSchemaParam{
Name: "next_step",
Description: openai.String("Agent reasoning and next action"),
Schema: nextStepSchema,
Strict: openai.Bool(true),
}
ctx := context.Background()
// we'll execute all tasks sequentially. You can add your tasks
// or prompt user to write their own
for _, task := range TASKS {
fmt.Print("\n\n")
printPanel("Launch agent with task", task)
// log will contain conversation context for the agent within task
log := []openai.ChatCompletionMessageParamUnion{
openai.SystemMessage(systemPrompt),
openai.UserMessage(task),
}
// let's limit number of reasoning steps by 20, just to be safe
for i := range 20 {
step := fmt.Sprintf("step_%d", i+1)
// This sample relies on OpenAI API. We specifically use 4o, since
// GPT-5 has bugs with constrained decoding as of August 14, 2025
completion, err := client.Chat.Completions.New(
ctx,
openai.ChatCompletionNewParams{
Messages: log,
ResponseFormat: openai.ChatCompletionNewParamsResponseFormatUnion{
OfJSONSchema: &openai.ResponseFormatJSONSchemaParam{
JSONSchema: schemaParam,
},
},
Model: "gpt-4o",
MaxCompletionTokens: openai.Int(10000),
},
)
if err != nil {
panic(err)
}
var nextStep NextStep
content := completion.Choices[0].Message.Content
err = json.Unmarshal([]byte(content), &nextStep)
if err != nil {
panic(err)
}
// now execute the tool by dispatching command to our handler
result, toolType, err := dispatch(nextStep.Function)
if err != nil {
panic(err)
}
// if SGR decided to finish, let's complete the task
// and quit this loop
if toolType == "report_completion" {
if completionResult, ok := result.(*ReportTaskCompletion); ok {
fmt.Printf("%sagent %s%s.\n", colorBlue, completionResult.Code, colorReset)
printRule("Summary")
for _, s := range completionResult.CompletedStepsLaconic {
fmt.Printf("%s- %s%s\n", colorGreen, s, colorReset)
}
printRule("")
break
}
}
// let's be nice and print the next remaining step (discard all others)
if len(nextStep.PlanRemainingStepsBrief) > 0 {
fmt.Printf("%sPlanning %s...%s %s\n", colorGray, step, colorReset, nextStep.PlanRemainingStepsBrief[0])
}
fmt.Printf(" %s\n", prettyJSON(nextStep.Function))
// Let's add tool request to conversation history as if OpenAI asked for it.
// a shorter way would be to just append `job.model_dump_json()` entirely
log = append(log, openai.AssistantMessage(content))
// Convert result to JSON string
var resultStr string
if str, ok := result.(string); ok {
resultStr = str
} else {
resultBytes, _ := json.Marshal(result)
resultStr = string(resultBytes)
}
// and now we add results back to the convesation history, so that agent
// we'll be able to act on the results in the next reasoning step.
log = append(log, openai.UserMessage(resultStr))
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment