Document playground

packages/compiler/src/compiler/chain.test.ts

@@ -1,11 +1,6 @@
 import { CHAIN_STEP_TAG } from '$compiler/constants'
 import CompileError from '$compiler/error/error'
-import {
-  AssistantMessage,
-  ContentType,
-  Conversation,
-  MessageRole,
-} from '$compiler/types'
+import { Conversation, MessageRole } from '$compiler/types'
 import { describe, expect, it, vi } from 'vitest'
 
 import { Chain } from './chain'
@@ -24,19 +19,8 @@ const getExpectedError = async <T>(
   throw new Error('Expected an error to be thrown')
 }
 
-const assistantMessage = (content?: string): AssistantMessage => ({
-  role: MessageRole.assistant,
-  content: [
-    {
-      type: ContentType.text,
-      value: content ?? '',
-    },
-  ],
-  toolCalls: [],
-})
-
-async function defaultCallback(): Promise<AssistantMessage> {
-  return assistantMessage('')
+async function defaultCallback(): Promise<string> {
+  return ''
 }
 
 async function complete({
@@ -45,11 +29,11 @@ async function complete({
   maxSteps = 50,
 }: {
   chain: Chain
-  callback?: (convo: Conversation) => Promise<AssistantMessage | undefined>
+  callback?: (convo: Conversation) => Promise<string>
   maxSteps?: number
 }): Promise<Conversation> {
   let steps = 0
-  let response: AssistantMessage | undefined = undefined
+  let response: string | undefined = undefined
   while (true) {
     const { completed, conversation } = await chain.step(response)
     if (completed) return conversation
@@ -152,7 +136,7 @@ describe('chain', async () => {
     expect(conversation1.messages[0]!.content[0]!.value).toBe('Message 1')
 
     const { completed: completed2, conversation: conversation2 } =
-      await chain.step(assistantMessage('response'))
+      await chain.step('response')
 
     expect(completed2).toBe(true)
     expect(conversation2.messages.length).toBe(3)
@@ -191,7 +175,7 @@ describe('chain', async () => {
       parameters: {},
     })
 
-    const action = () => chain.step(assistantMessage())
+    const action = () => chain.step('')
     const error = await getExpectedError(action, Error)
     expect(error.message).toBe(
       'A response is not allowed before the chain has started',
@@ -213,7 +197,7 @@ describe('chain', async () => {
 
     let { completed: stop } = await chain.step()
     while (!stop) {
-      const { completed } = await chain.step(assistantMessage())
+      const { completed } = await chain.step('')
       stop = completed
     }
 
@@ -432,18 +416,8 @@ describe('chain', async () => {
       parameters: {},
     })
 
-    const response = {
-      role: MessageRole.assistant,
-      content: [
-        {
-          type: ContentType.text,
-          value: 'foo',
-        },
-      ],
-    } as AssistantMessage
-
     await chain.step()
-    const { conversation } = await chain.step(response)
+    const { conversation } = await chain.step('foo')
 
     expect(conversation.messages.length).toBe(2)
     expect(conversation.messages[0]!.content[0]!.value).toBe('foo')
@@ -470,16 +444,15 @@ describe('chain', async () => {
     expect(step1.config.model).toBe('foo-1')
     expect(step1.config.temperature).toBe(0.5)
 
-    const { conversation: step2 } = await chain.step(assistantMessage())
+    const { conversation: step2 } = await chain.step('')
     expect(step2.config.model).toBe('foo-2')
     expect(step2.config.temperature).toBe(0.5)
 
-    const { conversation: step3 } = await chain.step(assistantMessage())
+    const { conversation: step3 } = await chain.step('')
     expect(step3.config.model).toBe('foo-1')
     expect(step3.config.temperature).toBe('1')
 
-    const { conversation: finalConversation } =
-      await chain.step(assistantMessage())
+    const { conversation: finalConversation } = await chain.step('')
     expect(finalConversation.config.model).toBe('foo-1')
     expect(finalConversation.config.temperature).toBe(0.5)
   })

packages/compiler/src/compiler/chain.ts

@@ -1,11 +1,6 @@
 import parse from '$compiler/parser'
 import { Fragment } from '$compiler/parser/interfaces'
-import {
-  AssistantMessage,
-  Config,
-  Conversation,
-  Message,
-} from '$compiler/types'
+import { Config, Conversation, Message } from '$compiler/types'
 
 import { Compile } from './compile'
 import Scope from './scope'
@@ -20,7 +15,7 @@ export class Chain {
   private ast: Fragment
   private scope: Scope
   private didStart: boolean = false
-  private completed: boolean = false
+  private _completed: boolean = false
 
   private messages: Message[] = []
   private config: Config | undefined
@@ -37,8 +32,8 @@ export class Chain {
     this.scope = new Scope(parameters)
   }
 
-  async step(response?: AssistantMessage): Promise<ChainStep> {
-    if (this.completed) {
+  async step(response?: string): Promise<ChainStep> {
+    if (this._completed) {
       throw new Error('The chain has already completed')
     }
     if (!this.didStart && response !== undefined) {
@@ -63,7 +58,7 @@ export class Chain {
     this.ast = ast
     this.messages.push(...messages)
     this.config = globalConfig ?? this.config
-    this.completed = completed || this.completed
+    this._completed = completed || this._completed
 
     const config = {
       ...this.config,
@@ -75,7 +70,11 @@ export class Chain {
         messages: this.messages,
         config,
       },
-      completed: this.completed,
+      completed: this._completed,
     }
   }
+
+  get completed(): boolean {
+    return this._completed
+  }
 }

packages/compiler/src/compiler/compile.ts

@@ -56,7 +56,7 @@ export class Compile {
 
   private messages: Message[] = []
   private config: Config | undefined
-  private stepResponse: AssistantMessage | undefined
+  private stepResponse: string | undefined
 
   private accumulatedText: string = ''
   private accumulatedContent: MessageContent[] = []
@@ -71,7 +71,7 @@ export class Compile {
     rawText: string
     globalScope: Scope
     ast: Fragment
-    stepResponse?: AssistantMessage
+    stepResponse?: string
   }) {
     this.rawText = rawText
     this.globalScope = globalScope
@@ -186,7 +186,16 @@ export class Compile {
   }
 
   private popStepResponse(): AssistantMessage | undefined {
-    const response = this.stepResponse
+    if (this.stepResponse === undefined) return undefined
+    const response = {
+      role: MessageRole.assistant,
+      content: [
+        {
+          type: ContentType.text,
+          value: this.stepResponse,
+        },
+      ],
+    } as AssistantMessage
     this.stepResponse = undefined
     return response
   }

packages/compiler/src/compiler/index.ts

@@ -19,6 +19,16 @@ export async function render({
   return conversation
 }
 
+export function createChain({
+  prompt,
+  parameters,
+}: {
+  prompt: string
+  parameters: Record<string, unknown>
+}): Chain {
+  return new Chain({ prompt, parameters })
+}
+
 export function readMetadata({
   prompt,
   referenceFn,
@@ -28,3 +38,5 @@ export function readMetadata({
 }): Promise<ConversationMetadata> {
   return new ReadMetadata({ prompt, referenceFn }).run()
 }
+
+export { Chain }

packages/compiler/src/compiler/readMetadata.test.ts

@@ -305,6 +305,21 @@ describe('parameters', async () => {
 
     expect(metadata.parameters).toEqual(new Set())
   })
+
+  it('adds the correct parameters to the scope context', async () => {
+    const prompt = `
+      ${CUSTOM_TAG_START}foo${CUSTOM_TAG_END}
+      ${CUSTOM_TAG_START}bar${CUSTOM_TAG_END}
+      ${CUSTOM_TAG_START}#each arr as val${CUSTOM_TAG_END}
+      ${CUSTOM_TAG_START}/each${CUSTOM_TAG_END}
+    `
+
+    const metadata = await readMetadata({
+      prompt: removeCommonIndent(prompt),
+    })
+
+    expect(metadata.parameters).toEqual(new Set(['foo', 'bar', 'arr']))
+  })
 })
 
 describe('referenced prompts', async () => {

packages/compiler/src/compiler/readMetadata.ts

@@ -23,7 +23,13 @@ import { ReferencePromptFn } from './compile'
 import { readConfig } from './config'
 import { updateScopeContextForNode } from './logic'
 import { ScopeContext } from './scope'
-import { isContentTag, isMessageTag, isRefTag, isToolCallTag } from './utils'
+import {
+  isChainStepTag,
+  isContentTag,
+  isMessageTag,
+  isRefTag,
+  isToolCallTag,
+} from './utils'
 
 function copyScopeContext(scopeContext: ScopeContext): ScopeContext {
   return {
@@ -223,6 +229,8 @@ export class ReadMetadata {
     }
 
     if (node.type === 'EachBlock') {
+      await this.updateScopeContext({ node: node.expression, scopeContext })
+
       const elseScope = copyScopeContext(scopeContext)
       for await (const childNode of node.else?.children ?? []) {
         await this.readBaseMetadata({
@@ -442,6 +450,25 @@ export class ReadMetadata {
         return
       }
 
+      if (isChainStepTag(node)) {
+        const attributes = await this.listTagAttributes({
+          tagNode: node,
+          scopeContext,
+          literalAttributes: ['as'],
+        })
+
+        if (attributes.has('as')) {
+          const asAttribute = node.attributes.find((a) => a.name === 'as')!
+          const asValue = (asAttribute.value as TemplateNode[])
+            .map((n) => n.value)
+            .join('')
+
+          scopeContext.definedVariables.add(asValue)
+        }
+
+        return
+      }
+
       this.baseNodeError(errors.unknownTag(node.name), node)
       return
     }

packages/web-ui/package.json

@@ -45,13 +45,15 @@
     "react": "18.3.0",
     "react-dom": "18.3.0",
     "react-resizable-panels": "^2.0.22",
+    "react-textarea-autosize": "^8.5.3",
     "tailwind-merge": "^2.4.0",
     "tailwindcss-animate": "^1.0.7",
     "use-debounce": "^10.0.1",
     "zod": "^3.23.8",
     "zustand": "^4.5.4"
   },
   "devDependencies": {
+    "@faker-js/faker": "^8.4.1",
     "@latitude-data/eslint-config": "workspace:*",
     "@latitude-data/typescript-config": "workspace:*",
     "@testing-library/dom": "^10.3.2",

packages/web-ui/src/ds/atoms/Badge/index.tsx

@@ -16,6 +16,8 @@ const badgeVariants = cva(
           'border-transparent bg-accent text-accent-foreground hover:bg-accent/80',
         destructive:
           'border-transparent bg-destructive text-destructive-foreground hover:bg-destructive/80',
+        muted:
+          'border-transparent bg-muted text-muted-foreground hover:bg-muted/80',
         outline: 'text-foreground',
       },
     },

packages/web-ui/src/ds/atoms/Button/index.tsx

@@ -2,8 +2,7 @@ import { ButtonHTMLAttributes, forwardRef, ReactNode } from 'react'
 import { cva, type VariantProps } from 'class-variance-authority'
 import { Slot, Slottable } from '@radix-ui/react-slot'
 
-import { IconProps, Icons } from '$ui/ds/atoms/Icons'
-import { colors } from '$ui/ds/tokens'
+import { Icon, IconProps } from '$ui/ds/atoms/Icons'
 import { cn } from '$ui/lib/utils'
 
 const buttonVariants = cva(
@@ -43,10 +42,7 @@ const buttonVariants = cva(
 export type ButtonProps = ButtonHTMLAttributes<HTMLButtonElement> &
   VariantProps<typeof buttonVariants> & {
     children: ReactNode
-    icon?: {
-      name: keyof typeof Icons
-      props?: IconProps
-    }
+    iconProps?: IconProps
     fullWidth?: boolean
     asChild?: boolean
     isLoading?: boolean
@@ -57,7 +53,7 @@ const Button = forwardRef<HTMLButtonElement, ButtonProps>(function Button(
     className,
     variant,
     size,
-    icon,
+    iconProps,
     fullWidth = false,
     asChild = false,
     isLoading,
@@ -67,8 +63,6 @@ const Button = forwardRef<HTMLButtonElement, ButtonProps>(function Button(
   ref,
 ) {
   const Comp = asChild ? Slot : 'button'
-  const ButtonIcon = icon ? Icons[icon.name] : null
-  const iconProps = icon?.props ?? {}
   return (
     <Comp
       disabled={isLoading}
@@ -80,15 +74,7 @@ const Button = forwardRef<HTMLButtonElement, ButtonProps>(function Button(
     >
       <Slottable>
         <div className='flex flex-row items-center gap-x-1'>
-          {ButtonIcon ? (
-            <ButtonIcon
-              className={cn({
-                [colors.textColors[iconProps.color!]]: iconProps.color,
-                'w-4': !iconProps.widthClass,
-                'h-4': !iconProps.heightClass,
-              })}
-            />
-          ) : null}
+          {iconProps ? <Icon {...iconProps} /> : null}
           {children}
         </div>
       </Slottable>