Merge branch 'messaging-part-4' into messaging-part-5

ee/hogai/assistant.py

@@ -79,6 +79,9 @@ def stream(self, conversation: Conversation) -> Generator[str, None, None]:
 
         chunks = AIMessageChunk(content="")
 
+        # Send a chunk to establish the connection avoiding the worker's timeout.
+        yield ""
+
         for update in generator:
             if is_value_update(update):
                 _, state_update = update

ee/hogai/trends/nodes.py

@@ -136,7 +136,11 @@ def _events_prompt(self) -> str:
         if not isinstance(response, CachedTeamTaxonomyQueryResponse):
             raise ValueError("Failed to generate events prompt.")
 
-        events = [item.event for item in response.results]
+        events: list[str] = []
+        for item in response.results:
+            if len(response.results) > 25 and item.count <= 3:
+                continue
+            events.append(item.event)
 
         # default for null in the
         tags: list[str] = ["all events"]

ee/hogai/trends/test/test_nodes.py

@@ -2,15 +2,13 @@
 
 from ee.hogai.trends.nodes import CreateTrendsPlanNode, GenerateTrendsNode
 from posthog.schema import AssistantMessage, ExperimentalAITrendsQuery, HumanMessage, VisualizationMessage
-from posthog.test.base import (
-    APIBaseTest,
-    ClickhouseTestMixin,
-)
+from posthog.test.base import APIBaseTest, ClickhouseTestMixin, _create_event, _create_person
 
 
 @override_settings(IN_UNIT_TESTING=True)
 class TestPlanAgentNode(ClickhouseTestMixin, APIBaseTest):
     def setUp(self):
+        super().setUp()
         self.schema = ExperimentalAITrendsQuery(series=[])
 
     def test_agent_reconstructs_conversation(self):
@@ -70,6 +68,24 @@ def test_agent_reconstructs_conversation_and_omits_unknown_messages(self):
         self.assertIn("Text", history[0].content)
         self.assertNotIn("{{question}}", history[0].content)
 
+    def test_agent_filters_out_low_count_events(self):
+        _create_person(distinct_ids=["test"], team=self.team)
+        for i in range(26):
+            _create_event(event=f"event{i}", distinct_id="test", team=self.team)
+            _create_event(event="distinctevent", distinct_id="test", team=self.team)
+        node = CreateTrendsPlanNode(self.team)
+        self.assertEqual(
+            node._events_prompt,
+            "<list of available events for filtering>\nall events\ndistinctevent\n</list of available events for filtering>",
+        )
+
+    def test_agent_preserves_low_count_events_for_smaller_teams(self):
+        _create_person(distinct_ids=["test"], team=self.team)
+        _create_event(event="distinctevent", distinct_id="test", team=self.team)
+        node = CreateTrendsPlanNode(self.team)
+        self.assertIn("distinctevent", node._events_prompt)
+        self.assertIn("all events", node._events_prompt)
+
 
 @override_settings(IN_UNIT_TESTING=True)
 class TestGenerateTrendsNode(ClickhouseTestMixin, APIBaseTest):

frontend/src/layout/navigation/EnvironmentSwitcher.tsx

@@ -157,6 +157,17 @@ function determineProjectSwitchUrl(pathname: string, newTeamId: number): string
     // and after switching is on a different page than before.
     let route = removeProjectIdIfPresent(pathname)
     route = removeFlagIdIfPresent(route)
+
+    // List of routes that should redirect to project home
+    // instead of keeping the current path.
+    const redirectToHomeRoutes = ['/products', '/onboarding']
+
+    const shouldRedirectToHome = redirectToHomeRoutes.some((redirectRoute) => route.includes(redirectRoute))
+
+    if (shouldRedirectToHome) {
+        return urls.project(newTeamId) // Go to project home
+    }
+
     return urls.project(newTeamId, route)
 }
 

frontend/src/lib/api.ts

