Skip to content

Instantly share code, notes, and snippets.

@paveljee

paveljee/sgr.ts

Forked from antonkuzmin/sgr.ts
Created August 15, 2025 09:01
Show Gist options

  • Select an option

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

Select an option

Learn more about clone URLs
Save paveljee/adb7d11f300be078047b87ecae82db1f to your computer and use it in GitHub Desktop.
Download ZIP
Raw
sgr.ts
/* package.json:
{
"name": "sgr_demo",
"version": "1.0.0",
"description": "TS-port of https://abdullin.com/schema-guided-reasoning/demo",
"scripts": {
"start": "bun src/sgr.ts"
},
"keywords": [],
"author": "[email protected]",
"license": "ISC",
"dependencies": {
"dotenv": "^17.2.1",
"openai": "^5.12.2",
"zod": "^3.25.76"
},
"devDependencies": {
"@types/node": "^22.13.13"
}
}
*/
/* .env:
OPENROUTER_KEY=xxxxx
*/
/* src/sgr.ts:
*Typescript port*
This Python code demonstrates Schema-Guided Reasoning (SGR) with OpenAI. It:
- 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
*/
// # 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 Invoice = {
id: string,
email: string,
file: string,
skus: string[],
discount_amount: number,
discount_percent: number,
total: number,
void: boolean,
}
type Product = { name: string, price: number }
type DBType = {
rules: Array<{
email: string,
rule: string,
}>
invoices: Record<string, Invoice>
emails: Array<{ to: string, subject: string, message: string }>
products: Record<string, Product>
}
const DB: DBType = {
rules: [],
invoices: {},
emails: [],
products: {
"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.
// from typing import List, Union, Literal, Annotated
// from annotated_types import MaxLen, Le, MinLen
// from pydantic import BaseModel, Field
import z from 'zod';
// # Tool: Sends an email with subject, message, attachments to a recipient
const SendEmail = z.object({
tool: z.literal("send_email"),
subject: z.string(),
message: z.string(),
files: z.array(z.string()),
recipient_email: z.string(),
});
type SendEmailType = z.infer<typeof SendEmail>;
// # Tool: Retrieves customer data such as rules, invoices, and emails from the database
const GetCustomerData = z.object({
tool: z.literal("get_customer_data"),
email: z.string(),
});
type GetCustomerDataType = z.infer<typeof GetCustomerData>;
// # Tool: Issues an invoice to a customer, allowing up to a 50% discount
const IssueInvoice = z.object({
tool: z.literal("issue_invoice"),
email: z.string(),
skus: z.array(z.string()),
discount_percent: z.number().int().positive().nonnegative().max(50),
});
type IssueInvoiceType = z.infer<typeof IssueInvoice>;
// # Tool: Cancels (voids) an existing invoice and records the reason
const VoidInvoice = z.object({
tool: z.literal("void_invoice"),
invoice_id: z.string(),
reason: z.string(),
});
type VoidInvoiceType = z.infer<typeof VoidInvoice>;
// # Tool: Saves a custom rule for interacting with a specific customer
const CreateRule = z.object({
tool: z.literal("remember"),
email: z.string(),
rule: z.string(),
});
type CreateRuleType = z.infer<typeof CreateRule>;
type AnyCommand =
| SendEmailType
| GetCustomerDataType
| IssueInvoiceType
| VoidInvoiceType
| CreateRuleType
// # 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.
const dispatch = (cmd: AnyCommand) => {
// # 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
if (cmd.tool === 'send_email') { //isinstance(cmd, SendEmail):
const email = {
"to": cmd.recipient_email,
"subject": cmd.subject,
"message": cmd.message,
}
DB.emails.push(email);
return email;
}
// # likewize rule creation just stores rule associated with customer
if (cmd.tool === 'remember') {
const rule = {
"email": cmd.email,
"rule": cmd.rule,
}
DB.rules.push(rule);
return rule;
}
// # customer data reading - doesn't change anything. It queries DB for all
// # records associated with the customer
if (cmd.tool === 'get_customer_data') {
const addr = cmd.email;
return {
"rules": DB.rules.filter(x => x.email === addr),
"invoices": Object.values(DB.invoices).filter(x => x.email === addr),
"emails": DB.emails.filter(x => x.to === addr),
}
}
// # 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
if (cmd.tool === 'issue_invoice') {
const skus = cmd.skus.map(id => DB.products[id]).filter(Boolean);
const total = skus.reduce((acc, x) => acc += x.price, 0);
if (skus.length < cmd.skus.length) {
return "Product {sku} not found";
}
const discount = +(total * 1.0 * cmd.discount_percent / 100.0).toFixed(2);
const invoice_id = `INV-${Object.keys(DB.invoices).length + 1}`;
const invoice: Invoice = {
"id": invoice_id,
"email": cmd.email,
"file": `/invoices/${invoice_id}.pdf`,
"skus": cmd.skus,
"discount_amount": discount,
"discount_percent": cmd.discount_percent,
"total": total,
"void": false,
}
DB.invoices[invoice_id] = invoice;
return invoice;
}
// # invoice cancellation marks a specific invoice as void
if (cmd.tool === 'void_invoice') {
const invoice = DB.invoices[cmd.invoice_id];
if (!invoice) {
return "Invoice {cmd.invoice_id} not found";
}
invoice.void = true;
return invoice;
}
}
// # Now, having such DB and tools, we could come up with a list of tasks
// # that we can carry out sequentially
const TASKS = [
// # 1. this one should create a new rule for sama
"Rule: address [email protected] as 'The SAMA', always give him 5% discount.",
// # 2. this should create a rule for elon
"Rule for [email protected]: Email his invoices to [email protected]",
// # 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 one of each product. Email him the invoice",
"[email protected] wants one of each product. Email him the 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]
"[email protected] wants 2x of what [email protected] got. Send invoice",
// # 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 elon @x.com but to finance @x.com
"redo last [email protected] invoice: use 3x discount of [email protected]",
]
// # 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.
const ReportTaskCompletion = z.object({
tool: z.literal("report_completion"),
completed_steps_laconic: z.array(z.string()),
code: z.enum(['completed', 'failed']),
});
// # now we have all sub - schemas in place, let's define SGR schema for the agent
const NextStep = z.object({
current_state: z.string(),
plan_remaining_steps_brief: z.array(z.string()).min(1).max(5),
task_completed: z.boolean(),
function: z.union([
ReportTaskCompletion,
SendEmail,
GetCustomerData,
IssueInvoice,
VoidInvoice,
CreateRule,
], { description: "execute first remaining step" }),
});
// # 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
const system_prompt = `
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: ${JSON.stringify(DB.products, null, '\t')}
`;
// # now we just need to implement the method to bring that all together
// # we will use rich for pretty printing in console
require('dotenv').config();
import OpenAI from 'OpenAI';
import { zodResponseFormat } from 'openai/helpers/zod';
import { ChatCompletionMessageParam } from 'OpenAI/resources/index.mjs';
const openai = new OpenAI({
baseURL: "https://openrouter.ai/api/v1",
apiKey: process.env.OPENROUTER_KEY,
});
const print = console.log;
// # Runs each defined task sequentially.The AI agent uses reasoning to determine
// # what steps are required to complete each task, executing tools as needed.
const execute_tasks = async () => {
// # we'll execute all tasks sequentially. You can add your tasks
// # of prompt user to write their own
for (let task of TASKS) {
print("\n\n");
print(`Launch agent with task: ${task}`);
// # log will contain conversation context for the agent within task
const log: ChatCompletionMessageParam[] = [
{ "role": "system", "content": system_prompt },
{ "role": "user", "content": task },
];
// # let's limit number of reasoning steps by 20, just to be safe
for (let i = 0; i < 20; i++) {
const step = `step_${i + 1}`;
print(`Planning ${step}...`);
// # This sample relies on OpenAI API. We specifically use 4o, since
// # GPT-5 has bugs with constrained decoding as of August 14, 2025
const completion = await openai.chat.completions.parse({
model: "gpt-4o",
messages: log,
max_completion_tokens: 10000,
response_format: zodResponseFormat(NextStep, "instruction"),
});
const job = completion.choices[0].message.parsed;
if (!job) {
print('ERROR: job cannot be parsed, skipping');
continue;
}
// # if SGR decided to finish, let's complete the task and quit this loop
if (job.function.tool === 'report_completion') {
print(`agent ${job.function.code}`);
print("### Summary ###");
job.function.completed_steps_laconic.forEach(s => print(' - ' + s));
print("### Database dump ###");
print(JSON.stringify(DB, null, '\t'));
print('###################################################\n\n');
break;
}
// # let's be nice and print the next remaining step (discard all others)
print(job.plan_remaining_steps_brief[0], `\n ${JSON.stringify(job.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.push({
"role": "assistant",
"content": job.plan_remaining_steps_brief[0],
"tool_calls": [{
"type": "function",
"id": step,
"function": {
"name": job.function.tool,
"arguments": JSON.stringify(job.function),
}
}]
});
// # now execute the tool by dispatching command to our handler
const result = dispatch(job.function)
const txt = typeof result === 'string' ? result : JSON.stringify(result);
print("OUTPUT", result);
// # 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.push({ "role": "tool", "content": txt, "tool_call_id": step });
}
}
}
execute_tasks();
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment