chore(wren-ai-service): add correlation ID to langfuse (#846)

Co-authored-by: Chih-Yu Yeh <chihyu.jimmy.yeh@gmail.com>

wren-ai-service/src/pipelines/common.py

@@ -90,14 +90,14 @@ async def _check_if_sql_executable(
         project_id: str | None = None,
     ):
         async with aiohttp.ClientSession() as session:
-            status, _, error = await self._engine.execute_sql(
+            status, _, addition = await self._engine.execute_sql(
                 sql,
                 session,
                 project_id=project_id,
             )
 
         if not status:
-            logger.exception(f"SQL is not executable: {error}")
+            logger.exception(f"SQL is not executable: {addition["error_message"]}")
 
         return status
 
@@ -165,22 +165,24 @@ async def _task(result: Dict[str, str]):
             quoted_sql, no_error = add_quotes(result["sql"])
 
             if no_error:
-                status, _, error = await self._engine.execute_sql(
+                status, _, addition = await self._engine.execute_sql(
                     quoted_sql, session, project_id=project_id
                 )
 
                 if status:
                     valid_generation_results.append(
                         {
                             "sql": quoted_sql,
+                            "correlation_id": addition["correlation_id"],
                         }
                     )
                 else:
                     invalid_generation_results.append(
                         {
                             "sql": quoted_sql,
                             "type": "DRY_RUN",
-                            "error": error,
+                            "error": addition["error_message"],
+                            "correlation_id": addition["correlation_id"],
                         }
                     )
             else:

wren-ai-service/src/providers/engine/wren.py

@@ -53,11 +53,16 @@ async def execute_sql(
             ) as response:
                 res = await response.json()
                 if data := res.get("data"):
-                    return True, data, None
+                    return True, data, {
+                        "correlation_id": res.get("correlationId"),
+                    }
                 return (
                     False,
                     None,
-                    res.get("errors", [{}])[0].get("message", "Unknown error"),
+                    {
+                        "error_message": res.get("errors", [{}])[0].get("message", "Unknown error"),
+                        "correlation_id": res.get("extensions", {}).get("other", {}).get("correlationId"),
+                    }
                 )
         except asyncio.TimeoutError:
             return False, None, f"Request timed out: {timeout} seconds"

wren-ui/src/apollo/server/resolvers/modelResolver.ts

@@ -933,14 +933,13 @@ export class ModelResolver {
       ? await ctx.projectService.getProjectById(projectId)
       : await ctx.projectService.getCurrentProject();
     const { manifest } = await ctx.deployService.getLastDeployment(project.id);
-    const previewRes = await ctx.queryService.preview(sql, {
+    return await ctx.queryService.preview(sql, {
       project,
       limit: limit,
       modelingOnly: false,
       manifest,
       dryRun,
     });
-    return dryRun ? { dryRun: 'success' } : previewRes;
   }
 
   public async getNativeSql(

wren-ui/src/apollo/server/services/queryService.ts

@@ -25,9 +25,7 @@ export interface ColumnMetadata {
   type: string;
 }
 
-export interface PreviewDataResponse {
-  correlationId?: string;
-  processTime?: string;
+export interface PreviewDataResponse extends IbisResponse {
   columns: ColumnMetadata[];
   data: any[][];
 }
@@ -60,7 +58,7 @@ export interface IQueryService {
   preview(
     sql: string,
     options: PreviewOptions,
-  ): Promise<PreviewDataResponse | boolean>;
+  ): Promise<IbisResponse| PreviewDataResponse | boolean>;
 
   describeStatement(
     sql: string,
@@ -97,7 +95,7 @@ export class QueryService implements IQueryService {
   public async preview(
     sql: string,
     options: PreviewOptions,
-  ): Promise<PreviewDataResponse | boolean> {
+  ): Promise<IbisResponse | PreviewDataResponse | boolean> {
     const { project, manifest: mdl, limit, dryRun } = options;
     const { type: dataSource, connectionInfo } = project;
     if (this.useEngine(dataSource)) {
@@ -117,8 +115,7 @@ export class QueryService implements IQueryService {
       this.checkDataSourceIsSupported(dataSource);
       logger.debug('Use ibis adaptor to preview');
       if (dryRun) {
-        await this.ibisDryRun(sql, dataSource, connectionInfo, mdl);
-        return true;
+        return await this.ibisDryRun(sql, dataSource, connectionInfo, mdl);
       } else {
         return await this.ibisQuery(sql, dataSource, connectionInfo, mdl, limit);
       }
@@ -178,7 +175,7 @@ export class QueryService implements IQueryService {
     dataSource: DataSourceName,
     connectionInfo: any,
     mdl: Manifest
-  ) {
+  ): Promise<IbisResponse> {
     const event = TelemetryEvent.IBIS_DRY_RUN;
     try {
       const res = await this.ibisAdaptor.dryRun(sql, {
@@ -187,6 +184,9 @@ export class QueryService implements IQueryService {
         mdl,
       });
       this.sendIbisEvent(event, res, { dataSource, sql });
+      return {
+        correlationId: res.correlationId,
+      }
     } catch (err: any) {
       this.sendIbisFailedEvent(event, err, { dataSource, sql });
       throw err;
@@ -209,7 +209,11 @@ export class QueryService implements IQueryService {
         limit,
       });
       this.sendIbisEvent(event, res, { dataSource, sql });
-      return this.transformDataType(res);
+      const data = this.transformDataType(res);
+      return { 
+        correlationId: res.correlationId,
+        ...data
+      };
     } catch (err: any) {
       this.sendIbisFailedEvent(event, err, { dataSource, sql });
       throw err;