Repository: PostHog/posthog
Pull request #31690: feat(sql-editor): Fix HogQL errors with AI
Author: login: Gilbert09
Created: 2025-04-29T09:17:48Z | Merged: 2025-04-29T13:49:26Z

## Changes
- Adds a button to the query window action bar to fix any current HogQL error in the current tab
- Uses LLM and the Max AI backend to run a prompt - we try the prompt up to 3 times and check whether it's parseable by our HogQL parser, if not, we run it again with the updated error message
- Only to be released internally behind a flag for more extensive testing and iterations
- Hooked up to the LLM observability tool for metrics + sending some of our own events 

**TODO post internal release:**
- Hook up clickhouse errors that don't get returned to the client
- Include hogql allowed functions 
- Update how we pass table structure into the LLM
- Pass the second/third/fourth query back into the LLM with the error message in the loop

https://github.com/user-attachments/assets/64ef1f77-16db-423d-ba7f-e16a75713d1f



## Does this work well for both Cloud and self-hosted?
Yup

## How did you test this code?
Lots of local testing


Diff:
diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts
index 6626bd378914..191e2b742b3f 100644
--- a/frontend/src/lib/api.ts
+++ b/frontend/src/lib/api.ts
@@ -1082,6 +1082,11 @@ class ApiRequest {
         return this.externalDataSchemas(teamId).addPathComponent(schemaId)
     }
 
+    // Fix HogQL errors
+    public fixHogQLErrors(teamId?: TeamType['id']): ApiRequest {
+        return this.environmentsDetail(teamId).addPathComponent('fix_hogql')
+    }
+
     // Insight Variables
     public insightVariables(teamId?: TeamType['id']): ApiRequest {
         return this.environmentsDetail(teamId).addPathComponent('insight_variables')
@@ -2895,6 +2900,11 @@ const api = {
                 .get()
         },
     },
+    fixHogQLErrors: {
+        async fix(query: string, error?: string): Promise<Record<string, any>> {
+            return await new ApiRequest().fixHogQLErrors().create({ data: { query, error } })
+        },
+    },
 
     dataWarehouseViewLinks: {
         async list(): Promise<PaginatedResponse<DataWarehouseViewLink>> {
diff --git a/frontend/src/lib/constants.tsx b/frontend/src/lib/constants.tsx
index 7262ba47f84c..a2fb78068715 100644
--- a/frontend/src/lib/constants.tsx
+++ b/frontend/src/lib/constants.tsx
@@ -259,6 +259,7 @@ export const FEATURE_FLAGS = {
     STARTUP_PROGRAM_INTENT: 'startup-program-intent', // owner: @pawel-cebula #team-billing
     SHOW_NEW_EXPERIMENTATION_ENGINE_BANNER: 'show-new-experimentation-engine-banner', // owner: @andehen #team-experiments
     WEB_ANALYTICS_FRUSTRATING_PAGES_TILE: 'web-analytics-frustrating-pages-tile', // owner: @lricoy #team-web-analytics
+    SQL_EDITOR_AI_ERROR_FIXER: 'sql-editor-ai-error-fixer', // owner: @Gilbert09 #team-data-warehouse
 } as const
 export type FeatureFlagKey = (typeof FEATURE_FLAGS)[keyof typeof FEATURE_FLAGS]
 
diff --git a/frontend/src/queries/schema.json b/frontend/src/queries/schema.json
index dbf038da4ccf..3a4274af2c25 100644
--- a/frontend/src/queries/schema.json
+++ b/frontend/src/queries/schema.json
@@ -615,7 +615,7 @@
             "type": "object"
         },
         "AssistantContextualTool": {
-            "enum": ["search_session_recordings", "generate_hogql_query"],
+            "enum": ["search_session_recordings", "generate_hogql_query", "fix_hogql_query"],
             "type": "string"
         },
         "AssistantDateRange": {
diff --git a/frontend/src/queries/schema/schema-assistant-messages.ts b/frontend/src/queries/schema/schema-assistant-messages.ts
index bbfbd07be869..e51bb7f1f28b 100644
--- a/frontend/src/queries/schema/schema-assistant-messages.ts
+++ b/frontend/src/queries/schema/schema-assistant-messages.ts
@@ -108,4 +108,4 @@ export interface AssistantToolCallMessage extends BaseAssistantMessage {
     tool_call_id: string
 }
 
-export type AssistantContextualTool = 'search_session_recordings' | 'generate_hogql_query'
+export type AssistantContextualTool = 'search_session_recordings' | 'generate_hogql_query' | 'fix_hogql_query'
diff --git a/frontend/src/scenes/data-warehouse/editor/QueryWindow.tsx b/frontend/src/scenes/data-warehouse/editor/QueryWindow.tsx
index 2b91c3f752c1..a70d0cf216cf 100644
--- a/frontend/src/scenes/data-warehouse/editor/QueryWindow.tsx
+++ b/frontend/src/scenes/data-warehouse/editor/QueryWindow.tsx
@@ -1,10 +1,21 @@
 import { Monaco } from '@monaco-editor/react'
-import { IconBolt, IconBook, IconBrackets, IconDownload, IconPlayFilled, IconSidebarClose } from '@posthog/icons'
-import { LemonDivider } from '@posthog/lemon-ui'
+import {
+    IconBolt,
+    IconBook,
+    IconBrackets,
+    IconDownload,
+    IconMagicWand,
+    IconPlayFilled,
+    IconSidebarClose,
+    IconWarning,
+} from '@posthog/icons'
+import { LemonDivider, Spinner } from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { router } from 'kea-router'
+import { FEATURE_FLAGS } from 'lib/constants'
 import { IconCancel } from 'lib/lemon-ui/icons'
 import { LemonButton } from 'lib/lemon-ui/LemonButton'
+import { featureFlagLogic } from 'lib/logic/featureFlagLogic'
 import type { editor as importedEditor } from 'monaco-editor'
 import { useMemo } from 'react'
 
@@ -12,6 +23,7 @@ import { dataNodeLogic } from '~/queries/nodes/DataNode/dataNodeLogic'
 
 import { dataWarehouseViewsLogic } from '../saved_queries/dataWarehouseViewsLogic'
 import { editorSizingLogic } from './editorSizingLogic'
+import { fixSQLErrorsLogic } from './fixSQLErrorsLogic'
 import { multitabEditorLogic } from './multitabEditorLogic'
 import { OutputPane } from './OutputPane'
 import { QueryHistoryModal } from './QueryHistoryModal'
@@ -48,6 +60,7 @@ export function QueryWindow({ onSetMonacoAndEditor }: QueryWindowProps): JSX.Ele
     const { sidebarWidth } = useValues(editorSizingLogic)
     const { resetDefaultSidebarWidth } = useActions(editorSizingLogic)
     const { setActiveTab } = useActions(editorSidebarLogic)
+    const { featureFlags } = useValues(featureFlagLogic)
 
     const isMaterializedView =
         !!editingView?.status &&
@@ -166,6 +179,7 @@ export function QueryWindow({ onSetMonacoAndEditor }: QueryWindowProps): JSX.Ele
                         {renderAddSQLVariablesButton()}
                     </>
                 )}
