Merge pull request #157 from bagofwords1/0.0.345-pre

0.0.345 pre

Author: yochze

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Release Notes
 
+## Version 0.0.345 (March 24, 2026)
+- Make test_connection and other data client utils async calls
+
 ## Version 0.0.344 (March 23, 2026)
 - Fix artifact sandbox: download React development builds in vendor script
 - Remove CDN fallbacks for airgapped deployments — missing vendored libs now fail loudly

backend/app/ai/context/builders/message_context_builder.py

@@ -314,6 +314,9 @@ async def build_context(
                                         digest_parts.append(f"mode: {rj.get('mode')}")
                                     if rj.get('artifact_id'):
                                         digest_parts.append(f"artifact_id: {rj.get('artifact_id')}")
+                                    viz_ids = rj.get('visualization_ids') or []
+                                    if viz_ids:
+                                        digest_parts.append(f"viz_ids: {', '.join(viz_ids)}")
                                     if rj.get('version'):
                                         digest_parts.append(f"v{rj.get('version')}")
                                     if digest_parts:
@@ -762,6 +765,9 @@ async def build(
                                     digest_parts.append(f"mode: {rj.get('mode')}")
                                 if rj.get('artifact_id'):
                                     digest_parts.append(f"artifact_id: {rj.get('artifact_id')}")
+                                viz_ids = rj.get('visualization_ids') or []
+                                if viz_ids:
+                                    digest_parts.append(f"viz_ids: {', '.join(viz_ids)}")
                                 if rj.get('version'):
                                     digest_parts.append(f"v{rj.get('version')}")
                                 if digest_parts:

backend/app/ai/tools/implementations/edit_artifact.py

@@ -297,6 +297,7 @@ def metadata(self) -> ToolMetadata:
             name="edit_artifact",
             description=(
                 "Edit an existing dashboard or artifact by applying targeted changes to its code. "
+                "Prioritize using read_artifact before editing an artifact"
                 "Use this instead of create_artifact when modifying an existing artifact's layout, styling, "
                 "filters, charts, or content. Preserves the existing design and only changes what is requested. "
                 "If the edit is adding a new visualization, you MUST ADD it as a parameter to the tool"

backend/app/ai/tools/implementations/read_artifact.py

@@ -31,8 +31,9 @@ def metadata(self) -> ToolMetadata:
             name="read_artifact",
             description=(
                 "Read an existing dashboard, slides, and artifact's code and metadata from the current report. "
-                "Use this to load previous artifact code into context before modifying with create_artifact or when the user wants to inspect or analyze an existing artifact. "
-                "Alwys use this before editing an artifact do have a full view of the existing code, visualizations, and layout. "
+                "Use this to load previous artifact code into context before modifying with edit_artifact (or create_artifact) or when the user wants to inspect or analyze an existing artifact. "
+                "ALWAYS use this before editing an artifact (edit_artifact) to have a full view of the existing code, visualizations, and layout. "
+                "If the user refers to a specific version of an artifact, ALWAYS load that version with this tool to have the correct code context for the edit. "
                 "IMPORTANT: The artifact_id is found in previous create_artifact results shown as 'artifact_id: <uuid>' in the conversation. "
                 "Do NOT ask the user for URLs or artifact IDs - extract the artifact_id from the conversation context."
             ),

backend/app/ai/tools/schemas/create_artifact.py

@@ -30,4 +30,5 @@ class CreateArtifactOutput(BaseModel):
     code: str = Field(..., description="The generated React/JSX code")
     mode: str = Field(..., description="Artifact mode, eiither 'page' for dashboards/reports or 'slides' for presentation, deck or powerpoint export")
     title: Optional[str] = Field(None, description="Artifact title")
+    visualization_ids: List[str] = Field(default_factory=list, description="All visualization IDs included in this artifact. Use these when making further edits with edit_artifact.")
     version: int = Field(default=1, description="Version number of the artifact")

backend/app/ai/tools/schemas/edit_artifact.py

@@ -13,7 +13,7 @@ class EditArtifactInput(BaseModel):
 
     artifact_id: str = Field(..., description="ID of the existing artifact to edit. Find this in previous create_artifact or read_artifact results as 'artifact_id: <uuid>' in the conversation.")
     edit_prompt: str = Field(..., description="Natural language description of the change to make. Be specific about what to change and how. E.g., 'Remove the filter bar', 'Make the revenue chart blue', 'Add a KPI card for total users'. Also use this to fix visual issues from a previous create_artifact (e.g., 'the bar chart is cut off on the right side', 'KPI cards are overlapping'). If adding new visualizations, describe where they should go in the layout.")
-    visualization_ids: Optional[List[str]] = Field(default=None, description="Optional list of NEW visualization IDs to add to the artifact (from create_data results). Existing visualization IDs are kept automatically. Provide this when the edit adds a new chart or data source that wasn't in the original artifact.")
+    visualization_ids: Optional[List[str]] = Field(default=None, description="List of NEW visualization IDs to include in the artifact. IMPORTANT: If you called create_data before this edit, you MUST pass the resulting visualization_id(s) here. Without them, the new visualizations will not appear in the dashboard. Existing visualization IDs from the original artifact are kept automatically — only pass new ones.")
     title: Optional[str] = Field(default=None, description="Updated title for the artifact. If not provided, the existing title is kept.")
 
 
@@ -32,5 +32,6 @@ class EditArtifactOutput(BaseModel):
     code: str = Field(..., description="The updated code after applying the edit")
     mode: str = Field(..., description="Artifact mode: 'page' or 'slides'")
     title: Optional[str] = Field(None, description="Artifact title")
+    visualization_ids: List[str] = Field(default_factory=list, description="All visualization IDs included in this artifact after the edit. Use these when making further edits.")
     version: int = Field(..., description="Bumped version number of the artifact")
     diff_applied: bool = Field(..., description="True if the edit was applied as a surgical search/replace diff. False if the tool fell back to a full code rewrite.")

backend/app/data_sources/clients/base.py

@@ -1,3 +1,4 @@
+import asyncio
 from abc import ABC, abstractmethod
 
 
@@ -33,3 +34,20 @@ def prompt_schema(self):
     @abstractmethod
     def execute_query(self, **kwargs):
         pass
+
+    # Async wrappers — offload blocking I/O to a thread so the event loop stays free.
+
+    async def atest_connection(self):
+        return await asyncio.to_thread(self.test_connection)
+
+    async def aget_schemas(self):
+        return await asyncio.to_thread(self.get_schemas)
+
+    async def aget_schema(self, table_name):
+        return await asyncio.to_thread(self.get_schema, table_name)
+
+    async def aprompt_schema(self):
+        return await asyncio.to_thread(self.prompt_schema)
+
+    async def aexecute_query(self, *args, **kwargs):
+        return await asyncio.to_thread(self.execute_query, *args, **kwargs)

backend/app/services/completion_service.py

@@ -2,6 +2,7 @@
 from fastapi.responses import StreamingResponse
 import json
 import logging
+import time
 from datetime import datetime
 from types import SimpleNamespace
 from uuid import uuid4
@@ -137,6 +138,9 @@ async def _get_instruction_suggestions_for_completion(
 tracer = get_tracer(__name__)
 
 
+logger = logging.getLogger(__name__)
+
+
 class CompletionService:
 
     def __init__(self):
@@ -1552,11 +1556,21 @@ async def _create_completion_stream_traced(
         build_id: str = None,
     ):
         try:
+            t0 = time.monotonic()
+            rid = str(report_id)[:8]
+
+            def _log(label):
+                elapsed = (time.monotonic() - t0) * 1000
+                logger.info(f"[stream:{rid}] {label} +{elapsed:.0f}ms")
+
+            _log("stream_start")
+
             # Validate report exists (same as regular create_completion)
             result = await db.execute(select(Report).filter(Report.id == report_id))
             report = result.scalar_one_or_none()
             if not report:
                 raise HTTPException(status_code=404, detail="Report not found")
+            _log("report_fetched")
 
             # Validate widget if provided
             if completion_data.prompt.widget_id:
@@ -1577,6 +1591,7 @@ async def _create_completion_stream_traced(
                 step = None
 
             span.add_event("validation_done")
+            _log("validation_done")
 
             # Get default model
             if completion_data.prompt and completion_data.prompt.model_id:
@@ -1589,11 +1604,13 @@ async def _create_completion_stream_traced(
                     status_code=400,
                     detail="No default LLM model configured. Please go to Settings > LLM and set a default model."
                 )
+            _log("model_resolved")
 
             small_model = await self.llm_service.get_default_model(db, organization, current_user, is_small=True)
             # Fallback: if no small model configured, use the main model
             if not small_model:
                 small_model = model
+            _log("small_model_resolved")
 
             span.set_attribute("llm.model_id", model.model_id)
 
@@ -1646,19 +1663,22 @@ async def _create_completion_stream_traced(
                     status_code=500,
                     detail=f"Failed to save completions: {str(e)}"
                 )
+            _log("db_commit_flush")
 
             span.set_attribute("completion.head_id", str(completion.id))
             span.set_attribute("completion.system_id", str(system_completion.id))
             span.add_event("completions_saved")
 
             # Mark image files with this completion_id (so they show attached to this message)
             await self._mark_images_with_completion(db, report.id, str(completion.id))
+            _log("images_marked")
 
             # Store mentions associated with the user completion (best-effort, non-blocking)
             try:
                 await self.mention_service.create_completion_mentions(db, completion)
             except Exception as e:
-                logging.error(f"Failed to create mentions for completion {completion.id}: {e}")
+                logger.error(f"Failed to create mentions for completion {completion.id}: {e}")
+            _log("mentions_created")
 
             # Audit log
             try:
@@ -1673,31 +1693,43 @@ async def _create_completion_stream_traced(
                 )
             except Exception:
                 pass
+            _log("audit_logged")
 
             org_settings = await organization.get_settings(db)
+            _log("org_settings_fetched")
             resolved_build_id = await self._resolve_build_id(db, organization, build_id)
+            _log("build_id_resolved")
 
             # Create event queue for streaming
             event_queue = CompletionEventQueue()
 
             async def run_agent_with_streaming():
                 """Run agent in background and stream events."""
+                at0 = time.monotonic()
+
+                def _alog(label):
+                    elapsed = (time.monotonic() - at0) * 1000
+                    logger.info(f"[stream:{rid}:agent] {label} +{elapsed:.0f}ms")
+
                 with tracer.start_as_current_span("completion.stream_agent_execution") as agent_span:
                     agent_span.set_attribute("report.id", str(report.id))
                     agent_span.set_attribute("completion.system_id", str(system_completion.id))
                     agent_span.set_attribute("llm.model_id", model.model_id)
                     async_session = create_async_session_factory()
                     async with async_session() as session:
                         try:
+                            _alog("session_opened")
+
                             # Re-fetch all database-dependent objects using the new session
                             report_obj = await session.get(Report, report.id)
                             completion_obj = await session.get(Completion, completion.id)
                             system_completion_obj = await session.get(Completion, system_completion.id)
                             widget_obj = await session.get(Widget, widget.id) if widget else None
                             step_obj = await session.get(Step, step.id) if step else None
+                            _alog("objects_refetched")
 
                             if not all([report_obj, completion_obj, system_completion_obj]):
-                                logging.error("Failed to fetch necessary objects for streaming agent.")
+                                logger.error("Failed to fetch necessary objects for streaming agent.")
                                 error_event = SSEEvent(
                                     event="completion.error",
                                     completion_id=str(system_completion.id),
@@ -1712,10 +1744,12 @@ async def run_agent_with_streaming():
                                     ds_clients = await self.data_source_service.construct_clients(session, data_source, current_user)
                                     clients.update(ds_clients)
                                 clients_span.set_attribute("data_sources.count", len(report_obj.data_sources))
+                            _alog(f"clients_constructed count={len(report_obj.data_sources)}")
 
                             # Pre-load files relationship in async context to avoid greenlet error in AgentV2.__init__
                             # (AgentV2.__init__ is synchronous, so lazy-loading files there would fail)
                             _ = report_obj.files
+                            _alog("files_preloaded")
 
                             # Create agent with event queue
                             agent = AgentV2(
@@ -1735,6 +1769,7 @@ async def run_agent_with_streaming():
                                 clients=clients,
                                 build_id=resolved_build_id,
                             )
+                            _alog("agent_initialized")
 
                             # Emit telemetry: stream started
                             try:
@@ -1754,9 +1789,11 @@ async def run_agent_with_streaming():
 
                             # Run agent execution
                             agent_span.add_event("agent_execution_started")
+                            _alog("agent_execution_start")
                             with tracer.start_as_current_span("completion.agent_execution"):
                                 await agent.main_execution()
                             agent_span.add_event("agent_execution_finished")
+                            _alog("agent_execution_done")
 
                             # Send completion finished event
                             finished_event = SSEEvent(
@@ -1765,6 +1802,7 @@ async def run_agent_with_streaming():
                                 data={"status": "success"}
                             )
                             await event_queue.put(finished_event)
+                            _alog("queue_finished")
 
                             # Emit telemetry: stream completed
                             try:
@@ -1783,7 +1821,8 @@ async def run_agent_with_streaming():
                         except Exception as e:
                             agent_span.set_status(StatusCode.ERROR, str(e))
                             agent_span.record_exception(e)
-                            logging.error(f"Agent streaming execution failed: {e}")
+                            _alog(f"agent_execution_error error={type(e).__name__}: {e}")
+                            logger.error(f"Agent streaming execution failed: {e}")
                             # Send error event
                             error_event = SSEEvent(
                                 event="completion.error",
@@ -1826,6 +1865,7 @@ async def run_agent_with_streaming():
 
             # Start agent execution in background
             asyncio.create_task(run_agent_with_streaming())
+            _log("task_spawned")
 
             # Stream events
             async def completion_stream_generator():
@@ -1873,11 +1913,11 @@ async def completion_stream_generator():
 
         except HTTPException as he:
             # Log the error and re-raise HTTP exceptions
-            logging.error(f"HTTP Exception in create_completion_stream: {str(he)}")
+            logger.error(f"HTTP Exception in create_completion_stream: {str(he)}")
             raise he
         except Exception as e:
             # Log and convert unexpected errors to HTTP exceptions
-            logging.error(f"Unexpected error in create_completion_stream: {str(e)}")
+            logger.error(f"Unexpected error in create_completion_stream: {str(e)}")
             raise HTTPException(
                 status_code=500,
                 detail=f"Unexpected error: {str(e)}"