Skip to content

Commit

Permalink
Create new AI Toolbar (#690)
Browse files Browse the repository at this point in the history
* Add new AI Toolbar and endpoints

* Pass model into processResult; update ai packages; repairText on result

* Add translations

* Track session activity on AI creation

* Speed up writer text

* Fix chart input E2E
  • Loading branch information
rob-gordon authored Jul 22, 2024
1 parent e407745 commit 015dea7
Show file tree
Hide file tree
Showing 41 changed files with 1,135 additions and 556 deletions.
4 changes: 2 additions & 2 deletions api/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,14 @@
"author": "",
"license": "ISC",
"dependencies": {
"@ai-sdk/openai": "^0.0.9",
"@ai-sdk/openai": "^0.0.37",
"@notionhq/client": "^0.4.13",
"@octokit/core": "^4.2.0",
"@sendgrid/mail": "^7.4.6",
"@supabase/supabase-js": "^2.31.0",
"@upstash/ratelimit": "^1.1.3",
"@vercel/kv": "^1.0.1",
"ai": "^3.2.19",
"ai": "^3.2.32",
"ajv": "^8.12.0",
"axios": "^0.27.2",
"csv-parse": "^5.3.6",
Expand Down
42 changes: 42 additions & 0 deletions api/prompt/_edit.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
import { VercelApiHandler } from "@vercel/node";
import { llmMany } from "../_lib/_llm";
import { z } from "zod";

const nodeSchema = z.object({
// id: z.string(),
// classes: z.string(),
label: z.string(),
});

const edgeSchema = z.object({
from: z.string(),
to: z.string(),
label: z.string().optional().default(""),
});

const graphSchema = z.object({
nodes: z.array(nodeSchema),
edges: z.array(edgeSchema),
});

const handler: VercelApiHandler = async (req, res) => {
const { graph, prompt } = req.body;
if (!graph || !prompt) {
throw new Error("Missing graph or prompt");
}

const result = await llmMany(
`${prompt}
Here is the current state of the flowchart:
${JSON.stringify(graph, null, 2)}
`,
{
updateGraph: graphSchema,
}
);

res.json(result);
};

export default handler;
103 changes: 103 additions & 0 deletions api/prompt/_shared.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
import { z } from "zod";
import { streamText } from "ai";
import { stripe } from "../_lib/_stripe";
import { kv } from "@vercel/kv";
import { Ratelimit } from "@upstash/ratelimit";
import { openai } from "@ai-sdk/openai";

export const reqSchema = z.object({
prompt: z.string().min(1),
document: z.string(),
});

export async function handleRateLimit(req: Request) {
const ip = getIp(req);
let isPro = false,
customerId: null | string = null;

const token = req.headers.get("Authorization");

if (token) {
const sid = token.split(" ")[1];
const sub = await stripe.subscriptions.retrieve(sid);
if (sub.status === "active" || sub.status === "trialing") {
isPro = true;
customerId = sub.customer as string;
}
}

const ratelimit = new Ratelimit({
redis: kv,
limiter: isPro
? Ratelimit.slidingWindow(3, "1m")
: Ratelimit.fixedWindow(3, "30d"),
});

const rateLimitKey = isPro ? `pro_${customerId}` : `unauth_${ip}`;
const { success, limit, reset, remaining } = await ratelimit.limit(
rateLimitKey
);

if (!success) {
return new Response("You have reached your request limit.", {
status: 429,
headers: {
"X-RateLimit-Limit": limit.toString(),
"X-RateLimit-Remaining": remaining.toString(),
"X-RateLimit-Reset": reset.toString(),
},
});
}

return null;
}

export async function processRequest(
req: Request,
systemMessage: string,
content: string,
model: Parameters<typeof openai.chat>[0] = "gpt-4-turbo"
) {
const rateLimitResponse = await handleRateLimit(req);
if (rateLimitResponse) return rateLimitResponse;

const result = await streamText({
model: openai.chat(model),
system: systemMessage,
temperature: 1,
messages: [
{
role: "user",
content,
},
],
});

return result.toTextStreamResponse();
}

function getIp(req: Request) {
return (
req.headers.get("x-real-ip") ||
req.headers.get("cf-connecting-ip") ||
req.headers.get("x-forwarded-for") ||
req.headers.get("x-client-ip") ||
req.headers.get("x-cluster-client-ip") ||
req.headers.get("forwarded-for") ||
req.headers.get("forwarded") ||
req.headers.get("via") ||
req.headers.get("x-forwarded") ||
req.headers.get
);
}

export const systemMessageStyle = `You can style nodes using classes at the end of a node. Available styles include:
- Colors: .color_blue, .color_red, .color_green, .color_yellow, .color_orange
- Shapes: .shape_circle, .shape_diamond, .shape_ellipse, .shape_right-rhomboid`;

export const systemMessageExample = `Node A
Node B .shape_circle
\\(Secret Node)
Node C
label from c to d: Node D .color_green.shape_diamond
label from d to a: (Node A)`;
102 changes: 11 additions & 91 deletions api/prompt/convert.ts
Original file line number Diff line number Diff line change
@@ -1,97 +1,32 @@
import { z } from "zod";
import { streamText } from "ai";
import { stripe } from "../_lib/_stripe";
import { kv } from "@vercel/kv";
import { Ratelimit } from "@upstash/ratelimit";
import { openai } from "@ai-sdk/openai";
import { processRequest, reqSchema } from "./_shared";

export const config = {
runtime: "edge",
};

const reqSchema = z.object({
prompt: z.string().min(1),
});
const systemMessage = `You are the Flowchart Fun creation assistant. When I give you a document respond with a diagram in Flowchart Fun syntax. The Flowchart Fun syntax you use indentation to express a tree shaped graph. You use text before a colon to labels to edges. You link back to earlier nodes by referring to their label in parentheses. The following characters must be escaped when used in a node or edge label: (,:,#, and .\n\nYou can style nodes using classes at the end of a node. Available styles include:
- Colors: .color_blue, .color_red, .color_green, .color_yellow
- Shapes: .shape_circle, .shape_diamond, .shape_hexagon
const systemMessage = `You are the Flowchart Fun creation assistant. When I give you a document respond with a diagram in Flowchart Fun syntax. The Flowchart Fun syntax you use indentation to express a tree shaped graph. You use text before a colon to labels to edges. You link back to earlier nodes by referring to their label in parentheses. The following characters must be escaped when used in a node or edge label: (,:,#, and .\n\nHere is a very simple graph illustrating the syntax:
Here is a very simple graph illustrating the syntax:
Node A
Node B
Node A .color_blue
Node B .shape_circle
\\(Secret Node)
Node C
label from c to d: Node D
Node C .color_green
label from c to d: Node D .shape_diamond
label from d to a: (Node A)
Note: Don't provide any explanation. Don't wrap your response in a code block.`;
export default async function handler(req: Request) {
const ip = getIp(req);

let isPro = false,
customerId: null | string = null;

// Check for auth token
const token = req.headers.get("Authorization");

if (token) {
// get sid from token
const sid = token.split(" ")[1];

// check if subscription is active or trialing
const sub = await stripe.subscriptions.retrieve(sid);
if (sub.status === "active" || sub.status === "trialing") {
isPro = true;
customerId = sub.customer as string;
}
}

// Implement rate-limiting based on IP for unauthorized users and customerId for authorized users
// Initialize Upstash Ratelimit
const ratelimit = new Ratelimit({
redis: kv,
limiter: isPro
? Ratelimit.slidingWindow(3, "1m") // Pro users: 3 requests per minute
: Ratelimit.fixedWindow(2, "30d"), // Unauthenticated users: 2 requests per month
});

// Determine the key for rate limiting
const rateLimitKey = isPro ? `pro_${customerId}` : `unauth_${ip}`;

// Check the rate limit
const { success, limit, reset, remaining } = await ratelimit.limit(
rateLimitKey
);

if (!success) {
return new Response("You have reached your request limit.", {
status: 429,
headers: {
"X-RateLimit-Limit": limit.toString(),
"X-RateLimit-Remaining": remaining.toString(),
"X-RateLimit-Reset": reset.toString(),
},
});
}

export default async function handler(req: Request) {
const body = await req.json();
const parsed = reqSchema.safeParse(body);

if (!parsed.success) {
return new Response(JSON.stringify(parsed.error), { status: 400 });
}

const result = await streamText({
model: openai.chat("gpt-4-turbo"),
system: systemMessage,
temperature: 0.15,
messages: [
{
role: "user",
content: getContent(parsed.data.prompt),
},
],
});

return result.toTextStreamResponse();
return processRequest(req, systemMessage, getContent(parsed.data.prompt));
}

function getContent(prompt: string): string {
Expand All @@ -100,18 +35,3 @@ function getContent(prompt: string): string {
prompt
);
}

function getIp(req: Request) {
return (
req.headers.get("x-real-ip") ||
req.headers.get("cf-connecting-ip") ||
req.headers.get("x-forwarded-for") ||
req.headers.get("x-client-ip") ||
req.headers.get("x-cluster-client-ip") ||
req.headers.get("forwarded-for") ||
req.headers.get("forwarded") ||
req.headers.get("via") ||
req.headers.get("x-forwarded") ||
req.headers.get
);
}
78 changes: 40 additions & 38 deletions api/prompt/edit.ts
Original file line number Diff line number Diff line change
@@ -1,42 +1,44 @@
import { VercelApiHandler } from "@vercel/node";
import { llmMany } from "../_lib/_llm";
import { z } from "zod";

const nodeSchema = z.object({
// id: z.string(),
// classes: z.string(),
label: z.string(),
});

const edgeSchema = z.object({
from: z.string(),
to: z.string(),
label: z.string().optional().default(""),
});

const graphSchema = z.object({
nodes: z.array(nodeSchema),
edges: z.array(edgeSchema),
});

const handler: VercelApiHandler = async (req, res) => {
const { graph, prompt } = req.body;
if (!graph || !prompt) {
throw new Error("Missing graph or prompt");
import { processRequest, reqSchema } from "./_shared";

export const config = {
runtime: "edge",
};

const systemMessage = `You are an AI document editor specializing in Flowchart Fun syntax. When given a document and editing instructions, return the same document with only the requested changes. Do not make any additional changes, including whitespace changes, beyond what is explicitly requested. Preserve all original formatting and content except where modifications are necessary.
Flowchart Fun Syntax:
- Use indentation to express a tree-shaped graph.
- Text before a colon represents labels for edges.
- Link back to earlier nodes by referring to their label in parentheses.
- Escape the following characters when used in a node or edge label: (,:,#, and \\.
- Use classes at the end of a node to apply styles. (e.g., .color_blue,.shape_circle)
Example:
Node A
Node B .color_blue
\\(Secret Node)
Node C
label from c to d: Node D
label from d to a: (Node A)
When editing, ensure that the Flowchart Fun syntax remains valid and consistent.`;

export default async function handler(req: Request) {
const body = await req.json();
const parsed = reqSchema.safeParse(body);

if (!parsed.success) {
return new Response(JSON.stringify(parsed.error), { status: 400 });
}

const result = await llmMany(
`${prompt}
Here is the current state of the flowchart:
${JSON.stringify(graph, null, 2)}
`,
{
updateGraph: graphSchema,
}
return processRequest(
req,
systemMessage,
getContent(parsed.data.prompt, parsed.data.document),
"gpt-4-turbo-2024-04-09"
);
}

res.json(result);
};

export default handler;
function getContent(prompt: string, document: string): string {
return `Edit the following document according to these instructions:\n\nInstructions: ${prompt}\n\nDocument:\n${document}`;
}
Loading

0 comments on commit 015dea7

Please sign in to comment.