@@ -1760,9 +1760,6 @@ const api = {
         async update(id: HogFunctionType['id'], data: Partial<HogFunctionType>): Promise<HogFunctionType> {
             return await new ApiRequest().hogFunction(id).update({ data })
         },
-        async sendBroadcast(id: HogFunctionType['id']): Promise<HogFunctionType> {
-            return await new ApiRequest().hogFunction(id).withAction('broadcast').create()
-        },
         async logs(
             id: HogFunctionType['id'],
             params: LogEntryRequestParams = {}

frontend/src/scenes/max/Max.stories.tsx

@@ -4,7 +4,7 @@ import { useEffect } from 'react'
 
 import { mswDecorator, useStorybookMocks } from '~/mocks/browser'
 
-import chatResponse from './__mocks__/chatResponse.json'
+import { chatResponseChunk } from './__mocks__/chatResponse.mocks'
 import { MaxInstance } from './Max'
 import { maxLogic } from './maxLogic'
 
@@ -13,7 +13,7 @@ const meta: Meta = {
     decorators: [
         mswDecorator({
             post: {
-                '/api/environments/:team_id/query/chat/': chatResponse,
+                '/api/environments/:team_id/query/chat/': (_, res, ctx) => res(ctx.text(chatResponseChunk)),
             },
         }),
     ],

frontend/src/scenes/max/__mocks__/chatResponse.mocks.ts

@@ -0,0 +1,3 @@
+import chatResponse from './chatResponse.json'
+
+export const chatResponseChunk = `data: ${JSON.stringify(chatResponse)}\n\n`

frontend/src/scenes/max/maxLogic.ts

@@ -1,4 +1,5 @@
 import { shuffle } from 'd3'
+import { createParser } from 'eventsource-parser'
 import { actions, kea, key, listeners, path, props, reducers, selectors } from 'kea'
 import { loaders } from 'kea-loaders'
 import api from 'lib/api'
@@ -118,20 +119,22 @@ export const maxLogic = kea<maxLogicType>([
                     messages: values.thread.map(({ status, ...message }) => message),
                 })
                 const reader = response.body?.getReader()
+
+                if (!reader) {
+                    return
+                }
+
                 const decoder = new TextDecoder()
 
-                if (reader) {
-                    let firstChunk = true
+                let firstChunk = true
 
-                    while (true) {
-                        const { done, value } = await reader.read()
-                        if (done) {
-                            actions.setMessageStatus(newIndex, 'completed')
-                            break
-                        }
+                const parser = createParser({
+                    onEvent: (event) => {
+                        const parsedResponse = parseResponse(event.data)
 
-                        const text = decoder.decode(value)
-                        const parsedResponse = parseResponse(text)
+                        if (!parsedResponse) {
+                            return
+                        }
 
                         if (firstChunk) {
                             firstChunk = false
@@ -145,6 +148,17 @@ export const maxLogic = kea<maxLogicType>([
                                 status: 'loading',
                             })
                         }
+                    },
+                })
+
+                while (true) {
+                    const { done, value } = await reader.read()
+
+                    parser.feed(decoder.decode(value))
+
+                    if (done) {
+                        actions.setMessageStatus(newIndex, 'completed')
+                        break
                     }
                 }
             } catch {
@@ -163,50 +177,11 @@ export const maxLogic = kea<maxLogicType>([
  * Parses the generation result from the API. Some generation chunks might be sent in batches.
  * @param response
  */
-function parseResponse(response: string, recursive = true): RootAssistantMessage | null {
+function parseResponse(response: string): RootAssistantMessage | null | undefined {
     try {
         const parsed = JSON.parse(response)
-        return parsed as RootAssistantMessage
+        return parsed as RootAssistantMessage | null | undefined
     } catch {
-        if (!recursive) {
-            return null
-        }
-
-        const results: [number, number][] = []
-        let pair: [number, number] = [0, 0]
-        let seq = 0
-
-        for (let i = 0; i < response.length; i++) {
-            const char = response[i]
-
-            if (char === '{') {
-                if (seq === 0) {
-                    pair[0] = i
-                }
-
-                seq += 1
-            }
-
-            if (char === '}') {
-                seq -= 1
-                if (seq === 0) {
-                    pair[1] = i
-                }
-            }
-
-            if (seq === 0) {
-                results.push(pair)
-                pair = [0, 0]
-            }
-        }
-
-        const lastPair = results.pop()
-
-        if (lastPair) {
-            const [left, right] = lastPair
-            return parseResponse(response.slice(left, right + 1), false)
-        }
-
         return null
     }
 }

frontend/src/scenes/pipeline/hogfunctions/HogFunctionConfiguration.tsx

@@ -69,7 +69,6 @@ export function HogFunctionConfiguration({ templateId, id }: HogFunctionConfigur
         templateHasChanged,
         forcedSubTemplateId,
         type,
-        broadcastLoading,
     } = useValues(logic)
     const {
         submitConfiguration,
@@ -81,7 +80,6 @@ export function HogFunctionConfiguration({ templateId, id }: HogFunctionConfigur
         setConfigurationValue,
         deleteHogFunction,
         setSubTemplateId,
-        sendBroadcast,
     } = useActions(logic)
 
     if (loading && !loaded) {
@@ -495,37 +493,6 @@ export function HogFunctionConfiguration({ templateId, id }: HogFunctionConfigur
                             )}
 
                             {!id || id === 'new' ? <HogFunctionTestPlaceholder /> : <HogFunctionTest id={id} />}
-
-                            {type === 'broadcast' ? (
-                                id && id !== 'new' ? (
-                                    <HogFunctionTestPlaceholder
-                                        title="Send broadcast"
-                                        description={
-                                            <div className="space-y-2">
-                                                <LemonButton
-                                                    type="primary"
-                                                    onClick={sendBroadcast}
-                                                    loading={personsCountLoading || broadcastLoading}
-                                                >
-                                                    Send to {personsCount} emails
-                                                </LemonButton>
-                                                <div>
-                                                    <strong>Please note:</strong> Clicking the button above will
-                                                    synchronously send to all the e-mails. While this is fine for
-                                                    testing with small lists, please don't use this for production
-                                                    usecases yet.
-                                                </div>
-                                            </div>
-                                        }
-                                    />
-                                ) : (
-                                    <HogFunctionTestPlaceholder
-                                        title="Send broadcast"
-                                        description="Save to send your broadcast"
-                                    />
-                                )
-                            ) : null}
-
                             <div className="flex gap-2 justify-end">{saveButtons}</div>
                         </div>
                     </div>

frontend/src/scenes/pipeline/hogfunctions/hogFunctionConfigurationLogic.tsx

@@ -188,7 +188,6 @@ export const hogFunctionConfigurationLogic = kea<hogFunctionConfigurationLogicTy
         setUnsavedConfiguration: (configuration: HogFunctionConfigurationType | null) => ({ configuration }),
         persistForUnload: true,
         setSampleGlobalsError: (error) => ({ error }),
-        sendBroadcast: true,
     }),
     reducers({
         showSource: [
@@ -403,22 +402,6 @@ export const hogFunctionConfigurationLogic = kea<hogFunctionConfigurationLogicTy
                 },
             },
         ],
-
-        broadcast: [
-            false,
-            {
-                sendBroadcast: async () => {
-                    const id = values.hogFunction?.id
-                    if (!id) {
-                        lemonToast.error('No broadcast to send')
-                        return false
-                    }
-                    await api.hogFunctions.sendBroadcast(id)
-                    lemonToast.success('Broadcast sent!')
-                    return true
-                },
-            },
-        ],
     })),
     forms(({ values, props, asyncActions }) => ({
         configuration: {

frontend/src/scenes/projectLogic.ts

@@ -19,7 +19,7 @@ export const projectLogic = kea<projectLogicType>([
         deleteProjectFailure: true,
     }),
     connect(() => ({
-        actions: [userLogic, ['loadUser']],
+        actions: [userLogic, ['loadUser', 'switchTeam']],
     })),
     reducers({
         projectBeingDeleted: [
@@ -101,6 +101,11 @@ export const projectLogic = kea<projectLogicType>([
         deleteProjectSuccess: () => {
             lemonToast.success('Project has been deleted')
         },
+        createProjectSuccess: ({ currentProject }) => {
+            if (currentProject) {
+                actions.switchTeam(currentProject.id)
+            }
+        },
     })),
     afterMount(({ actions }) => {
         const appContext = getAppContext()

hogvm/typescript/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@posthog/hogvm",
-    "version": "1.0.55",
+    "version": "1.0.56",
     "description": "PostHog Hog Virtual Machine",
     "types": "dist/index.d.ts",
     "source": "src/index.ts",

hogvm/typescript/src/execute.ts

@@ -122,7 +122,10 @@ export function exec(input: any[] | VMState | Bytecodes, options?: ExecOptions):
     let ops = vmState ? vmState.ops : 0
     const timeout = options?.timeout ?? DEFAULT_TIMEOUT_MS
     const maxAsyncSteps = options?.maxAsyncSteps ?? DEFAULT_MAX_ASYNC_STEPS
-    const rootGlobals: Record<string, any> = options?.globals ?? {}
+    const rootGlobals: Record<string, any> =
+        bytecodes.root?.globals && options?.globals
+            ? { ...bytecodes.root.globals, ...options.globals }
+            : bytecodes.root?.globals ?? options?.globals ?? {}
 
     if (callStack.length === 0) {
         callStack.push({

mypy-baseline.txt

@@ -808,7 +808,6 @@ posthog/temporal/tests/external_data/test_external_data_job.py:0: error: Invalid
 posthog/temporal/tests/external_data/test_external_data_job.py:0: error: Invalid index type "str" for "dict[Type, Sequence[str]]"; expected type "Type"  [index]
 posthog/temporal/tests/external_data/test_external_data_job.py:0: error: Invalid index type "str" for "dict[Type, Sequence[str]]"; expected type "Type"  [index]
 posthog/temporal/tests/external_data/test_external_data_job.py:0: error: Invalid index type "str" for "dict[Type, Sequence[str]]"; expected type "Type"  [index]
-posthog/temporal/tests/external_data/test_external_data_job.py:0: error: Invalid index type "str" for "dict[Type, Sequence[str]]"; expected type "Type"  [index]
 posthog/api/test/batch_exports/conftest.py:0: error: Signature of "run" incompatible with supertype "Worker"  [override]
 posthog/api/test/batch_exports/conftest.py:0: note: Superclass:
 posthog/api/test/batch_exports/conftest.py:0: note: def run(self) -> Coroutine[Any, Any, None]

package.json

@@ -76,7 +76,7 @@
         "@medv/finder": "^3.1.0",
         "@microlink/react-json-view": "^1.21.3",
         "@monaco-editor/react": "4.6.0",
-        "@posthog/hogvm": "^1.0.55",
+        "@posthog/hogvm": "^1.0.56",
         "@posthog/icons": "0.8.5",
         "@posthog/plugin-scaffold": "^1.4.4",
         "@react-hook/size": "^2.1.2",
@@ -124,6 +124,7 @@
         "esbuild-plugin-less": "^1.3.1",
         "esbuild-plugin-polyfill-node": "^0.3.0",
         "esbuild-sass-plugin": "^3.0.0",
+        "eventsource-parser": "^3.0.0",
         "expr-eval": "^2.0.2",
         "express": "^4.17.1",
         "fast-deep-equal": "^3.1.3",

plugin-server/package.json

@@ -54,7 +54,7 @@
         "@maxmind/geoip2-node": "^3.4.0",
         "@posthog/clickhouse": "^1.7.0",
         "@posthog/cyclotron": "file:../rust/cyclotron-node",
-        "@posthog/hogvm": "^1.0.55",
+        "@posthog/hogvm": "^1.0.56",
         "@posthog/plugin-scaffold": "1.4.4",
         "@sentry/node": "^7.49.0",
         "@sentry/profiling-node": "^0.3.0",