+                {featureFlags[FEATURE_FLAGS.SQL_EDITOR_AI_ERROR_FIXER] && <FixErrorButton />}
             </div>
             <QueryPane
                 originalValue={suggestedQueryInput && suggestedQueryInput != queryInput ? queryInput ?? ' ' : undefined}
@@ -258,3 +272,60 @@ function InternalQueryWindow(): JSX.Element | null {
 
     return <OutputPane />
 }
+
+function FixErrorButton(): JSX.Element {
+    const { queryInput, fixErrorsError, metadata } = useValues(multitabEditorLogic)
+    const { fixErrors: fixHogQLErrors } = useActions(multitabEditorLogic)
+    const { responseError } = useValues(dataNodeLogic)
+    const { responseLoading: fixHogQLErrorsLoading } = useValues(fixSQLErrorsLogic)
+
+    const queryError = responseError || metadata?.errors?.map((n) => n.message)?.join('. ') || undefined
+
+    const icon = useMemo(() => {
+        if (fixHogQLErrorsLoading) {
+            return <Spinner />
+        }
+
+        if (fixErrorsError) {
+            return <IconWarning className="text-warning" />
+        }
+
+        return <IconMagicWand />
+    }, [fixHogQLErrorsLoading, fixErrorsError])
+
+    const disabledReason = useMemo(() => {
+        if (!queryError) {
+            return 'No query error to fix'
+        }
+
+        if (fixErrorsError) {
+            return fixErrorsError
+        }
+
+        return false
+    }, [queryError, fixErrorsError])
+
+    const content = useMemo(() => {
+        if (fixHogQLErrorsLoading) {
+            return 'Fixing...'
+        }
+
+        if (fixErrorsError) {
+            return "Can't fix"
+        }
+
+        return 'Fix errors'
+    }, [fixErrorsError, fixHogQLErrorsLoading])
+
+    return (
+        <LemonButton
+            type="tertiary"
+            size="xsmall"
+            disabledReason={disabledReason}
+            icon={icon}
+            onClick={() => fixHogQLErrors(queryInput, queryError)}
+        >
+            {content}
+        </LemonButton>
+    )
+}
diff --git a/frontend/src/scenes/data-warehouse/editor/fixSQLErrorsLogic.ts b/frontend/src/scenes/data-warehouse/editor/fixSQLErrorsLogic.ts
new file mode 100644
index 000000000000..505bb94d784c
--- /dev/null
+++ b/frontend/src/scenes/data-warehouse/editor/fixSQLErrorsLogic.ts
@@ -0,0 +1,30 @@
+import { actions, kea, path } from 'kea'
+import { loaders } from 'kea-loaders'
+import api from 'lib/api'
+
+import type { fixSQLErrorsLogicType } from './fixSQLErrorsLogicType'
+
+export interface Response {
+    query: string
+    trace_id: string
+    error?: string
+}
+
+export const fixSQLErrorsLogic = kea<fixSQLErrorsLogicType>([
+    path(['scenes', 'data-warehouse', 'editor', 'fixSQLErrorsLogic']),
+    actions({
+        fixErrors: (query: string, error?: string) => ({ query, error }),
+    }),
+    loaders({
+        response: [
+            null as Response | null,
+            {
+                fixErrors: async ({ query, error }) => {
+                    const response = await api.fixHogQLErrors.fix(query, error)
+
+                    return response as Response
+                },
+            },
+        ],
+    }),
+])
diff --git a/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx b/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx
index 3dc6d899dd17..b650d09e01a5 100644
--- a/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx
+++ b/frontend/src/scenes/data-warehouse/editor/multitabEditorLogic.tsx
@@ -12,6 +12,7 @@ import { removeUndefinedAndNull } from 'lib/utils'
 import { copyToClipboard } from 'lib/utils/copyToClipboard'
 import isEqual from 'lodash.isequal'
 import { editor, Uri } from 'monaco-editor'
+import posthog from 'posthog-js'
 import { insightsApi } from 'scenes/insights/utils/api'
 import { urls } from 'scenes/urls'
 import { userLogic } from 'scenes/userLogic'
