-
-
Notifications
You must be signed in to change notification settings - Fork 229
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #692 from tone-row/dev
- Loading branch information
Showing
52 changed files
with
2,298 additions
and
845 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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)`; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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}`; | ||
} |
Oops, something went wrong.