@@ -31,6 +32,7 @@ import {
 import { dataWarehouseViewsLogic } from '../saved_queries/dataWarehouseViewsLogic'
 import { DATAWAREHOUSE_EDITOR_ITEM_ID, sizeOfInBytes } from '../utils'
 import { editorSceneLogic } from './editorSceneLogic'
+import { fixSQLErrorsLogic } from './fixSQLErrorsLogic'
 import type { multitabEditorLogicType } from './multitabEditorLogicType'
 import { outputPaneLogic, OutputTab } from './outputPaneLogic'
 import { ViewEmptyState } from './ViewLoadingState'
@@ -98,6 +100,8 @@ export const multitabEditorLogic = kea<multitabEditorLogicType>([
             ['setActiveTab'],
             editorSceneLogic,
             ['reportAIQueryPrompted', 'reportAIQueryAccepted', 'reportAIQueryRejected', 'reportAIQueryPromptOpen'],
+            fixSQLErrorsLogic,
+            ['fixErrors', 'fixErrorsSuccess', 'fixErrorsFailure'],
         ],
     })),
     actions(({ values }) => ({
@@ -296,8 +300,23 @@ export const multitabEditorLogic = kea<multitabEditorLogicType>([
                 closeHistoryModal: () => false,
             },
         ],
+        fixErrorsError: [
+            null as string | null,
+            {
+                setQueryInput: () => null,
+                fixErrorsFailure: (_, { error }) => error,
+            },
+        ],
     })),
     listeners(({ values, props, actions, asyncActions }) => ({
+        fixErrorsSuccess: ({ response }) => {
+            actions.setSuggestedQueryInput(response.query)
+
+            posthog.capture('ai-error-fixer-success', { trace_id: response.trace_id })
+        },
+        fixErrorsFailure: () => {
+            posthog.capture('ai-error-fixer-failure')
+        },
         shareTab: () => {
             const currentTab = values.activeModelUri
             if (!currentTab) {
diff --git a/posthog/api/__init__.py b/posthog/api/__init__.py
index d1809ca27543..f3607300d9da 100644
--- a/posthog/api/__init__.py
+++ b/posthog/api/__init__.py
@@ -1,6 +1,7 @@
 from rest_framework import decorators, exceptions, viewsets
 from rest_framework_extensions.routers import NestedRegistryItem
 
+import products.data_warehouse.backend.api.fix_hogql as fix_hogql
 import products.early_access_features.backend.api as early_access_feature
 from products.editor.backend.api import LLMProxyViewSet, MaxToolsViewSet
 from posthog.api import data_color_theme, metalytics, project, wizard
@@ -328,6 +329,12 @@ def register_grandfathered_environment_nested_viewset(
     "environment_external_data_schemas",
     ["team_id"],
 )
+environments_router.register(
+    r"fix_hogql",
+    fix_hogql.FixHogQLViewSet,
+    "project_fix_hogql",
+    ["team_id"],
+)
 
 # Organizations nested endpoints
 organizations_router = router.register(r"organizations", organization.OrganizationViewSet, "organizations")
diff --git a/posthog/schema.py b/posthog/schema.py
index a6163a544914..57b3f7ba094a 100644
--- a/posthog/schema.py
+++ b/posthog/schema.py
@@ -64,6 +64,7 @@ class AssistantBaseMultipleBreakdownFilter(BaseModel):
 class AssistantContextualTool(StrEnum):
     SEARCH_SESSION_RECORDINGS = "search_session_recordings"
     GENERATE_HOGQL_QUERY = "generate_hogql_query"
+    FIX_HOGQL_QUERY = "fix_hogql_query"
 
 
 class AssistantDateRange(BaseModel):
diff --git a/products/data_warehouse/backend/api/__init__.py b/products/data_warehouse/backend/api/__init__.py
new file mode 100644
index 000000000000..e7eb4641ad73
--- /dev/null
+++ b/products/data_warehouse/backend/api/__init__.py
@@ -0,0 +1,3 @@
+from .fix_hogql import FixHogQLViewSet
+
+__all__ = ["FixHogQLViewSet"]
diff --git a/products/data_warehouse/backend/api/fix_hogql.py b/products/data_warehouse/backend/api/fix_hogql.py
new file mode 100644
index 000000000000..224a1a72cbbc
--- /dev/null
+++ b/products/data_warehouse/backend/api/fix_hogql.py
@@ -0,0 +1,58 @@
+from typing import cast
+import uuid
+
+import posthoganalytics
+from posthoganalytics.ai.langchain.callbacks import CallbackHandler
+
+from posthog.api.routing import TeamAndOrgViewSetMixin
+from rest_framework import status, viewsets
+from rest_framework.request import Request
+from rest_framework.response import Response
+from langchain_core.runnables import RunnableConfig
+
+from posthog.models.user import User
+
+
+class FixHogQLViewSet(TeamAndOrgViewSetMixin, viewsets.ModelViewSet):
+    scope_object = "INTERNAL"
+
+    def create(self, request: Request, *args, **kwargs) -> Response:
+        from products.data_warehouse.backend.hogql_fixer_ai import HogQLQueryFixerTool
+
+        query = request.data.get("query", None)
+        error = request.data.get("error", "")
+
+        if query is None:
+            return Response(
+                status=status.HTTP_400_BAD_REQUEST,
+                data={"message": "No query provided"},
+            )
+
+        trace_id = f"fix_hogql_query_{uuid.uuid4()}"
+        user = cast(User, request.user)
+
+        config: RunnableConfig = {
+            "configurable": {
+                "contextual_tools": {
+                    "fix_hogql_query": {
+                        "hogql_query": query,
+                        "error_message": error,
+                    }
+                },
+                "team_id": self.team_id,
+                "trace_id": trace_id,
+                "distinct_id": user.distinct_id,
+            },
+            "callbacks": (
+                [CallbackHandler(posthoganalytics.default_client, distinct_id=user.distinct_id, trace_id=trace_id)]
+                if posthoganalytics.default_client
+                else None
+            ),
+        }
+
+        result = HogQLQueryFixerTool(_team_id=self.team_id, _context={}).invoke({}, config)
+
+        if result is None or len(result) == 0:
+            return Response({"trace_id": trace_id, "error": "Could not fix the query"}, status=400)
+
+        return Response({"query": result, "trace_id": trace_id}, status=200)
diff --git a/products/data_warehouse/backend/hogql_fixer_ai.py b/products/data_warehouse/backend/hogql_fixer_ai.py
new file mode 100644
index 000000000000..584e45da4f26
--- /dev/null
+++ b/products/data_warehouse/backend/hogql_fixer_ai.py
@@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
+
+USER_PROMPT = """
+Fix the errors in the HogQL query below and only return the new updated query in your response.
+
+- Don't update any other part of the query if it's not relevant to the error, including rewriting shorthand clickhouse SQL to the full version.
+- Don't change the capitalisation of the query if it's not relevant to the error, such as rewriting `select` as `SELECT` or `from` as `FROM`
+- There may also be more than one error in the syntax.
+
+Below is the current HogQL query and the error message
+"""
+
+
+class HogQLQueryFixerTool(MaxTool):
+    name: str = "fix_hogql_query"
+    description: str = "Fixes any error in the current HogQL query"
+    thinking_message: str = "Fixing errors in the SQL query"
+    root_system_prompt_template: str = SQL_ASSISTANT_ROOT_SYSTEM_PROMPT
+
+    def _run_impl(self) -> tuple[str, str | None]:
+        database = create_hogql_database(self._team_id)
+        hogql_context = HogQLContext(team_id=self._team_id, enable_select_queries=True, database=database)
+
+        serialized_database = serialize_database(hogql_context)
+        schema_description = "\n\n".join(
+            (
+                f"Table `{table_name}` with fields:\n"
+                + "\n".join(f"- {field.name} ({field.type})" for field in table.fields.values())
+                for table_name, table in serialized_database.items()
+                # Only the most important core tables, plus all warehouse tables
+                if table_name in ["events", "groups", "persons"]
+                or table_name in database.get_warehouse_tables()
+                or table_name in database.get_views()
+            )
+        )
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    SYSTEM_PROMPT.replace("{{schema_description}}", schema_description),
+                ),
+                (
+                    "user",
+                    USER_PROMPT
+                    + "\n\n<hogql_query>"
+                    + "{{{hogql_query}}}"
+                    + "</hogql_query>"
+                    + "\n\n<error>"
+                    + "{{{error_message}}}"
+                    + "</error>",
+                ),
+            ],
+            template_format="mustache",
+        )
+
+        for i in range(3):
+            try:
+                chain = prompt | self._model
+                result = chain.invoke(self.context)
+                parsed_result = self._parse_output(result, hogql_context)
+                break
+            except PydanticOutputParserException as e:
+                prompt += f"\n\nWe've ran this prompt {i+1} time{'s' if i+1 > 1 else ''} now. The newly updated query gave us this error: {e.validation_message}"
+        else:
+            return "", None
+
+        return parsed_result, parsed_result
+
+    @property
+    def _model(self):
+        return ChatOpenAI(model="gpt-4o", temperature=0, disable_streaming=True).with_structured_output(
+            SQL_SCHEMA,
+            method="function_calling",
+            include_raw=False,
+        )
+
+    def _parse_output(self, output, hogql_context: HogQLContext):
+        result = parse_pydantic_structured_output(SchemaGeneratorOutput[str])(output)  # type: ignore
+        # We also ensure the generated SQL is valid
+        assert result.query is not None
+        try:
+            print_ast(parse_select(result.query), context=hogql_context, dialect="clickhouse")
+        except (ExposedHogQLError, ResolutionError) as err:
+            err_msg = str(err)
+            if err_msg.startswith("no viable alternative"):
+                # The "no viable alternative" ANTLR error is horribly unhelpful, both for humans and LLMs
+                err_msg = f'ANTLR parsing error: "no viable alternative at input". This means that the query isn\' valid HogQL. The last 5 characters where we tripped up were "{result.query[-5:]}".'
+            raise PydanticOutputParserException(llm_output=result.query, validation_message=err_msg)
+        except Exception as e:
+            raise PydanticOutputParserException(llm_output=result.query, validation_message=str(e))
+
+        return result.query


Reviews:
user: login: skoob13
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T09:22:37Z
user: login: greptile-apps[bot]
body: 
## PR Summary

This PR introduces an AI-powered feature for fixing HogQL (PostHog's SQL variant) query errors, leveraging GPT-4 to analyze and suggest fixes for syntax issues.

- New `HogQLQueryFixerTool` in `hogql_fixer_ai.py` uses ANTLR grammar definitions to validate and fix SQL syntax errors
- Adds feature flag `SQL_EDITOR_AI_ERROR_FIXER` and new API endpoint `/api/environments/{team_id}/fix_hogql` for error fixing
- Implements `fixSQLErrorsLogic.ts` Kea module to handle frontend error fixing functionality with proper loading states
- Potential issues with thread safety in parser content globals and hardcoded model name 'gpt-4o'
- Missing rate limiting and timeout handling for LLM calls, could benefit from more robust error handling

<sub>10 file(s) reviewed, 5 comment(s)</sub>
<sub>[Edit PR Review Bot Settings](https://app.greptile.com/review/github) | [Greptile](https://greptile.com?utm_source=greptile_expert&utm_medium=github&utm_campaign=code_reviews&utm_content=posthog_31690)</sub>
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T11:09:39Z
user: login: graphite-app[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T11:09:59Z
user: login: rossgray
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T12:53:52Z
user: login: rossgray
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T12:57:24Z
user: login: rossgray
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T13:01:30Z
user: login: rossgray
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T13:02:20Z
user: login: rossgray
body: Changes overall seem good to me.

The main thing I am thinking about is whether there could be improved integration with the version of Max we're using already in the SQL Editor, both in the front end and the back end. For example, in the front end it would also make sense if Max could also correct queries using the same tool in the backend (assuming this is not possible). And in the backend, could we consolidate some of the code here, as both Max and this fixing tool share many similarities in the sense that they both need to know about HogQL syntax and its features, the difference being one is writing and editing SQL and one is fixing errors.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T13:07:22Z
user: login: Gilbert09
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-29T13:45:55Z
user: login: Gilbert09
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-29T13:46:32Z
user: login: Gilbert09
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-29T13:47:18Z

Review comments:
diff_hunk: @@ -0,0 +1,46 @@
+from typing import cast
+import uuid
+from posthog.api.routing import TeamAndOrgViewSetMixin
+from rest_framework import status, viewsets
+from rest_framework.request import Request
+from rest_framework.response import Response
+from langchain_core.runnables import RunnableConfig
+
+from posthog.models.user import User
+
+
+class FixHogQLViewSet(TeamAndOrgViewSetMixin, viewsets.ModelViewSet):
+    scope_object = "INTERNAL"
+
+    def create(self, request: Request, *args, **kwargs) -> Response:
+        from products.data_warehouse.backend.hogql_fixer_ai import HogQLQueryFixerTool
+
+        query = request.data.get("query", None)
+        error = request.data.get("error", "")
+
+        if query is None:
+            return Response(
+                status=status.HTTP_400_BAD_REQUEST,
+                data={"message": "No query provided"},
+            )
+
+        trace_id = str(uuid.uuid4())
+        user = cast(User, request.user)
+
+        config: RunnableConfig = {
+            "configurable": {
+                "contextual_tools": {
+                    "fix_hogql_query": {
+                        "hogql_query": query,
+                        "error_message": error,
+                    }
+                },
+                "team_id": self.team_id,
+                "trace_id": trace_id,
+                "distinct_id": user.distinct_id,
+            }
path: products/data_warehouse/backend/api/fix_hogql.py
user: login: skoob13
body: ```suggestion
            },
            "callbacks": (
		  [CallbackHandler(
	                posthoganalytics.default_client,
	                distinct_id=user.distinct_id,
	                trace_id=trace_id,
	            )]
	            if posthoganalytics.default_client
	            else None
	)
```

And the import is:
```python
import posthoganalytics
```
created_at: 2025-04-29T09:22:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -258,3 +272,60 @@ function InternalQueryWindow(): JSX.Element | null {
 
     return <OutputPane />
 }
+
+function FixErrorButton(): JSX.Element {
+    const { queryInput, fixErrorsError, metadata } = useValues(multitabEditorLogic)
+    const { fixErrors: fixHogQLErrors } = useActions(multitabEditorLogic)
+    const { responseError } = useValues(dataNodeLogic)
+    const { responseLoading: fixHogQLErrorsLoading } = useValues(fixSQLErrorsLogic)
+
+    const queryError = responseError || metadata?.errors?.map((n) => n.message)?.join('. ') || undefined
path: frontend/src/scenes/data-warehouse/editor/QueryWindow.tsx
user: login: greptile-apps[bot]
body: style: Consider handling empty string error messages to avoid displaying '.' as the only character when joining messages
created_at: 2025-04-29T11:06:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -328,6 +329,12 @@ def register_grandfathered_environment_nested_viewset(
     "environment_external_data_schemas",
     ["team_id"],
 )
+environments_router.register(
+    r"fix_hogql",
+    fix_hogql.FixHogQLViewSet,
+    "project_fix_hogql",
+    ["team_id"],
+)
path: posthog/api/__init__.py
user: login: greptile-apps[bot]
body: style: Inconsistent naming - the basename 'project_fix_hogql' should be 'environment_fix_hogql' since this is registered under environments_router and follows the environment pattern
created_at: 2025-04-29T11:06:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,58 @@
+from typing import cast
+import uuid
+
+import posthoganalytics
+from posthoganalytics.ai.langchain.callbacks import CallbackHandler
+
+from posthog.api.routing import TeamAndOrgViewSetMixin
+from rest_framework import status, viewsets
+from rest_framework.request import Request
+from rest_framework.response import Response
+from langchain_core.runnables import RunnableConfig
+
+from posthog.models.user import User
+
+
+class FixHogQLViewSet(TeamAndOrgViewSetMixin, viewsets.ModelViewSet):
+    scope_object = "INTERNAL"
+
+    def create(self, request: Request, *args, **kwargs) -> Response:
+        from products.data_warehouse.backend.hogql_fixer_ai import HogQLQueryFixerTool
+
+        query = request.data.get("query", None)
+        error = request.data.get("error", "")
+
+        if query is None:
+            return Response(
+                status=status.HTTP_400_BAD_REQUEST,
+                data={"message": "No query provided"},
+            )
+
+        trace_id = f"fix_hogql_query_{uuid.uuid4()}"
+        user = cast(User, request.user)
+
+        config: RunnableConfig = {
+            "configurable": {
+                "contextual_tools": {
+                    "fix_hogql_query": {
+                        "hogql_query": query,
+                        "error_message": error,
+                    }
+                },
+                "team_id": self.team_id,
+                "trace_id": trace_id,
+                "distinct_id": user.distinct_id,
+            },
+            "callbacks": (
+                [CallbackHandler(posthoganalytics.default_client, distinct_id=user.distinct_id, trace_id=trace_id)]
+                if posthoganalytics.default_client
+                else None
+            ),
+        }
+
+        result = HogQLQueryFixerTool(_team_id=self.team_id, _context={}).invoke({}, config)
path: products/data_warehouse/backend/api/fix_hogql.py
user: login: greptile-apps[bot]
body: logic: No error handling around invoke() call. AI service failures could cause 500 errors.
created_at: 2025-04-29T11:07:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,58 @@
+from typing import cast
+import uuid
+
+import posthoganalytics
+from posthoganalytics.ai.langchain.callbacks import CallbackHandler
+
+from posthog.api.routing import TeamAndOrgViewSetMixin
+from rest_framework import status, viewsets
+from rest_framework.request import Request
+from rest_framework.response import Response
+from langchain_core.runnables import RunnableConfig
+
+from posthog.models.user import User
+
+
+class FixHogQLViewSet(TeamAndOrgViewSetMixin, viewsets.ModelViewSet):
+    scope_object = "INTERNAL"
+
+    def create(self, request: Request, *args, **kwargs) -> Response:
+        from products.data_warehouse.backend.hogql_fixer_ai import HogQLQueryFixerTool
+
+        query = request.data.get("query", None)
+        error = request.data.get("error", "")
+
+        if query is None:
+            return Response(
+                status=status.HTTP_400_BAD_REQUEST,
+                data={"message": "No query provided"},
+            )
+
+        trace_id = f"fix_hogql_query_{uuid.uuid4()}"
+        user = cast(User, request.user)
+
+        config: RunnableConfig = {
+            "configurable": {
+                "contextual_tools": {
+                    "fix_hogql_query": {
+                        "hogql_query": query,
+                        "error_message": error,
+                    }
+                },
+                "team_id": self.team_id,
+                "trace_id": trace_id,
+                "distinct_id": user.distinct_id,
+            },
+            "callbacks": (
+                [CallbackHandler(posthoganalytics.default_client, distinct_id=user.distinct_id, trace_id=trace_id)]
+                if posthoganalytics.default_client
+                else None
+            ),
+        }
+
+        result = HogQLQueryFixerTool(_team_id=self.team_id, _context={}).invoke({}, config)
+
+        if result is None or len(result) == 0:
+            return Response({"trace_id": trace_id, "error": "Could not fix the query"}, status=400)
path: products/data_warehouse/backend/api/fix_hogql.py
user: login: greptile-apps[bot]
body: logic: Check if result is string type before checking len(). AI might return non-string.

```suggestion
        if result is None or not isinstance(result, str) or len(result) == 0:
            return Response({"trace_id": trace_id, "error": "Could not fix the query"}, status=400)
```
created_at: 2025-04-29T11:07:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
+
+USER_PROMPT = """
+Fix the errors in the HogQL query below and only return the new updated query in your response.
+
+- Don't update any other part of the query if it's not relevant to the error, including rewriting shorthand clickhouse SQL to the full version.
+- Don't change the capitalisation of the query if it's not relevant to the error, such as rewriting `select` as `SELECT` or `from` as `FROM`
+- There may also be more than one error in the syntax.
+
+Below is the current HogQL query and the error message
+"""
+
+
+class HogQLQueryFixerTool(MaxTool):
+    name: str = "fix_hogql_query"
+    description: str = "Fixes any error in the current HogQL query"
+    thinking_message: str = "Fixing errors in the SQL query"
+    root_system_prompt_template: str = SQL_ASSISTANT_ROOT_SYSTEM_PROMPT
+
+    def _run_impl(self) -> tuple[str, str | None]:
+        database = create_hogql_database(self._team_id)
+        hogql_context = HogQLContext(team_id=self._team_id, enable_select_queries=True, database=database)
+
+        serialized_database = serialize_database(hogql_context)
+        schema_description = "\n\n".join(
+            (
+                f"Table `{table_name}` with fields:\n"
+                + "\n".join(f"- {field.name} ({field.type})" for field in table.fields.values())
+                for table_name, table in serialized_database.items()
+                # Only the most important core tables, plus all warehouse tables
+                if table_name in ["events", "groups", "persons"]
+                or table_name in database.get_warehouse_tables()
+                or table_name in database.get_views()
+            )
+        )
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    SYSTEM_PROMPT.replace("{{schema_description}}", schema_description),
+                ),
+                (
+                    "user",
+                    USER_PROMPT
+                    + "\n\n<hogql_query>"
+                    + "{{{hogql_query}}}"
+                    + "</hogql_query>"
+                    + "\n\n<error>"
+                    + "{{{error_message}}}"
+                    + "</error>",
+                ),
+            ],
+            template_format="mustache",
+        )
+
+        for i in range(3):
+            try:
+                chain = prompt | self._model
+                result = chain.invoke(self.context)
+                parsed_result = self._parse_output(result, hogql_context)
+                break
+            except PydanticOutputParserException as e:
+                prompt += f"\n\nWe've ran this prompt {i+1} time{'s' if i+1 > 1 else ''} now. The newly updated query gave us this error: {e.validation_message}"
+        else:
+            return "", None
+
+        return parsed_result, parsed_result
+
+    @property
+    def _model(self):
+        return ChatOpenAI(model="gpt-4o", temperature=0, disable_streaming=True).with_structured_output(
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: greptile-apps[bot]
body: syntax: Model name 'gpt-4o' appears to be incorrect - should likely be 'gpt-4'

```suggestion
        return ChatOpenAI(model="gpt-4", temperature=0, disable_streaming=True).with_structured_output(
```
created_at: 2025-04-29T11:09:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
+
+USER_PROMPT = """
+Fix the errors in the HogQL query below and only return the new updated query in your response.
+
+- Don't update any other part of the query if it's not relevant to the error, including rewriting shorthand clickhouse SQL to the full version.
+- Don't change the capitalisation of the query if it's not relevant to the error, such as rewriting `select` as `SELECT` or `from` as `FROM`
+- There may also be more than one error in the syntax.
+
+Below is the current HogQL query and the error message
+"""
+
+
+class HogQLQueryFixerTool(MaxTool):
+    name: str = "fix_hogql_query"
+    description: str = "Fixes any error in the current HogQL query"
+    thinking_message: str = "Fixing errors in the SQL query"
+    root_system_prompt_template: str = SQL_ASSISTANT_ROOT_SYSTEM_PROMPT
+
+    def _run_impl(self) -> tuple[str, str | None]:
+        database = create_hogql_database(self._team_id)
+        hogql_context = HogQLContext(team_id=self._team_id, enable_select_queries=True, database=database)
+
+        serialized_database = serialize_database(hogql_context)
+        schema_description = "\n\n".join(
+            (
+                f"Table `{table_name}` with fields:\n"
+                + "\n".join(f"- {field.name} ({field.type})" for field in table.fields.values())
+                for table_name, table in serialized_database.items()
+                # Only the most important core tables, plus all warehouse tables
+                if table_name in ["events", "groups", "persons"]
+                or table_name in database.get_warehouse_tables()
+                or table_name in database.get_views()
+            )
+        )
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    SYSTEM_PROMPT.replace("{{schema_description}}", schema_description),
+                ),
+                (
+                    "user",
+                    USER_PROMPT
+                    + "\n\n<hogql_query>"
+                    + "{{{hogql_query}}}"
+                    + "</hogql_query>"
+                    + "\n\n<error>"
+                    + "{{{error_message}}}"
+                    + "</error>",
+                ),
+            ],
+            template_format="mustache",
+        )
+
+        for i in range(3):
+            try:
+                chain = prompt | self._model
+                result = chain.invoke(self.context)
+                parsed_result = self._parse_output(result, hogql_context)
+                break
+            except PydanticOutputParserException as e:
+                prompt += f"\n\nWe've ran this prompt {i+1} time{'s' if i+1 > 1 else ''} now. The newly updated query gave us this error: {e.validation_message}"
+        else:
+            return "", None
+
+        return parsed_result, parsed_result
+
+    @property
+    def _model(self):
+        return ChatOpenAI(model="gpt-4o", temperature=0, disable_streaming=True).with_structured_output(
+            SQL_SCHEMA,
+            method="function_calling",
+            include_raw=False,
+        )
+
+    def _parse_output(self, output, hogql_context: HogQLContext):
+        result = parse_pydantic_structured_output(SchemaGeneratorOutput[str])(output)  # type: ignore
+        # We also ensure the generated SQL is valid
+        assert result.query is not None
+        try:
+            print_ast(parse_select(result.query), context=hogql_context, dialect="clickhouse")
+        except (ExposedHogQLError, ResolutionError) as err:
+            err_msg = str(err)
+            if err_msg.startswith("no viable alternative"):
+                # The "no viable alternative" ANTLR error is horribly unhelpful, both for humans and LLMs
+                err_msg = f'ANTLR parsing error: "no viable alternative at input". This means that the query isn\' valid HogQL. The last 5 characters where we tripped up were "{result.query[-5:]}".'
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: graphite-app[bot]
body: There's a syntax error in the f-string on line 169. The single quote in `isn\'` is incorrectly escaped, which will cause a syntax error when the code runs. It should be either `isn't` or properly escaped as `isn\\'t`.
```suggestion
                err_msg = f'ANTLR parsing error: "no viable alternative at input". This means that the query isn\'t valid HogQL. The last 5 characters where we tripped up were "{result.query[-5:]}".'
```
  

*Spotted by [Diamond](https://app.graphite.dev/diamond/?org=PostHog&ref=ai-review-comment)*<i class='graphite__hidden'><br /><br />Is this helpful? React 👍 or 👎 to let us know.</i>
created_at: 2025-04-29T11:09:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: rossgray
body: I guess because this is being called in a module-level attribute, this will be executed on import. I don't know how long `HogQLLexer.g4` takes to read but it could slow down app startup slightly, and perhaps slow down tests. Might just be worth double checking, or perhaps read this file the first time it is used?
created_at: 2025-04-29T12:53:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: rossgray
body: I haven't checked how it's been done for Max inside the SQL editor but isn't there going to be some crossover here? Maybe it's worth consolidating the section of the prompt for HogQL knowledge so things like 'Important HogQL differences versus other SQL dialects' can be updated in a single place?
created_at: 2025-04-29T12:57:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
+
+USER_PROMPT = """
+Fix the errors in the HogQL query below and only return the new updated query in your response.
+
+- Don't update any other part of the query if it's not relevant to the error, including rewriting shorthand clickhouse SQL to the full version.
+- Don't change the capitalisation of the query if it's not relevant to the error, such as rewriting `select` as `SELECT` or `from` as `FROM`
+- There may also be more than one error in the syntax.
+
+Below is the current HogQL query and the error message
+"""
+
+
+class HogQLQueryFixerTool(MaxTool):
+    name: str = "fix_hogql_query"
+    description: str = "Fixes any error in the current HogQL query"
+    thinking_message: str = "Fixing errors in the SQL query"
+    root_system_prompt_template: str = SQL_ASSISTANT_ROOT_SYSTEM_PROMPT
+
+    def _run_impl(self) -> tuple[str, str | None]:
+        database = create_hogql_database(self._team_id)
+        hogql_context = HogQLContext(team_id=self._team_id, enable_select_queries=True, database=database)
+
+        serialized_database = serialize_database(hogql_context)
+        schema_description = "\n\n".join(
+            (
+                f"Table `{table_name}` with fields:\n"
+                + "\n".join(f"- {field.name} ({field.type})" for field in table.fields.values())
+                for table_name, table in serialized_database.items()
+                # Only the most important core tables, plus all warehouse tables
+                if table_name in ["events", "groups", "persons"]
+                or table_name in database.get_warehouse_tables()
+                or table_name in database.get_views()
+            )
+        )
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    SYSTEM_PROMPT.replace("{{schema_description}}", schema_description),
+                ),
+                (
+                    "user",
+                    USER_PROMPT
+                    + "\n\n<hogql_query>"
+                    + "{{{hogql_query}}}"
+                    + "</hogql_query>"
+                    + "\n\n<error>"
+                    + "{{{error_message}}}"
+                    + "</error>",
+                ),
+            ],
+            template_format="mustache",
+        )
+
+        for i in range(3):
+            try:
+                chain = prompt | self._model
+                result = chain.invoke(self.context)
+                parsed_result = self._parse_output(result, hogql_context)
+                break
+            except PydanticOutputParserException as e:
+                prompt += f"\n\nWe've ran this prompt {i+1} time{'s' if i+1 > 1 else ''} now. The newly updated query gave us this error: {e.validation_message}"
+        else:
+            return "", None
+
+        return parsed_result, parsed_result
+
+    @property
+    def _model(self):
+        return ChatOpenAI(model="gpt-4o", temperature=0, disable_streaming=True).with_structured_output(
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: rossgray
body: nit: I think this is the pattern we have been following elsewhere but feels like things like the model should be defined in a constant somewhere as I imagine we'll need to update it across the app within the next few months
created_at: 2025-04-29T13:01:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
+
+USER_PROMPT = """
+Fix the errors in the HogQL query below and only return the new updated query in your response.
+
+- Don't update any other part of the query if it's not relevant to the error, including rewriting shorthand clickhouse SQL to the full version.
+- Don't change the capitalisation of the query if it's not relevant to the error, such as rewriting `select` as `SELECT` or `from` as `FROM`
+- There may also be more than one error in the syntax.
+
+Below is the current HogQL query and the error message
+"""
+
+
+class HogQLQueryFixerTool(MaxTool):
+    name: str = "fix_hogql_query"
+    description: str = "Fixes any error in the current HogQL query"
+    thinking_message: str = "Fixing errors in the SQL query"
+    root_system_prompt_template: str = SQL_ASSISTANT_ROOT_SYSTEM_PROMPT
+
+    def _run_impl(self) -> tuple[str, str | None]:
+        database = create_hogql_database(self._team_id)
+        hogql_context = HogQLContext(team_id=self._team_id, enable_select_queries=True, database=database)
+
+        serialized_database = serialize_database(hogql_context)
+        schema_description = "\n\n".join(
+            (
+                f"Table `{table_name}` with fields:\n"
+                + "\n".join(f"- {field.name} ({field.type})" for field in table.fields.values())
+                for table_name, table in serialized_database.items()
+                # Only the most important core tables, plus all warehouse tables
+                if table_name in ["events", "groups", "persons"]
+                or table_name in database.get_warehouse_tables()
+                or table_name in database.get_views()
+            )
+        )
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    SYSTEM_PROMPT.replace("{{schema_description}}", schema_description),
+                ),
+                (
+                    "user",
+                    USER_PROMPT
+                    + "\n\n<hogql_query>"
+                    + "{{{hogql_query}}}"
+                    + "</hogql_query>"
+                    + "\n\n<error>"
+                    + "{{{error_message}}}"
+                    + "</error>",
+                ),
+            ],
+            template_format="mustache",
+        )
+
+        for i in range(3):
+            try:
+                chain = prompt | self._model
+                result = chain.invoke(self.context)
+                parsed_result = self._parse_output(result, hogql_context)
+                break
+            except PydanticOutputParserException as e:
+                prompt += f"\n\nWe've ran this prompt {i+1} time{'s' if i+1 > 1 else ''} now. The newly updated query gave us this error: {e.validation_message}"
+        else:
+            return "", None
+
+        return parsed_result, parsed_result
+
+    @property
+    def _model(self):
+        return ChatOpenAI(model="gpt-4o", temperature=0, disable_streaming=True).with_structured_output(
+            SQL_SCHEMA,
+            method="function_calling",
+            include_raw=False,
+        )
+
+    def _parse_output(self, output, hogql_context: HogQLContext):
+        result = parse_pydantic_structured_output(SchemaGeneratorOutput[str])(output)  # type: ignore
+        # We also ensure the generated SQL is valid
+        assert result.query is not None
+        try:
+            print_ast(parse_select(result.query), context=hogql_context, dialect="clickhouse")
+        except (ExposedHogQLError, ResolutionError) as err:
+            err_msg = str(err)
+            if err_msg.startswith("no viable alternative"):
+                # The "no viable alternative" ANTLR error is horribly unhelpful, both for humans and LLMs
+                err_msg = f'ANTLR parsing error: "no viable alternative at input". This means that the query isn\' valid HogQL. The last 5 characters where we tripped up were "{result.query[-5:]}".'
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: rossgray
body: sounds valid
created_at: 2025-04-29T13:02:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
+
+USER_PROMPT = """
+Fix the errors in the HogQL query below and only return the new updated query in your response.
+
+- Don't update any other part of the query if it's not relevant to the error, including rewriting shorthand clickhouse SQL to the full version.
+- Don't change the capitalisation of the query if it's not relevant to the error, such as rewriting `select` as `SELECT` or `from` as `FROM`
+- There may also be more than one error in the syntax.
+
+Below is the current HogQL query and the error message
+"""
+
+
+class HogQLQueryFixerTool(MaxTool):
+    name: str = "fix_hogql_query"
+    description: str = "Fixes any error in the current HogQL query"
+    thinking_message: str = "Fixing errors in the SQL query"
+    root_system_prompt_template: str = SQL_ASSISTANT_ROOT_SYSTEM_PROMPT
+
+    def _run_impl(self) -> tuple[str, str | None]:
+        database = create_hogql_database(self._team_id)
+        hogql_context = HogQLContext(team_id=self._team_id, enable_select_queries=True, database=database)
+
+        serialized_database = serialize_database(hogql_context)
+        schema_description = "\n\n".join(
+            (
+                f"Table `{table_name}` with fields:\n"
+                + "\n".join(f"- {field.name} ({field.type})" for field in table.fields.values())
+                for table_name, table in serialized_database.items()
+                # Only the most important core tables, plus all warehouse tables
+                if table_name in ["events", "groups", "persons"]
+                or table_name in database.get_warehouse_tables()
+                or table_name in database.get_views()
+            )
+        )
+
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                (
+                    "system",
+                    SYSTEM_PROMPT.replace("{{schema_description}}", schema_description),
+                ),
+                (
+                    "user",
+                    USER_PROMPT
+                    + "\n\n<hogql_query>"
+                    + "{{{hogql_query}}}"
+                    + "</hogql_query>"
+                    + "\n\n<error>"
+                    + "{{{error_message}}}"
+                    + "</error>",
+                ),
+            ],
+            template_format="mustache",
+        )
+
+        for i in range(3):
+            try:
+                chain = prompt | self._model
+                result = chain.invoke(self.context)
+                parsed_result = self._parse_output(result, hogql_context)
+                break
+            except PydanticOutputParserException as e:
+                prompt += f"\n\nWe've ran this prompt {i+1} time{'s' if i+1 > 1 else ''} now. The newly updated query gave us this error: {e.validation_message}"
+        else:
+            return "", None
+
+        return parsed_result, parsed_result
+
+    @property
+    def _model(self):
+        return ChatOpenAI(model="gpt-4o", temperature=0, disable_streaming=True).with_structured_output(
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: Gilbert09
body: Yep - good nit
created_at: 2025-04-29T13:45:55Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
+</hogql_parser>
+
+<hogql_lexer>
+{get_lexer_contents()}
+</hogql_lexer>
+
+You fix HogQL errors that may come from either HogQL resolver errors or clickhouse execution errors. You don't help with other knowledge.
+
+Important HogQL differences versus other SQL dialects:
+- JSON properties are accessed like `properties.foo.bar` instead of `properties->foo->bar`
+
+This is the schema of tables available in queries:
+
+{{schema_description}}
+
+Person or event metadata unspecified above (emails, names, etc.) is stored in `properties` fields, accessed like: `properties.foo.bar`.
+Note: "persons" means "users" here - instead of a "users" table, we have a "persons" table.
+
+Standardized events/properties such as pageview or screen start with `$`. Custom events/properties start with any other character.
+
+`virtual_table` and `lazy_table` fields are connections to linked tables, e.g. the virtual table field `person` allows accessing person properties like so: `person.properties.foo`.
+""".strip()
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: Gilbert09
body: Yeah, some of this is copied from the tool for generating queries - we should definitely centralize it. We don't really know how much is too much knowledge just yet though 
created_at: 2025-04-29T13:46:31Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,174 @@
+import os
+from ee.hogai.tool import MaxTool
+from posthog.hogql.database.database import create_hogql_database, serialize_database
+from posthog.hogql.context import HogQLContext
+from ee.hogai.graph.sql.toolkit import SQL_SCHEMA
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+from ee.hogai.graph.schema_generator.parsers import PydanticOutputParserException, parse_pydantic_structured_output
+from ee.hogai.graph.schema_generator.utils import SchemaGeneratorOutput
+
+from posthog.hogql.errors import ExposedHogQLError, ResolutionError
+from posthog.hogql.parser import parse_select
+from posthog.hogql.printer import print_ast
+
+
+_parser_contents: str | None = None
+_lexer_contents: str | None = None
+
+
+def get_parser_contents() -> str:
+    global _parser_contents
+
+    if _parser_contents is not None:
+        return _parser_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLParser.g4") as parser_file:
+        _parser_contents = parser_file.read()
+
+    return _parser_contents
+
+
+def get_lexer_contents() -> str:
+    global _lexer_contents
+
+    if _lexer_contents is not None:
+        return _lexer_contents
+
+    with open(f"{os.getcwd()}/posthog/hogql/grammar/HogQLLexer.g4") as parser_file:
+        _lexer_contents = parser_file.read()
+
+    return _lexer_contents
+
+
+SQL_ASSISTANT_ROOT_SYSTEM_PROMPT = """
+The user has written a HogQL query which contains error. They expect your help with tweaking the HogQL to fix these errors.
+
+IMPORTANT: This is currently your primary task. Therefore `fix_hogql_query` is currently your primary tool.
+Use `fix_hogql_query` when fixing any errors remotely related to HogQL.
+It's very important to disregard other tools for these purposes - the user expects `fix_hogql_query`.
+
+NOTE: When calling the `fix_hogql_query` tool, do not provide any response other than the tool call.
+"""
+
+SYSTEM_PROMPT = f"""
+HogQL is PostHog's variant of SQL. HogQL is a transpiler that outputs Clickhouse SQL. We use Antlr4 to define the HogQL language.
+Below is the antlr parser and lexer definitions - when writing HogQL, ensure you follow the grammar rules.
+
+<hogql_parser>
+{get_parser_contents()}
path: products/data_warehouse/backend/hogql_fixer_ai.py
user: login: Gilbert09
body: I think this will be a nominal amount of extra start-up time. The files being read are pretty small and are on the local filesystem
created_at: 2025-04-29T13:47:18Z
author_association: MEMBER

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-29T09:29:56Z
body: **Size Change:** +895 B (+0.01%) 

**Total Size:** 13.6 MB



<details><summary>ℹ️ <strong>View Unchanged</strong></summary>

| Filename | Size | Change |
| :--- | :---: | :---: |
| `frontend/dist/toolbar.js` | 13.6 MB | +895 B (+0.01%) |

</details>



<a href="https://github.com/preactjs/compressed-size-action"><sub>compressed-size-action</sub></a>
author_association: CONTRIBUTOR
user: login: posthog-bot
created_at: 2025-04-29T11:09:34Z
body: ## 📸 UI snapshots have been updated

**1** snapshot changes in total. **0** added, **1** modified, **0** deleted:

- **`chromium`**: **0** added, **1** modified, **0** deleted ([diff for shard 5](https://github.com/PostHog/posthog/pull/31690/commits/247835884cbdd6ed09a1d1281d810c84d1a2f2fa))
- `webkit`: **0** added, **0** modified, **0** deleted

Triggered by [this commit](https://github.com/PostHog/posthog/pull/31690/commits/74aefd6be2fb5a2a8e8aa4c3068c2339ed116f38).

👉 **[Review this PR's diff of snapshots.](https://github.com/PostHog/posthog/pull/31690/files#:~:text=frontend/__snapshots__/)**
author_association: CONTRIBUTOR
user: login: posthog-bot
created_at: 2025-04-29T11:18:46Z
body: ## 📸 UI snapshots have been updated

**1** snapshot changes in total. **0** added, **1** modified, **0** deleted:

- **`chromium`**: **0** added, **1** modified, **0** deleted ([diff for shard 5](https://github.com/PostHog/posthog/pull/31690/commits/d95717a980ebab0d53a6fe4d309df9196ef275c1))
- `webkit`: **0** added, **0** modified, **0** deleted

Triggered by [this commit](https://github.com/PostHog/posthog/pull/31690/commits/2ebc5a8185a543535aeb5795ffaabd14d52e72d6).

👉 **[Review this PR's diff of snapshots.](https://github.com/PostHog/posthog/pull/31690/files#:~:text=frontend/__snapshots__/)**
author_association: CONTRIBUTOR
user: login: Gilbert09
created_at: 2025-04-29T13:49:09Z
body: Thanks, Ross, I do agree with everything here. There's some intentional separation going on here as directed by the Max AI team as they roll out more support for doing things like "single tool calls" - this is why we have a dedicated API endpoint for this and a separate frontend logic right now. Hopefully these can be removed soon once the extra support has been rolled out
author_association: MEMBER