diff --git a/lumen/ai/__init__.py b/lumen/ai/__init__.py index 20a0f4b48..3ccaf49d9 100644 --- a/lumen/ai/__init__.py +++ b/lumen/ai/__init__.py @@ -3,7 +3,6 @@ from . import agents, embeddings, llm # noqa from .analysis import Analysis # noqa from .coordinator import Coordinator, DependencyResolver, Planner # noqa -from .memory import memory # noqa from .tools import IterativeTableLookup # noqa from .ui import ChatUI, ExplorerUI # noqa from .vector_store import DuckDBVectorStore, NumpyVectorStore # noqa diff --git a/lumen/ai/actor.py b/lumen/ai/actor.py index 691964dfe..0b0baf92b 100644 --- a/lumen/ai/actor.py +++ b/lumen/ai/actor.py @@ -13,8 +13,8 @@ from pydantic import BaseModel from .config import PROMPTS_DIR, SOURCE_TABLE_SEPARATOR +from .context import ContextModel, TContext from .llm import Llm, Message -from .memory import _Memory, memory from .utils import ( class_name_to_llm_spec_key, log_debug, render_template, warn_on_unused_variables, wrap_logfire_on_method, @@ -130,16 +130,18 @@ def _add_step(self, title: str = "", **kwargs): kwargs['steps_layout'] = self.steps_layout return nullcontext(self._null_step) if self.interface is None else self.interface.add_step(title=title, **kwargs) - async def _gather_prompt_context(self, prompt_name: str, messages: list[Message], **context): + async def _gather_prompt_context(self, prompt_name: str, messages: list[Message], context: TContext, **kwargs): """Gather context for the prompt template.""" - context["current_datetime"] = datetime.datetime.now() - return context + prompt_context = dict(kwargs) + prompt_context["memory"] = context + prompt_context["current_datetime"] = datetime.datetime.now() + return prompt_context - async def _render_prompt(self, prompt_name: str, messages: list[Message], **context) -> str: + async def _render_prompt(self, prompt_name: str, messages: list[Message], context: TContext, **kwargs) -> str: """Render a prompt template with context.""" prompt_template = self._lookup_prompt_key(prompt_name, "template") overrides = self.template_overrides.get(prompt_name, {}) - context = await self._gather_prompt_context(prompt_name, messages, **context) + prompt_context = await self._gather_prompt_context(prompt_name, messages, context, **kwargs) prompt_label = f"\033[92m{self.name}.prompts['{prompt_name}']['template']\033[0m" try: @@ -148,7 +150,7 @@ async def _render_prompt(self, prompt_name: str, messages: list[Message], **cont path_exists = False if isinstance(prompt_template, str) and not path_exists: # check if all the format_kwargs keys are contained in prompt_template - format_kwargs = dict(**overrides, **context) + format_kwargs = dict(**overrides, **prompt_context) warn_on_unused_variables(prompt_template, format_kwargs, prompt_label) try: prompt = prompt_template.format(**format_kwargs) @@ -161,7 +163,7 @@ async def _render_prompt(self, prompt_name: str, messages: list[Message], **cont prompt = render_template( prompt_template, overrides=overrides, - **context + **prompt_context ) prompt = prompt.strip() log_debug(f"{prompt_label}:\n\033[90m{prompt}\033[0m", show_length=True) @@ -171,9 +173,10 @@ async def _invoke_prompt( self, prompt_name: str, messages: list[Message], + context: TContext, response_model: type[BaseModel] | None = None, model_spec: str | None = None, - **context + **kwargs ) -> Any: """ Render a prompt and invoke the LLM. @@ -184,21 +187,23 @@ async def _invoke_prompt( Name of the prompt template to use messages : list[Message] The conversation messages + context : TContext + The context dictionary response_model : type[BaseModel], optional Pydantic model to structure the response model_spec : str, optional Specification for which LLM to use - **context : dict + **kwargs : dict Additional context variables for the prompt template Returns ------- The structured response from the LLM """ - system = await self._render_prompt(prompt_name, messages, **context) + system = await self._render_prompt(prompt_name, messages, context, **kwargs) if response_model is None: try: - response_model = self._get_model(prompt_name, **context) + response_model = self._get_model(prompt_name, **kwargs) except (KeyError, AttributeError): pass @@ -220,10 +225,11 @@ async def _stream_prompt( self, prompt_name: str, messages: list[Message], + context: TContext, response_model: type[BaseModel] | None = None, model_spec: str | None = None, field: str | None = None, - **context + **kwargs ): """ Render a prompt and stream responses from the LLM. @@ -234,13 +240,15 @@ async def _stream_prompt( Name of the prompt template to use messages : list[Message] The conversation messages + context : TContext + The context dictionary response_model : type[BaseModel], optional Pydantic model to structure the response model_spec : str, optional Specification for which LLM to use field : str, optional Specific field to extract from the response model - **context : dict + **kwargs : dict Additional context variables for the prompt template Yields @@ -248,12 +256,12 @@ async def _stream_prompt( Chunks of the response from the LLM as they are generated """ # Render the prompt - system = await self._render_prompt(prompt_name, messages, **context) + system = await self._render_prompt(prompt_name, messages, context, **kwargs) # Determine the response model if response_model is None: try: - response_model = self._get_model(prompt_name, **context) + response_model = self._get_model(prompt_name, **kwargs) except (KeyError, AttributeError): pass @@ -283,15 +291,12 @@ def llm_spec_key(self): return class_name_to_llm_spec_key(type(self).__name__) + class Actor(LLMUser): interface = param.ClassSelector(class_=ChatFeed, doc=""" The interface for the Coordinator to interact with.""") - memory = param.ClassSelector(class_=_Memory, default=None, doc=""" - Local memory which will be used to provide the agent context. - If None the global memory will be used.""") - def __init__(self, **params): super().__init__(**params) self._null_step = NullStep() @@ -311,22 +316,21 @@ def _add_step(self, title: str = "", **kwargs): """ return nullcontext(self._null_step) if self.interface is None else self.interface.add_step(title=title, **kwargs) - @property - def _memory(self) -> _Memory: - return memory if self.memory is None else self.memory - - async def _gather_prompt_context(self, prompt_name: str, messages: list[Message], **context): + async def _gather_prompt_context(self, prompt_name: str, messages: list[Message], context: TContext, **kwargs): """Gather context for the prompt template.""" - context = await super()._gather_prompt_context(prompt_name, messages, **context) - context["memory"] = self._memory + context = await super()._gather_prompt_context(prompt_name, messages, context, **kwargs) context["actor_name"] = self.name context["source_table_sep"] = SOURCE_TABLE_SEPARATOR return context @abstractmethod - async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> Any: + async def respond( + self, messages: list[Message], context: TContext, **kwargs: dict[str, Any] + ) -> tuple[list[Any], ContextModel]: """ - Responds to the provided messages. + Responds to the provided messages and context, returning + a list of visual outputs and a ContextModel containing + the context for subsequent Actors. """ @@ -353,25 +357,25 @@ class ContextProvider(param.Parameterized): not_with = param.List(default=[], doc=""" List of actors that this actor should not be invoked with.""") - provides = param.List(default=[], readonly=True, doc=""" - List of context values it provides to current working memory.""") - purpose = param.String(default="", doc=""" A descriptive statement of this actor's functionality and capabilities. Serves as a high-level explanation for other actors to understand what this actor does and when it might be useful to invoke it.""") - requires = param.List(default=[], readonly=True, doc=""" - List of context values it requires to be in memory.""") + input_schema: type[ContextModel] = ContextModel + output_schema: type[ContextModel] = ContextModel + + async def prepare(self, context: TContext): + pass async def requirements(self, messages: list[Message]) -> list[str]: - return self.requires + return list(self.input_schema.__annotations__) def __str__(self): string = ( f"- {self.name[:-5]}: {' '.join(self.purpose.strip().split())}\n" - f" Requires: `{'`, `'.join(self.requires)}`\n" - f" Provides: `{'`, `'.join(self.provides)}`\n" + f" Requires: `{'`, `'.join(self.input_schema.__annotations__)}`\n" + f" Provides: `{'`, `'.join(self.output_schema.__annotations__)}`\n" ) if self.conditions: string += " Conditions:\n" + "\n".join(f" - {condition}" for condition in self.conditions) + "\n" diff --git a/lumen/ai/agents.py b/lumen/ai/agents.py index 6dafca8c1..9b4a78d04 100644 --- a/lumen/ai/agents.py +++ b/lumen/ai/agents.py @@ -1,11 +1,11 @@ -from __future__ import annotations - import asyncio import json import traceback -from functools import partial -from typing import Any, ClassVar, Literal +from collections.abc import Callable +from typing import ( + Annotated, Any, ClassVar, Literal, NotRequired, +) import pandas as pd import panel as pn @@ -20,7 +20,6 @@ from pydantic import BaseModel, create_model from pydantic.fields import FieldInfo -from ..base import Component from ..dashboard import Config from ..pipeline import Pipeline from ..sources.base import BaseSQLSource, Source @@ -36,15 +35,17 @@ VEGA_LITE_EXAMPLES_OPENAI_DB_FILE, VEGA_MAP_LAYER, VEGA_ZOOMABLE_MAP_ITEMS, MissingContextError, RetriesExceededError, ) -from .controls import AnnotationControls, RetryControls, SourceControls +from .context import ContextModel, TContext +from .controls import SourceControls from .llm import Llm, Message, OpenAI -from .memory import _Memory from .models import ( DbtslQueryParams, DiscoveryQueries, DiscoverySufficiency, DistinctQuery, PartialBaseModel, QueryCompletionValidation, RetrySpec, SampleQuery, - SqlQuery, VegaLiteSpec, VegaLiteSpecUpdate, make_sql_model, + VegaLiteSpec, VegaLiteSpecUpdate, make_sql_model, +) +from .schemas import ( + DbtslMetadata, DbtslMetaset, SQLMetaset, get_metaset, ) -from .schemas import get_metaset from .services import DbtslMixin from .tools import ToolUser from .translate import param_to_pydantic @@ -54,19 +55,18 @@ retry_llm_output, set_nested, stream_details, ) from .vector_store import DuckDBVectorStore -from .views import AnalysisOutput, LumenOutput, VegaLiteOutput +from .views import ( + AnalysisOutput, LumenOutput, SQLOutput, VegaLiteOutput, +) class Agent(Viewer, ToolUser, ContextProvider): """ Agents are actors responsible for taking a user query and - performing a particular task and responding by adding context to - the current memory and creating outputs. - - Each Agent can require certain context that is needed to perform - the task and should declare any context it itself provides. + performing a particular task, either by adding context or + generating outputs. - Agents have access to an LLM and the current memory and can + Agents have access to an LLM and are given context and can solve tasks by executing a series of prompts or by rendering contents such as forms or widgets to gather user input. """ @@ -130,7 +130,8 @@ def __panel__(self): async def _stream(self, messages: list[Message], system_prompt: str) -> Any: message = None model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) - async for output_chunk in self.llm.stream(messages, system=system_prompt, model_spec=model_spec, field="output"): + output = self.llm.stream(messages, system=system_prompt, model_spec=model_spec, field="output") + async for output_chunk in output: if self.interface is None: if message is None: message = ChatMessage(output_chunk, user=self.user) @@ -140,29 +141,27 @@ async def _stream(self, messages: list[Message], system_prompt: str) -> Any: message = self.interface.stream(output_chunk, replace=True, message=message, user=self.user, max_width=self._max_width) return message - async def _gather_prompt_context(self, prompt_name: str, messages: list, **context): - context = await super()._gather_prompt_context(prompt_name, messages, **context) + async def _gather_prompt_context(self, prompt_name: str, messages: list, context: TContext, **kwargs): + context = await super()._gather_prompt_context(prompt_name, messages, context, **kwargs) if "tool_context" not in context: - context["tool_context"] = await self._use_tools(prompt_name, messages) + context["tool_context"] = await self._use_tools(prompt_name, messages, context) return context # Public API @classmethod - async def applies(cls, memory: _Memory) -> bool: + async def applies(cls, context: TContext) -> bool: """ Additional checks to determine if the agent should be used. """ return True - async def requirements(self, messages: list[Message]) -> list[str]: - return self.requires - async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: + ) -> tuple[list[Any], TContext]: """ Provides a response to the user query. @@ -173,12 +172,21 @@ async def respond( messages: list[Message] The list of messages corresponding to the user query and any other system messages to be included. + context: TContext + A mapping containing context for the agent to perform its task. step_title: str | None If the Agent response is part of a longer query this describes the step currently being processed. """ - system_prompt = await self._render_prompt("main", messages) - return await self._stream(messages, system_prompt) + system_prompt = await self._render_prompt("main", messages, context) + return [await self._stream(messages, system_prompt)], context + + +class SourceOutputs(ContextModel): + + source: Source + + document_sources: list[Any] class SourceAgent(Agent): @@ -194,23 +202,21 @@ class SourceAgent(Agent): purpose = param.String(default="Allows a user to upload new datasets, data, or documents.") - requires = param.List(default=[], readonly=True) - - provides = param.List(default=["sources", "source", "document_sources"], readonly=True) - source_controls: ClassVar[SourceControls] = SourceControls _extensions = ("filedropper",) + output_schema = SourceOutputs + async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: - source_controls = self.source_controls(memory=self._memory, cancellable=True, replace_controls=True) - + ) -> tuple[list[Any], SourceOutputs]: + source_controls = self.source_controls(context=context, cancellable=True, replace_controls=True) output = pn.Column(source_controls) - if "source" not in self._memory: + if "source" not in context: help_message = "No datasets or documents were found, **please upload at least one to continue**..." else: help_message = "**Please upload new dataset(s)/document(s) to continue**, or click cancel if this was unintended..." @@ -221,8 +227,8 @@ async def respond( if source_controls._cancel_button.clicks > 0: self.interface.undo() self.interface.disabled = False - return None - return source_controls + return [], {} + return [], source_controls.outputs class ChatAgent(Agent): @@ -254,21 +260,31 @@ class ChatAgent(Agent): } ) - requires = param.List(default=[], readonly=True) - async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: - context = {"tool_context": await self._use_tools("main", messages)} - if "vector_metaset" not in self._memory and "source" in self._memory and "table" in self._memory: - source = self._memory["source"] - self._memory["vector_metaset"] = await get_metaset( - [source], [f"{source.name}{SOURCE_TABLE_SEPARATOR}{self._memory['table']}"], + ) -> tuple[list[Any], dict[str, Any]]: + prompt_context = {"tool_context": await self._use_tools("main", messages, context)} + if "vector_metaset" not in context and "source" in context and "table" in context: + source = context["source"] + context["vector_metaset"] = await get_metaset( + [source], [f"{source.name}{SOURCE_TABLE_SEPARATOR}{context['table']}"], ) - system_prompt = await self._render_prompt("main", messages, **context) - return await self._stream(messages, system_prompt) + system_prompt = await self._render_prompt("main", messages, context, **prompt_context) + return [await self._stream(messages, system_prompt)], {} + + +class AnalystInputs(ContextModel): + + data: NotRequired[Any] + + source: Source + + pipeline: Pipeline + + sql: NotRequired[str] class AnalystAgent(ChatAgent): @@ -293,17 +309,18 @@ class AnalystAgent(ChatAgent): } ) - requires = param.List(default=["source", "pipeline"], readonly=True) + input_schema = AnalystInputs async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: - messages = await super().respond(messages, step_title) - if len(self._memory.get("data", [])) == 0 and self._memory.get("sql"): - self._memory["sql"] = f"{self._memory['sql']}\n-- No data was returned from the query." - return messages + ) -> tuple[list[Any], TContext]: + messages, out_context = await super().respond(messages, context, step_title=step_title) + if len(context.get("data", [])) == 0 and context.get("sql"): + context["sql"] = f"{context['sql']}\n-- No data was returned from the query." + return messages, out_context class ListAgent(Agent): @@ -314,8 +331,6 @@ class ListAgent(Agent): purpose = param.String(default=""" Renders a list of items to the user and lets the user pick one.""") - requires = param.List(default=[], readonly=True) - _extensions = ("tabulator",) _column_name = None @@ -324,7 +339,7 @@ class ListAgent(Agent): __abstract = True - def _get_items(self) -> dict[str, list[str]]: + def _get_items(self, context: TContext) -> dict[str, list[str]]: """Return dict of items grouped by source/category""" def _use_item(self, event): @@ -344,9 +359,10 @@ def _use_item(self, event): async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: - items = self._get_items() + ) -> tuple[list[Any], dict[str, Any]]: + items = self._get_items(context) # Create tabs with one tabulator per source tabs = [] @@ -385,7 +401,14 @@ async def respond( self._tabs ), user="Assistant" ) - return self._tabs + return [self._tabs], {} + + +class TableListInputs(ContextModel): + + source: Source + + visible_slugs: NotRequired[set[str]] class TableListAgent(ListAgent): @@ -401,25 +424,25 @@ class TableListAgent(ListAgent): not_with = param.List(default=["DbtslAgent", "SQLAgent"]) purpose = param.String(default=""" - Displays a list of all available data & datasets in memory. Not useful for identifying which dataset to use for analysis.""") - - requires = param.List(default=["source"], readonly=True) + Displays a list of all available data & datasets. Not useful for identifying which dataset to use for analysis.""") _column_name = "Data" _message_format = "Show the data: {item}" + input_schema = TableListInputs + @classmethod - async def applies(cls, memory: _Memory) -> bool: - return len(memory.get('visible_slugs', set())) > 1 + async def applies(cls, context: TContext) -> bool: + return len(context.get('visible_slugs', set())) > 1 - def _get_items(self) -> dict[str, list[str]]: - if "closest_tables" in self._memory: + def _get_items(self, context: TContext) -> dict[str, list[str]]: + if "closest_tables" in context: # If we have closest_tables from search, return as a single group - return {"Search Results": self._memory["closest_tables"]} + return {"Search Results": context["closest_tables"]} # Group tables by source - visible_slugs = self._memory.get('visible_slugs', set()) + visible_slugs = context.get('visible_slugs', set()) if not visible_slugs: return {} @@ -443,6 +466,11 @@ def _get_items(self) -> dict[str, list[str]]: return tables_by_source +class DocumentListInputs(ContextModel): + + document_sources: dict[str, Any] + + class DocumentListAgent(ListAgent): """ The DocumentListAgent lists all available documents provided by the user. @@ -456,25 +484,22 @@ class DocumentListAgent(ListAgent): ) purpose = param.String(default=""" - Displays a list of all available documents in memory.""") - - requires = param.List(default=["document_sources"], readonly=True) + Displays a list of all available documents.""") _column_name = "Documents" _message_format = "Tell me about: {item}" @classmethod - async def applies(cls, memory: _Memory) -> bool: - sources = memory.get("document_sources") + async def applies(cls, context: TContext) -> bool: + sources = context.get("document_sources") if not sources: - return False # source not loaded yet; always apply + return False return len(sources) > 1 - def _get_items(self) -> dict[str, list[str]]: + def _get_items(self, context: TContext) -> dict[str, list[str]]: # extract the filename, following this pattern `Filename: 'filename'`` - documents = [doc["metadata"].get("filename", "untitled") for doc in self._memory.get("document_sources", [])] - # Return all documents under a single "Documents" category + documents = [doc["metadata"].get("filename", "untitled") for doc in context.get("document_sources", [])] return {"Documents": documents} if documents else {} @@ -492,13 +517,15 @@ class LumenBaseAgent(Agent): _output_type = LumenOutput _retry_target_keys = [] - def _update_spec(self, memory: _Memory, event: param.parameterized.Event): + def _update_spec(self, context: TContext, event: param.parameterized.Event): """ - Update the specification in memory. + Update the specification in the context dictionary. """ @staticmethod - def _prepare_lines_for_retry(original_output: str, retry_target_keys: list[str] | None = None) -> tuple[list[str], bool, dict | None]: + def _prepare_lines_for_retry( + original_output: str, retry_target_keys: list[str] | None = None + ) -> tuple[list[str], bool, dict | None]: """ Prepare lines for retry by optionally extracting targeted sections from YAML. @@ -572,77 +599,68 @@ def _apply_line_changes_to_output( else: return apply_changes(lines, line_changes) - async def _retry_output_by_line( + async def revise( self, - feedback: str, + instruction: str, messages: list[Message], - memory: _Memory, + context: TContext, original_output: str, language: str | None = None, - **context + **kwargs ) -> str: """ - Retry the output by line, allowing the user to provide feedback on why the output was not satisfactory, or an error. + Retry the output by line, allowing the user to provide instruction on why the output was not satisfactory, or an error. """ # Prepare lines for retry processing lines, targeted, original_spec = self._prepare_lines_for_retry( original_output, self._retry_target_keys ) - - with self.param.update(memory=memory): - numbered_text = "\n".join(f"{i:2d}: {line}" for i, line in enumerate(lines, 1)) - system = await self._render_prompt( - "retry_output", - messages=messages, - numbered_text=numbered_text, - language=language, - feedback=feedback, - **context - ) + numbered_text = "\n".join(f"{i:2d}: {line}" for i, line in enumerate(lines, 1)) + system = await self._render_prompt( + "retry_output", + messages, + context, + numbered_text=numbered_text, + language=language, + feedback=instruction, + **kwargs + ) retry_model = self._lookup_prompt_key("retry_output", "response_model") - invoke_kwargs = dict( - messages=messages, + result = await self.llm.invoke( + messages, system=system, response_model=retry_model, - model_spec="edit", + model_spec="edit" ) - result = await self.llm.invoke(**invoke_kwargs) - - # Apply line changes and return result return self._apply_line_changes_to_output( lines, result.lines_changes, targeted, original_spec, self._retry_target_keys ) - def _render_lumen( - self, - component: Component, - messages: list | None = None, - title: str | None = None, - **kwargs, - ): - async def _retry_invoke(event: param.parameterized.Event): - with out.param.update(loading=True): - out.spec = await self._retry_output_by_line(event.new, messages, memory, out.spec, language=out.language) - - memory = self._memory - retry_controls = RetryControls() - retry_controls.param.watch(_retry_invoke, "reason") - out = self._output_type( - component=component, - footer=[retry_controls], - title=title, - **kwargs - ) - out.param.watch(partial(self._update_spec, self._memory), "spec") - if "outputs" in self._memory: - # We have to create a new list to trigger an event - # since inplace updates will not trigger updates - # and won't allow diffing between old and new values - self._memory["outputs"] = self._memory["outputs"] + [out] - if self.interface is not None: - message_kwargs = dict(value=out, user=self.user) - self.interface.stream(replace=True, max_width=self._max_width, **message_kwargs) - return out + +class SQLInputs(ContextModel): + + data: NotRequired[Any] + + source: Source + + sources: Annotated[list[Source], ("accumulate", "source")] + + sql: NotRequired[str] + + sql_metaset: SQLMetaset + + visible_slugs: NotRequired[set[str]] + + +class SQLOutputs(ContextModel): + + data: Any + + table: str + + sql: str + + pipeline: Pipeline class SQLAgent(LumenBaseAgent): @@ -691,69 +709,20 @@ class SQLAgent(LumenBaseAgent): } ) - provides = param.List(default=["table", "sql", "pipeline", "data"], readonly=True) - - requires = param.List(default=["sources", "source", "sql_metaset"], readonly=True) - user = param.String(default="SQL") _extensions = ("codeeditor", "tabulator") - _output_type = LumenOutput - - def _update_spec(self, memory: _Memory, event: param.parameterized.Event): - memory["sql"] = event.new + _output_type = SQLOutput - @retry_llm_output() - async def _generate_sql_queries( - self, messages: list[Message], dialect: str, step_number: int, - is_final: bool, context_entries: list[dict] | None, - sql_plan_context: str | None = None, errors: list | None = None - ) -> dict[str, str]: - """Generate SQL queries using LLM.""" - # Build SQL history from context - sql_query_history = {} - if context_entries: - for entry in context_entries: - for query in entry.get("queries", []): - sql_query_history[query["sql"]] = query["table_status"] - - # Render prompt - system_prompt = await self._render_prompt( - "main", - messages, - dialect=dialect, - step_number=step_number, - is_final_step=is_final, - current_step=messages[0]["content"] if not is_final else "", - sql_query_history=sql_query_history, - current_iteration=getattr(self, '_current_iteration', 1), - sql_plan_context=sql_plan_context, - errors=errors, - ) + input_schema = SQLInputs + output_schema = SQLOutputs - # Generate SQL - model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) - - sql_response_model = self._get_model("main", is_final=is_final) - output = await self.llm.invoke( - messages, - system=system_prompt, - model_spec=model_spec, - response_model=sql_response_model, - ) - if not output: - raise ValueError("No output was generated.") - - sql_queries = {} - for query_obj in ([output] if isinstance(output, SqlQuery) else output.queries): - if query_obj.query and query_obj.expr_slug: - sql_queries[query_obj.expr_slug.strip()] = query_obj.query.strip() - - return sql_queries + def _update_spec(self, context: TContext, event: param.parameterized.Event): + context["sql"] = event.new async def _validate_sql( - self, sql_query: str, expr_slug: str, dialect: str, + self, context: TContext, sql_query: str, expr_slug: str, dialect: str, source, messages: list[Message], step, max_retries: int = 2, discovery_context: str | None = None, ) -> str: @@ -782,14 +751,14 @@ async def _validate_sql( if "KeyError" in feedback: feedback += " The data does not exist; select from available data sources." - retry_result = await self._retry_output_by_line( - feedback, messages, self._memory, sql_query, language=f"sql.{dialect}", discovery_context=discovery_context + retry_result = await self.revise( + feedback, messages, context, sql_query, language=f"sql.{dialect}", discovery_context=discovery_context ) sql_query = clean_sql(retry_result, dialect) return sql_query async def _execute_query( - self, source, expr_slug: str, sql_query: str, + self, source, context: TContext, expr_slug: str, sql_query: str, is_final: bool, should_materialize: bool, step ) -> tuple[Pipeline, Source, str]: """Execute SQL query and return pipeline and summary.""" @@ -799,7 +768,7 @@ async def _execute_query( ) if should_materialize: - self._memory["source"] = sql_expr_source + context["source"] = sql_expr_source # Create pipeline if is_final: @@ -825,14 +794,15 @@ async def _finalize_execution( results: dict, context_entries: list[dict], messages: list[Message], + context: TContext, step_title: str | None, raise_if_empty: bool = False - ) -> None: + ) -> tuple[LumenOutput, SQLOutputs]: """Finalize execution for final step.""" # Get first result (typically only one for final step) expr_slug, result = next(iter(results.items())) - # Update memory + # Update context pipeline = result["pipeline"] df = await get_data(pipeline) @@ -840,31 +810,29 @@ async def _finalize_execution( if df.empty and raise_if_empty: raise ValueError(f"\nQuery `{result['sql']}` returned empty results; ensure all the WHERE filter values exist in the dataset.") - self._memory["data"] = await describe_data(df) - self._memory["sql"] = result["sql"] - self._memory["pipeline"] = pipeline - self._memory["table"] = pipeline.table - self._memory["source"] = pipeline.source - self._memory["sql_plan_context"] = context_entries - - # Render output - self._render_lumen( - pipeline, - messages=messages, - title=step_title, - spec=result["sql"] + view = self._output_type( + component=pipeline, title=step_title, spec=result["sql"] ) + return view, { + "data": await describe_data(df), + "sql": result["sql"], + "pipeline": pipeline, + "table": pipeline.table, + "source": pipeline.source, + "sql_plan_context": context_entries + } async def _render_execute_query( self, messages: list[Message], + context: TContext, source: Source, step_title: str, success_message: str, discovery_context: str | None = None, raise_if_empty: bool = False, output_title: str | None = None - ) -> Pipeline: + ) -> tuple[LumenOutput, SQLOutputs]: """ Helper method that generates, validates, and executes final SQL queries. @@ -895,6 +863,7 @@ async def _render_execute_query( system_prompt = await self._render_prompt( "main", messages, + context, dialect=source.dialect, step_number=1, is_final_step=True, @@ -925,12 +894,13 @@ async def _render_execute_query( # Validate SQL validated_sql = await self._validate_sql( - sql_query, expr_slug, source.dialect, source, messages, step, discovery_context=discovery_context + context, sql_query, expr_slug, source.dialect, source, messages, + step, discovery_context=discovery_context ) # Execute and get results pipeline, sql_expr_source, summary = await self._execute_query( - source, expr_slug, validated_sql, is_final=True, + source, context, expr_slug, validated_sql, is_final=True, should_materialize=True, step=step ) @@ -944,24 +914,26 @@ async def _render_execute_query( } } - await self._finalize_execution( - results, [], messages, output_title, raise_if_empty=raise_if_empty + view, out_context = await self._finalize_execution( + results, [], messages, context, output_title, raise_if_empty=raise_if_empty ) step.status = "success" step.success_title = success_message - return pipeline + return view, out_context async def _select_discoveries( self, messages: list[Message], + context: TContext, error_context: str, ) -> DiscoveryQueries: """Let LLM choose which discoveries to run based on error and available tables.""" system_prompt = await self._render_prompt( "select_discoveries", messages, + context, error_context=error_context, ) @@ -977,6 +949,7 @@ async def _select_discoveries( async def _check_discovery_sufficiency( self, messages: list[Message], + context: TContext, error_context: str, discovery_results: list[tuple[str, str]], ) -> DiscoverySufficiency: @@ -984,6 +957,7 @@ async def _check_discovery_sufficiency( system_prompt = await self._render_prompt( "check_sufficiency", messages, + context, error_context=error_context, discovery_results=discovery_results, ) @@ -1044,14 +1018,15 @@ async def _run_discoveries_parallel( async def _explore_tables( self, messages: list[Message], + context: TContext, source: Source, error_context: str, step_title: str | None = None, - ) -> Any: + ) -> tuple[LumenOutput, SQLOutputs]: """Run adaptive exploration: initial discoveries → sufficiency check → optional follow-ups → final answer.""" # Step 1: LLM selects initial discoveries with self._add_step(title="Selecting initial discoveries", steps_layout=self._steps_layout) as step: - selection = await self._select_discoveries(messages, error_context) + selection = await self._select_discoveries(messages, context, error_context) step.stream(f"Strategy: {selection.reasoning}\n\nSelected {len(selection.queries)} initial discoveries") # Step 2: Run initial discoveries in parallel @@ -1059,7 +1034,7 @@ async def _explore_tables( # Step 3: Check if discoveries are sufficient with self._add_step(title="Evaluating discovery sufficiency", steps_layout=self._steps_layout) as step: - sufficiency = await self._check_discovery_sufficiency(messages, error_context, initial_results) + sufficiency = await self._check_discovery_sufficiency(messages, context, error_context, initial_results) step.stream(f"Assessment: {sufficiency.reasoning}") if sufficiency.sufficient: @@ -1081,7 +1056,8 @@ async def _explore_tables( ]) return await self._render_execute_query( - messages=messages, + messages, + context, source=source, step_title="Generating final answer with discoveries", success_message="Final answer generated with discovery context", @@ -1093,15 +1069,17 @@ async def _explore_tables( async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: + ) -> tuple[list[Any], SQLOutputs]: """Execute SQL generation with one-shot attempt first, then exploration if needed.""" # Setup sources - source = self._memory["source"] + source = context["source"] try: # Try one-shot approach first - pipeline = await self._render_execute_query( - messages=messages, + out, out_context = await self._render_execute_query( + messages, + context, source=source, step_title="Attempting one-shot SQL generation...", success_message="One-shot SQL generation successful", @@ -1114,8 +1092,16 @@ async def respond( # If exploration is disabled, re-raise the error instead of falling back raise e # Fall back to exploration mode if enabled - pipeline = await self._explore_tables(messages, source, str(e), step_title) - return pipeline + out_context = await self._explore_tables(messages, context, source, str(e), step_title) + return [out], out_context + + + +class DbtslOutputs(SQLOutputs): + + dbtsl_vector_metaset: DbtslMetadata + + dbtsl_sql_metaset: DbtslMetaset class DbtslAgent(LumenBaseAgent, DbtslMixin): @@ -1147,8 +1133,6 @@ class DbtslAgent(LumenBaseAgent, DbtslMixin): } ) - provides = param.List(default=["table", "sql", "pipeline", "data", "dbtsl_vector_metaset", "dbtsl_sql_metaset"], readonly=True) - requires = param.List(default=["source", "dbtsl_metaset"], readonly=True) source = param.ClassSelector( @@ -1163,17 +1147,21 @@ class DbtslAgent(LumenBaseAgent, DbtslMixin): _output_type = LumenOutput + output_schema = DbtslOutputs + def __init__(self, source: Source, **params): super().__init__(source=source, **params) - def _update_spec(self, memory: _Memory, event: param.parameterized.Event): + def _update_spec(self, context: TContext, event: param.parameterized.Event): """ - Update the SQL specification in memory. + Update the SQL specification. """ - memory["sql"] = event.new + context["sql"] = event.new @retry_llm_output() - async def _create_valid_query(self, messages: list[Message], title: str | None = None, errors: list | None = None): + async def _create_valid_query( + self, messages: list[Message], title: str | None = None, errors: list | None = None + ) -> DbtslOutputs: """ Create a valid dbt Semantic Layer query based on user messages. """ @@ -1183,6 +1171,7 @@ async def _create_valid_query(self, messages: list[Message], title: str | None = errors=errors, ) + out_context = {} with self._add_step(title=title or "dbt Semantic Layer query", steps_layout=self._steps_layout) as step: model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) response = self.llm.stream( @@ -1218,7 +1207,7 @@ async def _create_valid_query(self, messages: list[Message], title: str | None = formatted_params = json.dumps(query_params, indent=2) step.stream(f"\n\n`{expr_slug}`\n```json\n{formatted_params}\n```") - self._memory["dbtsl_query_params"] = query_params + out_context["dbtsl_query_params"] = query_params except asyncio.CancelledError as e: step.failed_title = "Cancelled dbt Semantic Layer query generation" raise e @@ -1244,7 +1233,7 @@ async def _create_valid_query(self, messages: list[Message], title: str | None = step.stream(f"\nCompiled SQL:\n```sql\n{sql_query}\n```", replace=False) sql_expr_source = self.source.create_sql_expr_source({expr_slug: sql_query}) - self._memory["sql"] = sql_query + out_context["sql"] = sql_query # Apply transforms sql_transforms = [SQLLimit(limit=1_000_000, write=self.source.dialect, pretty=True, identify=False)] @@ -1261,42 +1250,59 @@ async def _create_valid_query(self, messages: list[Message], title: str | None = sql_metaset = await get_metaset([sql_expr_source], [expr_slug]) vector_metaset = sql_metaset.vector_metaset - # Update memory - self._memory["data"] = await describe_data(df) - self._memory["source"] = sql_expr_source - self._memory["pipeline"] = pipeline - self._memory["table"] = pipeline.table - self._memory["dbtsl_vector_metaset"] = vector_metaset - self._memory["dbtsl_sql_metaset"] = sql_metaset - return sql_query, pipeline + # Update context + out_context["data"] = await describe_data(df) + out_context["source"] = sql_expr_source + out_context["pipeline"] = pipeline + out_context["table"] = pipeline.table + out_context["dbtsl_vector_metaset"] = vector_metaset + out_context["dbtsl_sql_metaset"] = sql_metaset except Exception as e: report_error(e, step) raise e + return out_context async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: + ) -> tuple[list[Any], DbtslOutputs]: """ Responds to user messages by generating and executing a dbt Semantic Layer query. """ try: - sql_query, pipeline = await self._create_valid_query(messages, step_title) + out_context = await self._create_valid_query(messages, step_title) except RetriesExceededError as e: traceback.print_exception(e) - self._memory["__error__"] = str(e) + context["__error__"] = str(e) return None - self._render_lumen(pipeline, messages=messages, title=step_title, spec=sql_query) - return pipeline + pipeline = out_context["pipeline"] + view = self._output_type( + component=pipeline, title=step_title, spec=out_context["sql"] + ) + return [view], out_context -class BaseViewAgent(LumenBaseAgent): - requires = param.List(default=["pipeline", "table", "data"], readonly=True) +class ViewInputs(ContextModel): + + data: Any + + pipeline: Pipeline + + sql_metaset: NotRequired[SQLMetaset] + + table: str + - provides = param.List(default=["view"], readonly=True) +class ViewOutputs(ContextModel): + + view = Any + + +class BaseViewAgent(LumenBaseAgent): prompts = param.Dict( default={ @@ -1304,11 +1310,14 @@ class BaseViewAgent(LumenBaseAgent): } ) + input_schema = ViewInputs + output_schema = ViewOutputs + def __init__(self, **params): self._last_output = None super().__init__(**params) - def _build_errors_context(self, pipeline: Pipeline, errors: list[str] | None) -> dict: + def _build_errors_context(self, pipeline: Pipeline, context: TContext, errors: list[str] | None) -> dict: errors_context = {} if errors: errors = ("\n".join(f"{i + 1}. {error}" for i, error in enumerate(errors))).strip() @@ -1318,10 +1327,10 @@ def _build_errors_context(self, pipeline: Pipeline, errors: list[str] | None) -> last_output = "" vector_metadata_map = None - if "sql_metaset" in self._memory: - vector_metadata_map = self._memory["sql_metaset"].vector_metaset.vector_metadata_map - elif "dbtsl_sql_metaset" in self._memory: - vector_metadata_map = self._memory["dbtsl_sql_metaset"].vector_metaset.vector_metadata_map + if "sql_metaset" in context: + vector_metadata_map = context["sql_metaset"].vector_metaset.vector_metadata_map + elif "dbtsl_sql_metaset" in context: + vector_metadata_map = context["dbtsl_sql_metaset"].vector_metaset.vector_metadata_map columns_context = "" if vector_metadata_map is not None: @@ -1341,16 +1350,18 @@ def _build_errors_context(self, pipeline: Pipeline, errors: list[str] | None) -> async def _create_valid_spec( self, messages: list[Message], + context: TContext, pipeline: Pipeline, schema: dict[str, Any], step_title: str | None = None, errors: list[str] | None = None, ) -> dict[str, Any]: - errors_context = self._build_errors_context(pipeline, errors) + errors_context = self._build_errors_context(pipeline, context, errors) doc = self.view_type.__doc__.split("\n\n")[0] if self.view_type.__doc__ else self.view_type.__name__ system = await self._render_prompt( "main", messages, + context, table=pipeline.table, doc=doc, **errors_context, @@ -1367,7 +1378,9 @@ async def _create_valid_spec( e = None error = "" spec = "" - with self._add_step(title=step_title or "Generating view...", steps_layout=self._steps_layout) as step: + with self._add_step( + title=step_title or "Generating view...", steps_layout=self._steps_layout + ) as step: async for output in response: chain_of_thought = output.chain_of_thought or "" step.stream(chain_of_thought, replace=True) @@ -1391,7 +1404,7 @@ async def _create_valid_spec( title="Re-attempted view generation", steps_layout=self._steps_layout, ) as retry_step: - view = await self._retry_output_by_line(e, messages, self._memory, yaml.safe_dump(spec), language="") + view = await self.revise(e, messages, context, yaml.safe_dump(spec), language="") if "yaml_spec: " in view: view = view.split("yaml_spec: ")[-1].rstrip('"').rstrip("'") retry_step.stream(f"\n\n```yaml\n{view}\n```") @@ -1409,33 +1422,35 @@ async def _create_valid_spec( async def _extract_spec(self, spec: dict[str, Any]): return dict(spec) - async def _update_spec(self, memory: _Memory, event: param.parameterized.Event): - memory["view"] = dict(await self._extract_spec(event.new), type=self.view_type) + async def _update_spec(self, context: TContext, event: param.parameterized.Event): + context["view"] = dict(await self._extract_spec(event.new), type=self.view_type) async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: + ) -> tuple[list[Any], ViewOutputs]: """ Generates a visualization based on user messages and the current data pipeline. """ - pipeline = self._memory.get("pipeline") + pipeline = context.get("pipeline") if not pipeline: - raise ValueError("No current pipeline found in memory.") + raise ValueError("Context did not contain a pipeline.") schema = await get_schema(pipeline) if not schema: raise ValueError("Failed to retrieve schema for the current pipeline.") - spec = await self._create_valid_spec(messages, pipeline, schema, step_title) - self._memory["view"] = dict(spec, type=self.view_type) + spec = await self._create_valid_spec(messages, context, pipeline, schema, step_title) + context["view"] = dict(spec, type=self.view_type) view = self.view_type(pipeline=pipeline, **spec) - self._render_lumen(view, messages=messages, title=step_title) - return view + out = self._output_type(component=view, title=step_title) + return [out], {"view": view} class hvPlotAgent(BaseViewAgent): + conditions = param.List( default=[ "Use for exploratory data analysis, interactive plots, and dynamic filtering", @@ -1476,12 +1491,12 @@ def _get_model(self, prompt_name: str, schema: dict[str, Any]) -> type[BaseModel ) return model[self.view_type.__name__] - async def _update_spec(self, memory: _Memory, event: param.parameterized.Event): + async def _update_spec(self, context: TContext, event: param.parameterized.Event): spec = yaml.load(event.new, Loader=yaml.SafeLoader) - memory["view"] = dict(await self._extract_spec(spec), type=self.view_type) + context["view"] = dict(await self._extract_spec(context, spec), type=self.view_type) - async def _extract_spec(self, spec: dict[str, Any]): - pipeline = self._memory["pipeline"] + async def _extract_spec(self, context: TContext, spec: dict[str, Any]): + pipeline = context["pipeline"] spec = {key: val for key, val in spec.items() if val is not None} spec["type"] = "hvplot_ui" self.view_type.validate(spec) @@ -1516,6 +1531,8 @@ class VegaLiteAgent(BaseViewAgent): } ) + user = param.String(default="Vega") + vector_store_path = param.Path(default=None, check_exists=False, doc=""" Path to a custom vector store for storing and retrieving Vega-Lite examples; if not provided a default store will be used depending on the LLM-- @@ -1547,7 +1564,8 @@ def _get_vector_store(self): response = requests.get(f"{VECTOR_STORE_ASSETS_URL}{db_file}", timeout=5) response.raise_for_status() uri.write_bytes(response.content) - self._vector_store = DuckDBVectorStore(uri=str(uri)) + # Use a read-only connection to avoid lock conflicts + self._vector_store = DuckDBVectorStore(uri=str(uri), read_only=True) return self._vector_store def _deep_merge_dicts(self, base_dict: dict[str, Any], update_dict: dict[str, Any]) -> dict[str, Any]: @@ -1661,16 +1679,18 @@ async def _update_spec_step( vega_spec: dict[str, Any], prompt_name: str, messages: list[Message], + context: TContext, doc: str | None = None, ) -> tuple[str, dict[str, Any]]: """Update a Vega-Lite spec with incremental changes for a specific step.""" - with self.interface.param.update(callback_exception="raise"), self._add_step(title=step_desc, steps_layout=self._steps_layout) as step: + with self._add_step(title=step_desc, steps_layout=self._steps_layout) as step: system_prompt = await self._render_prompt( prompt_name, messages, + context, vega_spec=yaml.dump(vega_spec, default_flow_style=False), doc=doc, - table=self._memory["pipeline"].table, + table=context["pipeline"].table, ) model_spec = self.prompts.get(prompt_name, {}).get("llm_spec", self.llm_spec_key) @@ -1686,13 +1706,13 @@ async def _update_spec_step( update_dict = yaml.safe_load(result.yaml_update) return step_name, update_dict - async def _update_spec(self, memory: _Memory, event: param.parameterized.Event): + async def _update_spec(self, context: TContext, event: param.parameterized.Event): try: spec = await self._extract_spec({"yaml_spec": event.new}) except Exception as e: traceback.print_exception(e) return - memory["view"] = dict(spec, type=self.view_type) + context["view"] = dict(spec, type=self.view_type) def _add_zoom_params(self, vega_spec: dict) -> None: """Add zoom parameters to vega spec.""" @@ -1765,20 +1785,28 @@ def _extract_as_keys(cls, transforms: list[dict]) -> list[str]: return list(dict.fromkeys(as_fields)) @retry_llm_output() - async def _generate_basic_spec(self, messages: list[Message], pipeline: Pipeline, doc_examples: list, doc: str, errors: list | None = None) -> dict[str, Any]: + async def _generate_basic_spec( + self, + messages: list[Message], + context: TContext, + pipeline: Pipeline, + doc_examples: list, + doc: str, + errors: list | None = None + ) -> dict[str, Any]: """Generate the basic VegaLite spec structure.""" - errors_context = self._build_errors_context(pipeline, errors) + errors_context = self._build_errors_context(pipeline, context, errors) with self._add_step(title="Creating basic plot structure", steps_layout=self._steps_layout) as step: system_prompt = await self._render_prompt( "main", messages, + context, table=pipeline.table, doc=doc, doc_examples=doc_examples, **errors_context, ) model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) - response = self.llm.stream( messages, system=system_prompt, @@ -1818,9 +1846,9 @@ async def _extract_spec(self, spec: dict[str, Any]): # # add pan/zoom controls to all plots except geographic ones and points overlaid on line plots # # because those result in an blank plot without error # vega_spec["params"] = [{"bind": "scales", "name": "grid", "select": "interval"}] - return {"spec": vega_spec, "sizing_mode": "stretch_both", "min_height": 300, "max_width": 1200} + return {"spec": vega_spec, "sizing_mode": "stretch_both", "min_height": 300} - async def _get_doc_examples(self, user_query: str) -> list: + async def _get_doc_examples(self, user_query: str) -> list[str]: # Query vector store for relevant examples doc_examples = [] vector_store = self._get_vector_store() @@ -1841,50 +1869,52 @@ async def _get_doc_examples(self, user_query: str) -> list: k += 1 if k >= 3: # Limit to top 3 examples break + return doc_examples - async def _retry_output_by_line( + async def revise( self, feedback: str, messages: list[Message], - memory: _Memory, + context: TContext, original_output: str, language: str | None = None, - **context + **kwargs ) -> str: doc_examples = await self._get_doc_examples(feedback) context["doc_examples"] = doc_examples - return await super()._retry_output_by_line(feedback, messages, memory, original_output, language, **context) + return await super().revise(feedback, messages, context, original_output, language, **kwargs) async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None, - ) -> Any: + ) -> tuple[list[Any], TContext]: """ Generates a VegaLite visualization using progressive building approach with real-time updates. """ - pipeline = self._memory.get("pipeline") + pipeline = context.get("pipeline") if not pipeline: - raise ValueError("No current pipeline found in memory.") + raise ValueError("Context did not contain a pipeline.") schema = await get_schema(pipeline) if not schema: raise ValueError("Failed to retrieve schema for the current pipeline.") user_query = messages[-1].get("content", "") if messages[-1].get("role") == "user" else "" - doc_examples = await self._get_doc_examples(user_query) + try: + doc_examples = await self._get_doc_examples(user_query) + except Exception: + doc_examples = [] # Step 1: Generate basic spec doc = self.view_type.__doc__.split("\n\n")[0] if self.view_type.__doc__ else self.view_type.__name__ # Produces {"spec": {$schema: ..., ...}, "sizing_mode": ..., ...} - full_dict = await self._generate_basic_spec(messages, pipeline, doc_examples, doc) + full_dict = await self._generate_basic_spec(messages, context, pipeline, doc_examples, doc) # Step 2: Show complete plot immediately - self._memory["view"] = dict(full_dict, type=self.view_type) view = self.view_type(pipeline=pipeline, **full_dict) - out = self._render_lumen(view, messages=messages, title=step_title) - # Get the latest spec from the rendered view, which includes type: vega-lite - full_dict = yaml.safe_load(out.spec) + out = self._output_type(component=view, title=step_title) # Step 3: enhancements (LLM-driven creative decisions) steps = { @@ -1895,85 +1925,42 @@ async def respond( for step_name, step_desc in steps.items(): # Only pass the vega lite 'spec' portion to prevent ballooning context step_name, update_dict = await self._update_spec_step( - step_name, step_desc, full_dict["spec"], step_name, messages, doc=doc + step_name, step_desc, out.spec, step_name, messages, context, doc=doc ) try: - test_spec = self._deep_merge_dicts(full_dict["spec"], update_dict) + test_spec = self._deep_merge_dicts(out.spec, update_dict) await self._extract_spec({"yaml_spec": yaml.dump(test_spec)}) # Validation except Exception as e: log_debug(f"Skipping invalid {step_name} update due to error: {e}") continue - full_dict["spec"] = self._deep_merge_dicts(full_dict["spec"], update_dict) - out.spec = yaml.dump(full_dict) + spec = self._deep_merge_dicts(out.spec, update_dict) + out.spec = spec log_debug(f"📊 Applied {step_name} updates and refreshed visualization") - self._memory["view"] = full_dict - return view - - def _render_lumen( - self, - component: Component, - messages: list | None = None, - title: str | None = None, - **kwargs, - ): - """Override to add annotation controls alongside retry controls.""" - async def _retry_invoke(event: param.parameterized.Event): - with out.param.update(loading=True): - out.spec = await self._retry_output_by_line(event.new, messages, self._memory, out.spec, language=out.language) - - async def _annotation_invoke(event: param.parameterized.Event): - """Handle annotation request.""" - with out.param.update(loading=True): - current_dict = yaml.safe_load(out.spec) - updated_dict = await self._apply_annotation( - annotation_request=event.new, - current_dict=current_dict, - messages=messages, - ) - # Already converted to full spec - out.spec = yaml.dump(updated_dict) - - retry_controls = RetryControls() - retry_controls.param.watch(_retry_invoke, "reason") + # Update final context state + out_context = {"view": view} + return [out], out_context - annotation_controls = AnnotationControls() - annotation_controls.param.watch(_annotation_invoke, "annotation_request") - - out = self._output_type( - component=component, - footer=[retry_controls, annotation_controls], - title=title, - **kwargs - ) - out.param.watch(partial(self._update_spec, self._memory), "spec") - if "outputs" in self._memory: - # We have to create a new list to trigger an event - # since inplace updates will not trigger updates - # and won't allow diffing between old and new values - self._memory["outputs"] = self._memory["outputs"] + [out] - if self.interface is not None: - message_kwargs = dict(value=out, user=self.user) - self.interface.stream(replace=True, max_width=self._max_width, **message_kwargs) - return out - - async def _apply_annotation( + async def annotate( self, - annotation_request: str, - current_dict: dict, + instruction: str, messages: list[Message], + context: TContext, + spec: dict ) -> dict: """ Apply annotations based on user request. Parameters ---------- - annotation_request : str + instruction: str User's description of what to annotate - current_spec : dict - The current VegaLite specification (full dict with 'spec' key) - messages : list[Message] + messages: list[Message] Chat history for context + context: TContext + Session context + spec : dict + The current VegaLite specification (full dict with 'spec' key) Returns ------- @@ -1983,14 +1970,15 @@ async def _apply_annotation( # Add user's annotation request to messages context annotation_messages = messages + [{ "role": "user", - "content": f"Add annotations: {annotation_request}" + "content": f"Add annotations: {instruction}" }] with self.interface.param.update(callback_exception="raise"): system_prompt = await self._render_prompt( "annotate_plot", annotation_messages, - vega_spec=yaml.dump(current_dict["spec"], default_flow_style=False), + context, + vega_spec=yaml.dump(spec["spec"], default_flow_style=False), ) model_spec = self.prompts.get("annotate_plot", {}).get("llm_spec", self.llm_spec_key) @@ -2003,18 +1991,29 @@ async def _apply_annotation( update_dict = yaml.safe_load(result.yaml_update) # Merge and validate - final_dict = current_dict.copy() + final_dict = spec.copy() try: final_dict["spec"] = self._deep_merge_dicts(final_dict["spec"], update_dict) await self._extract_spec({"yaml_spec": yaml.dump(final_dict["spec"])}) except Exception as e: log_debug(f"Skipping invalid annotation update due to error: {e}") - # Return original spec if annotation fails - return final_dict - return final_dict +class AnalysisInputs(ContextModel): + + data: NotRequired[Any] + + pipeline: Pipeline + + +class AnalysisOutputs(ContextModel): + + analysis: Callable + + view: Any + + class AnalysisAgent(LumenBaseAgent): analyses = param.List([]) @@ -2035,21 +2034,18 @@ class AnalysisAgent(LumenBaseAgent): } ) - provides = param.List(default=["view"]) - - requires = param.List(default=["pipeline"]) - _output_type = AnalysisOutput - def _update_spec(self, memory: _Memory, event: param.parameterized.Event): + def _update_spec(self, context: TContext, event: param.parameterized.Event): pass async def respond( self, messages: list[Message], + context: TContext, step_title: str | None = None - ) -> Any: - pipeline = self._memory["pipeline"] + ) -> tuple[list[Any], TContext]: + pipeline = context.get("pipeline") analyses = {a.name: a for a in self.analyses if await a.applies(pipeline)} if not analyses: log_debug("No analyses apply to the current data.") @@ -2071,7 +2067,7 @@ async def respond( "main", messages, analyses=analyses, - data=self._memory.get("data"), + data=context.get("data"), ) model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) analysis_name = ( @@ -2089,51 +2085,59 @@ async def respond( analysis_name = next(iter(analyses)) view = None + out_context = {} with self.interface.param.update(callback_exception="raise"): with self._add_step(title=step_title or "Creating view...", steps_layout=self._steps_layout) as step: await asyncio.sleep(0.1) # necessary to give it time to render before calling sync function... - analysis_callable = analyses[analysis_name].instance(agents=self.agents, memory=self._memory, interface=self.interface) + analysis_callable = analyses[analysis_name].instance(agents=self.agents, context=context, interface=self.interface) data = await get_data(pipeline) for field in analysis_callable._field_params: analysis_callable.param[field].objects = list(data.columns) - self._memory["analysis"] = analysis_callable + context["analysis"] = analysis_callable if analysis_callable.autorun: if asyncio.iscoroutinefunction(analysis_callable.__call__): - view = await analysis_callable(pipeline) + view = await analysis_callable(pipeline, context) else: - view = await asyncio.to_thread(analysis_callable, pipeline) + view = await asyncio.to_thread(analysis_callable, pipeline, context) if isinstance(view, Viewable): - view = Panel(object=view, pipeline=self._memory.get("pipeline")) + view = Panel(object=view, pipeline=context.get("pipeline")) spec = view.to_spec() if isinstance(view, View): view_type = view.view_type - self._memory["view"] = dict(spec, type=view_type) + out_context["view"] = dict(spec, type=view_type) elif isinstance(view, Pipeline): - self._memory["pipeline"] = view + out_context["pipeline"] = view # Ensure data reflects processed pipeline - if pipeline is not self._memory["pipeline"]: - pipeline = self._memory["pipeline"] + if pipeline is not out_context["pipeline"]: + pipeline = out_context["pipeline"] data = await get_data(pipeline) if len(data) > 0: - self._memory["data"] = await describe_data(data) + out_context["data"] = await describe_data(data) yaml_spec = yaml.dump(spec) step.stream(f"Generated view\n```yaml\n{yaml_spec}\n```") step.success_title = "Generated view" else: step.success_title = "Configure the analysis" - analysis = self._memory["analysis"] - pipeline = self._memory["pipeline"] + analysis = out_context["analysis"] + pipeline = out_context["pipeline"] if view is None and analysis.autorun: self.interface.stream("Failed to find an analysis that applies to this data") else: - self._render_lumen(view, analysis=analysis, pipeline=pipeline, title=step_title) + out = self._output_type( + component=view, title=step_title, analysis=analysis, pipeline=pipeline + ) self.interface.stream( analysis.message or f"Successfully created view with {analysis_name} analysis.", user="Assistant" ) - return view + return [] if view is None else [out], out_context + + +class ValidationOutputs(ContextModel): + + validation_result: str class ValidationAgent(Agent): @@ -2164,16 +2168,15 @@ class ValidationAgent(Agent): } ) - requires = param.List(default=[], readonly=True) - - provides = param.List(default=["validation_result"], readonly=True) + output_schema = ValidationOutputs async def respond( self, messages: list[Message], + context: TContext, render_output: bool = False, step_title: str | None = None, - ) -> Any: + ) -> tuple[list[Any], ValidationOutputs]: def on_click(event): if messages: user_messages = [msg for msg in reversed(messages) if msg.get("role") == "user"] @@ -2182,10 +2185,10 @@ def on_click(event): self.interface.send(f"Follow these suggestions to fulfill the original intent {original_query}\n\n{suggestions_list}") executed_steps = None - if "plan" in self._memory and hasattr(self._memory["plan"], "steps"): - executed_steps = [f"{step.actor}: {step.instruction}" for step in self._memory["plan"].steps] + if "plan" in context and hasattr(context["plan"], "steps"): + executed_steps = [f"{step.actor}: {step.instruction}" for step in context["plan"].steps] - system_prompt = await self._render_prompt("main", messages, executed_steps=executed_steps) + system_prompt = await self._render_prompt("main", messages, context, executed_steps=executed_steps) model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) result = await self.llm.invoke( @@ -2194,12 +2197,9 @@ def on_click(event): model_spec=model_spec, response_model=QueryCompletionValidation, ) - - self._memory["validation_result"] = result - response_parts = [] if result.correct: - return result + return [result], {"validation_result": result} response_parts.append(f"**Query Validation: ✗ Incomplete** - {result.chain_of_thought}") if result.missing_elements: @@ -2213,4 +2213,4 @@ def on_click(event): footer_objects = [button] formatted_response = "\n\n".join(response_parts) self.interface.stream(formatted_response, user=self.user, max_width=self._max_width, footer_objects=footer_objects) - return result + return [result], {"validation_result": result} diff --git a/lumen/ai/analysis.py b/lumen/ai/analysis.py index a66b400da..d3336ab1b 100644 --- a/lumen/ai/analysis.py +++ b/lumen/ai/analysis.py @@ -5,11 +5,11 @@ from panel.chat import ChatFeed from panel.viewable import Viewable +from panel_material_ui import AutocompleteInput, TextInput from ..base import Component from .config import SOURCE_TABLE_SEPARATOR from .controls import SourceControls -from .memory import _Memory, memory from .utils import get_data @@ -30,13 +30,11 @@ class Analysis(param.ParameterizedFunction): The columns required for the analysis. May use tuples to declare that one of the columns must be present.""") + context = param.Dict() + interface = param.ClassSelector(class_=ChatFeed, doc=""" The ChatInterface instance that will be used to stream messages.""") - memory = param.ClassSelector(class_=_Memory, default=None, doc=""" - Local memory which will be used to provide the agent context. - If None the global memory will be used.""") - message = param.String(default="", doc="The message to display on interface when the analysis is run.") _run_button = param.Parameter(default=None) @@ -48,10 +46,6 @@ class Analysis(param.ParameterizedFunction): _field_params = [] - @property - def _memory(self): - return memory if self.memory is None else self.memory - @classmethod async def applies(cls, pipeline) -> bool: applies = True @@ -80,7 +74,7 @@ class Join(Analysis): index_col = param.String(doc="The column to join on in the left table.") - context = param.String(doc="Additional context to provide to the LLM.") + guidance = param.String(doc="Additional context to provide to the LLM.") _callable_by_llm = False @@ -93,22 +87,22 @@ def _update_table_name(self, event): def controls(self): self._source_controls = SourceControls( - multiple=True, replace_controls=False, memory=self._memory + multiple=True, replace_controls=False, context=self.context ) self._run_button = self._source_controls._add_button self._source_controls.param.watch(self._update_table_name, "_last_table") - source = self._memory.get("source") - table = self._memory.get("table") + source = self.context.get("source") + table = self.context.get("table") self._previous_source = source self._previous_table = table columns = list(source.get_schema(table).keys()) - index_col = pn.widgets.AutocompleteInput.from_param( - self.param.index_col, options=columns, name="Join on", + index_col = AutocompleteInput.from_param( + self.param.index_col, options=columns, label="Join on", placeholder="Start typing column name", search_strategy="includes", case_sensitive=False, restrict=False ) - context = pn.widgets.TextInput.from_param(self.param.context, name="Context") + context = TextInput.from_param(self.param.guidance) controls = pn.FlexBox( index_col, context, @@ -128,10 +122,10 @@ async def __call__(self, pipeline): content += f" left join on {self.index_col}" else: content += " based on the closest matching columns" - if self.context: - content += f"\nadditional context:\n{self.context!r}" + if self.guidance: + content += f"\nadditional context:\n{self.guidance!r}" await agent.answer(messages=[{"role": "user", "content": content}]) - pipeline = self._memory["pipeline"] + pipeline = self.context["pipeline"] self.message = f"Joined {self._previous_table} with {self.table_name}." return pipeline diff --git a/lumen/ai/assets/lumen_template.docx b/lumen/ai/assets/lumen_template.docx new file mode 100644 index 000000000..756a927c7 Binary files /dev/null and b/lumen/ai/assets/lumen_template.docx differ diff --git a/lumen/ai/components.py b/lumen/ai/components.py index ba6881f73..ec191669b 100644 --- a/lumen/ai/components.py +++ b/lumen/ai/components.py @@ -1,18 +1,11 @@ +from __future__ import annotations + from pathlib import Path from typing import Any import param from panel.custom import Child, JSComponent -from panel.layout import Column, HSpacer, Row -from panel.pane import Markdown -from panel.viewable import Viewer -from panel_material_ui import ( - Card, CheckBoxGroup, IconButton, Switch, Typography, -) - -from .config import SOURCE_TABLE_SEPARATOR -from .memory import memory CSS = """ /* Base styles for the split container */ @@ -287,299 +280,3 @@ def _handle_msg(self, msg): collapsed = msg['collapsed'] with param.discard_events(self): self.collapsed = collapsed - - -class TableSourceCard(Viewer): - """ - A component that displays a single data source as a card with table selection controls. - - The card includes: - - A header with the source name and a checkbox to toggle all tables - - A delete button (if multiple sources exist) - - Individual checkboxes for each table in the source - - Metadata display showing source information like filenames and other key-value pairs - """ - - all_selected = param.Boolean(default=True, doc=""" - Whether all tables should be selected by default.""") - - collapsed = param.Boolean(default=False, doc=""" - Whether the card should start collapsed.""") - - deletable = param.Boolean(default=True, doc=""" - Whether to show the delete button.""") - - delete = param.Event(doc="""Action to delete this source from memory.""") - - selected = param.List(default=None, doc=""" - List of currently selected table names.""") - - source = param.Parameter(doc=""" - The data source to display in this card.""") - - def __init__(self, **params): - super().__init__(**params) - self.all_tables = self.source.get_tables() - - # Determine which tables are currently visible - if self.selected is None: - visible_tables = [] - for table in self.all_tables: - visible_tables.append(table) - self.selected = visible_tables - - # Create widgets once in init - self.source_toggle = Switch.from_param( - self.param.all_selected, - name=f"{self.source.name}", - margin=(5, -5, 0, 3), - sizing_mode='fixed', - ) - - self.delete_button = IconButton.from_param( - self.param.delete, - icon='delete', - icon_size='1em', - color="danger", - margin=(5, 0, 0, 0), - sizing_mode='fixed', - width=40, - height=40, - visible=self.param.deletable - ) - - # Create table checkboxes with metadata - self.table_controls = self._create_table_controls() - - # Create source-level metadata display (if any non-table metadata exists) - self.metadata_display = self._create_source_metadata_display() - - def _create_table_controls(self): - """Create table checkboxes with per-table metadata displayed next to each checkbox.""" - table_controls = [] - - for table in self.all_tables: - # Create checkbox for this table - checkbox = CheckBoxGroup( - value=[table] if table in self.selected else [], - options=[table], - sizing_mode='stretch_width', - margin=(2, 10), - name="", - ) - - # Get metadata for this table - table_metadata = self.source.metadata.get(table, {}) if self.source.metadata else {} - metadata_parts = [] - - for key, value in table_metadata.items(): - if isinstance(value, list): - value_str = ', '.join(str(v) for v in value) - else: - value_str = str(value) - metadata_parts.append(f"{key}: {value_str}") - - if metadata_parts: - metadata_text = '; '.join(metadata_parts) - metadata_display = Typography( - metadata_text, - variant="caption", - color="text.secondary", - margin=(-10, 10, 0, 42), - sizing_mode='stretch_width', - styles={"min-height": "unset"} # Override ChatMessage CSS - ) - table_controls.extend([checkbox, metadata_display]) - else: - table_controls.append(checkbox) - - # Watch checkbox changes - checkbox.param.watch(self._on_table_checkbox_change, 'value') - - return Column(*table_controls, margin=0, sizing_mode='stretch_width') - - def _on_table_checkbox_change(self, event): - """Handle individual table checkbox changes.""" - # Collect all selected tables from all checkboxes - selected_tables = [] - for obj in self.table_controls.objects: - if obj.value: - selected_tables.extend(obj.value) - - # Update selected parameter - self.selected = selected_tables - - def _create_source_metadata_display(self): - """Create a metadata display widget for source-level metadata (non-table metadata).""" - metadata_parts = [] - - if self.source.metadata: - # Only show metadata that's not table-specific - for key, value in self.source.metadata.items(): - if key not in self.all_tables: # Skip table-specific metadata - if isinstance(value, list): - value_str = ', '.join(str(v) for v in value) - else: - value_str = str(value) - metadata_parts.append(f"{key}: {value_str}") - - if metadata_parts: - metadata_text = '; '.join(metadata_parts) - return Typography( - metadata_text, - variant="caption", - color="text.secondary", - margin=(0, 10, 5, 10), - sizing_mode='stretch_width', - ) - else: - return Typography( - "", - margin=0, - sizing_mode='stretch_width', - visible=False - ) - - @param.depends('all_selected', watch=True) - def _on_source_toggle(self): - """Handle source checkbox toggle (all tables on/off).""" - if not self.all_selected and len(self.selected) == len(self.all_tables): - # Important to check to see if all tables are selected for intuitive behavior - self.selected = [] - elif self.all_selected: - self.selected = self.all_tables - - @param.depends('selected', watch=True) - def _update_visible_slugs(self): - """Update visible_slugs in memory based on selected tables.""" - self.all_selected = len(self.selected) == len(self.all_tables) - for table in self.all_tables: - table_slug = f"{self.source.name}{SOURCE_TABLE_SEPARATOR}{table}" - if table in self.selected: - memory['visible_slugs'].add(table_slug) - else: - memory['visible_slugs'].discard(table_slug) - memory.trigger('visible_slugs') - - @param.depends('delete', watch=True) - def _delete_source(self): - """Handle source deletion via param.Action.""" - if self.source in memory.get("sources", []): - # Remove all tables from this source from visible_slugs - for table in self.all_tables: - table_slug = f"{self.source.name}{SOURCE_TABLE_SEPARATOR}{table}" - memory['visible_slugs'].discard(table_slug) - - memory["sources"] = [ - source for source in memory.get("sources", []) - if source is not self.source - ] - - def __panel__(self): - card_header = Row( - self.source_toggle, - HSpacer(), - self.delete_button, - sizing_mode='stretch_width', - align='start', - height=35, - margin=0 - ) - - # Create the card content with metadata display - card_content = Column( - self.metadata_display, - self.table_controls, - margin=0, - sizing_mode='stretch_width' - ) - - # Create the card - return Card( - card_content, - header=card_header, - collapsible=True, - collapsed=self.param.collapsed, - sizing_mode='stretch_width', - margin=0, - name="TableSourceCard" - ) - - -class SourceCatalog(Viewer): - """ - A component that displays all data sources with table selection controls. - - This component shows each source as a collapsible card with: - - A header checkbox to toggle all tables in the source - - Individual checkboxes for each table - - A delete button to remove the source (if multiple sources exist) - - Tables can be selectively shown/hidden using the checkboxes, which updates - the 'visible_slugs' set in memory. - """ - - sources = param.List(default=[], doc=""" - List of data sources to display in the catalog.""") - - def __init__(self, **params): - self._title = Markdown(margin=0) - self._cards_column = Column( - margin=0, - ) - self._layout = Column( - self._title, - self._cards_column, - margin=0, - sizing_mode='stretch_width' - ) - super().__init__(**params) - - @param.depends("sources", watch=True, on_init=True) - def _refresh(self, sources=None): - """ - Trigger the catalog with new sources. - - Args: - sources: Optional list of sources. If None, uses sources from memory. - """ - sources = self.sources or memory.get('sources', []) - - # Create a lookup of existing cards by source - existing_cards = { - card.source: card for card in self._cards_column.objects - if isinstance(card, TableSourceCard) and card.source in sources - } - - # Build the new cards list - source_cards = [] - multiple_sources = len(sources) > 1 - for source in sources: - if source in existing_cards: - # Reuse existing card and update its deletable property - card = existing_cards[source] - card.deletable = multiple_sources - source_cards.append(card) - else: - # Create new card for new source - source_card = TableSourceCard( - source=source, - deletable=multiple_sources, - collapsed=multiple_sources, - ) - source_cards.append(source_card) - - self._cards_column.objects = source_cards - - if len(self.sources) == 0: - self._title.object = "No sources available. Add a source to get started." - else: - self._title.object = "Select the table and document sources you want visible to the LLM." - - def __panel__(self): - """ - Create the source catalog UI. - - Returns a Column containing all source cards or a message if no sources exist. - """ - return self._layout diff --git a/lumen/ai/context.py b/lumen/ai/context.py new file mode 100644 index 000000000..f8448e105 --- /dev/null +++ b/lumen/ai/context.py @@ -0,0 +1,687 @@ +from __future__ import annotations + +import sys + +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, Annotated, Any, Literal, NotRequired, Required, TypedDict, + Union, get_args, get_origin, get_type_hints, +) + +if TYPE_CHECKING: + from .actor import Actor + from .report import Action + +TContext = Mapping[str, Any] + + +class ContextModel(TypedDict): + """ + Baseclass for context models, responsible for defining the inputs and outputs + of an Actor. + """ + +@dataclass(frozen=True) +class AccumulateSpec: + from_key: str + # how to extend: either "accumulate" a list or a callable that takes a list of all from_key + # values and may process them in some way. + func: str = Callable[[list[Any]], Any] | Literal["extend"] + # provide to remove duplicates; can be "value" or a callable(item)->hashable + dedupe_by: str | Callable[[Any], Any] | None = None + # if contexts also supply the accumulator field directly, do we include those too? + include_target_field: bool = True + + +def _parse_accumulate_meta(annotation: Any) -> AccumulateSpec | None: + """ + Parse the accumulate metadata from an Annotated field. + + Parameters + ---------- + annotation: Annotated[Any, ...] + The Annotated field to parse. + + Returns + ------- + AccumulateSpec | None + The accumulate specification, or None if the field is not an Annotated field. + + """ + if get_origin(annotation) is not Annotated: + return None + base, *meta = get_args(annotation) + # accept either a tuple ("accumulate", "source"), a Callabe or an AccumulateSpec + for m in meta: + if isinstance(m, AccumulateSpec): + return m + elif isinstance(m, tuple) and len(m) >= 2: + # ("accumulate", "source", {optional kwargs...}) + if m[0] == "accumulate": + func = list + else: + func = m[0] + from_key = m[1] + kwargs = m[2] if len(m) > 2 else {} + return AccumulateSpec(from_key=from_key, func=func, **kwargs) + +def _dedupe(seq: Iterable[Any], key: str | Callable[[Any], Any] | None) -> list[Any]: + """ + Deduplicate a sequence of items based on a key. + + Parameters + ---------- + seq: Iterable[Any] + The sequence to deduplicate. + key: str | Callable[[Any], Any] | None + The key to deduplicate by. + + Returns + ------- + list[Any] + The deduplicated sequence. + + Examples + -------- + >>> _dedupe([1, 2, 3, 2, 1], "value") + [1, 2, 3] + """ + if key is None: + return list(seq) + key_fn: Callable[[Any], Any] + if key == "value": + key_fn = lambda x: x + elif callable(key): + key_fn = key # type: ignore[assignment] + else: + raise ValueError(f"Unknown dedupe_by: {key!r}") + seen = set() + out = [] + for item in seq: + k = key_fn(item) + if k not in seen: + seen.add(k) + out.append(item) + return out + + +class LWW(TypedDict): + """ + A dictionary that overrides the last-write-wins (LWW) behavior. + """ + __override_lww__: Literal[True] + + +def merge_contexts( + schema: type[TypedDict], + contexts: list[Mapping[str, Any]], + *, + lww_keys: Iterable[str] | None = None, # override keys for last-write-wins +) -> dict[str, Any]: + """ + Merge a list of context dicts according to the schema's Annotated metadata. + - Accumulated fields: extend/append from 'from_key' occurrences across contexts + - Regular fields: last-write-wins (LWW) unless overridden by lww_keys + + Parameters + ---------- + schema: type[TypedDict] + The schema to merge the contexts according to. + contexts: list[Mapping[str, Any]] + The list of contexts to merge. + + Returns + ------- + dict[str, Any] + The merged context. + """ + annotations: dict[str, Any] = get_type_hints(schema, include_extras=True) + result: dict[str, Any] = {} + + # Pre-scan accumulators from schema + accumulators: dict[str, AccumulateSpec] = {} + for field, ann in annotations.items(): + spec = _parse_accumulate_meta(ann) + if spec: + accumulators[field] = spec + + for ctx in contexts: + for k, v in ctx.items(): + if k in accumulators: + continue + if v is not None: + result[k] = v + + for target_field, spec in accumulators.items(): + payloads = [] + for ctx in contexts: + if spec.from_key not in ctx or ctx[spec.from_key] is None: + continue + payload = ctx[spec.from_key] + if spec.include_target_field and target_field in ctx and (tgt_value:= ctx[target_field]) is not None: + if isinstance(tgt_value, list): + payloads.extend(tgt_value) + else: + payloads.append(tgt_value) + payloads.append(payload) + + bucket = spec.func(payloads) + bucket = _dedupe(bucket, spec.dedupe_by) + if bucket: + result[target_field] = bucket + elif spec.func == "accumulate": + result.setdefault(target_field, []) + + return result + + +@dataclass +class ValidationIssue: + path: str | tuple[str, ...] | list[str] + key: str + expected: Any + actual: Any | None + message: str + severity: str = "error" + + +def resolved_annotations(tp: type) -> dict[str, Any]: + """Resolve ForwardRefs and preserve Annotated extras.""" + mod = sys.modules[tp.__module__] + return get_type_hints(tp, globalns=vars(mod), localns=None, include_extras=True) + + +def unwrap_annotated(tp: Any) -> tuple[Any, list[Any]]: + """Return (inner_type, metadata[]) for Annotated, else (tp, []).""" + if get_origin(tp) is Annotated: + base, *meta = get_args(tp) + return base, list(meta) + return tp, [] + + +_NOTREQ_ORIGIN = getattr(NotRequired, "__origin__", NotRequired) # defensive +_REQ_ORIGIN = getattr(Required, "__origin__", Required) + +def unwrap_field_type(tp: Any) -> Any: + """Unwrap Annotated, NotRequired, and Required to get the inner type.""" + # Unwrap Annotated + if get_origin(tp) is Annotated: + tp = get_args(tp)[0] + # Unwrap NotRequired / Required (PEP 655) + origin = get_origin(tp) + if origin in (_NOTREQ_ORIGIN, _REQ_ORIGIN): + tp = get_args(tp)[0] + return tp + +def is_typed_dict(tp: Any) -> bool: + return isinstance(tp, type) and issubclass(tp, dict) and hasattr(tp, "__required_keys__") + +def schema_fields(schema: type[TypedDict]) -> dict[str, dict[str, Any]]: + """ + { key: { 'type': , 'required': bool, 'meta': [...] } } + """ + ann = resolved_annotations(schema) + req = getattr(schema, "__required_keys__", set()) + out: dict[str, dict[str, Any]] = {} + for k, tp in ann.items(): + base, meta = unwrap_annotated(tp) + base = unwrap_field_type(base) + out[k] = {"type": base, "required": (k in req), "meta": meta} + return out + +def _accumulate_from_key(meta: list[Any]) -> str | None: + for m in meta: + if isinstance(m, tuple) and len(m) >= 2 and m[0] == "accumulate": + return m[1] + return None + +def input_dependency_keys(schema: type[TypedDict]) -> set[str]: + """ + Keys that, if invalidated, require rerunning a task with this inputs schema. + Includes the field names themselves PLUS any ('accumulate','') sources. + """ + deps: set[str] = set() + # include_extras=True keeps Annotated metadata + ann = get_type_hints(schema, include_extras=True) + for key, tp in ann.items(): + deps.add(key) + if get_origin(tp) is Annotated: + _, *meta = get_args(tp) + src = _accumulate_from_key(meta) + if src: + deps.add(src) + return deps + +def isinstance_like(value: Any, tp: Any) -> bool: + """A pragmatic structural checker for common typing forms.""" + tp = unwrap_field_type(tp) + + origin = get_origin(tp) + args = get_args(tp) + + if tp is Any: + return True + if origin is Union: # includes Optional + return any(isinstance_like(value, alt) for alt in args) + if is_typed_dict(tp): + if not isinstance(value, Mapping): + return False + ann = resolved_annotations(tp) + req = getattr(tp, "__required_keys__", set()) + # required keys present? + for k in req: + if k not in value: + return False + if not isinstance_like(value[k], ann[k]): + return False + # optional keys (if present) must match + opt = getattr(tp, "__optional_keys__", set()) + for k in (value.keys() & opt): + if not isinstance_like(value[k], ann[k]): + return False + # ignore extra keys + return True + if origin in (list, tuple): + if not isinstance(value, (list, tuple)): + return False + if not args: + return True + (elem_type,) = args if len(args) == 1 else (args[0],) + return all(isinstance_like(v, elem_type) for v in value) + if origin in (dict, Mapping): + if not isinstance(value, Mapping): + return False + if len(args) == 2: + kt, vt = args + return all(isinstance_like(k, kt) and isinstance_like(v, vt) for k, v in value.items()) + return True + if isinstance(tp, type): + return isinstance(value, tp) + # Fallback: accept + return True + + +def collect_task_outputs(task: Action | Actor) -> dict[str, Any]: + """ + Returns {key: type} for what this task guarantees to add/enrich in context, + based on task.output_schema (a TypedDict schema). + """ + out_schema = task.output_schema + fields = schema_fields(out_schema) + return {k: f["type"] for k, f in fields.items()} + + +def _base_type(tp: Any) -> Any: + tp = unwrap_field_type(tp) + return get_args(tp)[0] if get_origin(tp) is Annotated else tp + + +def types_compatible(expected: Any, produced: Any) -> bool: + """ + Is a value of 'produced' type usable where 'expected' is required? + (Permissive but practical; handles Annotated, Union/Optional, list/dict/Mapping, TypedDict.) + """ + expected = unwrap_field_type(_base_type(expected)) + produced = unwrap_field_type(_base_type(produced)) + + if expected is Any or produced is Any: + return True + + if get_origin(expected) is Union: + return any(types_compatible(alt, produced) for alt in get_args(expected)) + + e_origin, p_origin = get_origin(expected), get_origin(produced) + e_args, p_args = get_args(expected), get_args(produced) + + if e_origin and p_origin: + if e_origin in (list, tuple) and p_origin in (list, tuple): + e_elem = e_args[0] if e_args else Any + p_elem = p_args[0] if p_args else Any + return types_compatible(e_elem, p_elem) + + if e_origin in (dict, Mapping) and p_origin in (dict, Mapping): + ek = e_args[0] if len(e_args) == 2 else Any + ev = e_args[1] if len(e_args) == 2 else Any + pk = p_args[0] if len(p_args) == 2 else Any + pv = p_args[1] if len(p_args) == 2 else Any + return types_compatible(ek, pk) and types_compatible(ev, pv) + + return False + + try: + if isinstance(expected, type) and issubclass(expected, dict) and hasattr(expected, "__required_keys__"): + if not (isinstance(produced, type) and issubclass(produced, dict) and hasattr(produced, "__required_keys__")): + return True + exp_ann = get_type_hints(expected, include_extras=True) + prod_ann = get_type_hints(produced, include_extras=True) + for k in expected.__required_keys__: # type: ignore[attr-defined] + if k not in prod_ann or not types_compatible(exp_ann[k], prod_ann[k]): + return False + opt = getattr(expected, "__optional_keys__", set()) + for k in (opt & prod_ann.keys()): + if k in exp_ann and not types_compatible(exp_ann[k], prod_ann[k]): + return False + return True + except Exception: + pass + + if isinstance(expected, type) and isinstance(produced, type): + return issubclass(produced, expected) + + return True + +def _list_elem_type(tp: Any) -> Any | None: + """If tp is list[X] (or Annotated[list[X], ...]), return X; else None.""" + tp = unwrap_field_type(tp) + if get_origin(tp) is list: + args = get_args(tp) + return args[0] if args else Any + return None + +def _tname(tp: Any) -> str: + try: + return getattr(tp, "__name__", str(tp)) + except Exception: + return str(tp) + +def _accumulate_alt_key(meta: list[Any]) -> str | None: + for m in meta: + if isinstance(m, tuple) and len(m) >= 2 and m[0] == "accumulate": + return m[1] + return None + +def validate_task_inputs( + task: Action | Actor, value_ctx: Mapping[str, Any], available_types: dict[str, Any], path: str +) -> list[ValidationIssue]: + """ + Validates that task.input_schema are satisfied by value_ctx or prior available_types. + """ + issues: list[ValidationIssue] = [] + schema = task.input_schema + + if not schema: + return issues + + fields = schema_fields(schema) + for key, spec in fields.items(): + expected = spec["type"] + required = spec["required"] + meta = spec["meta"] + alt_key = _accumulate_alt_key(meta) + + # Concrete value present in context + if key in value_ctx: + val = value_ctx[key] + if not isinstance_like(val, expected): + issues.append(ValidationIssue( + path=path, key=key, expected=expected, actual=type(val), + message=f"Type mismatch for '{key}'. Expected {_tname(expected)}, got {_tname(type(val))}." + )) + continue + + # No value; does upstream promise a *compatible* type + if key in available_types: + produced = available_types[key] + if not types_compatible(expected, produced): + issues.append(ValidationIssue( + path=path, key=key, expected=expected, actual=produced, + message=(f"Incompatible upstream type for '{key}': " + f"expected {_tname(expected)} but pipeline provides {_tname(produced)}.") + )) + else: + continue # satisfied by upstream type + + # Does the schema provide an accumulator + if alt_key: + elem = _list_elem_type(expected) # only special-case when expected is list[Elem] + + # 3a) Concrete value under alt_key? + if alt_key in value_ctx: + val = value_ctx[alt_key] + if elem is not None: + # Accept single Elem when accumulating into list[Elem] + if isinstance_like(val, elem): + continue + # Otherwise fall back to normal check against full expected + if isinstance_like(val, expected): + continue + issues.append(ValidationIssue( + path=path, key=key, expected=expected, actual=type(val), + message=(f"Type mismatch via accumulator '{alt_key}' for '{key}': " + f"expected {_tname(expected)}, got {_tname(type(val))}.") + )) + continue + + # 3b) Upstream type promised under alt_key? + if alt_key in available_types: + produced = available_types[alt_key] + if elem is not None: + # Accept Elem (or compatible) when target expects list[Elem] + if types_compatible(elem, produced): + continue + # Otherwise require full compatibility with expected + if types_compatible(expected, produced): + continue + issues.append(ValidationIssue( + path=path, key=key, expected=expected, actual=produced, + message=(f"Incompatible upstream type via accumulator '{alt_key}' for '{key}': " + f"expected {_tname(expected)} but pipeline provides {_tname(produced)}.") + )) + continue + + # 4) Still unsatisfied + if required: + issues.append(ValidationIssue( + path=path, key=key, expected=expected, actual=None, + message=f"Missing required key '{key}'." + )) + return issues + +def _normalize_path(path: str | tuple[str, ...] | list[str]) -> list[str]: + """ + Normalize a path to a list of segments. + + Supports: + - tuple/list of strings: ('TaskGroup', 'Step[0]', 'Summarize') + - string with ' -> ' delimiters: 'TaskGroup[0] -> Summarize' + """ + if isinstance(path, (tuple, list)): + return [str(p) for p in path] + # String form: split on '->' + parts = [seg.strip() for seg in str(path).split("->")] + # Keep '[idx]' attached if present; callers can pass nicer tuples to avoid parsing. + return parts + +def _pretty_type(tp: Any) -> str: + """Conservative type pretty-printer.""" + try: + return getattr(tp, "__name__", str(tp)) + except Exception: + return str(tp) + + +class Node: + __slots__ = ("children", "leaf_issues") + def __init__(self): + self.children: dict[str, Node] = {} + self.leaf_issues: list[ValidationIssue] = [] + + +def render_issues_tree(issues: list[ValidationIssue], *, title: str = "Validation errors") -> str: + """ + Build a tree like: + + Validation errors + ├─ TaskGroup + │ └─ Step[0] -> Summarize + │ ├─ summary [error] Type mismatch: expected str, got int + │ └─ sources [error] Missing required key 'sources' + └─ TaskGroup[1] -> Report + └─ report [error] Missing required key 'report' + """ + # Build a nested tree: dict[str, node]; each node has children dict + issues list + + root = Node() + + for issue in issues: + segs = _normalize_path(issue.path) + node = root + for seg in segs: + node = node.children.setdefault(seg, Node()) + node.leaf_issues.append(issue) + + lines: list[str] = [] + + def walk(node: Node, prefix: str = "", is_last: bool = True, label: str | None = None): + branch = "└─ " if is_last else "├─ " + child_prefix = prefix + (" " if is_last else "│ ") + + if label is not None: + lines.append(prefix + branch + label) + + # Collect children sorted for stable output + items = list(node.children.items()) + for i, (name, child) in enumerate(items): + walk(child, child_prefix, i == len(items) - 1, name) + + # Leaf issues under this node + if node.leaf_issues: + # If this node also had children, we still print issues as leaves + for j, iss in enumerate(node.leaf_issues): + leaf_branch = "└─ " if j == len(node.leaf_issues) - 1 else "├─ " + # one-line summary per issue + exp = _pretty_type(iss.expected) + act = _pretty_type(iss.actual) if iss.actual is not None else "∅" + lines.append( + (child_prefix if label is not None else prefix) + + leaf_branch + + f"{iss.key} [{iss.severity}] {iss.message}" + + (f" (expected {exp}, got {act})" if "expected" not in iss.message.lower() else "") + ) + + top_items = list(root.children.items()) + for i, (name, child) in enumerate(top_items): + walk(child, "", i == len(top_items) - 1, name) + + if root.leaf_issues: + lines.append("└─ (root)") + for j, iss in enumerate(root.leaf_issues): + leaf_branch = "└─ " if j == len(root.leaf_issues) - 1 else "├─ " + exp = _pretty_type(iss.expected) + act = _pretty_type(iss.actual) if iss.actual is not None else "∅" + lines.append(f" {leaf_branch}{iss.key} [{iss.severity}] {iss.message} (expected {exp}, got {act})") + tree = "\n".join(lines) + return f"{title}\n```\n{tree}\n```" + + +class ContextError(RuntimeError): + """ + Raised when context validation fails. + Wraps a list of ValidationIssue objects and renders them + as a readable ASCII tree in str(). + """ + + def __init__(self, issues: list[ValidationIssue], *, title: str = "Context validation failed"): + self.issues = issues + self.title = title + # build the rendered message once for efficiency + self._message = render_issues_tree(issues, title=title) + super().__init__(self._message) + + def __str__(self) -> str: + return self._message + + def __repr__(self) -> str: + return f"" + + def summary(self) -> str: + """ + Short one-line summary suitable for logs. + """ + n_errors = sum(1 for i in self.issues if i.severity == "error") + n_warnings = sum(1 for i in self.issues if i.severity == "warning") + return f"{self.title}: {n_errors} error(s), {n_warnings} warning(s)" + + +def _class_name(obj: Any) -> str: + return obj.__class__.__name__ + + +def _normalize_not_with(obj: Any) -> set[str]: + """ + Accepts: + - missing/None -> empty set + - iterable of strings and/or types -> convert to class-name strings + """ + raw = getattr(obj, "not_with", None) + if not raw: + return set() + names: set[str] = set() + for item in raw: + if isinstance(item, str): + names.add(item) + elif isinstance(item, type): + names.add(item.__name__) + else: + names.add(getattr(item, "__name__", str(item))) + return names + + +def validate_taskgroup_exclusions(group, *, path: str = "TaskGroup") -> list[ValidationIssue]: + """ + Validate that within this TaskGroup there are no tasks that exclude each other. + A conflict occurs if task_i.not_with contains task_j's CLASS NAME, or vice versa. + + Recurses into nested TaskGroups, but only checks conflicts *within* each group, + not across siblings or parents. + """ + issues: list[ValidationIssue] = [] + + tasks = list(group) + + table: list[tuple[int, Any, str, set[str]]] = [ + (idx, t, _class_name(t), _normalize_not_with(t)) + for idx, t in enumerate(tasks) + ] + + seen_pairs: set[tuple[str, str]] = set() + for i in range(len(table)): + idx_i, task_i, name_i, not_with_i = table[i] + for j in range(i + 1, len(table)): + idx_j, task_j, name_j, not_with_j = table[j] + + i_blocks_j = name_j in not_with_i + j_blocks_i = name_i in not_with_j + if not (i_blocks_j or j_blocks_i): + continue + + pair_key = tuple(sorted((name_i, name_j))) + if pair_key in seen_pairs: + continue + seen_pairs.add(pair_key) + + tname_i = getattr(task_i, "name", name_i) + tname_j = getattr(task_j, "name", name_j) + + reasons = [] + if i_blocks_j: + reasons.append(f"{name_i}.not_with contains '{name_j}'") + if j_blocks_i: + reasons.append(f"{name_j}.not_with contains '{name_i}'") + reason_txt = " and ".join(reasons) + key = f"{name_i} × {name_j}" # noqa: RUF001 + issues.append(ValidationIssue( + path=path, + key=key, + expected=None, + actual=None, + message=( + f"Mutually exclusive tasks selected: " + f"'{tname_i}' (index {idx_i}) and '{tname_j}' (index {idx_j}). " + f"Conflict because {reason_txt}. Remove one of them." + ), + severity="error", + )) + + return issues diff --git a/lumen/ai/controls.py b/lumen/ai/controls.py index a958d1f0e..149d5b529 100644 --- a/lumen/ai/controls.py +++ b/lumen/ai/controls.py @@ -10,17 +10,23 @@ import pandas as pd import param -from panel.pane.markup import HTML +from panel.io import state +from panel.layout import Column, HSpacer, Row +from panel.pane.markup import HTML, Markdown from panel.viewable import Viewer from panel.widgets import FileDropper, Tabulator, Tqdm +from panel_gwalker import GraphicWalker from panel_material_ui import ( - Button, ChatAreaInput, Column, FlexBox, Row, Select, Tabs, TextInput, - ToggleIcon, + Button, Card, ChatAreaInput, CheckBoxGroup, FlexBox, IconButton, + MultiChoice, Select, Switch, Tabs, TextInput, ToggleIcon, Typography, ) +from ..pipeline import Pipeline from ..sources.duckdb import DuckDBSource +from ..transforms.sql import SQLLimit from ..util import detect_file_encoding -from .memory import _Memory, memory +from .config import SOURCE_TABLE_SEPARATOR +from .context import TContext TABLE_EXTENSIONS = ("csv", "parquet", "parq", "json", "xlsx", "geojson", "wkt", "zip") @@ -125,6 +131,8 @@ class SourceControls(Viewer): clear_uploads = param.Boolean(default=True, doc="Clear uploaded file tabs") + context = param.Dict(default={}) + disabled = param.Boolean(default=False, doc="Disable controls") downloaded_files = param.Dict(default={}, doc="Downloaded files to add as tabs") @@ -137,16 +145,14 @@ class SourceControls(Viewer): cancel = param.Event(doc="Cancel") - memory = param.ClassSelector(class_=_Memory, default=None, doc=""" - Local memory which will be used to provide the agent context. - If None the global memory will be used.""") - multiple = param.Boolean(default=True, doc="Allow multiple files") show_input = param.Boolean(default=True, doc="Whether to show the input controls") replace_controls = param.Boolean(default=False, doc="Replace controls on add") + outputs = param.Dict(default={}) + table_upload_callbacks = {} _last_table = param.String(default="", doc="Last table added") @@ -155,7 +161,6 @@ class SourceControls(Viewer): def __init__(self, **params): super().__init__(**params) - self.tables_tabs = Tabs(sizing_mode="stretch_width") self._markitdown = None self._file_input = FileDropper( @@ -184,7 +189,7 @@ def __init__(self, **params): self._input_tabs = Tabs( ("Upload Files", self._file_input), ("Download from URL", self._url_input), - sizing_mode="stretch_both", + sizing_mode="stretch_width", dynamic=True, active=self.param.active, ) @@ -231,10 +236,6 @@ def __init__(self, **params): self._active_download_task = None # Track active download task for cancellation self._add_downloaded_files_as_tabs() - @property - def _memory(self): - return memory if self.memory is None else self.memory - def _handle_cancel(self, event): """Handle cancel button click by cancelling active download task""" if self._active_download_task and not self._active_download_task.done(): @@ -608,8 +609,13 @@ def _add_table( else: df_rel.to_view(table) duckdb_source.tables[table] = sql_expr - self._memory["source"] = duckdb_source - self._memory["table"] = table + self.outputs["source"] = duckdb_source + if "sources" not in self.outputs: + self.outputs["sources"] = [duckdb_source] + else: + self.outputs["sources"].append(duckdb_source) + self.outputs["table"] = table + self.param.trigger('outputs') self._last_table = table return 1 @@ -637,15 +643,16 @@ def _add_document( "comments": document_controls._metadata_input.value, } document = {"text": text, "metadata": metadata} - if "document_sources" in self._memory: - for i, source in enumerate(self._memory["document_sources"]): + if "document_sources" in self.outputs: + for i, source in enumerate(self.outputs["document_sources"]): if source.get("metadata", {})["filename"] == metadata["filename"]: - self._memory["document_sources"][i] = document + self.outputs["document_sources"][i] = document break else: - self._memory["document_sources"].append(document) + self.outputs["document_sources"].append(document) else: - self._memory["document_sources"] = [document] + self.outputs["document_sources"] = [document] + self.param.trigger('outputs') return 1 @param.depends("add", watch=True) @@ -658,7 +665,6 @@ def add_medias(self): return with self.menu.param.update(loading=True): - source = None n_tables = 0 n_docs = 0 @@ -687,10 +693,10 @@ def add_medias(self): n_docs += self._add_document(media_controls.file_obj, media_controls) if self.replace_controls: - src = self._memory.get("source") + src = self.output.get("source") if src: self.tables_tabs[:] = [ - (t, Tabulator(src.get(t), sizing_mode="stretch_both", pagination="remote")) + (t, Tabulator(src.get(t), sizing_mode="stretch_width", pagination="remote")) for t in src.get_tables() ] self.menu[0].visible = False @@ -713,10 +719,6 @@ def add_medias(self): self._file_input.value = {} self._url_input.value = "" - if n_docs > 0: - # Rather than triggering document sources on every upload, trigger it once - self._memory.trigger("document_sources") - # Clear uploaded files and URLs from memory if (n_tables + n_docs) > 0: self._message_placeholder.param.update( @@ -731,11 +733,15 @@ def __panel__(self): return self.menu -class RetryControls(Viewer): +class RevisionControls(Viewer): + + active = param.Boolean(False, doc="Click to revise") - active = param.Boolean(False, doc="Click to retry") + instruction = param.String(doc="Instruction to LLM to revise output") - reason = param.String(doc="Reason for retry") + toggle_kwargs = {} + + input_kwargs = {} def __init__(self, **params): super().__init__(**params) @@ -743,13 +749,12 @@ def __init__(self, **params): self.param.active, active_icon="cancel", color="default", - description="Prompt LLM to retry", - icon="edit", icon_size="1em", label="", margin=(5, 0), size="small", - sx={".MuiIcon-root": {"color": "var(--mui-palette-default-dark)"}} + sx={".MuiIcon-root": {"color": "var(--mui-palette-default-dark)"}}, + **self.toggle_kwargs ) self._text_input = TextInput( placeholder="Enter feedback and press the to retry.", @@ -758,14 +763,12 @@ def __init__(self, **params): margin=(5, 0), size="small" ) - row = Row(icon, self._text_input) - self._row = row - self._text_input.param.watch(self._enter_reason, "enter_pressed") + self._row = Row(icon, self._text_input) def _enter_reason(self, _): self.param.update( - reason=self._text_input.value_input, + instruction=self._text_input.value_input, active=False, ) @@ -773,47 +776,427 @@ def __panel__(self): return self._row -class AnnotationControls(Viewer): +class RetryControls(RevisionControls): + + instruction = param.String(doc="Reason for retry") + + input_kwargs = {"placeholder": "Enter feedback and press the to retry."} + + toggle_kwargs = {"icon": "edit", "description": "Prompt LLM to retry"} + + + +class AnnotationControls(RevisionControls): """Controls for adding annotations to visualizations.""" - active = param.Boolean(False, doc="Click to add annotations") + input_kwargs = { + "placeholder": "Describe what to annotate (e.g., 'highlight peak values', 'mark outliers')..." + } + + toggle_kwargs = { + "description": "Add annotations to highlight key insights", + "icon": "chat-bubble", + } + + +class TableSourceCard(Viewer): + """ + A component that displays a single data source as a card with table selection controls. + + The card includes: + - A header with the source name and a checkbox to toggle all tables + - A delete button (if multiple sources exist) + - Individual checkboxes for each table in the source + - Metadata display showing source information like filenames and other key-value pairs + """ + + all_selected = param.Boolean(default=True, doc=""" + Whether all tables should be selected by default.""") + + collapsed = param.Boolean(default=False, doc=""" + Whether the card should start collapsed.""") + + context = param.Dict() + + deletable = param.Boolean(default=True, doc=""" + Whether to show the delete button.""") - annotation_request = param.String(doc="User's annotation request") + delete = param.Event(doc="""Action to delete this source from memory.""") + + selected = param.List(default=None, doc=""" + List of currently selected table names.""") + + source = param.Parameter(doc=""" + The data source to display in this card.""") def __init__(self, **params): super().__init__(**params) - icon = ToggleIcon.from_param( - self.param.active, - active_icon="cancel", - color="default", - description="Add annotations to highlight key insights", - icon="chat-bubble", - icon_size="1em", - label="", - margin=(5, 0), - size="small", - sx={".MuiIcon-root": {"color": "var(--mui-palette-default-dark)"}} + self.all_tables = self.source.get_tables() + + # Determine which tables are currently visible + if self.selected is None: + visible_tables = [] + for table in self.all_tables: + visible_tables.append(table) + self.selected = visible_tables + + # Create widgets once in init + self.source_toggle = Switch.from_param( + self.param.all_selected, + name=f"{self.source.name}", + margin=(5, -5, 0, 3), + sizing_mode='fixed', ) - self._text_input = TextInput( - placeholder="Describe what to annotate (e.g., 'highlight peak values', 'mark outliers')...", - visible=icon.param.value, - max_length=200, - margin=(5, 0), - size="small" + + self.delete_button = IconButton.from_param( + self.param.delete, + icon='delete', + icon_size='1em', + color="danger", + margin=(5, 0, 0, 0), + sizing_mode='fixed', + width=40, + height=40, + visible=self.param.deletable ) - row = Row(icon, self._text_input) - self._row = row - self._text_input.param.watch(self._enter_request, "enter_pressed") + # Create table checkboxes with metadata + self.table_controls = self._create_table_controls() - def _enter_request(self, _): - """Handle Enter key press in text input.""" - self.param.update( - annotation_request=self._text_input.value_input, - active=False, + # Create source-level metadata display (if any non-table metadata exists) + self.metadata_display = self._create_source_metadata_display() + + def _create_table_controls(self): + """Create table checkboxes with per-table metadata displayed next to each checkbox.""" + table_controls = [] + + for table in self.all_tables: + # Create checkbox for this table + checkbox = CheckBoxGroup( + value=[table] if table in self.selected else [], + options=[table], + sizing_mode='stretch_width', + margin=(2, 10), + name="", + ) + + # Get metadata for this table + table_metadata = self.source.metadata.get(table, {}) if self.source.metadata else {} + metadata_parts = [] + + for key, value in table_metadata.items(): + if isinstance(value, list): + value_str = ', '.join(str(v) for v in value) + else: + value_str = str(value) + metadata_parts.append(f"{key}: {value_str}") + + if metadata_parts: + metadata_text = '; '.join(metadata_parts) + metadata_display = Typography( + metadata_text, + variant="caption", + color="text.secondary", + margin=(-10, 10, 0, 42), + sizing_mode='stretch_width', + styles={"min-height": "unset"} # Override ChatMessage CSS + ) + table_controls.extend([checkbox, metadata_display]) + else: + table_controls.append(checkbox) + + # Watch checkbox changes + checkbox.param.watch(self._on_table_checkbox_change, 'value') + + return Column(*table_controls, margin=0, sizing_mode='stretch_width') + + def _on_table_checkbox_change(self, event): + """Handle individual table checkbox changes.""" + # Collect all selected tables from all checkboxes + selected_tables = [] + for obj in self.table_controls.objects: + if obj.value: + selected_tables.extend(obj.value) + + # Update selected parameter + self.selected = selected_tables + + def _create_source_metadata_display(self): + """Create a metadata display widget for source-level metadata (non-table metadata).""" + metadata_parts = [] + + if self.source.metadata: + # Only show metadata that's not table-specific + for key, value in self.source.metadata.items(): + if key not in self.all_tables: # Skip table-specific metadata + if isinstance(value, list): + value_str = ', '.join(str(v) for v in value) + else: + value_str = str(value) + metadata_parts.append(f"{key}: {value_str}") + + if metadata_parts: + metadata_text = '; '.join(metadata_parts) + return Typography( + metadata_text, + variant="caption", + color="text.secondary", + margin=(0, 10, 5, 10), + sizing_mode='stretch_width', + ) + else: + return Typography( + "", + margin=0, + sizing_mode='stretch_width', + visible=False + ) + + @param.depends('all_selected', watch=True) + def _on_source_toggle(self): + """Handle source checkbox toggle (all tables on/off).""" + if not self.all_selected and len(self.selected) == len(self.all_tables): + # Important to check to see if all tables are selected for intuitive behavior + self.selected = [] + elif self.all_selected: + self.selected = self.all_tables + + @param.depends('selected', watch=True) + def _update_visible_slugs(self): + """Update visible_slugs in memory based on selected tables.""" + self.all_selected = len(self.selected) == len(self.all_tables) + for table in self.all_tables: + table_slug = f"{self.source.name}{SOURCE_TABLE_SEPARATOR}{table}" + if table in self.selected: + self.context['visible_slugs'].add(table_slug) + else: + self.context['visible_slugs'].discard(table_slug) + + @param.depends('delete', watch=True) + def _delete_source(self): + """Handle source deletion via param.Action.""" + if self.source in self.context.get("sources", []): + # Remove all tables from this source from visible_slugs + for table in self.all_tables: + table_slug = f"{self.source.name}{SOURCE_TABLE_SEPARATOR}{table}" + self.context['visible_slugs'].discard(table_slug) + + self.context["sources"] = [ + source for source in self.context.get("sources", []) + if source is not self.source + ] + + def __panel__(self): + card_header = Row( + self.source_toggle, + HSpacer(), + self.delete_button, + sizing_mode='stretch_width', + align='start', + height=35, + margin=0 + ) + + # Create the card content with metadata display + card_content = Column( + self.metadata_display, + self.table_controls, + margin=0, + sizing_mode='stretch_width' + ) + + # Create the card + return Card( + card_content, + header=card_header, + collapsible=True, + collapsed=self.param.collapsed, + sizing_mode='stretch_width', + margin=0, + name="TableSourceCard" + ) + + +class SourceCatalog(Viewer): + """ + A component that displays all data sources with table selection controls. + + This component shows each source as a collapsible card with: + - A header checkbox to toggle all tables in the source + - Individual checkboxes for each table + - A delete button to remove the source (if multiple sources exist) + + Tables can be selectively shown/hidden using the checkboxes, which updates + the 'visible_slugs' set in memory. + """ + + context = param.Dict(default={}) + + sources = param.List(default=[], doc=""" + List of data sources to display in the catalog.""") + + def __init__(self, /, context: TContext | None = None, **params): + self._title = Markdown(margin=0) + self._cards_column = Column() + self._layout = Column( + self._title, + self._cards_column, + sizing_mode='stretch_width' + ) + if context is None: + raise ValueError("SourceCatalog must be given a context dictionary.") + if "source" in context and "sources" not in context: + context["sources"] = [context["source"]] + if "visible_slugs" not in context: + context["visible_slugs"] = set() + super().__init__(context=context, **params) + + @param.depends("sources", watch=True, on_init=True) + async def sync(self, context: TContext | None = None): + """ + Trigger the catalog with new sources. + + Args: + sources: Optional list of sources. If None, uses sources from memory. + """ + context = context or self.context + sources = self.sources or context.get('sources', []) + + # Create a lookup of existing cards by source + existing_cards = { + card.source: card for card in self._cards_column.objects + if isinstance(card, TableSourceCard) and card.source in sources + } + + # Build the new cards list + source_cards = [] + multiple_sources = len(sources) > 1 + for source in sources: + if source in existing_cards: + # Reuse existing card and update its deletable property + card = existing_cards[source] + card.deletable = multiple_sources + source_cards.append(card) + else: + # Create new card for new source + source_card = TableSourceCard( + context=context, + source=source, + deletable=multiple_sources, + collapsed=multiple_sources, + ) + source_cards.append(source_card) + + self._cards_column.objects = source_cards + + if len(self.sources) == 0: + self._title.object = "No sources available. Add a source to get started." + else: + self._title.object = "Select the table and document sources you want visible to the LLM." + + def __panel__(self): + """ + Create the source catalog UI. + + Returns a Column containing all source cards or a message if no sources exist. + """ + return self._layout + + +class TableExplorer(Viewer): + """ + TableExplorer provides a high-level entrypoint to explore tables in a split UI. + It allows users to load tables, explore them using Graphic Walker, and then + interrogate the data via a chat interface. + """ + + context = param.Dict(default={}) + + def __init__(self, **params): + self._initialized = False + super().__init__(**params) + self._table_select = MultiChoice( + label="Select table(s) to preview", sizing_mode='stretch_width', + max_height=200, max_items=5, margin=0 + ) + self._explore_button = Button( + name='Explore table(s)', icon='add_chart', button_type='primary', icon_size="2em", + disabled=self._table_select.param.value.rx().rx.not_(), on_click=self._update_explorers, + margin=(0, 0, 0, 10), width=200, align='end' + ) + self._input_row = Row(self._table_select, self._explore_button) + self._source_map = {} + self._tabs = Tabs(dynamic=True, sizing_mode='stretch_both') + self._layout = Column( + self._input_row, self._tabs, sizing_mode='stretch_both', ) - # Clear the input after submission - self._text_input.value = "" + + @param.depends("context", watch=True, on_init=True) + async def sync(self, context: TContext | None = None): + init = not self._initialized + self._initialized = True + context = context or self.context + if "sources" in context: + sources = context["sources"] + elif "source" in context: + sources = [context["source"]] + else: + return + selected = list(self._table_select.value) + deduplicate = len(sources) > 1 + new = {} + + # Build the source map for UI display + for source in sources: + tables = source.get_tables() + for table in tables: + if deduplicate: + table = f'{source.name}{SOURCE_TABLE_SEPARATOR}{table}' + + if (table.split(SOURCE_TABLE_SEPARATOR, maxsplit=1)[-1] not in self._source_map and + not init and not len(selected) > self._table_select.max_items and state.loaded): + selected.append(table) + new[table] = source + + self._source_map.clear() + self._source_map.update(new) + selected = selected if len(selected) == 1 else [] + self._table_select.param.update(options=list(self._source_map), value=selected) + self._input_row.visible = bool(self._source_map) + self._initialized = True + + def _explore_table_if_single(self, event): + """ + If only one table is uploaded, help the user load it + without requiring them to click twice. This step + only triggers when the Upload in the Overview tab is used, + i.e. does not trigger with uploads through the SourceAgent + """ + if len(self._table_select.options) == 1: + self._explore_button.param.trigger("value") + + def _update_explorers(self, event): + if not event.new: + return + + with self._explore_button.param.update(loading=True): + explorers = [] + for table in self._table_select.value: + source = self._source_map[table] + if SOURCE_TABLE_SEPARATOR in table: + _, table = table.split(SOURCE_TABLE_SEPARATOR, maxsplit=1) + pipeline = Pipeline( + source=source, table=table, sql_transforms=[SQLLimit(limit=100_000, read=source.dialect)] + ) + table_label = f"{table[:25]}..." if len(table) > 25 else table + walker = GraphicWalker( + pipeline.param.data, sizing_mode='stretch_both', min_height=800, + kernel_computation=True, name=table_label, tab='data' + ) + explorers.append(walker) + + self._tabs.objects = explorers + self._table_select.value = [] def __panel__(self): - return self._row + return self._layout diff --git a/lumen/ai/coordinator.py b/lumen/ai/coordinator.py index 128129181..8cafd1721 100644 --- a/lumen/ai/coordinator.py +++ b/lumen/ai/coordinator.py @@ -22,12 +22,12 @@ from .actor import Actor from .agents import Agent, AnalysisAgent, ChatAgent -from .components import TableSourceCard from .config import ( DEMO_MESSAGES, GETTING_STARTED_SUGGESTIONS, PROMPTS_DIR, - SOURCE_TABLE_SEPARATOR, MissingContextError, + MissingContextError, ) -from .controls import SourceControls +from .context import ContextError, TContext +from .controls import SourceControls, TableSourceCard from .llm import LlamaCpp, Llm, Message from .logs import ChatLogs from .models import ( @@ -39,7 +39,7 @@ ) from .utils import ( fuse_messages, get_root_exception, log_debug, mutate_user_message, - normalized_name, stream_details, wrap_logfire, + normalized_name, wrap_logfire, ) from .views import AnalysisOutput @@ -59,6 +59,8 @@ class Plan(Section): agents = param.List(item_type=Actor, default=[]) + coordinator = param.ClassSelector(class_=param.Parameterized) + interface = param.ClassSelector(class_=ChatFeed) def _render_task_history(self, i: int) -> tuple[list[Message], str]: @@ -76,39 +78,46 @@ def _render_task_history(self, i: int) -> tuple[list[Message], str]: for msg in self.history ], todos - async def _run_task(self, i: int, task: Self | Actor, **kwargs): + async def _run_task(self, i: int, task: Self | Actor, context: TContext, **kwargs): outputs = [] - with self.interface.add_step(title=f"{task.title}...", user="Runner", layout_params={"title": "🏗️ Plan Execution Steps"}, steps_layout=self.steps_layout) as step: - self._coordinator._todos_title.object = f"⚙️ Working on task {task.title!r}..." + with self.interface.add_step( + title=f"{task.title}...", user="Runner", layout_params={"title": "🏗️ Plan Execution Steps"}, + steps_layout=self.coordinator.steps_layout + ) as step: + self.coordinator._todos_title.object = f"⚙️ Working on task {task.title!r}..." step.stream(f"`Working on task {task.title}`:\n\n{task.instruction}") history, todos = self._render_task_history(i) - todos_obj = self._coordinator._todos - todos_obj.object = todos + self.coordinator._todos.object = todos try: kwargs = {"agents": self.agents} if 'agents' in task.param else {} with task.param.update( - memory=self.memory, interface=self.interface, steps_layout=self.steps_layout, + interface=self.interface, steps_layout=self.coordinator.steps_layout, history=history, **kwargs ): - outputs += await task.execute(**kwargs) + new, task_context = await task.execute(context, **kwargs) + outputs += new except Exception as e: # Handle the exception using the dedicated error handler - error_outputs = await self._handle_task_execution_error(e, task, step, i) + task_context = {} + error_outputs = await self._handle_task_execution_error(e, task, task_context, step, i) if error_outputs is not None: - return error_outputs + return error_outputs, task_context if isinstance(task, TaskGroup): - unprovided = [p for actor in task for p in actor.provides if p not in self.memory] + unprovided = [ + p for actor in task for p in actor.output_schema.__annotations__ + if p not in task_context + ] else: unprovided = [] if unprovided: step.failed_title = f"{task.title} did not provide {', '.join(unprovided)}. Aborting the plan." - raise RuntimeError(f"{task.title} failed to provide declared context.") + raise RuntimeError(f"{task.title!r} task failed to provide declared context.") log_debug(f"\033[96mCompleted: {task.title}\033[0m", show_length=False) step.stream(f"\n\nSuccessfully completed task {task.title}:\n\n> {task.instruction}", replace=True) step.success_title = f"{task.title} successfully completed" - return outputs + return outputs, task_context - async def _handle_task_execution_error(self, e: Exception, task: Self | Actor, step: ChatStep, i: int) -> list | None: + async def _handle_task_execution_error(self, e: Exception, task: Self | Actor, context: TContext, step: ChatStep, i: int) -> list | None: """ Handle exceptions that occur during task execution. Returns outputs if the error was handled successfully, None if the error should be re-raised. @@ -122,7 +131,7 @@ async def _handle_task_execution_error(self, e: Exception, task: Self | Actor, s log_debug(f"\033[93mMissing context detected: {root_exception!s}\033[0m") # Find which agent provided the pipeline or relevant context - provider_index = self._find_context_provider(i) + provider_index = self._find_context_provider(i, context) if provider_index is not None: # Re-run from the provider with the error as feedback outputs = await self._retry_from_provider(provider_index, i, str(root_exception)) @@ -139,15 +148,15 @@ async def _handle_task_execution_error(self, e: Exception, task: Self | Actor, s raise e else: traceback.print_exception(e) - self.memory['__error__'] = str(e) + context['__error__'] = str(e) raise e - def _find_context_provider(self, failed_index: int) -> int | None: + def _find_context_provider(self, failed_index: int, context: TContext) -> int | None: """ Find the task that provided the pipeline or other relevant context. Search backwards from the failed task. """ - pipeline_exists = 'pipeline' in self.memory + pipeline_exists = 'pipeline' in context if not pipeline_exists: return @@ -156,9 +165,8 @@ def _find_context_provider(self, failed_index: int) -> int | None: # Check if this task provides pipeline or other relevant context if isinstance(task, TaskGroup): for actor in task: - if hasattr(actor, 'provides'): - if 'pipeline' in actor.provides: - return idx + if 'pipeline' in actor.output_schema.__annotations__: + return idx async def _retry_from_provider(self, provider_index: int, failed_index: int, error_message: str) -> list: """ @@ -195,23 +203,37 @@ async def _retry_from_provider(self, provider_index: int, failed_index: int, err f"- [{'x' if tidx < idx else '🔄' if tidx == idx else ' '}] {'' + t.instruction + '' if tidx == idx else t.instruction}" for tidx, t in enumerate(self) ) - self._coordinator._todos.object = todos + self.coordinator._todos.object = todos # Run with mutated history kwargs = {"agents": self.agents} if 'agents' in task.param else {} with task.param.update( - memory=self.memory, interface=self.interface, steps_layout=self.steps_layout, + interface=self.interface, steps_layout=self.steps_layout, history=retry_history, **kwargs ): outputs += await task.execute(**kwargs) retry_step.success_title = f"✅ {task.title} successfully completed on retry" return outputs - async def execute(self, **kwargs): - ret = await super().execute(**kwargs) + async def execute(self, context: TContext = None, **kwargs) -> list[Any]: + context = context or self.context + if '__error__' in context: + del context['__error__'] + outputs, out_context = await super().execute(context, **kwargs) _, todos = self._render_task_history(len(self)) - self._coordinator._todos.object = todos - return ret + self.coordinator._todos.object = todos + if self.status == 'success': + self.coordinator._todos_title.object = f"✅ Sucessfully completed {self.title!r}" + else: + self.coordinator._todos_title.object = f"❌ Failed to execute {self.title!r}" + if "pipeline" in out_context: + await self.coordinator._add_analysis_suggestions(out_context) + log_debug("\033[92mCompleted: Plan\033[0m", show_sep="below") + if self.interface is not None: + for message_obj in self.interface.objects[::-1]: + if isinstance(message_obj.object, Card): + message_obj.object.collapsed = True + return outputs, out_context class Coordinator(Viewer, VectorLookupToolUser): @@ -237,6 +259,8 @@ class Coordinator(Viewer, VectorLookupToolUser): agents = param.List(default=[ChatAgent], doc=""" List of agents to coordinate.""") + context = param.Dict(default={}) + demo_inputs = param.List(default=DEMO_MESSAGES, doc=""" List of instructions to demo the Coordinator.""") @@ -266,6 +290,7 @@ def __init__( interface: ChatFeed | None = None, agents: list[Agent | type[Agent]] | None = None, tools: list[Tool | type[Tool]] | None = None, + context: TContext | None = None, vector_store: VectorStore | None = None, document_vector_store: VectorStore | None = None, logs_db_path: str = "", @@ -309,7 +334,7 @@ def on_rerun(instance, _): self._logs.update_status(message_id=id(message), removed=True) def on_clear(instance, _): - self._memory.cleanup() + pass def on_submit(event=None, instance=None): chat_input = self.interface.active_widget @@ -323,12 +348,12 @@ def on_submit(event=None, instance=None): # Reset value input because reset has no time to propagate self._main[:] = [self.interface] - old_sources = self._memory.get("sources", []) + old_sources = context.get("sources", []) if uploaded: # Process uploaded files through SourceControls if any exist source_controls = SourceControls( downloaded_files={key: value["value"] for key, value in uploaded.items()}, - memory=self._memory, + context=context, replace_controls=False, show_input=False, clear_uploads=True # Clear the uploads after processing @@ -337,7 +362,7 @@ def on_submit(event=None, instance=None): chat_input.value_uploaded = {} source_cards = [ TableSourceCard(source=source, name=source.name) - for source in self._memory.get("sources", []) if source not in old_sources + for source in context.get("sources", []) if source not in old_sources ] if len(source_cards) > 1: source_view = Accordion(*source_cards, sizing_mode="stretch_width", name="TableSourceCard") @@ -372,7 +397,9 @@ def on_submit(event=None, instance=None): variant="h1" ) - num_sources = len(self._memory.get("sources", [])) + if context is None: + context = {} + num_sources = len(context.get("sources", [])) prefix_text = "Add your dataset to begin, then" if num_sources == 0 else f"{num_sources} source{'s' if num_sources > 1 else ''} connected;" welcome_text = Typography( f"{prefix_text} ask any question, or select a quick action below." @@ -454,7 +481,8 @@ def on_submit(event=None, instance=None): params["prompts"]["main"]["tools"] += [tool for tool in tools] super().__init__( llm=llm, agents=instantiated, interface=interface, logs_db_path=logs_db_path, - vector_store=vector_store, document_vector_store=document_vector_store, **params + vector_store=vector_store, document_vector_store=document_vector_store, context=context, + **params ) interface.button_properties = { @@ -463,17 +491,6 @@ def on_submit(event=None, instance=None): "clear": {"callback": on_clear}, } - # Set up automatic synchronization between source and sources FIRST - # so the initial setup below triggers automatic sync - self._memory.on_change("source", self._sync_source_to_sources) - self._memory.on_change("sources", self._sync_sources_to_source) - self._memory.on_change("sources", self._update_visible_slugs) - - # Initialize memory - self._memory["sources"] = self._memory.get("sources", []) - self._sync_source_to_sources(None, None, self._memory.get("source", None)) - self._sync_sources_to_source(None, None, self._memory["sources"]) - self._update_visible_slugs(None, None, self._memory["sources"]) # Use existing method to create suggestions self._add_suggestions_to_footer( @@ -485,72 +502,6 @@ def on_submit(event=None, instance=None): hide_after_use=False ) - def _update_visible_slugs(self, key=None, old_sources=None, new_sources=None): - """ - Update visible_slugs when sources change. - This is the central place where table visibility is managed. - """ - if not new_sources: - self._memory['visible_slugs'] = set() - return - - # Calculate all available table slugs from sources - all_slugs = set() - for source in new_sources: - tables = source.get_tables() - for table in tables: - table_slug = f'{source.name}{SOURCE_TABLE_SEPARATOR}{table}' - all_slugs.add(table_slug) - - # Update visible_slugs, preserving existing visibility where possible - # This ensures removed tables are filtered out, new tables are added - current_visible = self._memory.get('visible_slugs', set()) - if current_visible: - # Keep intersection of current visible and available slugs - # Plus add any new slugs that weren't previously available - self._memory['visible_slugs'] = current_visible.intersection(all_slugs) | (all_slugs - current_visible) - else: - # If no visible_slugs set, make all tables visible - self._memory['visible_slugs'] = all_slugs - - def _sync_source_to_sources(self, key, old_source, new_source): - """ - When source is set/changed, automatically add it to sources if it doesn't exist. - This eliminates the need for manual dual updates. - """ - if new_source is None: - return - - current_sources = self._memory.get("sources", []) - # Check if the new source already exists in sources (by name) - existing_source = next( - (source for source in current_sources if source.name == new_source.name), - None - ) - - if existing_source is None: - # Add new source to sources list - self._memory["sources"] = current_sources + [new_source] - elif existing_source is not new_source: - # Replace existing source with new one (in case it's updated) - updated_sources = [ - new_source if source.name == new_source.name else source - for source in current_sources - ] - self._memory["sources"] = updated_sources - - def _sync_sources_to_source(self, key, old_sources, new_sources): - """ - When sources changes, ensure source is set to the first source. - """ - if not new_sources: - return - - current_source = self._memory.get('source') - # If current source is not in the new sources list, update it - if current_source is None or current_source not in new_sources: - self._memory['source'] = new_sources[0] - def __panel__(self): return self._main @@ -562,7 +513,7 @@ def _add_suggestions_to_footer( analysis: bool = False, append_demo: bool = True, hide_after_use: bool = True, - memory = None + context: TContext | None = None ): if not suggestions: return @@ -571,9 +522,6 @@ async def hide_suggestions(_=None): if len(self.interface.objects) > num_objects: suggestion_buttons.visible = False - if memory is None: - memory = self._memory - async def use_suggestion(event): if self._main[0] is not self.interface: self._main[:] = [self.interface] @@ -599,15 +547,10 @@ async def use_suggestion(event): log_debug("No analysis agent found.") return messages = [{"role": "user", "content": contents}] - original_memory = agent.memory - try: - with agent.param.update(memory=memory, agents=self.agents): - await agent.respond(messages) - # Pass the same memory to _add_analysis_suggestions - await self._add_analysis_suggestions(memory=memory) - finally: - # Reset agent memory to original state - agent.memory = original_memory + with agent.param.update(agents=self.agents): + await agent.respond(messages, context) + # Pass the same memory to _add_analysis_suggestions + await self._add_analysis_suggestions(context) else: self.interface.send(contents) @@ -662,11 +605,9 @@ async def run_demo(event): self.interface.param.watch(hide_suggestions, "objects") - async def _add_analysis_suggestions(self, memory=None): - if memory is None: - memory = self._memory - pipeline = memory["pipeline"] - current_analysis = memory.get("analysis") + async def _add_analysis_suggestions(self, context: TContext): + pipeline = context["pipeline"] + current_analysis = context.get("analysis") # Clear current_analysis unless the last message is the same AnalysisOutput if current_analysis and self.interface.objects: @@ -683,15 +624,15 @@ async def _add_analysis_suggestions(self, memory=None): [f"Apply {analysis.__name__}" for analysis in applicable_analyses], append_demo=False, analysis=True, + context=context, hide_after_use=False, num_objects=len(self.interface.objects), - memory=memory ) @wrap_logfire(span_name="Chat Invoke") - async def _chat_invoke(self, contents: list | str, user: str, instance: ChatInterface) -> Plan: - log_debug(f"New Message: \033[91m{contents!r}\033[0m", show_sep="above") - return await self.respond(contents) + async def _chat_invoke(self, messages: list[Message], context: TContext, user: str, instance: ChatInterface) -> Plan: + log_debug(f"New Message: \033[91m{messages!r}\033[0m", show_sep="above") + return await self.respond(messages, context) async def _fill_model(self, messages, system, agent_model): model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) @@ -704,25 +645,32 @@ async def _fill_model(self, messages, system, agent_model): return out async def _pre_plan( - self, messages: list[Message], agents: dict[str, Agent], tools: dict[str, Tool] + self, messages: list[Message], context: TContext, agents: dict[str, Agent], tools: dict[str, Tool] ) -> tuple[dict[str, Agent], dict[str, Tool], dict[str, Any]]: """ Pre-plan step to prepare the agents and tools for the execution graph. This is where we can modify the agents and tools based on the messages. """ # Filter agents by exclusions and applies - agents = {agent_name: agent for agent_name, agent in agents.items() if not any(excluded_key in self._memory for excluded_key in agent.exclusions)} - applies = await asyncio.gather(*[agent.applies(self._memory) for agent in agents.values()]) + agents = {agent_name: agent for agent_name, agent in agents.items() if not any(excluded_key in context for excluded_key in agent.exclusions)} + applies = await asyncio.gather(*[agent.applies(context) for agent in agents.values()]) agents = {agent_name: agent for (agent_name, agent), aapply in zip(agents.items(), applies, strict=False) if aapply} # Filter tools by exclusions and applies - tools = {tool_name: tool for tool_name, tool in tools.items() if not any(excluded_key in self._memory for excluded_key in tool.exclusions)} - applies = await asyncio.gather(*[tool.applies(self._memory) for tool in tools.values()]) + tools = {tool_name: tool for tool_name, tool in tools.items() if not any(excluded_key in context for excluded_key in tool.exclusions)} + applies = await asyncio.gather(*[tool.applies(context) for tool in tools.values()]) tools = {tool_name: tool for (tool_name, tool), tapply in zip(tools.items(), applies, strict=False) if tapply} return agents, tools, {} - async def _compute_plan(self, messages: list[Message], agents: dict[str, Agent], tools: dict[str, Tool], pre_plan_output: dict) -> Plan: + async def _compute_plan( + self, + messages: list[Message], + context: TContext, + agents: dict[str, Agent], + tools: dict[str, Tool], + pre_plan_output: dict + ) -> Plan: """ Compute the execution graph for the given messages and agents. The graph is a list of ExecutionNode objects that represent @@ -750,8 +698,8 @@ def _serialize(self, obj: Any, exclude_passwords: bool = True) -> str: obj = obj.value return str(obj) - async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> str: - self._memory["agent_tool_contexts"] = {} + async def respond(self, messages: list[Message], context: TContext, **kwargs: dict[str, Any]) -> Plan | None: + context = {"agent_tool_contexts": [], **context} with self.interface.param.update(loading=True): if isinstance(self.llm, LlamaCpp): with self.interface.add_step(title="Loading LlamaCpp model...", success_title="Using the cached LlamaCpp model", user="Assistant") as step: @@ -773,45 +721,25 @@ async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> st agents = {normalized_name(agent): agent for agent in self.agents} tools = {normalized_name(tool): tool for tool in self._tools["main"]} - agents, tools, pre_plan_output = await self._pre_plan(messages, agents, tools) - plan = await self._compute_plan(messages, agents, tools, pre_plan_output) + agents, tools, pre_plan_output = await self._pre_plan(messages, context, agents, tools) + context["plan"] = plan = await self._compute_plan(messages, context, agents, tools, pre_plan_output) if plan is None: msg = ( "Assistant could not settle on a plan of action to perform the requested query. " "Please restate your request." ) self.interface.stream(msg, user='Lumen') - return msg - - if '__error__' in self._memory: - del self._memory['__error__'] - with self.interface.param.update(callback_exception="raise"): - with plan.param.update( - history=messages, memory=self._memory, interface=self.interface, - steps_layout=self.steps_layout, agents=list(agents.values()) - ): - # Pass coordinator reference to plan for todo updates - plan._coordinator = self - await plan.execute() - - if plan.status == 'success': - self._todos_title.object = f"✅ Sucessfully completed {plan.title!r}" - else: - self._todos_title.object = f"❌ Failed to execute {plan.title!r}" - - if "pipeline" in self._memory: - await self._add_analysis_suggestions() - log_debug("\033[92mCompleted: Coordinator\033[0m", show_sep="below") + return - for message_obj in self.interface.objects[::-1]: - if isinstance(message_obj.object, Card): - message_obj.object.collapsed = True return plan - async def _check_tool_relevance(self, tool: Tool, tool_output: str, actor: Actor, actor_task: str, messages: list[Message]) -> bool: + async def _check_tool_relevance( + self, tool: Tool, tool_output: str, actor: Actor, actor_task: str, messages: list[Message], context: TContext + ) -> bool: result = await self._invoke_prompt( "tool_relevance", messages, + context, tool_name=tool.name, tool_purpose=getattr(tool, "purpose", ""), tool_output=tool_output, @@ -822,66 +750,11 @@ async def _check_tool_relevance(self, tool: Tool, tool_output: str, actor: Actor return result.yes - async def _handle_tool_result(self, tool: Tool, result: str, step: ChatStep, plan: Plan | None = None, messages: list[Message] | None = None): - """Handle string tool results and determine relevance for future agents.""" - # Caller should ensure result is a non-empty string before calling this method - - # Display the result - stream_details(result, step, title="Results", auto=False) - - # Early exit if no execution graph provided - if not plan: - return - - # Find agents that will be used in future nodes - future_agents = {} - for agent in self.agents: - # Skip tools, only interested in non-tool agents - if isinstance(agent, Tool): - continue - - # Find if this agent appears in future nodes - for task in plan.subtasks: - for actor in task.subtasks: - if actor is agent: - future_agents[agent] = task.instruction - break - - # Early exit if no future agents found - if not future_agents: - return - - # Check relevance and store context for each future agent - for agent, task in future_agents.items(): - # Get tool and agent provides/requires lists - tool_provides = getattr(tool, "provides", []) - agent_requires = getattr(agent, "requires", []) - - # If tool provides at least one thing the agent requires, consider it relevant - # without performing the more expensive relevance check - direct_dependency = any(provided in agent_requires for provided in tool_provides) - - is_relevant = False - if direct_dependency: - log_debug(f"Direct dependency detected: {tool.name} provides at least one requirement for {agent.name}") - # The agent already has it formatted in its template - continue - elif len(result) < 1000: - is_relevant = True - else: - # Otherwise, check semantic relevance - is_relevant = await self._check_tool_relevance( - tool, result, agent, task, messages - ) - - if is_relevant: - # Initialize agent_tool_contexts if needed - if agent.name not in self._memory["agent_tool_contexts"]: - self._memory["agent_tool_contexts"][agent.name] = {} - - # Store the tool output in agent's context - self._memory["agent_tool_contexts"][agent.name][tool.name] = result - log_debug(f"Added {tool.name} output to {agent.name}'s context") + async def sync(self, context: TContext | None): + context = context or self.context + for tools in self._tools.values(): + for tool in tools: + await tool.sync(context) class DependencyResolver(Coordinator): @@ -903,15 +776,16 @@ class DependencyResolver(Coordinator): async def _choose_agent( self, messages: list[Message], + context: TContext, agents: list[Agent] | None = None, primary: bool = False, unmet_dependencies: tuple[str] | None = None ): if agents is None: agents = self.agents - applies = await asyncio.gather(*[agent.applies(self._memory) for agent in agents]) + applies = await asyncio.gather(*[agent.applies(context) for agent in agents]) agents = [agent for agent, aapply in zip(agents, applies, strict=False) if aapply] - applies = await asyncio.gather(*[tool.applies(self._memory) for tool in self._tools['main']]) + applies = await asyncio.gather(*[tool.applies(context) for tool in self._tools['main']]) tools = [tool for tool, tapply in zip(self._tools['main'], applies, strict=False) if tapply] agent_names = tuple(sagent.name[:-5] for sagent in agents) + tuple(tool.name for tool in tools) @@ -921,19 +795,31 @@ async def _choose_agent( if len(agent_names) == 1: return agent_model(agent=agent_names[0], chain_of_thought='') system = await self._render_prompt( - "main", messages, agents=agents, tools=tools, primary=primary, + "main", + messages, + context, + agents=agents, + tools=tools, + primary=primary, unmet_dependencies=unmet_dependencies ) return await self._fill_model(messages, system, agent_model) - async def _compute_plan(self, messages, agents: dict[str, Agent], tools: dict[str, Tool], pre_plan_output: dict[str, Any]) -> Plan | None: + async def _compute_plan( + self, + messages: list[Message], + context: TContext, + agents: dict[str, Agent], + tools: dict[str, Tool], + pre_plan_output: dict[str, Any] + ) -> Plan | None: if len(agents) == 1: agent = next(iter(agents.values())) else: agent = None with self.interface.add_step(title="Selecting primary agent...", user="Assistant") as step: try: - output = await self._choose_agent(messages, self.agents, primary=True) + output = await self._choose_agent(messages, context, self.agents, primary=True) except Exception as e: if self.interface.callback_exception not in ('raise', 'verbose'): step.failed_title = 'Failed to select main agent...' @@ -980,7 +866,7 @@ async def _compute_plan(self, messages, agents: dict[str, Agent], tools: dict[st ) ) step.success_title = f"Solved a dependency with {output.agent_or_tool}" - return Plan(*(tasks[::-1] + [TaskGroup(agent, instruction=cot)]), history=messages) + return Plan(*(tasks[::-1] + [TaskGroup(agent, instruction=cot)]), history=messages, context=context, coordinator=self) class Planner(Coordinator): @@ -1011,25 +897,26 @@ def __init__(self, **params): params["planner_tools"] = self._initialize_tools_for_prompt(params["planner_tools"], **params) super().__init__(**params) - async def _check_follow_up_question(self, messages: list[Message]) -> bool: + async def _check_follow_up_question(self, messages: list[Message], context: TContext) -> bool: """Check if the user's query is a follow-up question about the previous dataset.""" # Only check if data is in memory - if "data" not in self._memory: + if "data" not in context: return False # Use the follow_up prompt to check result = await self._invoke_prompt( "follow_up", messages, + context ) is_follow_up = result.yes if not is_follow_up: - self._memory.pop("pipeline", None) + context.pop("pipeline", None) return is_follow_up - async def _execute_planner_tools(self, messages: list[Message]): + async def _execute_planner_tools(self, messages: list[Message], context: TContext): """Execute planner tools to gather context before planning.""" if not self.planner_tools: return @@ -1050,13 +937,10 @@ async def _execute_planner_tools(self, messages: list[Message]): with self.interface.add_step(title="Gathering context for planning...", user="Assistant", steps_layout=steps_layout) as step: for tool in self.planner_tools: is_relevant = await self._check_tool_relevance( - tool, "", self, f"Gather context for planning to answer {user_query}", messages + tool, "", self, f"Gather context for planning to answer {user_query}", messages, context ) if not is_relevant: - # remove the keys if they're irrelevant - for key in tool.provides: - self._memory.pop(key, None) continue tool_name = getattr(tool, "name", type(tool).__name__) @@ -1064,20 +948,25 @@ async def _execute_planner_tools(self, messages: list[Message]): task = TaskGroup( tool, interface=self.interface, - memory=self._memory, instruction=user_query, title=f"Gathering context with {tool_name}", steps_layout=steps_layout, ) - await task.execute() + await task.execute(context) if task.status != "error": step.stream(f"\n\n✗ Failed to gather context from {tool_name}") continue - async def _pre_plan(self, messages: list[Message], agents: dict[str, Agent], tools: dict[str, Tool]) -> tuple[dict[str, Agent], dict[str, Tool], dict[str, Any]]: - is_follow_up = await self._check_follow_up_question(messages) + async def _pre_plan( + self, + messages: list[Message], + context: TContext, + agents: dict[str, Agent], + tools: dict[str, Tool], + ) -> tuple[dict[str, Agent], dict[str, Tool], dict[str, Any]]: + is_follow_up = await self._check_follow_up_question(messages, context) if not is_follow_up: - await self._execute_planner_tools(messages) + await self._execute_planner_tools(messages, context) else: log_debug("\033[92mDetected follow-up question, using existing context\033[0m") with self.interface.add_step( @@ -1086,7 +975,7 @@ async def _pre_plan(self, messages: list[Message], agents: dict[str, Agent], too step.stream("Detected that this is a follow-up question related to the previous dataset.") step.stream("\n\nUsing the existing data in memory to answer without re-executing data retrieval.") step.success_title = "Using existing data for follow-up question" - agents, tools, pre_plan_output = await super()._pre_plan(messages, agents, tools) + agents, tools, pre_plan_output = await super()._pre_plan(messages, context, agents, tools) pre_plan_output["is_follow_up"] = is_follow_up return agents, tools, pre_plan_output @@ -1094,6 +983,7 @@ async def _pre_plan(self, messages: list[Message], agents: dict[str, Agent], too async def _make_plan( self, messages: list[Message], + context: TContext, agents: dict[str, Agent], tools: dict[str, Tool], unmet_dependencies: set[str], @@ -1107,33 +997,34 @@ async def _make_plan( tools = list(tools.values()) all_provides = set() for provider in agents + tools: - all_provides |= set(provider.provides) - all_provides |= set(self._memory.keys()) + all_provides |= set(provider.output_schema.__annotations__) + all_provides |= set(context) # filter agents using applies - agents = [agent for agent in agents if await agent.applies(self._memory)] + agents = [agent for agent in agents if await agent.applies(context)] # ensure these candidates are satisfiable # e.g. DbtslAgent is unsatisfiable if DbtslLookup was used in planning # but did not provide dbtsl_metaset # also filter out agents where excluded keys exist in memory - agents = [agent for agent in agents if len(set(agent.requires) - all_provides) == 0 and type(agent).__name__ != "ValidationAgent"] - tools = [tool for tool in tools if len(set(tool.requires) - all_provides) == 0] + agents = [agent for agent in agents if len(set(agent.input_schema.__annotations__) - all_provides) == 0 and type(agent).__name__ != "ValidationAgent"] + tools = [tool for tool in tools if len(set(tool.input_schema.__annotations__) - all_provides) == 0] reasoning = None while reasoning is None: # candidates = agents and tools that can provide # the unmet dependencies agent_candidates = [ agent for agent in agents - if not unmet_dependencies or set(agent.provides) & unmet_dependencies + if not unmet_dependencies or set(agent.output_schema.__annotations__) & unmet_dependencies ] tool_candidates = [ tool for tool in tools - if not unmet_dependencies or set(tool.provides) & unmet_dependencies + if not unmet_dependencies or set(tool.output_schema.__annotations__) & unmet_dependencies ] system = await self._render_prompt( "main", messages, + context, agents=agents, tools=tools, unmet_dependencies=unmet_dependencies, @@ -1152,7 +1043,7 @@ async def _make_plan( ): self.steps_layout.title = "🧠 Reasoning about the plan..." if reasoning.chain_of_thought: # do not replace with empty string - self._memory["reasoning"] = reasoning.chain_of_thought + context["reasoning"] = reasoning.chain_of_thought step.stream(reasoning.chain_of_thought, replace=True) previous_plans.append(reasoning.chain_of_thought) @@ -1164,11 +1055,12 @@ async def _resolve_plan( agents: dict[str, Agent], tools: dict[str, Tool], messages: list[Message], + context: TContext, previous_actors: list[str], ) -> tuple[Plan, set[str], list[str]]: table_provided = False tasks = [] - provided = set(self._memory) + provided = set(context) unmet_dependencies = set() steps = [] actors = [] @@ -1221,22 +1113,15 @@ async def _resolve_plan( log_debug(f"Warning: Agent or tool '{key}' not found in available agents/tools") continue - # Check not_with constraints - not_with = getattr(subagent, 'not_with', []) - conflicts = [actor for actor in actors_in_graph if actor in not_with] - if conflicts: - # just to prompt the LLM - unmet_dependencies.add(f"{key} is incompatible with {', '.join(conflicts)}") - requires = set(await subagent.requirements(messages)) - provided |= set(subagent.provides) + provided |= set(subagent.output_schema.__annotations__) unmet_dependencies = (unmet_dependencies | requires) - provided has_table_lookup = any( any(isinstance(st, TableLookup) for st in task) for task in tasks ) if "table" in unmet_dependencies and not table_provided and "SQLAgent" in agents and has_table_lookup: - provided |= set(agents['SQLAgent'].provides) + provided |= set(agents['SQLAgent'].output_schema.__annotations__) sql_step = type(step)( actor='SQLAgent', instruction='Load the table', @@ -1265,7 +1150,7 @@ async def _resolve_plan( last_task = tasks[-1] if isinstance(last_task[0], Tool): - if "AnalystAgent" in agents and all(r in provided for r in agents["AnalystAgent"].requires): + if "AnalystAgent" in agents and all(r in provided for r in agents["AnalystAgent"].input_schema.__annotations__): actor = "AnalystAgent" else: actor = "ChatAgent" @@ -1278,7 +1163,7 @@ async def _resolve_plan( log_debug(f"Skipping summarization with {actor} due to conflicts: {conflicts}") raw_plan.steps = steps previous_actors = actors - return Plan(*tasks, title=raw_plan.title, history=messages), unmet_dependencies, previous_actors + return Plan(*tasks, title=raw_plan.title, history=messages, context=context, coordinator=self), previous_actors summarize_step = type(step)( actor=actor, @@ -1312,9 +1197,16 @@ async def _resolve_plan( actors_in_graph.add("ValidationAgent") raw_plan.steps = steps - return Plan(*tasks, title=raw_plan.title, history=messages), unmet_dependencies, actors + return Plan(*tasks, title=raw_plan.title, history=messages, context=context, coordinator=self), actors - async def _compute_plan(self, messages: list[Message], agents: dict[str, Agent], tools: dict[str, Tool], pre_plan_output: dict[str, Any]) -> Plan: + async def _compute_plan( + self, + messages: list[Message], + context: TContext, + agents: dict[str, Agent], + tools: dict[str, Tool], + pre_plan_output: dict[str, Any] + ) -> Plan: tool_names = list(tools) agent_names = list(agents) plan_model = self._get_model("main", agents=agent_names, tools=tool_names) @@ -1338,7 +1230,8 @@ async def _compute_plan(self, messages: list[Message], agents: dict[str, Agent], log_debug(f"\033[91m!! Attempt {attempts}\033[0m") plan = None try: - raw_plan = await self._make_plan(messages, agents, tools, unmet_dependencies, previous_actors, previous_plans, + raw_plan = await self._make_plan( + messages, context, agents, tools, unmet_dependencies, previous_actors, previous_plans, plan_model, istep, is_follow_up=pre_plan_output["is_follow_up"] ) except asyncio.CancelledError as e: @@ -1349,11 +1242,13 @@ async def _compute_plan(self, messages: list[Message], agents: dict[str, Agent], self._todos_title.object = istep.failed_title = 'Failed to make plan. Ensure LLM is configured correctly and/or try again.' traceback.print_exception(e) raise e - plan, unmet_dependencies, previous_actors = await self._resolve_plan( - raw_plan, agents, tools, messages, previous_actors + plan, previous_actors = await self._resolve_plan( + raw_plan, agents, tools, messages, context, previous_actors ) - if unmet_dependencies: - istep.stream(f"The plan didn't account for {unmet_dependencies!r}", replace=True) + try: + plan.validate() + except ContextError as e: + istep.stream(str(e), replace=True) attempts += 1 else: planned = True @@ -1363,7 +1258,6 @@ async def _compute_plan(self, messages: list[Message], agents: dict[str, Agent], e = RuntimeError("Planner failed to come up with viable plan after 5 attempts.") traceback.print_exception(e) raise e - self._memory["plan"] = raw_plan # Store the todo message reference for later updates self._todo_step = istep diff --git a/lumen/ai/llm.py b/lumen/ai/llm.py index 2dec1e294..fd1d09dfb 100644 --- a/lumen/ai/llm.py +++ b/lumen/ai/llm.py @@ -1,9 +1,11 @@ from __future__ import annotations import asyncio +import base64 import os from functools import partial +from pathlib import Path from types import SimpleNamespace from typing import Any, Literal, TypedDict @@ -13,6 +15,7 @@ from instructor import Mode, patch from instructor.dsl.partial import Partial +from instructor.processing.multimodal import Image from pydantic import BaseModel from .interceptor import Interceptor @@ -25,6 +28,13 @@ class Message(TypedDict): content: str name: str | None + +class ImageResponse(BaseModel): + # To easily analyze images, we need instructor patch activated, + # so we use a pass-thru dummy string basemodel + output: str + + BASE_MODES = list(Mode) @@ -122,6 +132,36 @@ def _add_system_message( messages = [{"role": "system", "content": system}] + messages return messages, input_kwargs + def _serialize_image_pane(self, image: pn.pane.image.ImageBase | Image) -> Image: + if isinstance(image, Image): + return image + + image_object = image.object + if isinstance(image_object, bytes): + # convert bytes to base64 string + base64_str = base64.b64encode(image_object).decode('utf-8') + image = Image.from_raw_base64(base64_str) + elif isinstance(image_object, (Path, str)) and Path(image_object).is_file(): + image = Image.from_path(image_object) + elif isinstance(image_object, str): + image = Image.from_url(image_object) + return image + + def _check_for_image(self, messages: list[Message]) -> tuple[list[Message], bool]: + contains_image = False + for i, message in enumerate(messages): + content = message.get("content") + if isinstance(content, (Image, pn.pane.image.ImageBase)): + messages[i]["content"] = self._serialize_image_pane(content) + contains_image = True + + elif isinstance(content, list): + for item in content: + if isinstance(item, (Image, pn.pane.image.ImageBase)): + messages[i]["content"] = self._serialize_image_pane(item) + contains_image = True + return messages, contains_image + @classmethod def warmup(cls, model_kwargs: dict | None): """ @@ -167,10 +207,19 @@ async def invoke( kwargs = dict(self._client_kwargs) kwargs.update(input_kwargs) + messages, contains_image = self._check_for_image(messages) + if contains_image: + # Currently instructor does not support streaming with multimodal + # https://github.com/567-labs/instructor/issues/1872 + kwargs["stream"] = False + if response_model is not None: - if allow_partial: + if allow_partial and isinstance(response_model, BaseModel): response_model = Partial[response_model] kwargs["response_model"] = response_model + # check if any of the messages contain images + elif response_model is None and contains_image: + kwargs["response_model"] = ImageResponse output = await self.run_client(model_spec, messages, **kwargs) if output is None or output == "": @@ -261,6 +310,10 @@ async def stream( model_spec=model_spec, **kwargs, ) + if isinstance(chunks, BaseModel): + yield getattr(chunks, field) if field is not None else chunks + return + try: async for chunk in chunks: if response_model is None: @@ -326,9 +379,12 @@ class LlamaCpp(Llm, LlamaCppMixin): select_models = param.List(default=[ "unsloth/Qwen3-8B-GGUF", - "microsoft/Phi-3-mini-4k-instruct-gguf", - "meta-llama/Llama-2-7b-chat-hf", - "TheBloke/CodeLlama-7B-Instruct-GGUF" + "unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF", + "unsloth/DeepSeek-V3.1-GGUF", + "unsloth/gpt-oss-20b-GGUF", + "unsloth/GLM-4.6-GGUF", + "microsoft/Phi-4-GGUF", + "meta-llama/Llama-3.3-70B-Instruct-GGUF" ], constant=True, doc="Available models for selection dropdowns") temperature = param.Number(default=0.4, bounds=(0, None), constant=True) @@ -555,15 +611,19 @@ class MistralAI(Llm): model_kwargs = param.Dict(default={ "default": {"model": "mistral-small-latest"}, - "edit": {"model": "mistral-large-latest"}, + "edit": {"model": "mistral-medium-latest"}, }) select_models = param.List(default=[ - "codestral-latest", - "mistral-7b-instruct", + "mistral-medium-latest", + "magistral-medium-latest", "mistral-large-latest", + "magistral-small-latest", "mistral-small-latest", - "mixtral-8x7b-instruct" + "codestral-latest", + "ministral-8b-latest", + "ministral-3b-latest", + "devstral-small-latest" ], constant=True, doc="Available models for selection dropdowns") temperature = param.Number(default=0.7, bounds=(0, 1), constant=True) @@ -670,16 +730,19 @@ class AnthropicAI(Llm): mode = param.Selector(default=Mode.ANTHROPIC_TOOLS, objects=[Mode.ANTHROPIC_JSON, Mode.ANTHROPIC_TOOLS]) model_kwargs = param.Dict(default={ - "default": {"model": "claude-3-5-haiku-latest"}, - "edit": {"model": "claude-3-5-sonnet-latest"}, + "default": {"model": "claude-haiku-4-5"}, + "edit": {"model": "claude-sonnet-4-5"}, }) select_models = param.List(default=[ + "claude-sonnet-4-5", + "claude-sonnet-4-0", + "claude-3-7-sonnet-latest", + "claude-opus-4-1", + "claude-opus-4-0", + "claude-haiku-4-5", "claude-3-5-haiku-latest", - "claude-3-5-sonnet-latest", - "claude-3-haiku-20240307", - "claude-3-sonnet-20240229", - "claude-4-sonnet-latest" + "claude-3-haiku-20240307" ], constant=True, doc="Available models for selection dropdowns") temperature = param.Number(default=0.7, bounds=(0, 1), constant=True) @@ -732,16 +795,18 @@ class GoogleAI(Llm): mode = param.Selector(default=Mode.GENAI_TOOLS, objects=[Mode.GENAI_TOOLS, Mode.GENAI_STRUCTURED_OUTPUTS]) model_kwargs = param.Dict(default={ - "default": {"model": "gemini-2.0-flash"}, # Cost-optimized, low latency - "edit": {"model": "gemini-2.5-flash-preview-05-20"}, # Thinking model, balanced price/performance + "default": {"model": "gemini-2.5-flash"}, # Best price-performance with thinking + "edit": {"model": "gemini-2.5-pro"}, # State-of-the-art thinking model }) select_models = param.List(default=[ - "gemini-1.5-flash", - "gemini-1.5-pro", - "gemini-2.0-flash", + "gemini-2.5-pro", "gemini-2.5-flash", - "gemini-2.5-flash-lite" + "gemini-2.5-flash-lite", + "gemini-2.0-flash", + "gemini-2.0-flash-lite", + "gemini-1.5-flash", + "gemini-1.5-pro" ], constant=True, doc="Available models for selection dropdowns") temperature = param.Number(default=1, bounds=(0, 1), constant=True) @@ -774,7 +839,7 @@ async def get_client(self, model_spec: str | dict, response_model: BaseModel | N return partial(client.chat.completions.create, model=model, **self._get_create_kwargs(response_model)) else: chat = llm.aio.models - if kwargs.pop("stream"): + if kwargs.pop("stream", None): return partial(chat.generate_content_stream, model=model, **self._get_create_kwargs(response_model)) else: return partial(chat.generate_content, model=model, **self._get_create_kwargs(response_model)) @@ -795,10 +860,8 @@ async def run_client(self, model_spec: str | dict, messages: list[Message], **kw config = GenerateContentConfig(temperature=self.temperature) return await client(messages=messages, config=config, **kwargs) else: - kwargs.pop("stream") - system_instruction = next( - message["content"] for message in messages if message["role"] == "system" - ) + kwargs.pop("stream", None) + system_instruction = next((message["content"] for message in messages if message["role"] == "system"), "Be helpful.") config = GenerateContentConfig(temperature=self.temperature, system_instruction=system_instruction) prompt = messages.pop(-1)["content"] return await client(contents=[prompt], **kwargs) @@ -838,15 +901,19 @@ class Ollama(OpenAI): mode = param.Selector(default=Mode.JSON) model_kwargs = param.Dict(default={ - "default": {"model": "qwen2.5-coder:7b"}, + "default": {"model": "qwen3:8b"}, }) select_models = param.List(default=[ - "codellama:7b", - "llama2:13b", - "llama3.2:latest", - "mistral:latest", - "qwen2.5-coder:7b" + "qwen3:8b", + "qwen3-coder:30b", + "deepseek-r1:7b", + "llama3.3:70b", + "llama4:latest", + "gemma3:12b", + "mistral-small3.2:24b", + "qwen2.5-coder:7b", + "phi4:14b" ], constant=True, doc="Available models for selection dropdowns") temperature = param.Number(default=0.25, bounds=(0, None), constant=True) @@ -976,14 +1043,15 @@ class LiteLLM(Llm): mode = param.Selector(default=Mode.TOOLS, objects=BASE_MODES) model_kwargs = param.Dict(default={ - "default": {"model": "gpt-4o-mini"}, - "edit": {"model": "claude-3-5-sonnet-latest"}, - "sql": {"model": "gpt-4o-mini"}, + "default": {"model": "gpt-4.1-mini"}, + "edit": {"model": "anthropic/claude-sonnet-4-5"}, + "sql": {"model": "gpt-4.1-mini"}, }, doc=""" Model configurations by type. LiteLLM supports model strings like: - - OpenAI: "gpt-4", "gpt-4o-mini" - - Anthropic: "claude-3-5-sonnet-latest", "claude-3-haiku" - - Google: "gemini/gemini-pro", "gemini/gemini-1.5-flash" + - OpenAI: "gpt-4.1-mini", "gpt-4.1-nano", "gpt-5-mini" + - Anthropic: "anthropic/claude-sonnet-4-5", "anthropic/claude-haiku-4-5" + - Google: "gemini/gemini-2.0-flash", "gemini/gemini-2.5-flash" + - Mistral: "mistral/mistral-medium-latest", "mistral/mistral-small-latest" - And many more with format: "provider/model" or just "model" for defaults """) @@ -992,12 +1060,17 @@ class LiteLLM(Llm): Example: {"routing_strategy": "least-busy", "num_retries": 3}""") select_models = param.List(default=[ - "anthropic/claude-3-haiku", - "azure/gpt-4", - "claude-3-5-sonnet-latest", - "gemini/gemini-pro", - "gpt-4o-mini", - "openai/gpt-4" + "gpt-4.1-mini", + "gpt-4.1-nano", + "gpt-5-mini", + "anthropic/claude-sonnet-4-5", + "anthropic/claude-haiku-4-5", + "anthropic/claude-opus-4-1", + "gemini/gemini-2.0-flash", + "gemini/gemini-2.5-flash", + "mistral/mistral-medium-latest", + "mistral/mistral-small-latest", + "mistral/codestral-latest" ], constant=True, doc="Available models for selection dropdowns") temperature = param.Number(default=0.7, bounds=(0, 2), constant=True) diff --git a/lumen/ai/memory.py b/lumen/ai/memory.py deleted file mode 100644 index 8b4c02252..000000000 --- a/lumen/ai/memory.py +++ /dev/null @@ -1,81 +0,0 @@ -from __future__ import annotations - -import weakref - -from collections import defaultdict -from functools import partial - -import param - -from panel.io.state import state - -from ..config import SessionCache - - -class _Memory(SessionCache): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._session_callbacks = weakref.WeakKeyDictionary() - self._global_callbacks = defaultdict(list) - self._session_rx = weakref.WeakKeyDictionary() - self._global_rx = {} - - @property - def _callbacks(self): - if state.curdoc: - if state.curdoc not in self._session_callbacks: - self._session_callbacks[state.curdoc] = defaultdict(list) - return self._session_callbacks[state.curdoc] - return self._global_callbacks - - @property - def _rx(self): - if state.curdoc: - if state.curdoc not in self._session_rx: - self._session_rx[state.curdoc] = {} - return self._session_rx[state.curdoc] - return self._global_rx - - def __setitem__(self, key, new): - if key in self: - old = self[key] - else: - old = None - super().__setitem__(key, new) - self._trigger_update(key, old, new) - - def cleanup(self): - if state.curdoc: - self._session_callbacks[state.curdoc].clear() - self._session_rx[state.curdoc].clear() - else: - self._global_callbacks.clear() - self._global_rx.clear() - - def on_change(self, key, callback): - self._callbacks[key].append(callback) - - def remove_on_change(self, key, callback): - self._callbacks[key].remove(callback) - - def rx(self, key): - if key in self._rx: - return self._rx[key] - self._rx[key] = rxp = param.rx(self[key]) - return rxp - - def trigger(self, key): - self._trigger_update(key, self[key], self[key]) - - def _trigger_update(self, key, old, new): - for cb in self._callbacks[key]: - if param.parameterized.iscoroutinefunction(cb): - state.execute(partial(cb, key, old, new)) - else: - cb(key, old, new) - if key in self._rx: - self._rx[key].rx.value = new - - -memory = _Memory() diff --git a/lumen/ai/prompts/AnalystAgent/main.jinja2 b/lumen/ai/prompts/AnalystAgent/main.jinja2 index a762f8bf2..32b55f68f 100644 --- a/lumen/ai/prompts/AnalystAgent/main.jinja2 +++ b/lumen/ai/prompts/AnalystAgent/main.jinja2 @@ -92,7 +92,7 @@ collection issue or seasonal effect that requires validation before making strat Here was the plan that was executed: {% if memory.get('reasoning') %} """ -{{ memory.reasoning }} +{{ memory['reasoning'] }} """ {% endif %} @@ -102,14 +102,14 @@ Here is the current dataset: The data is empty. Critique the SQL query and suggest what other columns or values should be used instead. Then prompt the user to click the rerun button below if they'd like to try again. {% else %} -{{ memory.data }} +{{ memory['data'] }} {%- endif %} {%- endif -%} {% if memory.get('sql') %} Here is the current SQL query: ```sql -{{ memory.sql }} +{{ memory['sql'] }} ``` {%- endif -%} {%- endblock -%} diff --git a/lumen/ai/prompts/BaseViewAgent/main.jinja2 b/lumen/ai/prompts/BaseViewAgent/main.jinja2 index 0db71f5bd..919d17eb4 100644 --- a/lumen/ai/prompts/BaseViewAgent/main.jinja2 +++ b/lumen/ai/prompts/BaseViewAgent/main.jinja2 @@ -4,10 +4,10 @@ Available visualization types: {{ doc }} -The current SQL table name: {{ memory['table'] }}. +The current SQL table name: {{ memory['pipeline'].table }}. Here is the current dataset to use: -{{ memory.data }} +{{ memory['data'] }} {%- if "view" in memory %} The previous view specification was: @@ -16,7 +16,7 @@ The previous view specification was: ``` {%- endif %} -{% if memory.source.dialect == "snowflake" %} +{% if memory['pipeline'].source.dialect == "snowflake" %} Please use all upper case for field names. {% endif %} {%- endblock %} diff --git a/lumen/ai/prompts/ChatAgent/main.jinja2 b/lumen/ai/prompts/ChatAgent/main.jinja2 index 942efe8c8..fa3071a72 100644 --- a/lumen/ai/prompts/ChatAgent/main.jinja2 +++ b/lumen/ai/prompts/ChatAgent/main.jinja2 @@ -10,7 +10,7 @@ If asked who developed you (Lumen), mention the HoloViz Team and link to https:/ Here was the plan that was executed: {% if 'reasoning' in memory %} """ -{{ memory.reasoning }} +{{ memory['reasoning'] }} """ {% endif %} {%- if 'data' in memory %} diff --git a/lumen/ai/prompts/FunctionTool/main.jinja2 b/lumen/ai/prompts/FunctionTool/main.jinja2 index 70e0e06fb..6fdd833a2 100644 --- a/lumen/ai/prompts/FunctionTool/main.jinja2 +++ b/lumen/ai/prompts/FunctionTool/main.jinja2 @@ -6,11 +6,11 @@ You are invoking a function, provide the arguments to the function given the cur {%- block context %} {%- if 'table' in memory %} -The data table currently being worked on is: {{ memory.table }} +The data table currently being worked on is: {{ memory['table'] }} {%- endif %} {%- if 'data' in memory %} Here is a summary of the current data: -{{ memory.data }} +{{ memory['data'] }} {%- endif %} {% endblock %} diff --git a/lumen/ai/prompts/Planner/main.jinja2 b/lumen/ai/prompts/Planner/main.jinja2 index 228ff24cc..17b8edb0e 100644 --- a/lumen/ai/prompts/Planner/main.jinja2 +++ b/lumen/ai/prompts/Planner/main.jinja2 @@ -29,7 +29,7 @@ Ground Rules: ## 🛠️ Tools {% for tool in tools %} {%- set missing_reqs = [] %} -{%- for req in tool.requires %} +{%- for req in tool.input_schema.__annotations__ %} {%- if req not in memory.keys() or memory[req] is none %} {%- set _ = missing_reqs.append(req) %} {%- endif %} @@ -38,7 +38,7 @@ Ground Rules: {%- else %} ✅ READY {%- endif %} {{ ' '.join(dedent(tool.purpose).strip().split()) }} -{% if tool.provides -%}Provides: `{{ tool.provides | join('`, `') }}`{%- endif %} +{% if tool.output_schema.__annotations__ -%}Provides: `{{ tool.output_schema.__annotations__ | join('`, `') }}`{%- endif %} Conditions for use: {%- for condition in tool.conditions %} - {{ dedent(condition) }} @@ -49,7 +49,7 @@ Conditions for use: ## 🧑‍💼 Agents {% for agent in agents %} {%- set missing_reqs = [] %} -{%- for req in agent.requires %} +{%- for req in agent.input_schema.__annotations__ %} {%- if req not in memory.keys() or memory[req] is none %} {%- set _ = missing_reqs.append(req) %} {%- endif %} @@ -58,12 +58,12 @@ Conditions for use: {%- set providers = [] %} {%- for req in missing_reqs %} {%- for tool in tools %} - {%- if req in tool.provides %} + {%- if req in tool.output_schema.__annotations__ %} {%- set _ = providers.append(tool.__class__.__name__) %} {%- endif %} {%- endfor %} {%- for other_agent in agents %} - {%- if req in other_agent.provides %} + {%- if req in other_agent.output_schema.__annotations__ %} {%- set _ = providers.append(other_agent.name[:-5]) %} {%- endif %} {%- endfor %} @@ -71,14 +71,14 @@ Conditions for use: {%- else %} ✅ READY {%- endif %} {{ ' '.join(dedent(agent.purpose).strip().split()) }} -{% if agent.provides -%}Provides: `{{ agent.provides | join('`, `') or 'None' }}`{%- endif %} +{% if agent.output_schema.__annotations__ -%}Provides: `{{ agent.output_schema.__annotations__ | join('`, `') or 'None' }}`{%- endif %} Guidelines: {%- for condition in agent.conditions %} - {{ dedent(condition) }} {%- endfor -%} {% if agent.not_with %}Should never be used together with: `{{ agent.not_with | join('`, `') }}`{%- endif %} {% endfor %} -# Current Data Context +## Current Data Context {%- if memory.get('document_sources') %} 📂 Documents available: @@ -139,7 +139,7 @@ Failed because: The plan didn't satisfy dependencies: `{{ unmet_dependencies }}` {% if candidates %} Available solutions: {%- for candidate in candidates %} -- `{{ candidate.name[:-5] if candidate.name.endswith('Agent') else candidate.__class__.__name__ }}` provides: `{{ candidate.provides | join('`, `') }}` +- `{{ candidate.name[:-5] if candidate.name.endswith('Agent') else candidate.__class__.__name__ }}` provides: `{{ candidate.output_schema.__annotations__ | join('`, `') }}` {%- endfor %} {% endif %} diff --git a/lumen/ai/prompts/SQLAgent/retry_output.jinja2 b/lumen/ai/prompts/SQLAgent/retry_output.jinja2 index a8497abeb..4ec1f94e0 100644 --- a/lumen/ai/prompts/SQLAgent/retry_output.jinja2 +++ b/lumen/ai/prompts/SQLAgent/retry_output.jinja2 @@ -3,7 +3,7 @@ {% block instructions %} {{ super() }} -{% if discovery_context %} +{% if discovery_context is defined %} ## Discovery Results: {{ discovery_context }} {% endif %} diff --git a/lumen/ai/prompts/ValidationAgent/main.jinja2 b/lumen/ai/prompts/ValidationAgent/main.jinja2 index fba0c9e9d..d83ffb4f2 100644 --- a/lumen/ai/prompts/ValidationAgent/main.jinja2 +++ b/lumen/ai/prompts/ValidationAgent/main.jinja2 @@ -41,9 +41,10 @@ SQL: SELECT state, SUM(amount) FROM data GROUP BY state Steps: {% for step in executed_steps %}{{ step }}{% if not loop.last %}, {% endif %}{% endfor %} {% endif %} {%- if memory.get('sql') %} -SQL: `{{ memory.sql }}` +SQL: `{{ memory['sql'] }}` {%- endif -%} {%- if memory.get('data') is not none %} -Data Overview: {{ memory["data"] }} +Data Overview: +{{ memory["data"] }} {%- endif -%} {%- endblock -%} diff --git a/lumen/ai/report.py b/lumen/ai/report.py index 05ace2c35..1fc246da6 100644 --- a/lumen/ai/report.py +++ b/lumen/ai/report.py @@ -2,13 +2,17 @@ import asyncio import io +import os +import tempfile import traceback as tb from abc import abstractmethod -from collections.abc import Iterator +from collections.abc import Iterable, Iterator +from datetime import datetime from functools import partial +from pathlib import Path from types import FunctionType -from typing import Any, final +from typing import Any, TypedDict, final import panel as pn import param @@ -20,7 +24,7 @@ from panel.viewable import Viewable, Viewer from panel_material_ui import ( Accordion, Alert, Button, ChatFeed, ChatMessage, Container, Dialog, - Divider, FileDownload, IconButton, Progress, Select, Tabs, TextAreaInput, + Divider, FileDownload, IconButton, Progress, Select, TextAreaInput, TextInput, Typography, ) from typing_extensions import Self @@ -28,21 +32,26 @@ from ..pipeline import Pipeline from ..sources.base import BaseSQLSource from ..views.base import Panel, View -from .actor import Actor -from .agents import AnalystAgent +from .actor import Actor, ContextProvider, TContext +from .agents import AnalystAgent, LumenBaseAgent from .config import MissingContextError +from .context import ( + LWW, ContextError, ValidationIssue, collect_task_outputs, + input_dependency_keys, merge_contexts, validate_task_inputs, + validate_taskgroup_exclusions, +) +from .controls import RetryControls from .export import ( format_output, make_md_cell, make_preamble, write_notebook, ) from .llm import Llm -from .memory import _Memory -from .schemas import get_metaset +from .schemas import SQLMetaset, get_metaset from .tools import FunctionTool, Tool from .utils import ( describe_data, extract_block_source, get_block_names, wrap_logfire_on_method, ) -from .views import LumenOutput +from .views import LumenOutput, VegaLiteOutput class Task(Viewer): @@ -53,6 +62,8 @@ class Task(Viewer): abort_on_error = param.Boolean(default=False, doc=""" If True, the report will abort if an error occurs.""") + context = param.Dict() + history = param.List(doc=""" Conversation history to include as context for the task.""") @@ -65,14 +76,7 @@ class Task(Viewer): llm = param.ClassSelector(class_=Llm, doc=""" The LLM to use for the task.""") - memory = param.ClassSelector(class_=_Memory, doc=""" - The memory to use for the task.""") - - outputs = param.List(doc=""" - The outputs of the task.""") - - title = param.String(doc=""" - The title of the task.""") + parent = param.Parameter() running = param.Boolean(doc=""" Whether the task is currently running.""") @@ -83,6 +87,12 @@ class Task(Viewer): steps_layout = param.ClassSelector(default=None, class_=(ListLike, NamedListLike), allow_None=True, doc=""" The layout progress updates will be streamed to.""") + title = param.String(doc=""" + The title of the task.""") + + views = param.List(doc=""" + The generated viewable outputs of the task.""") + def __init_subclass__(cls, **kwargs): """ Apply wrap_logfire to all the subclasses' execute automatically @@ -92,6 +102,7 @@ def __init_subclass__(cls, **kwargs): def __init__(self, **params): super().__init__(**params) + self._prepared = False self._init_view() def __repr__(self): @@ -103,7 +114,9 @@ def __repr__(self): return f"{self.__class__.__name__}({', '.join(params)})" def _init_view(self): - self._view = self._container = Column(sizing_mode='stretch_width', styles={'min-height': 'unset'}, height_policy='fit') + self._view = self._container = Column( + sizing_mode='stretch_width', styles={'min-height': 'unset'}, height_policy='fit' + ) def _populate_view(self): self._view[:] = [] @@ -111,7 +124,7 @@ def _populate_view(self): def reset(self): """Resets the view, removing generated outputs.""" self._view[:] = [] - self.outputs.clear() + self.views.clear() def _render_controls(self): return [ @@ -128,14 +141,8 @@ def _render_output(self, out): return Typography(out, margin=(20, 10)) elif isinstance(out, ChatMessage): return Typography(out.object, margin=(20, 10)) - elif isinstance(out, (Viewable, View)): + elif isinstance(out, (Viewable, View, LumenOutput)): return out - elif isinstance(out, LumenOutput): - return Tabs( - ('Specification', out), - ('Output', pn.param.ParamMethod(out.render, inplace=True, sizing_mode='stretch_width')), - active=1, sizing_mode='stretch_width', min_height=0, height_policy='fit' - ) def editor(self, show_title: bool = True) -> Viewable: """ @@ -156,10 +163,13 @@ def __panel__(self): return self._container @abstractmethod - async def _execute(self, **kwargs): + async def _execute(self, context: TContext, **kwargs) -> tuple[list[Any], TContext]: raise NotImplementedError(f"{self.__class__.__name__} does not implement the execute.") - async def execute(self, **kwargs) -> list[Any]: + async def prepare(self, context: TContext | None = None): + self._prepared = True + + async def execute(self, context: TContext | None = None, **kwargs) -> tuple[list[Any], TContext]: """ Executes the task. @@ -172,7 +182,10 @@ async def execute(self, **kwargs) -> list[Any]: ------- The outputs of the task. """ - return await self._execute(**kwargs) + context = dict(self.context or {}, **(context or {})) + if not self._prepared: + await self.prepare(context) + return await self._execute(context, **kwargs) class TaskGroup(Task): @@ -191,18 +204,73 @@ def __init__(self, *tasks, **params): tasks = params.pop('tasks', []) else: tasks = list(tasks) - outputs, _tasks = [], [] + views, _tasks, _contexts = [], [], [] for task in tasks: if isinstance(task, FunctionType): task = FunctionTool(task) elif isinstance(task, Task): - outputs += task.outputs + views += task.views + task.parent = self _tasks.append(task) - super().__init__(_tasks=_tasks, outputs=outputs, **params) - self._current = 0 + _contexts.append({}) + super().__init__(_tasks=_tasks, views=views, **params) + self._watchers = {} + self._task_outputs = {} + self._task_contexts = {} + self._task_rendered = {} self._init_view() self._populate_view() + def validate( + self, context: TContext | None = None, + available_types: dict[str, Any] | None = None, + path: str | None = None, + raise_on_error: bool = True + ): + """ + Validate the task group and its subtasks. + + Parameters + ---------- + context : TContext | None, optional + The context to validate against, by default None + available_types : dict[str, Any] | None, optional + Dictionary of available types for validation, by default None + path : str | None, optional + Path to the current task group for error reporting, by default None + raise_on_error : bool, optional + Whether to raise an error if validation issues are found, by default True + + Returns + ------- + tuple[list[ValidationIssue], dict[str, Any]] + A tuple containing: + - List of validation issues found + - Dictionary of output types from all tasks + """ + cur_path = path or self.name + issues: list[ValidationIssue] = validate_taskgroup_exclusions(self, path=cur_path) + value_ctx = dict((self.context or {}), **(context or {})) + types_out: dict[str, Any] = dict(available_types or {}) + for idx, t in enumerate(self): + subpath = f"{cur_path}[{idx}] -> {t.name}" + if isinstance(t, TaskGroup): + sub_issues, sub_types = t.validate( + value_ctx, + available_types=types_out, + path=subpath, + ) + else: + sub_issues = validate_task_inputs( + t, value_ctx, types_out, subpath + ) + sub_types = collect_task_outputs(t) + issues.extend(sub_issues) + types_out.update(sub_types) + if path is None and issues and raise_on_error: + raise ContextError(issues) + return issues, types_out + def __repr__(self): params = [] if self.instruction: @@ -229,20 +297,79 @@ def _populate_view(self): using _add_outputs. """ - def _add_outputs(self, i: int, task: Task | Actor, outputs: list, **kwargs): + async def _retry_invoke( + self, i: int, task: Task | Actor, context: TContext, view: LumenOutput, config: dict[str, Any], event: param.parameterized.Event + ): + invalidation_keys = set(task.output_schema.__annotations__) + self.invalidate(invalidation_keys, start=i+1) + if isinstance(task, LumenBaseAgent): + with view.editor.param.update(loading=True): + messages = list(self.history) + task_context = self._get_context(i, context, task) + view.spec = await task.revise( + event.new, messages, task_context, view.spec, language=view.language + ) + root = self + while root.parent is not None: + root = root.parent + with root.param.update(config): + await root.execute() + + def _add_outputs( + self, i: int, task: Task | Actor, views: list, context: TContext, out_context: TContext | None, **kwargs + ): + # Attach retry controls + for view in views: + if not isinstance(view, LumenOutput): + continue + retry_controls = RetryControls() + view.footer = [retry_controls] + self._watchers[i] = retry_controls.param.watch( + partial(self._retry_invoke, i, task, context, view, {'interface': self.interface}), + "instruction" + ) + + # Track context and outputs + if i >= 0: + self._task_contexts[task] = out_context + self._task_outputs[task] = views + + # Find view and output to insert the new outputs after + if i > 0: + prev_task = self._tasks[i-1] + prev_out = self._task_rendered[task] + if isinstance(prev_task, Task): + prev_view = prev_task.views[-1] + else: + prev_view = self._task_views[prev_task][-1] + else: + prev_view = prev_out = None + + idx = 0 if prev_out is None else (self._view.index(prev_out) + 1) if isinstance(task, Task): - self._view.append(task) - self.outputs += task.outputs + self._view.insert(idx, task) + self._task_rendered[task] = self._view[idx] + views = task.views else: - views = [] - for out in outputs: + rendered = [] + for out in views: view = self._render_output(out) if view is not None: - views.append(view) - self._view.extend(views) - self.outputs += outputs - - def _watch_child_outputs(self, previous, event): + rendered.append(view) + if rendered: + rendered_col = Column(*rendered) + if task is not None: + self._task_rendered[task] = rendered_col + self._view.insert(idx, rendered_col) + new_views = list(self.views) + view_idx = 0 if prev_view is None else self.views.index(prev_view) + for vi, view in enumerate(views): + new_views.insert(view_idx+vi, view) + self.views = new_views + + def _watch_child_outputs( + self, i: int, previous: list, context: TContext, event: param.Event, **kwargs + ): pass def append(self, task: Task | Actor): @@ -271,21 +398,51 @@ def insert(self, index, task: Task | Actor): self._tasks.insert(index, task) self._populate_view() + def merge(self, other: TaskGroup): + """ + Merges another task group into the current task group. + + Parameters + --------- + other: TaskGroup + The other task group to merge with. + + Returns + ------- + self: TaskGroup + The current task group. + """ + for task in other: + other.parent = self + self._tasks.append(task) + self._task_contexts.update(other._task_contexts) + self._task_rendered.update(other._task_rendered) + self._task_outputs.update(other._task_outputs) + self._view[:] = list(self._view) + list(other._view) + self.views = self.views + other.views + return self + + async def prepare(self, context: TContext | None = None): + context = context or (self.context or {}) + for task in self._tasks: + await task.prepare(context) + self._prepared = True + def reset(self): """ Resets the view, removing generated outputs. """ - self._current = 0 - self.outputs.clear() + self._task_outputs.clear() + self.views.clear() self._populate_view() for task in self._tasks: if isinstance(task, Task): task.reset() - async def _run_task(self, i: int, task: Self | Actor, **kwargs) -> list[Any]: - pre = 0 if self.memory is None else len(self.memory['outputs']) + async def _run_task( + self, i: int, task: Self | Actor, context: TContext, **kwargs + ) -> tuple[list[Any], TContext]: outputs = [] - memory = task.memory or self.memory messages = list(self.history) if self.instruction: user_msg = None @@ -301,12 +458,11 @@ async def _run_task(self, i: int, task: Self | Actor, **kwargs) -> list[Any]: with task.param.update( interface=self.interface, llm=task.llm or self.llm, - memory=memory, steps_layout=self.steps_layout ): if isinstance(task, Actor): try: - out = await task.respond(messages, **kwargs) + out, out_context = await task.respond(messages, context, **kwargs) except MissingContextError: # Re-raise MissingContextError to allow retry logic at Plan level raise @@ -316,29 +472,31 @@ async def _run_task(self, i: int, task: Self | Actor, **kwargs) -> list[Any]: f'Executing task {type(task).__name__} failed.', alert_type='error', sizing_mode="stretch_width" ) - return [alert] + return [alert], {} # Handle Tool specific behaviors if isinstance(task, Tool): # Handle View/Viewable results regardless of agent type - if isinstance(out, (View, Viewable)): - if isinstance(out, Viewable): - pipeline = None if self.memory is None else self.memory.get('pipeline') - out = Panel(object=out, pipeline=pipeline) - out = LumenOutput( - component=out, title=self.title + rendered = [] + for o in out: + if not isinstance(o, (View, Viewable)): + continue + if isinstance(o, Viewable): + pipeline = None if context is None else context.get('pipeline') + o = Panel(object=o, pipeline=pipeline) + o = LumenOutput( + component=o, title=self.title ) - message_kwargs = dict(value=out, user=task.name) + message_kwargs = dict(value=o, user=task.name) if self.interface: self.interface.stream(**message_kwargs) - self.memory['outputs'] = self.memory['outputs'] + [out] - new = self.memory['outputs'][pre:] - if not new and isinstance(out, (Viewable, View, LumenOutput)): - new = [out] - outputs += new + rendered.append(o) + out = rendered else: with task.param.update(running=True, history=messages): - outputs += await task.execute(**kwargs) - return outputs + out, out_context = await task.execute(context, **kwargs) + self._task_outputs[task] = out + outputs += out + return outputs, out_context def _render_tasks(self) -> Viewable: tasks = [] @@ -418,29 +576,46 @@ def editor(self, show_title=True): ) ) - async def _execute(self, **kwargs): + def _get_context(self, i: int, context: TContext | None, task: Task | Actor) -> TContext | None: + contexts = ([self.context] if self.context else []) + ([context] if context else []) + contexts += [self._task_contexts[task] for task in self._tasks[:i]] + if isinstance(task, Actor): + return merge_contexts(task.input_schema, contexts) + elif isinstance(task, TaskGroup): + subcontexts = [] + for subtask in task: + subtask_context = self._get_context(i, context, subtask) + if subtask_context is not None: + subcontexts.append(subtask_context) + return merge_contexts(LWW, subcontexts) + elif isinstance(task, Action): + return merge_contexts(task.input_schema, contexts) + else: + raise TypeError("Abstract Task does not implement _get_context.") + + async def _execute(self, context: TContext, **kwargs): """ Executes the tasks. Arguments --------- + context: TContext + The context given to the task **kwargs: dict Additional keyword arguments to pass to the tasks. """ - if self.memory is not None and 'outputs' not in self.memory: - self.memory['outputs'] = [] - if self._current != 0: - outputs = list(self.outputs) - else: - outputs = [Typography(f"{'#'*self.level} {self.title}", margin=(10, 10, 0, 10))] if self.title else [] - if outputs: - self._add_outputs(-1, None, outputs, **kwargs) + title = Typography(f"{'#'*self.level} {self.title}", margin=(10, 10, 0, 10)) + views = [title] if self.title else [] + if views and not self._task_outputs: + self._add_outputs(-1, None, views, context, None, **kwargs) for i, task in enumerate(self._tasks): - if i < self._current: + if task in self._task_outputs: + views += self._task_outputs[task] continue + subcontext = self._get_context(i, context, task) new = [] try: - new = await self._run_task(i, task, **kwargs) + new, new_context = await self._run_task(i, task, subcontext, **kwargs) except MissingContextError: # Re-raise MissingContextError to allow retry logic at Plan level raise @@ -451,21 +626,98 @@ async def _execute(self, **kwargs): break else: self.status = "success" - outputs += new - self._add_outputs(i, task, new, **kwargs) - self._current = i + 1 - return outputs + views += new + self._add_outputs( + i, task, new, context, new_context, **kwargs + ) + contexts = ([self.context] if self.context else []) + contexts += [self._task_contexts[task] for task in self] + return views, merge_contexts(LWW, contexts) + + def invalidate(self, keys: Iterable[str], start: int = 0, propagate: bool = True) -> tuple[bool, set[str]]: + """ + Invalidates tasks and propagates context dependencies within a TaskGroup. + + Parameters + ---------- + keys : Iterable[str] + A set of context keys that have been modified or invalidated by the user. + These represent outputs whose dependent tasks should be re-evaluated. + start : int + Index of the task to start invalidating. + propagate: bool + Whether to propagate the invalidation to the parent task group. + + Returns + ------- + tuple[bool, set[str]] + A tuple ``(invalidated, keys)`` where: + - ``invalidated`` is ``True`` if any tasks in this group (or nested groups) + were marked invalid due to overlapping input dependencies. + - ``keys`` is the full, cumulative set of invalidated context keys after + propagating dependencies through all affected tasks. + + Notes + ----- + - If a task's input keys intersect with the provided invalidation keys, + that task is considered stale and will be rerun. + - When a task is invalidated, its recorded outputs are removed from + ``self._task_outputs`` and its output keys are added to the invalidation set, + ensuring that downstream tasks depending on those outputs are also invalidated. + - Nested TaskGroups are traversed recursively, and their invalidations + propagate upward to the parent group. + """ + if self.parent is not None and propagate: + parent_idx = self.parent._tasks.index(self) + if start >= len(self): + # If the last task is being invalidated, + # only subsequent tasks on the parent + # have to be invalidated + parent_idx += 1 + self.parent.invalidate(keys, start=parent_idx) + return + keys = set(keys) + invalidated = False + views = list(self.views) + rendered_views = list(self._view) + for i, task in enumerate(self): + if i < start: + continue + if isinstance(task, ContextProvider): + deps = input_dependency_keys(task.input_schema) + if not (deps & keys): + continue + invalidated = True + self._task_contexts.pop(task, None) + outputs = self._task_outputs.pop(task, []) + rendered = self._task_rendered.pop(task, None) + keys |= set(task.output_schema.__annotations__) + if rendered is not None: + rendered_views.remove(rendered) + views = [view for view in views if view not in outputs] + if isinstance(task, Task): + task.reset() + else: + old = list(task.views) + subtask_invalidated, subtask_keys = task.invalidate(keys, propagate=not propagate) + new = list(task.views) + invalidated = invalidated or subtask_invalidated + keys |= subtask_keys + views = [view for view in views if not (view in old and view not in new)] + self.views = views + self._view[:] = rendered_views + return invalidated, keys def to_notebook(self): """ Returns the notebook representation of the tasks. """ - if len(self) and not len(self.outputs): + if len(self) and not len(self.output_schema): raise RuntimeError( "Report has not been executed, run report before exporting to_notebook." ) cells, extensions = [], ['tabulator'] - for out in self.outputs: + for out in self.views: ext = None if isinstance(out, Typography): level = int(out.variant[1:]) if out.variant and out.variant.startswith('h') else 0 @@ -502,8 +754,12 @@ def __repr__(self) -> str: tasks = [f"\n {task!r}" for task in self._tasks] return f"{self.__class__.__name__}({', '.join(params)}{''.join(tasks)})" - def _add_outputs(self, i: int, task: Task | Actor, outputs: list, **kwargs): - self.outputs += outputs + def _add_outputs(self, i: int, task: Task | Actor, views: list, context: TContext, out_context: TContext | None, **kwargs): + if out_context is not None: + self._task_contexts[task] = out_context + self._task_outputs[task] = views + self._task_rendered[task] = task + self.views = self.views + views def _render_controls(self): return [ @@ -543,26 +799,23 @@ def _populate_view(self): def _open_settings(self, event): self._dialog.open = True - def _watch_child_outputs(self, i: int, previous: list, event: param.Event, **kwargs): + def _watch_child_outputs(self, i: int, previous: list, context: TContext, event: param.Event, **kwargs): for out in (event.old or []): if out not in event.new: out.param.unwatch(self._watchers[out]) + state = dict(interface=self.interface, llm=self.llm, steps_layout=self.steps_layout) for out in event.new: - if out not in self._watchers and isinstance(out, LumenOutput): - context = dict( - interface=self.interface, llm=self.llm, memory=self.memory.clone(), - steps_layout=self.steps_layout - ) - self._watchers[out] = out.param.watch(partial(self._rerun, i+1, context), 'spec') + if isinstance(out, LumenOutput) and out not in self._watchers: + self._watchers[out] = out.param.watch(partial(self._rerun, i+1, dict(context), state), 'spec') - async def _rerun(self, i: int, context: dict, _: param.Event, **kwargs): + async def _rerun(self, i: int, context: TContext, state: dict, _: param.Event, **kwargs): for task in self._tasks[i:]: task.reset() with self.param.update(running=True): for j, task in enumerate(self._tasks[i:]): try: - with self.param.update(context): - await self._run_task(i+j, task, **kwargs) + with self.param.update(state): + await self._run_task(i+j, task, context, **kwargs) except Exception as e: tb.print_exception(e) self.status = "error" @@ -571,19 +824,24 @@ async def _rerun(self, i: int, context: dict, _: param.Event, **kwargs): else: self.status = "success" - async def _run_task(self, i: int, task: Task | Actor, **kwargs) -> list[Any]: - if self.memory: - self.memory['outputs'] = [] - instructions = "\n".join(f"{i+1}. {task.instruction}" if hasattr(task, 'instruction') else f"{i+1}. " for i, task in enumerate(self._tasks)) - self.memory['reasoning'] = f"{self.title}\n\n{instructions}" + async def _run_task(self, i: int, task: Task | Actor, context: TContext | None, **kwargs) -> list[Any]: + if context is not None: + instructions = "\n".join( + f"{i+1}. {task.instruction}" if hasattr(task, 'instruction') else f"{i+1}. " + for i, task in enumerate(self._tasks) + ) + context['reasoning'] = f"{self.title}\n\n{instructions}" if isinstance(task, Task): - watcher = task.param.watch(partial(self._watch_child_outputs, i, list(self.outputs), **kwargs), 'outputs') + watcher = task.param.watch( + partial(self._watch_child_outputs, i, list(self.views), context, **kwargs), + 'views' + ) try: - outputs = await super()._run_task(i, task, **kwargs) + outputs, out = await super()._run_task(i, task, context, **kwargs) finally: if isinstance(task, Task): task.param.unwatch(watcher) - return outputs + return outputs, out @param.depends('running', watch=True) async def _running(self): @@ -607,8 +865,42 @@ class Report(TaskGroup): _tasks = param.List(item_type=Section) + docx_context = param.Dict(default={}, doc=""" + Context dictionary for docx template rendering. If keys are not provided, + the following defaults will be used: + + - 'title': self.title or 'Lumen Report' + - 'subtitle': 'Generated on {date}' (e.g., 'Generated on October 27, 2025') + - 'cover_page_header': '' (empty string) + - 'cover_page_footer': '' (empty string) + - 'content_page_header': '' (empty string) + + The following keys are always auto-generated and cannot be overridden: + - 'sections': List of section dicts with 'title', 'image', and 'caption' + - 'page_break': R('\f') for page breaks + + Example: + report.docx_context = { + 'title': 'Q4 Sales Report', + 'subtitle': 'Quarterly Analysis', + 'cover_page_header': 'ACME Corporation', + 'cover_page_footer': 'Confidential' + }""") + + docx_template_path = param.String( + default=str(Path(__file__).parent / "assets" / "lumen_template.docx"), + doc="""Path to the docx template file.""") + level = 1 + def __init__(self, *tasks, **params): + if not tasks: + tasks = params.pop('tasks', []) + else: + tasks = list(tasks) + super().__init__(*tasks, **params) + pn.state.execute(self.prepare) + def _init_view(self): self._title = Typography( self.param.title, variant="h1", margin=(0, 10, 0, 10) @@ -633,12 +925,18 @@ def _init_view(self): icon="settings", on_click=self._open_settings, size="large", color="default", margin=0, description="Configure Report" ) - self._export = FileDownload( + self._notebook_export_btn = FileDownload( callback=self._notebook_export, label="\u200b", variant='text', icon='get_app', icon_size="2.4em", color="default", margin=(8, 0, 10, 0), sx={".MuiButton-startIcon": {"mr": 0, "color": "var(--mui-palette-default-dark)"}}, description="Export Report to .ipynb", filename=f"{self.title or 'Report'}.ipynb" ) + self._docx_export_btn = FileDownload( + callback=self._docx_export, label="\u200b", variant='text', icon='description', + icon_size="2.4em", color="default", margin=(8, 0, 10, 0), + sx={".MuiButton-startIcon": {"mr": 0, "color": "var(--mui-palette-default-dark)"}}, + description="Export Report to .docx", filename=f"{self.title or 'Report'}.docx" + ) self._dialog = Dialog( TextInput.from_param(self.param.title, margin=(10, 0, 0, 0), sizing_mode="stretch_width"), show_close_button=True, @@ -650,7 +948,8 @@ def _init_view(self): self._run, self._clear, self._collapse, - self._export, + self._notebook_export_btn, + self._docx_export_btn, self._settings, sizing_mode="stretch_width" ) @@ -663,17 +962,26 @@ def _init_view(self): @param.depends('title', watch=True) def _update_filename(self): - self._export.filename = f"{self.title or 'Report'}.ipynb" + self._notebook_export_btn.filename = f"{self.title or 'Report'}.ipynb" + self._docx_export_btn.filename = f"{self.title or 'Report'}.docx" - def _add_outputs(self, i: int, task: Task | Actor, outputs: list, **kwargs): - self.outputs += outputs + def _add_outputs(self, i: int, task: Task | Actor, views: list, context: TContext, out_context: dict | None, **kwargs): + if out_context is not None: + self._task_contexts[task] = out_context + self._task_outputs[task] = views + self._task_rendered[task] = task + self.views = self.views + views def _notebook_export(self): return io.StringIO(self.to_notebook()) - async def _execute(self, *args): + def _docx_export(self): + """Callback for FileDownload to export report as docx.""" + return self.to_docx() + + async def _execute(self, context: TContext, *args): with self._run.param.update(loading=True): - return await super()._execute() + return await super()._execute(context) def _expand_all(self, event): if self._collapse.icon == "unfold_less": @@ -690,11 +998,11 @@ def _populate_view(self): self._view[:] = objects = [(task.title, task) for task in self._tasks] self._view.active = list(range(len(objects))) - async def _run_task(self, i: int, task: Section, **kwargs): + async def _run_task(self, i: int, task: Section, context: TContext, **kwargs): self._view.active = self._view.active + [i] - watcher = task.param.watch(partial(self._watch_child_outputs, self.outputs), "outputs") + watcher = task.param.watch(partial(self._watch_child_outputs, i, self.views, dict(context)), "views") try: - outputs = await super()._run_task(i, task, **kwargs) + outputs = await super()._run_task(i, task, context, **kwargs) finally: task.param.unwatch(watcher) return outputs @@ -708,8 +1016,160 @@ def __panel__(self): ) ) + def to_docx(self) -> io.BytesIO: + """ + Export the report to a Word document (.docx) format. + + Returns + ------- + BytesIO + A BytesIO buffer containing the rendered docx document. + + Raises + ------ + RuntimeError + If the report has not been executed yet. + FileNotFoundError + If the template file is not found. + """ + from docxtpl import DocxTemplate, R + + # Validate execution + if len(self) and not len(self.outputs): + raise RuntimeError( + "Report has not been executed, run report before exporting to_docx." + ) + + # Load template + template_path = Path(self.docx_template_path) + if not template_path.exists(): + raise FileNotFoundError(f"Template file not found: {template_path}") + + doc = DocxTemplate(str(template_path)) + + # Start with copy of docx_context + context = dict(self.docx_context) + + # Set defaults for missing keys + if 'title' not in context: + context['title'] = self.title or "Lumen Report" + + if 'subtitle' not in context: + date_string = datetime.now().strftime("%B %d, %Y") + context['subtitle'] = f"Generated on {date_string}" + + if 'cover_page_header' not in context: + context['cover_page_header'] = "" + + if 'cover_page_footer' not in context: + context['cover_page_footer'] = "" + + if 'content_page_header' not in context: + context['content_page_header'] = "" + + # Always generate sections from report outputs + context['sections'] = self._generate_sections(doc) + + # Always set page_break + context['page_break'] = R("\f") + + # Render template + doc.render(context) + + # Return as BytesIO + buffer = io.BytesIO() + doc.save(buffer) + buffer.seek(0) + return buffer + + def _generate_sections(self, doc) -> list[dict]: + """ + Generate sections list from report tasks for docx template. + + Arguments + --------- + doc : DocxTemplate + The document template instance (needed for InlineImage creation) + + Returns + ------- + list[dict] + List of section dictionaries with title, image, and caption + """ + from docx.shared import Mm + from docxtpl import InlineImage, RichText + + sections = [] + + for section in self._tasks: + if not isinstance(section, Section): + continue + + section_dict = { + "title": section.title or "Untitled Section", + "image": None, + "caption": RichText("") + } + + # Process section outputs to find visualizations and captions + image_found = False + for i, out in enumerate(section.outputs): + if isinstance(out, VegaLiteOutput) and not image_found: + # Convert LumenOutput to image + image_path = self._output_to_image(out) + if image_path: + section_dict["image"] = InlineImage(doc, image_path, width=Mm(160)) + image_found = True + + # Check if next output is a Typography for caption + if i + 1 < len(section.outputs): + next_out = section.outputs[i + 1] + if isinstance(next_out, Typography): + section_dict["caption"] = RichText(next_out.object) + break + + if section_dict["image"]: # Only add section if it has an image + sections.append(section_dict) + + return sections + + def _output_to_image(self, output: LumenOutput) -> str | None: + """ + Convert a LumenOutput to an image file path. + + Arguments + --------- + output : LumenOutput + The output to convert -class Action(Task): + Returns + ------- + str | None + Path to temporary image file, or None if conversion failed + """ + # Create a temporary file for the image + tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False) + tmp_path = tmp.name + tmp.close() + try: + # Render the component and save as image + component = output.component + with open(tmp_path, 'wb') as f: + vega_pane = component.__panel__()._pane + vega_pane.param.update( + width=650, + height=400, + ) + image_bytes = vega_pane.export("png", scale=2, ppi=144) + f.write(image_bytes) + return tmp_path + except Exception as e: + self.param.warning(f"Failed to convert output to image: {e}") + os.unlink(tmp_path) + return None + + +class Action(Task, ContextProvider): """ An `Action` implements an execute method that performs some unit of work and optionally generates outputs to be rendered. @@ -719,12 +1179,25 @@ class Action(Task): Whether the outputs should be rendered.""") @final - async def execute(self, **kwargs): - outputs = await super().execute(**kwargs) + async def execute(self, context: TContext | None = None, **kwargs): + views, out_context = await super().execute(context, **kwargs) if self.render_outputs: - self._view[:] = [self._render_output(out) for out in outputs] - self.outputs += outputs - return outputs + self._view[:] = [self._render_output(out) for out in views] + self.views = self.views + views + return views, out_context + + +class SQLQueryInputs(TypedDict): + + source: BaseSQLSource + + +class SQLQueryOutputs(TypedDict): + source: BaseSQLSource + pipeline: Pipeline + data: dict + sql_metaset: SQLMetaset + table: str class SQLQuery(Action): @@ -733,8 +1206,8 @@ class SQLQuery(Action): and generates an LumenOutput to be rendered. """ - generate_caption = param.Boolean(default=True, doc=""" - Whether to generate a caption for the data.""") + schema = param.Dict(default=None, doc=""" + Optional schema to use to not infer schema from data.""") source = param.ClassSelector(class_=BaseSQLSource, doc=""" The Source to execute the SQL expression on.""") @@ -742,11 +1215,22 @@ class SQLQuery(Action): sql_expr = param.String(default="", doc=""" The SQL expression to use for the action.""") + table_params = param.List(default=[], doc=""" + List of parameters to pass to the SQL expression. + Parameters are used with placeholders (?) in the SQL expression.""") + table = param.String(doc=""" The name of the table generated from the SQL expression.""") - user_content = param.String(default="Generate a short caption for the data", doc=""" - Additional instructions to provide to the analyst agent, i.e. what to focus on.""") + template_overrides = param.Dict(default={}, doc=""" + Template overrides to provide to the AnalystAgent.""") + + analyst_instructions = param.String(default=None, doc=""" + Instructions to provide to the analyst agent, i.e. what to focus on; + if unset no additional instructions are provided.""") + + inputs = SQLQueryInputs + outputs = SQLQueryOutputs def _render_controls(self): return [ @@ -768,7 +1252,7 @@ def __repr__(self): params.append(f"title='{self.title}'") return f"{self.__class__.__name__}({', '.join(params)})" - async def _execute(self, **kwargs): + async def _execute(self, context: TContext, **kwargs) -> tuple[list[Any], SQLQueryOutputs]: """ Executes the action. @@ -783,31 +1267,31 @@ async def _execute(self, **kwargs): """ source = self.source if source is None: - if self.memory is None or 'source' not in self.memory: + if context is None or 'source' not in context: raise ValueError( "SQLQuery could not resolve a source. Either provide " "an explicit source or ensure another action or actor " "provides a source." ) - source = self.memory['source'] + source = context['source'] if not self.table: raise ValueError("SQLQuery must declare a table name.") source = source.create_sql_expr_source({self.table: self.sql_expr}) pipeline = Pipeline(source=source, table=self.table) - if self.memory is not None: - self.memory["source"] = source - if "sources" not in self.memory: - self.memory["sources"] = [] - self.memory["sources"].append(source) - self.memory["pipeline"] = pipeline - self.memory["data"] = await describe_data(pipeline.data) - self.memory["sql_metaset"] = await get_metaset([source], [self.table]) - self.memory["table"] = self.table + out_context = { + "source": source, + "pipeline": pipeline, + "data": await describe_data(pipeline.data), + "sql_metaset": await get_metaset([source], [self.table]), + "table": self.table, + } out = LumenOutput(component=pipeline) - outputs = [Typography(f"### {self.title}", variant='h4', margin=(10, 10, 0, 10)), out] if self.title else [out] - if self.generate_caption: - caption = await AnalystAgent(llm=self.llm).respond( - [{"role": "user", "content": self.user_content}] + title = Typography(f"### {self.title}", variant='h4', margin=(10, 10, 0, 10)) + outputs = [title, out] if self.title else [out] + if self.analyst_instructions: + caption_out, _ = await AnalystAgent(llm=self.llm).respond( + [{"role": "user", "content": self.analyst_instructions}], context ) + caption = caption_out[0] outputs.append(Typography(caption.object)) - return outputs + return outputs, out_context diff --git a/lumen/ai/schemas.py b/lumen/ai/schemas.py index 8cd54c4c5..b497282ed 100644 --- a/lumen/ai/schemas.py +++ b/lumen/ai/schemas.py @@ -3,6 +3,8 @@ from dataclasses import dataclass, field from typing import Any +import yaml + from ..sources import Source from .config import SOURCE_TABLE_SEPARATOR from .utils import ( @@ -40,20 +42,22 @@ class VectorMetaset: def _generate_context(self, include_columns: bool = False, truncate: bool = False) -> str: """ - Generate formatted text representation of the context. + Generate YAML formatted representation of the context. Args: include_columns: Whether to include column details in the context truncate: Whether to truncate strings and columns for brevity """ - context = "" + tables_data = {} + for table_slug, vector_metadata in self.vector_metadata_map.items(): base_sql = truncate_string(vector_metadata.base_sql, max_length=200) if truncate else vector_metadata.base_sql - context += f"{table_slug!r} (access this table with: {base_sql})\n" + + table_data = {'read_with': base_sql} if vector_metadata.description: desc = truncate_string(vector_metadata.description, max_length=100) if truncate else vector_metadata.description - context += f"Info: {desc}\n" + table_data['info'] = desc # Only include columns if explicitly requested if include_columns and vector_metadata.columns: @@ -61,27 +65,29 @@ def _generate_context(self, include_columns: bool = False, truncate: bool = Fals cols_to_show = vector_metadata.columns show_ellipsis = False - original_indices = [] if truncate: cols_to_show, original_indices, show_ellipsis = truncate_iterable(cols_to_show, max_length) else: cols_to_show = list(cols_to_show) - original_indices = list(range(len(cols_to_show))) - for i, (col, orig_idx) in enumerate(zip(cols_to_show, original_indices, strict=False)): + columns_data = {} + for i, col in enumerate(cols_to_show): if show_ellipsis and i == len(cols_to_show) // 2: - context += "...\n" - - if i == 0: - context += "Cols:\n" + columns_data['...'] = '...' col_name = truncate_string(col.name) if truncate else col.name - context += f"{orig_idx}. {col_name!r}" if col.description: col_desc = truncate_string(col.description, max_length=100) if truncate else col.description - context += f": {col_desc}" - context += "\n" - return context + columns_data[col_name] = col_desc + else: + columns_data[col_name] = None + + if columns_data: + table_data['columns'] = columns_data + + tables_data[table_slug] = table_data + + return yaml.dump(tables_data, default_flow_style=False, allow_unicode=True, sort_keys=False) @property def table_context(self) -> str: @@ -120,16 +126,16 @@ class SQLMetaset: def _generate_context(self, include_columns: bool = False, truncate: bool = False) -> str: """ - Generate formatted context with both vector and SQL data. + Generate YAML formatted context with both vector and SQL data. Args: include_columns: Whether to include column details in the context truncate: Whether to truncate strings and columns for brevity Returns: - Formatted context string + YAML formatted context string """ - context = "" + tables_data = {} for table_slug in self.sql_metadata_map.keys(): vector_metadata = self.vector_metaset.vector_metadata_map.get(table_slug) @@ -137,22 +143,24 @@ def _generate_context(self, include_columns: bool = False, truncate: bool = Fals continue base_sql = truncate_string(vector_metadata.base_sql, max_length=200) if truncate else vector_metadata.base_sql - context += f"\n{table_slug!r} (access this table with: {base_sql})\n" + + table_data = {'read_with': base_sql} if vector_metadata.description: desc = truncate_string(vector_metadata.description, max_length=100) if truncate else vector_metadata.description - context += f"Info: {desc}\n" + table_data['info'] = desc sql_data: SQLMetadata = self.sql_metadata_map.get(table_slug) if sql_data: # Get the count from schema if sql_data.schema.get("__len__"): - context += f"Row count: {len(sql_data.schema)}\n" + table_data['row_count'] = len(sql_data.schema) # Only include columns if explicitly requested if include_columns and vector_metadata.columns: cols_to_show = vector_metadata.columns - context += "Columns:" + columns_data = {} + for col in cols_to_show: schema_data = None if sql_data and col.name in sql_data.schema: @@ -160,23 +168,34 @@ def _generate_context(self, include_columns: bool = False, truncate: bool = Fals if truncate and schema_data == "": continue - # Get column name - context += f"\n- {col.name}" + col_info = {} # Get column description with optional truncation if col.description: col_desc = truncate_string(col.description, max_length=100) if truncate else col.description - context += f": {col_desc}" - else: - context += ": " + col_info['description'] = col_desc # Add schema info for the column if available - if schema_data: - if truncate and schema_data.get('type') == 'enum': - schema_data = truncate_string(str(schema_data), max_length=50) - context += f" `{schema_data}`" - context += "\n" - return context.replace("'type': 'str', ", "") # Remove type info for lower token + if schema_data and schema_data != "": + if isinstance(schema_data, dict): + # Remove 'type': 'str' for token efficiency + schema_copy = {k: v for k, v in schema_data.items() if not (k == 'type' and v == 'str')} + if truncate and schema_copy.get('type') == 'enum': + schema_str = str(schema_copy) + if len(schema_str) > 50: + schema_copy = truncate_string(schema_str, max_length=50) + col_info.update(schema_copy) + else: + col_info['value'] = schema_data + + columns_data[col.name] = col_info if col_info else None + + if columns_data: + table_data['columns'] = columns_data + + tables_data[table_slug] = table_data + + return yaml.dump(tables_data, default_flow_style=False, allow_unicode=True, sort_keys=False) @property def table_context(self) -> str: @@ -203,7 +222,7 @@ def __str__(self) -> str: return self.table_context -async def get_metaset(sources: list[Source], tables: list[str]) -> SQLMetaset: +async def get_metaset(sources: list[Source], tables: list[str], prev: SQLMetaset | None = None, schema: dict | None = None) -> SQLMetaset: """ Get the metaset for the given sources and tables. @@ -213,6 +232,8 @@ async def get_metaset(sources: list[Source], tables: list[str]) -> SQLMetaset: The sources to get the metaset for. tables: list[str] The tables to get the metaset for. + schema: dict | None + Optional schema to use instead of fetching from sources. Returns ------- @@ -233,27 +254,35 @@ async def get_metaset(sources: list[Source], tables: list[str]) -> SQLMetaset: else: source_name = next(iter(sources)).name table_name = table_slug - source = next((s for s in sources if s.name == source_name), None) - schema = await get_schema(source, table_name, include_count=True) - tables_info[table_slug] = SQLMetadata( - table_slug=table_slug, - schema=schema, - ) - try: - metadata = source.get_metadata(table_name) - except Exception as e: - log_debug(f"Failed to get metadata for table {table_name} in source {source_name}: {e}") - metadata = {} - tables_metadata[table_slug] = VectorMetadata( - table_slug=table_slug, - similarity=1, - base_sql=source.get_sql_expr(source.normalize_table(table_name)), - description=metadata.get("description"), - columns=[ - Column(name=col_name, description=col_values.pop("description", None), metadata=col_values) - for col_name, col_values in metadata.get("columns", {}).items() - ], - ) + + if prev and table_slug in prev.sql_metadata_map: + sql_metadata = prev.sql_metadata_map[table_slug] + else: + source = next((s for s in sources if s.name == source_name), None) + if schema is None: + schema = await get_schema(source, table_name, include_count=True) + sql_metadata = SQLMetadata(table_slug=table_slug, schema=schema) + tables_info[table_slug] = sql_metadata + + if prev and table_slug in prev.vector_metaset.vector_metadata_map: + vector_metadata = prev.vector_metaset.vector_metadata_map[table_slug] + else: + try: + metadata = source.get_metadata(table_name) + except Exception as e: + log_debug(f"Failed to get metadata for table {table_name} in source {source_name}: {e}") + metadata = {} + vector_metadata = VectorMetadata( + table_slug=table_slug, + similarity=1, + base_sql=source.get_sql_expr(source.normalize_table(table_name)), + description=metadata.get("description"), + columns=[ + Column(name=col_name, description=col_values.pop("description", None), metadata=col_values) + for col_name, col_values in metadata.get("columns", {}).items() + ], + ) + tables_metadata[table_slug] = vector_metadata vector_metaset = VectorMetaset(vector_metadata_map=tables_metadata, query=None) return SQLMetaset( vector_metaset=vector_metaset, diff --git a/lumen/ai/tools.py b/lumen/ai/tools.py index 7d5e74fb7..7cd8bc894 100644 --- a/lumen/ai/tools.py +++ b/lumen/ai/tools.py @@ -1,24 +1,23 @@ import asyncio import traceback -from functools import partial from types import FunctionType -from typing import Any +from typing import Annotated, Any, TypedDict import param from panel.io import cache as pn_cache -from panel.io.state import state from panel.pane import HoloViews as HoloViewsPanel, panel as as_panel from panel.viewable import Viewable +from ..sources.base import Source from ..sources.duckdb import DuckDBSource from ..views.base import HoloViews, View from .actor import Actor, ContextProvider from .config import PROMPTS_DIR, SOURCE_TABLE_SEPARATOR +from .context import ContextModel, TContext from .embeddings import NumpyEmbeddings from .llm import Message -from .memory import _Memory from .models import ( ThinkingYesNo, make_iterative_selection_model, make_refined_query_model, ) @@ -116,15 +115,14 @@ def _initialize_tools_for_prompt(self, tools_or_key: str | list, **params) -> li instantiated_tools.append(tool(**tool_kwargs)) return instantiated_tools - async def _use_tools(self, prompt_name: str, messages: list[Message]) -> str: + async def _use_tools(self, prompt_name: str, messages: list[Message], context: TContext) -> str: tools_context = "" # TODO: INVESTIGATE WHY or self.tools is needed for tool in self._tools.get(prompt_name, []) or self.tools: - if all(requirement in self._memory for requirement in tool.requires): - with tool.param.update(memory=self.memory): - tool_context = await tool.respond(messages) - if tool_context: - tools_context += f"\n{tool_context}" + if all(requirement in context for requirement in tool.input_schema.__annotations__): + tool_context = await tool.respond(messages, context) + if tool_context: + tools_context += f"\n{tool_context}" return tools_context @@ -225,12 +223,22 @@ class Tool(Actor, ContextProvider): ]) @classmethod - async def applies(cls, memory: _Memory) -> bool: + async def applies(cls, context: TContext) -> bool: """ Additional checks to determine if the tool should be used. """ return True + async def sync(self, context: TContext): + """ + Allows the tool to update when the provided context changes. + """ + + + +class VectorLookupOutputs(ContextModel): + document_chunks: list[str] + class VectorLookupTool(Tool): """ @@ -238,14 +246,18 @@ class VectorLookupTool(Tool): chunks. """ - # Class variable to track which sources are currently being processed - _sources_in_progress = {} enable_query_refinement = param.Boolean(default=True, doc=""" Whether to enable query refinement for improving search results.""") + max_refinement_iterations = param.Integer(default=3, bounds=(1, 10), doc=""" + Maximum number of refinement iterations to perform.""") + min_similarity = param.Number(default=0.3, doc=""" The minimum similarity to include a document.""") + min_refinement_improvement = param.Number(default=0.05, bounds=(0, 1), doc=""" + Minimum improvement in similarity score required to keep refining.""") + n = param.Integer(default=5, bounds=(1, None), doc=""" The number of document results to return.""") @@ -262,25 +274,28 @@ class VectorLookupTool(Tool): refinement_similarity_threshold = param.Number(default=0.3, bounds=(0, 1), doc=""" Similarity threshold below which query refinement is triggered.""") - max_refinement_iterations = param.Integer(default=3, bounds=(1, 10), doc=""" - Maximum number of refinement iterations to perform.""") - - min_refinement_improvement = param.Number(default=0.05, bounds=(0, 1), doc=""" - Minimum improvement in similarity score required to keep refining.""") - vector_store = param.ClassSelector(class_=VectorStore, constant=True, doc=""" Vector store object which is queried to provide additional context before responding.""") _item_type_name: str = None + # Class variable to track which sources are currently being processed + _sources_in_progress = {} + __abstract = True + outputs = VectorLookupOutputs + def __init__(self, **params): if 'vector_store' not in params: params['vector_store'] = NumpyVectorStore(embeddings=NumpyEmbeddings()) super().__init__(**params) + async def prepare(self, context: TContext): + context["tables_metadata"] = {} + await self._update_vector_store(context) + def _handle_ready_task_done(self, task): """Properly handle exceptions from async ready tasks.""" try: @@ -315,7 +330,8 @@ def _format_results_for_refinement(self, results: list[dict[str, Any]]) -> str: async def _refine_query( self, original_query: str, - results: list[dict[str, Any]] + results: list[dict[str, Any]], + context: TContext ) -> str: """ Refines the search query based on initial search results. @@ -332,7 +348,7 @@ async def _refine_query( str A refined search query """ - results_description = self._format_results_for_refinement(results) + results_description = self._format_results_for_refinement(results, context) messages = [{"role": "user", "content": original_query}] @@ -341,6 +357,7 @@ async def _refine_query( system_prompt = await self._render_prompt( "refine_query", messages, + context, results=results, results_description=results_description, original_query=original_query, @@ -362,7 +379,7 @@ async def _refine_query( step.status = "failed" return original_query - async def _perform_search_with_refinement(self, query: str, **kwargs) -> list[dict[str, Any]]: + async def _perform_search_with_refinement(self, query: str, context: TContext, **kwargs) -> list[dict[str, Any]]: """ Performs a vector search with optional query refinement. @@ -411,7 +428,7 @@ async def _perform_search_with_refinement(self, query: str, **kwargs) -> list[di iteration += 1 step.stream(f"Processing refinement iteration {iteration}/{self.max_refinement_iterations}\n\n") - refined_query = await self._refine_query(current_query, results) + refined_query = await self._refine_query(current_query, results, context) if refined_query == current_query: step.stream("Refinement returned unchanged query, stopping iterations.") @@ -454,7 +471,7 @@ async def _perform_search_with_refinement(self, query: str, **kwargs) -> list[di return best_results - async def respond(self, messages: list[Message], **kwargs: Any) -> str: + async def respond(self, messages: list[Message], context: TContext, **kwargs: Any) -> tuple[list[Any], ]: """ Respond to a user query using the vector store. @@ -473,7 +490,7 @@ async def respond(self, messages: list[Message], **kwargs: Any) -> str: query = messages[-1]["content"] # Perform search with refinement - results = await self._perform_search_with_refinement(query) + results = await self._perform_search_with_refinement(query, context) closest_doc_chunks = [ f"{result['text']} (Relevance: {result['similarity']:.1f} - " f"Metadata: {result['metadata']})" @@ -486,7 +503,11 @@ async def respond(self, messages: list[Message], **kwargs: Any) -> str: message = "Please augment your response with the following context if relevant:\n" message += "\n".join(f"- {doc}" for doc in closest_doc_chunks) - return message + return [message], {"document_chunks": closest_doc_chunks} + + +class VectorLookupInputs(ContextModel): + document_sources: dict[str, dict[str, str]] class DocumentLookup(VectorLookupTool): @@ -499,21 +520,13 @@ class DocumentLookup(VectorLookupTool): purpose = param.String(default=""" Looks up relevant documents based on the user query.""") - requires = param.List(default=["document_sources"], readonly=True, doc=""" - List of context that this Tool requires to be run.""") - sync_sources = param.Boolean(default=True, doc=""" Whether to automatically sync newly added document sources to the vector store.""") # Override the item type name _item_type_name = "documents" - def __init__(self, **params): - super().__init__(**params) - if self.sync_sources: - self._memory.on_change('document_sources', self._update_vector_store) - if "document_sources" in self._memory: - state.execute(partial(self._update_vector_store, None, None, self._memory['document_sources'])) + inputs = VectorLookupInputs def _format_results_for_refinement(self, results: list[dict[str, Any]]) -> str: """ @@ -543,10 +556,10 @@ def _format_results_for_refinement(self, results: list[dict[str, Any]]) -> str: return "\n".join(formatted_results) - async def _update_vector_store(self, _, __, sources): + async def _update_vector_store(self, context: TContext): # Build a list of document items for a single upsert operation items_to_upsert = [] - for source in sources: + for source in context.get("sources"): metadata = source.get("metadata", {}) metadata["type"] = self._item_type_name items_to_upsert.append({"text": source["text"], "metadata": metadata}) @@ -556,6 +569,16 @@ async def _update_vector_store(self, _, __, sources): await self.vector_store.upsert(items_to_upsert) +class TableLookupInputs(ContextModel): + + sources: Annotated[list[Source], ("accumulate", "source")] + + +class TableLookupOutputs(ContextModel): + + vector_metaset: VectorMetaset + + class TableLookup(VectorLookupTool): """ TableLookup tool that creates a vector store of all available tables @@ -568,10 +591,6 @@ class TableLookup(VectorLookupTool): "Not useful for data related queries", ]) - # Class variable to track which sources are currently being processed - _sources_in_progress = {} - _progress_lock = asyncio.Lock() # Lock for thread-safe access - prompts = param.Dict( default={ "refine_query": { @@ -591,12 +610,6 @@ class TableLookup(VectorLookupTool): Not to be used for finding tables for further analysis (e.g. SQL), because it does not provide a schema.""") - requires = param.List(default=["sources"], readonly=True, doc=""" - List of context that this Tool requires to be run.""") - - provides = param.List(default=["tables_metadata", "vector_metaset"], readonly=True, doc=""" - List of context values this Tool provides to current working memory.""") - include_metadata = param.Boolean(default=True, doc=""" Whether to include table descriptions in the embeddings and responses.""") @@ -619,23 +632,26 @@ class TableLookup(VectorLookupTool): sync_sources = param.Boolean(default=True, doc=""" Whether to automatically sync newly added data sources to the vector store.""") - _item_type_name = "tables" - _ready = param.Boolean(default=False, allow_None=True, doc=""" Whether the vector store is ready.""") + _item_type_name = "tables" + + # Class variable to track which sources are currently being processed + _sources_in_progress = {} + _progress_lock = asyncio.Lock() # Lock for thread-safe access + + inputs = TableLookupInputs + outputs = TableLookupOutputs + def __init__(self, **params): super().__init__(**params) - if "tables_metadata" not in self._memory: - self._memory["tables_metadata"] = {} # used for storing table metadata for LLM self._raw_metadata = {} self._semaphore = asyncio.Semaphore(self.max_concurrent) - if self.sync_sources: - self._memory.on_change("sources", self._update_vector_store) - if "sources" in self._memory: - state.execute(partial(self._update_vector_store, None, None, self._memory["sources"])) - def _format_results_for_refinement(self, results: list[dict[str, Any]]) -> str: + def _format_results_for_refinement( + self, results: list[dict[str, Any]], context: TContext + ) -> str: """ Format table search results for inclusion in the refinement prompt. @@ -649,7 +665,7 @@ def _format_results_for_refinement(self, results: list[dict[str, Any]]) -> str: str Formatted description of table results """ - visible_slugs = self._memory.get('visible_slugs', set()) + visible_slugs = context.get('visible_slugs', set()) formatted_results = [] for result in results: source_name = result['metadata'].get("source", "unknown") @@ -660,13 +676,13 @@ def _format_results_for_refinement(self, results: list[dict[str, Any]]) -> str: text = result["text"] description = f"- {text} {table_slug} (Similarity: {result.get('similarity', 0):.3f})" - if tables_vector_data := self._memory["tables_metadata"].get(table_slug): + if tables_vector_data := context["tables_metadata"].get(table_slug): if table_description := tables_vector_data.get("description"): description += f"\n Info: {table_description}" formatted_results.append(description) return "\n".join(formatted_results) - async def _enrich_metadata(self, source, table_name: str): + async def _enrich_metadata(self, source, table_name: str, context: TContext): """Fetch metadata for a table and return enriched entry for batch processing.""" async with self._semaphore: source_metadata = self._raw_metadata[source.name] @@ -687,8 +703,8 @@ async def _enrich_metadata(self, source, table_name: str): # IMPORTANT: re-insert using table slug # we need to store a copy of tables_vector_data so it can be used to inject context into the LLM table_slug = f"{source.name}{SOURCE_TABLE_SEPARATOR}{table_name}" - self._memory["tables_metadata"][table_slug] = vector_info.copy() - self._memory["tables_metadata"][table_slug]["source_name"] = source.name # need this to rebuild the slug + context["tables_metadata"][table_slug] = vector_info.copy() + context["tables_metadata"][table_slug]["source_name"] = source.name # need this to rebuild the slug # Create column schema objects columns = [] @@ -725,7 +741,12 @@ async def _enrich_metadata(self, source, table_name: str): # Return the entry instead of upserting directly return {"text": enriched_text, "metadata": vector_metadata} - async def _update_vector_store(self, _, __, sources): + async def _update_vector_store(self, context: TContext): + sources = context.get("sources") + if sources is None and "source" in context: + sources = [context["source"]] + else: + raise ValueError("Context does not contain a \"source\" or \"sources\".") self._ready = False await asyncio.sleep(0.5) # allow main thread time to load UI first tasks = [] @@ -738,7 +759,6 @@ async def _update_vector_store(self, _, __, sources): if vector_store_id not in self._sources_in_progress: self._sources_in_progress[vector_store_id] = set() - log_debug(f"[TableLookup] Starting _update_vector_store for {len(sources)} sources") for source in sources: @@ -767,7 +787,7 @@ async def _update_vector_store(self, _, __, sources): if self.include_metadata: for table in tables: - task = asyncio.create_task(self._enrich_metadata(source, table)) + task = asyncio.create_task(self._enrich_metadata(source, table, context)) tasks.append(task) else: await self.vector_store.upsert([ @@ -832,13 +852,24 @@ async def _gather_with_semaphore(): if enriched_entries: log_debug(f"[TableLookup] Upserting {len(enriched_entries)} enriched entries") - await self.vector_store.upsert(enriched_entries) + # Add timeout to the upsert operation to prevent deadlock + try: + await asyncio.wait_for( + self.vector_store.upsert(enriched_entries), + timeout=60 + ) + log_debug(f"[TableLookup] Successfully upserted {len(enriched_entries)} entries") + except asyncio.TimeoutError: + log_debug("[TableLookup] Timeout during upsert operation (60 seconds)") + self._ready = None # Set to error state + return else: log_debug("[TableLookup] No enriched entries to upsert") log_debug("[TableLookup] All table metadata tasks completed.") self._ready = True except Exception as e: log_debug(f"[TableLookup] Error in _mark_ready_when_done: {e!s}") + traceback.print_exc() self._ready = None # Set to error state raise finally: @@ -846,13 +877,14 @@ async def _gather_with_semaphore(): if vector_store_id is not None and processed_sources: self._cleanup_sources_in_progress(vector_store_id, processed_sources) - async def _gather_info(self, messages: list[dict[str, str]]) -> dict: + async def _gather_info(self, messages: list[dict[str, str]], context: TContext) -> TableLookupOutputs: """Gather relevant information about the tables based on the user query.""" + query = messages[-1]["content"] # Count total number of tables available across all sources total_tables = 0 - for source in self._memory.get("sources", []): + for source in context.get("sources", []): total_tables += len(source.get_tables()) # Skip query refinement if there are fewer than 5 tables @@ -863,33 +895,33 @@ async def _gather_info(self, messages: list[dict[str, str]]) -> dict: filters["type"] = self._item_type_name results = await self.vector_store.query(query, top_k=self.n, filters=filters) else: - results = await self._perform_search_with_refinement(query) + results = await self._perform_search_with_refinement(query, context) - any_matches = any(result['similarity'] >= self.min_similarity for result in results) + any_matches = any(result["similarity"] >= self.min_similarity for result in results) same_table = len([result["metadata"]["table_name"] for result in results]) vector_metadata_map = {} for result in results: - source_name = result['metadata']["source"] - table_name = result['metadata']["table_name"] - for source in self._memory.get("sources", []): + source_name = result["metadata"]["source"] + table_name = result["metadata"]["table_name"] + for source in context.get("sources", []): if source.name == source_name: sql = source.get_sql_expr(source.normalize_table(table_name)) break table_slug = f"{source_name}{SOURCE_TABLE_SEPARATOR}{table_name}" - similarity_score = result['similarity'] + similarity_score = result["similarity"] # Filter by visible_slugs if specified - visible_slugs = self._memory.get('visible_slugs', set()) + visible_slugs = context.get("visible_slugs", set()) if visible_slugs and table_slug not in visible_slugs: continue - if any_matches and result['similarity'] < self.min_similarity and not same_table: + if any_matches and result["similarity"] < self.min_similarity and not same_table: continue columns = [] table_description = None - if table_metadata := self._memory["tables_metadata"].get(table_slug): + if table_metadata := context["tables_metadata"].get(table_slug): table_description = table_metadata.get("description") column_metadata = table_metadata.get("columns", {}) @@ -908,13 +940,13 @@ async def _gather_info(self, messages: list[dict[str, str]]) -> dict: description=table_description, base_sql=sql, columns=columns, - metadata=self._memory["tables_metadata"].get(table_slug, {}).copy() + metadata=context["tables_metadata"].get(table_slug, {}).copy() ) vector_metadata_map[table_slug] = vector_metadata # If query contains an exact table name, mark it as max similarity - visible_slugs = self._memory.get('visible_slugs', set()) - tables_to_check = self._memory["tables_metadata"] + visible_slugs = context.get("visible_slugs", set()) + tables_to_check = context["tables_metadata"] if visible_slugs: tables_to_check = {slug: data for slug, data in tables_to_check.items() if slug in visible_slugs} @@ -923,23 +955,30 @@ async def _gather_info(self, messages: list[dict[str, str]]) -> dict: if table_slug in vector_metadata_map: vector_metadata_map[table_slug].similarity = 1 - self._memory["vector_metaset"] = VectorMetaset( + vector_metaset = VectorMetaset( vector_metadata_map=vector_metadata_map, query=query) - return vector_metadata_map + return {"vector_metaset": vector_metaset} - def _format_context(self) -> str: + def _format_context(self, outputs: TableLookupOutputs) -> str: """Generate formatted text representation from schema objects.""" # Get schema objects from memory - vector_metaset = self._memory.get("vector_metaset") + vector_metaset = outputs.get("vector_metaset") return str(vector_metaset) - async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> str: + async def respond( + self, messages: list[Message], context: TContext, **kwargs: dict[str, Any] + ) -> tuple[list[Any], TableLookupOutputs]: """ Fetches tables based on the user query and returns formatted context. """ # Run the process to build schema objects - await self._gather_info(messages) - return self._format_context() + out_model = await self._gather_info(messages, context) + return [self._format_context(out_model)], out_model + + +class IterativeTableLookupOutputs(TableLookupOutputs): + + sql_metaset: SQLMetaset class IterativeTableLookup(TableLookup): @@ -959,9 +998,6 @@ class IterativeTableLookup(TableLookup): purpose = param.String(default=""" Looks up the most relevant tables and provides SQL schemas of those tables.""") - provides = param.List(default=["tables_metadata", "vector_metaset", "sql_metaset"], readonly=True, doc=""" - List of context values this Tool provides to current working memory.""") - # Override sync_sources to False by default to prevent duplicate processing # when used together with TableLookup sync_sources = param.Boolean(default=False, doc=""" @@ -989,40 +1025,32 @@ class IterativeTableLookup(TableLookup): }, ) - def __init__(self, **params): - super().__init__(**params) - self._memory.on_change('visible_slugs', self._update_sql_metaset_for_visible_tables) + outputs = IterativeTableLookupOutputs - async def _update_sql_metaset_for_visible_tables(self, key, old_slugs, new_slugs): + async def sync(self, context: TContext): """ Update sql_metaset when visible_slugs changes. This ensures SQL operations only work with visible tables by regenerating the sql_metaset using get_metaset. """ - if new_slugs is None: - new_slugs = set() + slugs = context.get("visible_slugs", set()) # If no visible slugs or no sources, clear sql_metaset - if not new_slugs or not self._memory.get("sources"): - if "sql_metaset" in self._memory: - del self._memory["sql_metaset"] - return - - # Only update if we have an existing sql_metaset to work with - if "sql_metaset" not in self._memory: + if not slugs or not context.get("sources"): + if "sql_metaset" in context: + del context["sql_metaset"] return try: - sources = self._memory["sources"] - visible_tables = list(new_slugs) - if visible_tables: - new_sql_metaset = await get_metaset(sources, visible_tables) - self._memory["sql_metaset"] = new_sql_metaset - log_debug(f"[IterativeTableLookup] Updated sql_metaset with {len(visible_tables)} visible tables") + sources = context["sources"] + visible_tables = list(slugs) + new_sql_metaset = await get_metaset(sources, visible_tables, prev=context.get("sql_metaset")) + context["sql_metaset"] = new_sql_metaset + log_debug(f"[IterativeTableLookup] Updated sql_metaset with {len(visible_tables)} visible tables") except Exception as e: log_debug(f"[IterativeTableLookup] Error updating sql_metaset: {e}") - async def _gather_info(self, messages: list[dict[str, str]]) -> dict: + async def _gather_info(self, messages: list[dict[str, str]], context: TContext) -> IterativeTableLookupOutputs: """ Performs an iterative table selection process to gather context. This function: @@ -1031,14 +1059,15 @@ async def _gather_info(self, messages: list[dict[str, str]]) -> dict: 3. Gets complete schemas for these tables 4. Repeats until the LLM is satisfied with the context """ - vector_metadata_map = await super()._gather_info(messages) - vector_metaset = self._memory.get("vector_metaset") + out_model = await super()._gather_info(messages, context) + vector_metaset = out_model["vector_metaset"] + vector_metadata_map = vector_metaset.vector_metadata_map sql_metadata_map = {} examined_slugs = set(sql_metadata_map.keys()) # Filter to only include visible tables all_slugs = list(vector_metadata_map.keys()) - visible_slugs = self._memory.get('visible_slugs', set()) + visible_slugs = context.get('visible_slugs', set()) if visible_slugs: all_slugs = [slug for slug in all_slugs if slug in visible_slugs] @@ -1046,7 +1075,7 @@ async def _gather_info(self, messages: list[dict[str, str]]) -> dict: selected_slugs = [] chain_of_thought = "" max_iterations = self.max_selection_iterations - sources = {source.name: source for source in self._memory["sources"]} + sources = {source.name: source for source in context["sources"]} for iteration in range(1, max_iterations + 1): with self._add_step( @@ -1133,7 +1162,7 @@ async def _gather_info(self, messages: list[dict[str, str]]) -> dict: stream_details(str(vector_metaset.vector_metadata_map[table_slug]), step, title="Table details", auto=False) try: view_definition = truncate_string( - self._memory["tables_metadata"].get(table_slug, {}).get("view_definition", ""), + context["tables_metadata"].get(table_slug, {}).get("view_definition", ""), max_length=300 ) @@ -1175,21 +1204,31 @@ async def _gather_info(self, messages: list[dict[str, str]]) -> dict: vector_metaset=vector_metaset, sql_metadata_map=sql_metadata_map, ) - self._memory["sql_metaset"] = sql_metaset - return sql_metadata_map + return {"sql_metaset": sql_metaset} - def _format_context(self) -> str: + def _format_context(self, outputs: IterativeTableLookupOutputs) -> str: """Generate formatted text representation from schema objects.""" - sql_metaset = self._memory.get("sql_metaset") + sql_metaset = outputs.get("sql_metaset") return str(sql_metaset) @classmethod - async def applies(cls, memory: _Memory) -> bool: - visible_slugs = memory.get('visible_slugs', set()) - is_necessary = len(visible_slugs) > 1 - if not is_necessary: - memory["sql_metaset"] = await get_metaset(sources=memory["sources"], tables=visible_slugs) - return is_necessary + async def applies(cls, context: TContext) -> bool: + visible_slugs = context.get('visible_slugs', set()) + return len(visible_slugs) > 1 + + async def respond( + self, messages: list[Message], context: TContext, **kwargs: dict[str, Any] + ) -> tuple[list[Any], IterativeTableLookupOutputs]: + """ + Fetches tables based on the user query and returns formatted context. + """ + out_model = await self._gather_info(messages, context) + return [self._format_context(out_model)], out_model + + + +class DbtslLookupOutputs(ContextModel): + dbtsl_metaset: DbtslMetaset class DbtslLookup(VectorLookupTool, DbtslMixin): @@ -1229,21 +1268,16 @@ class DbtslLookup(VectorLookupTool, DbtslMixin): Useful for quickly gathering information about dbt semantic layers and their metrics to plan the steps. Not useful for looking up what datasets are available. Likely useful for all queries.""") - requires = param.List(default=["source"], readonly=True, doc=""" - List of context that this Tool requires to be run.""") - - provides = param.List(default=["dbtsl_metaset"], readonly=True, doc=""" - List of context values this Tool provides to current working memory.""") - _item_type_name = "metrics" + outputs = DbtslLookupOutputs + def __init__(self, environment_id: int, **params): params["environment_id"] = environment_id super().__init__(**params) self._metric_objs = {} - state.execute(partial(self._update_vector_store, None, None)) - async def _update_vector_store(self, _, __): + async def _update_vector_store(self, context: TContext): """ Updates the vector store with metrics from the dbt semantic layer client. @@ -1273,7 +1307,9 @@ async def _update_vector_store(self, _, __): if items_to_upsert: await self.vector_store.upsert(items_to_upsert) - async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> str: + async def respond( + self, messages: list[Message], context: TContext, **kwargs: dict[str, Any] + ) -> tuple[list[Any], TContext]: """ Fetches metrics based on the user query, populates the DbtslMetaset, and returns formatted context. @@ -1283,14 +1319,14 @@ async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> st # Search for relevant metrics with self._add_step(title="dbt Semantic Layer Search") as search_step: search_step.stream(f"Searching for metrics relevant to: '{query}'") - results = await self._perform_search_with_refinement(query) + results = await self._perform_search_with_refinement(query, context) closest_metrics = [r for r in results if r['similarity'] >= self.min_similarity] if not closest_metrics: search_step.stream("\n⚠️ No metrics found with sufficient relevance to the query.") search_step.status = "failed" - self._memory.pop("dbtsl_metaset", None) - return "" + context.pop("dbtsl_metaset", None) + return [], context metrics_info = [f"- {r['metadata'].get('name')}" for r in closest_metrics] stream_details("\n".join(metrics_info), search_step, title=f"Found {len(closest_metrics)} relevant chunks", auto=False) @@ -1378,11 +1414,11 @@ async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> st step.status = "failed" if not can_answer_query: - self._memory.pop("dbtsl_metaset", None) - return "" + context.pop("dbtsl_metaset", None) + return [], context - self._memory["dbtsl_metaset"] = metaset - return str(metaset) + context["dbtsl_metaset"] = metaset + return [str(metaset)], context @pn_cache async def _fetch_dimension_values(self, client, metric_name, dim, num_cols, step): @@ -1460,8 +1496,14 @@ def __init__(self, function, **params): ) self._model = model - async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> str: - prompt = await self._render_prompt("main", messages) + @property + def inputs(self): + return TypedDict(f"{self.function.__name__}Inputs", {f: Any for f in self.requires}) + + async def respond( + self, messages: list[Message], context: TContext, **kwargs: dict[str, Any] + ) -> tuple[list[Any], ContextModel]: + prompt = await self._render_prompt("main", messages, context) kwargs = {} if any(field not in self.requires for field in self._model.model_fields): model_spec = self.prompts["main"].get("llm_spec", self.llm_spec_key) @@ -1473,25 +1515,26 @@ async def respond(self, messages: list[Message], **kwargs: dict[str, Any]) -> st allow_partial=False, max_retries=3, ) - arguments = dict(kwargs, **{k: self._memory[k] for k in self.requires}) + arguments = dict(kwargs, **{k: context[k] for k in self.requires}) if param.parameterized.iscoroutinefunction(self.function): result = await self.function(**arguments) else: result = self.function(**arguments) if isinstance(result, (View, Viewable)): - return result + return [result], {} elif self.render_output: p = as_panel(result) if isinstance(p, HoloViewsPanel): - return HoloViews(object=p.object) - return p + p = HoloViews(object=p.object) + return [p], {} + out_model = {} if self.provides: if len(self.provides) == 1 and not isinstance(result, dict): - self._memory[self.provides[0]] = result + out_model[self.provides[0]] = result else: - self._memory.update({result[key] for key in self.provides}) - return self.formatter.format( + out_model.update({result[key] for key in self.provides}) + return [self.formatter.format( function=self.function.__name__, arguments=', '.join(f'{k}={v!r}' for k, v in arguments.items()), output=result - ) + )], out_model diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py index ccfddaf8d..72f0a9334 100644 --- a/lumen/ai/ui.py +++ b/lumen/ai/ui.py @@ -7,7 +7,7 @@ from functools import partial from io import StringIO from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any import param @@ -16,43 +16,41 @@ from panel.io.document import hold from panel.io.resources import CSS_URLS from panel.io.state import state -from panel.layout import Column as PnColumn, HSpacer +from panel.layout import Column, HSpacer from panel.pane import SVG, Markdown from panel.param import ParamMethod from panel.util import edit_readonly from panel.viewable import Child, Children, Viewer from panel_gwalker import GraphicWalker from panel_material_ui import ( - Accordion, Button, ChatFeed, ChatInterface, ChatMessage, Column, Dialog, - Divider, FileDownload, IconButton, MenuList, MenuToggle, MultiChoice, Page, - Paper, Row, Switch, Tabs, ToggleIcon, + Accordion, Button, ChatFeed, ChatInterface, ChatMessage, + Column as MuiColumn, Dialog, Divider, FileDownload, IconButton, MenuList, + MenuToggle, Page, Paper, Row, Switch, Tabs, ToggleIcon, ) +from panel_splitjs import VSplit from ..pipeline import Pipeline from ..sources import Source from ..sources.duckdb import DuckDBSource -from ..transforms.sql import SQLLimit from ..util import log from .agents import ( AnalysisAgent, AnalystAgent, ChatAgent, DocumentListAgent, SourceAgent, SQLAgent, TableListAgent, ValidationAgent, VegaLiteAgent, ) -from .components import SourceCatalog, SplitJS +from .components import SplitJS from .config import PROVIDED_SOURCE_NAME, SOURCE_TABLE_SEPARATOR -from .controls import SourceControls +from .context import TContext +from .controls import SourceCatalog, SourceControls, TableExplorer from .coordinator import Coordinator, Plan, Planner from .export import ( export_notebook, make_md_cell, make_preamble, render_cells, write_notebook, ) -from .llm import Llm, OpenAI +from .llm import Llm, Message, OpenAI from .llm_dialog import LLMConfigDialog -from .memory import _Memory, memory from .report import Report from .utils import wrap_logfire from .vector_store import VectorStore - -if TYPE_CHECKING: - from .views import LumenOutput +from .views import LumenOutput, SQLOutput DataT = str | Path | Source | Pipeline @@ -93,99 +91,6 @@ """ -class TableExplorer(Viewer): - """ - TableExplorer provides a high-level entrypoint to explore tables in a split UI. - It allows users to load tables, explore them using Graphic Walker, and then - interrogate the data via a chat interface. - """ - - interface = param.ClassSelector(class_=ChatFeed, doc=""" - The interface for the Coordinator to interact with.""") - - def __init__(self, **params): - super().__init__(**params) - self._table_select = MultiChoice( - label="Select table(s) to preview", sizing_mode='stretch_width', - max_height=200, max_items=5, margin=0 - ) - self._explore_button = Button( - name='Explore table(s)', icon='add_chart', button_type='primary', icon_size="2em", - disabled=self._table_select.param.value.rx().rx.not_(), on_click=self._update_explorers, - margin=(0, 0, 0, 10), width=200, align='end' - ) - self._input_row = Row(self._table_select, self._explore_button) - self._source_map = {} - memory.on_change("sources", self._update_source_map) - self._update_source_map(init=True) - - self._tabs = Tabs(dynamic=True, sizing_mode='stretch_both') - self._layout = Column( - self._input_row, self._tabs, sizing_mode='stretch_both', - ) - - def _update_source_map(self, key=None, old=None, sources=None, init=False): - if sources is None: - sources = memory["sources"] - selected = list(self._table_select.value) - deduplicate = len(sources) > 1 - new = {} - - # Build the source map for UI display - for source in sources: - tables = source.get_tables() - for table in tables: - if deduplicate: - table = f'{source.name}{SOURCE_TABLE_SEPARATOR}{table}' - - if (table.split(SOURCE_TABLE_SEPARATOR, maxsplit=1)[-1] not in self._source_map and - not init and not len(selected) > self._table_select.max_items and state.loaded): - selected.append(table) - new[table] = source - - self._source_map.clear() - self._source_map.update(new) - selected = selected if len(selected) == 1 else [] - self._table_select.param.update(options=list(self._source_map), value=selected) - self._input_row.visible = bool(self._source_map) - - def _explore_table_if_single(self, event): - """ - If only one table is uploaded, help the user load it - without requiring them to click twice. This step - only triggers when the Upload in the Overview tab is used, - i.e. does not trigger with uploads through the SourceAgent - """ - if len(self._table_select.options) == 1: - self._explore_button.param.trigger("value") - - def _update_explorers(self, event): - if not event.new: - return - - with self._explore_button.param.update(loading=True), self.interface.param.update(loading=True): - explorers = [] - for table in self._table_select.value: - source = self._source_map[table] - if SOURCE_TABLE_SEPARATOR in table: - _, table = table.split(SOURCE_TABLE_SEPARATOR, maxsplit=1) - pipeline = Pipeline( - source=source, table=table, sql_transforms=[SQLLimit(limit=100_000, read=source.dialect)] - ) - table_label = f"{table[:25]}..." if len(table) > 25 else table - walker = GraphicWalker( - pipeline.param.data, sizing_mode='stretch_both', min_height=800, - kernel_computation=True, name=table_label, tab='data' - ) - explorers.append(walker) - - self._tabs.objects = explorers - self._table_select.value = [] - - def __panel__(self): - return self._layout - - class UI(Viewer): """ UI provides a baseclass and high-level entrypoint to start chatting with your data. @@ -199,6 +104,8 @@ class UI(Viewer): List of custom analyses. If provided the AnalysesAgent will be added.""" ) + context = param.Dict(default={}) + coordinator = param.ClassSelector( class_=Coordinator, default=Planner, is_instance=False, doc=""" The Coordinator class that will be responsible for coordinating the Agents.""" @@ -281,7 +188,8 @@ def __init__( super().__init__(**params) SourceAgent.source_controls = self.source_controls SourceControls.table_upload_callbacks = self.table_upload_callbacks - self._source_controls = self.source_controls(memory=memory) + self._source_controls = self.source_controls(context=self.context) + self._source_controls.param.watch(self._sync_sources, 'outputs') log.setLevel(self.log_level) agents = self.agents @@ -331,6 +239,7 @@ def __init__( self._settings_menu.param.watch(self._toggle_validation_agent, 'toggled') self._coordinator = self.coordinator( agents=agents, + context=self.context, interface=self.interface, llm=self.llm, tools=self.tools, @@ -398,10 +307,10 @@ def __init__( # Set up actions for the ChatAreaInput speed dial self._setup_actions() self._table_lookup_tool = None # Will be set after coordinator is initialized - self._source_catalog = SourceCatalog() + self._source_catalog = SourceCatalog(context=self.context) self._source_accordion = Accordion( ("Add Sources", self._source_controls), ("View Sources", self._source_catalog), - margin=(-30, 10, 0, 10), sizing_mode="stretch_both", toggle=True, active=[0] + margin=(-30, 10, 0, 10), sizing_mode="stretch_width", toggle=True, active=[0] ) self._sources_dialog_content = Dialog( self._source_accordion, close_on_click=True, show_close_button=True, @@ -421,14 +330,56 @@ def __init__( self._report_toggle = None self._contextbar = None - memory.on_change("sources", self._update_source_catalog) - # LLM status is now managed through actions - if state.curdoc and state.curdoc.session_context: state.on_session_destroyed(self._destroy) state.onload(self._initialize_new_llm) + @param.depends('context', on_init=True, watch=True) + async def _sync_sources(self, event=None): + context = event.new if event else self.context + if 'sources' in context: + old_sources = self.context.get("sources", [self.context["source"]] if "source" in self.context else []) + new_sources = [src for src in context["sources"] if src not in old_sources] + self.context["sources"] = old_sources + new_sources + + all_slugs = set() + for source in new_sources: + tables = source.get_tables() + for table in tables: + table_slug = f'{source.name}{SOURCE_TABLE_SEPARATOR}{table}' + all_slugs.add(table_slug) + + # Update visible_slugs, preserving existing visibility where possible + # This ensures removed tables are filtered out, new tables are added + current_visible = self.context.get('visible_slugs', set()) + if current_visible: + # Keep intersection of current visible and available slugs + # Plus add any new slugs that weren't previously available + self.context['visible_slugs'] = current_visible.intersection(all_slugs) | (all_slugs - current_visible) + else: + # If no visible_slugs set, make all tables visible + self.context['visible_slugs'] = all_slugs + if "source" in context: + if "source" in self.context: + old_source = self.context["source"] + if "sources" not in self.context: + self.context["sources"] = [old_source] + elif old_source not in self.context["sources"]: + self.context["sources"].append(old_source) + self.context["source"] = context["source"] + if "table" in context: + self.context["table"] = context["table"] + if "document_sources" in context: + new_docs = context["document_sources"] + if "document_sources" not in self.context: + self.context["document_sources"] = new_docs + else: + self.context["document_sources"].extend(new_docs) + await self._explorer.sync() + await self._source_catalog.sync() + await self._coordinator.sync(self.context) + def _setup_actions(self): """Set up actions for the ChatAreaInput speed dial.""" existing_actions = self.interface.active_widget.actions @@ -534,7 +485,8 @@ def _resolve_data(self, data: DataT | list[DataT] | dict[DataT] | None): uri=':memory:' ) sources.append(source) - memory["sources"] = sources + self.context["sources"] = sources + self.context["source"] = sources[-1] def show(self, **kwargs): return self._create_view(server=True).show(**kwargs) @@ -574,7 +526,7 @@ def _update_source_catalog(self, *args, **kwargs): """ Update the sources dialog content when memory sources change. """ - self._source_catalog.sources = memory["sources"] + self._source_catalog.sources = self.context["sources"] self._source_accordion.active = [1] def servable(self, title: str | None = None, **kwargs): @@ -621,7 +573,7 @@ class ChatUI(UI): class Exploration(param.Parameterized): - context = param.ClassSelector(class_=_Memory) + context = param.Dict() conversation = Children() @@ -681,7 +633,7 @@ def __init__( self._explorations.on_action('up', self._move_up) self._explorations.on_action('down', self._move_down) self._explorations.on_action('remove', self._delete_exploration) - self._explorer = TableExplorer(interface=self.interface) + self._explorer = TableExplorer(context=self.context) self._explorations_intro = Markdown( EXPLORATIONS_INTRO, margin=(0, 0, 10, 10), @@ -700,10 +652,10 @@ def __init__( self._report_toggle.param.watch(self._toggle_report_mode, ['value']) self._last_synced = self._home = Exploration( - context=memory, + context=self.context, title='Home', conversation=self.interface.objects, - view=Column(self._explorations_intro, self._explorer) + view=MuiColumn(self._explorations_intro, self._explorer) ) home_item = {'label': 'Home', 'icon': 'home', 'view': self._home} self._explorations.param.update(items=[home_item], value=home_item) @@ -729,7 +681,8 @@ def __init__( self._split = SplitJS( left=self._coordinator, right=self._output, - invert=self.chat_ui_position != 'left', + sizes=(40, 60), + expanded_sizes=(40, 60), sizing_mode='stretch_both' ) self._report = Column() @@ -809,7 +762,7 @@ def _destroy(self, session_context): Cleanup on session destroy """ for c in self._explorations.items[1:]: - c['view'].context.cleanup() + c['view'].context.clear() def _global_export_notebook(self): cells, extensions = [], [] @@ -884,7 +837,7 @@ def _snapshot_messages(self, new=False): messages.append(msg) return messages - async def _add_exploration(self, title: str, memory: _Memory): + async def _add_exploration(self, plan: Plan): # Sanpshot previous conversation last_exploration = self._explorations.value['view'] is_home = last_exploration is self._home @@ -895,15 +848,16 @@ async def _add_exploration(self, title: str, memory: _Memory): conversation = list(self.interface.objects) # Create new exploration - output = Column(sizing_mode='stretch_both', loading=True) + output = MuiColumn(sizing_mode='stretch_both', loading=True) exploration = Exploration( - context=memory, + context=plan.context, conversation=conversation, - title=title, + plan=plan, + title=plan.title, view=output ) view_item = { - 'label': title, + 'label': plan.title, 'view': exploration, 'icon': None, 'actions': [{'action': 'remove', 'label': 'Remove', 'icon': 'delete'}] @@ -917,32 +871,16 @@ async def _add_exploration(self, title: str, memory: _Memory): ) self._idle.clear() self._output[:] = [output] - self._notebook_export.filename = f"{title.replace(' ', '_')}.ipynb" + self._notebook_export.filename = f"{plan.title.replace(' ', '_')}.ipynb" await self._update_conversation() self._last_synced = exploration + return exploration - def _add_outputs(self, exploration: Exploration, outputs: list[LumenOutput] | str): - memory = exploration.context + def _add_outputs(self, exploration: Exploration, outputs: list[Any], context: TContext): view = exploration.view - if "sql" in memory: - sql = memory.rx("sql") - sql_pane = Markdown( - param.rx('```sql\n{sql}\n```').format(sql=sql), - margin=(-15, 0, 0, 0), sizing_mode='stretch_width', name='SQL' - ) - if sql.count('\n') > 10: - sql_pane = PnColumn( - sql_pane, max_height=325, scroll='y-auto', name='SQL' - ) - if len(view) and view[0].name == 'SQL': - view[0] = sql_pane - else: - view.insert(0, sql_pane) - content = [] - if view.loading: - from panel_gwalker import GraphicWalker - pipeline = memory['pipeline'] + if view.loading and 'pipeline' in context: + pipeline = context['pipeline'] content.append( ('Overview', GraphicWalker( pipeline.param.data, @@ -952,88 +890,126 @@ def _add_outputs(self, exploration: Exploration, outputs: list[LumenOutput] | st )) ) for out in outputs: + if not isinstance(out, LumenOutput): + continue title = out.title or type(out).__name__.replace('Output', '') if len(title) > 25: title = f"{title[:25]}..." - content.append((title, ParamMethod(out.render, inplace=True, sizing_mode='stretch_both'))) + output = ParamMethod(out.render, inplace=True, sizing_mode='stretch_both') + vsplit = VSplit(out.editor, output, sizes=(20, 80), expanded_sizes=(20, 80), sizing_mode="stretch_both") + content.append((title, vsplit)) if view.loading: - tabs = Tabs(*content, dynamic=True, active=len(outputs), sizing_mode='stretch_both') + tabs = Tabs(*content, dynamic=True, active=len(content), sizing_mode='stretch_both') view.append(tabs) else: tabs = view[-1] tabs.extend(content) tabs.active = len(tabs)-1 - def _wrap_callback(self, callback): - async def wrapper(contents: list | str, user: str, instance: ChatInterface): - prev = self._explorations.value - prev_memory = prev['view'].context - new_exploration = False - local_memory = prev_memory.clone() - local_memory["outputs"] = outputs = [] - - async def render_plan(_, old, new): - nonlocal new_exploration - plan = local_memory["plan"] - if any(step.actor in ('SQLAgent', 'DbtslAgent') for step in plan.steps): - # Expand the contextbar when the first exploration is created - await self._add_exploration(plan.title, local_memory) - new_exploration = True + def _render_view(self, out: LumenOutput) -> VSplit: + title = out.title or type(out).__name__.replace('Output', '') + if len(title) > 25: + title = f"{title[:25]}..." + output = ParamMethod(out.render, inplace=True, sizing_mode='stretch_both') + return (title, VSplit(out.editor, output, sizes=(20, 80), expanded_sizes=(20, 80), sizing_mode="stretch_both")) - def sync_available_sources_memory(_, __, sources): - """ - For cases when the user uploads a dataset through SourceAgent - this will update the available_sources in the global memory - """ - memory["sources"] += [ - source for source in sources if source not in memory["sources"] - ] + def _add_views(self, exploration: Exploration, event: param.parameterized.Event): + outputs = [view for view in event.new if isinstance(view, LumenOutput) and view not in event.old] - local_memory.on_change('plan', render_plan) - local_memory.on_change('sources', sync_available_sources_memory) - - async def render_output(_, old, new): - added = [out for out in new if out not in old] - exploration = self._explorations.value['view'] - self._add_outputs(exploration, added) - exploration.view.loading = False - outputs[:] = new - if self._split.collapsed and prev['label'] == 'Home': - self._split.param.update( - collapsed=False, - sizes=self._split.expanded_sizes, - ) - local_memory.on_change('outputs', render_output) - - # Remove exploration on error if no outputs have been - # added yet and we launched a new exploration - async def remove_output(_, __, ___): - nonlocal new_exploration - if "__error__" in local_memory: - del memory['__error__'] - if outputs or not new_exploration: - return - exploration = self._explorations.value['view'] - prev['view'].conversation = exploration.conversation - with hold(): - self._explorations.param.update( - items=self._explorations.items[:-1], - value=prev - ) - await self._update_conversation() - new_exploration = False - local_memory.on_change('__error__', remove_output) + view = exploration.view + if len(view) == 0: + tabs = None + content = [('Overview', "Waiting on data...")] + else: + tabs = view[0] + content = [] + + for out in outputs: + title, vsplit = self._render_view(out) + content.append((title, vsplit)) + if tabs and isinstance(out, SQLOutput) and not isinstance(tabs[0], GraphicWalker): + tabs[0] = ("Overview", GraphicWalker( + out.component.param.data, + kernel_computation=True, + tab='data', + sizing_mode='stretch_both' + )) + if tabs is None: + tabs = Tabs(*content, dynamic=True, active=len(content), sizing_mode='stretch_both') + view.append(tabs) + else: + tabs.extend(content) + tabs.active = len(tabs)-1 + if self._split.collapsed: + self._split.param.update( + collapsed=False, + sizes=self._split.expanded_sizes, + ) + + def _update_views(self, exploration: Exploration, event: param.parameterized.Event): + current = [view for view in event.new if isinstance(view, LumenOutput)] + old = [view for view in event.old if isinstance(view, LumenOutput)] + + idx = None + tabs = exploration.view[0] + content = list(zip(tabs._names, tabs, strict=False)) + for out in old: + if out in current: + continue + matches = [ + i for i, (_, vsplit) in enumerate(content) + if isinstance(vsplit, VSplit) and out.editor in vsplit + ] + if matches: + idx = matches[0] + content.pop(idx) + + for out in current: + if out in old: + idx = next(i for i, tab in enumerate(tabs[1:]) if out.editor in tab) + 1 + continue + title, vsplit = self._render_view(out) + content.insert(idx+1, (title, vsplit)) + tabs[:] = content + + def _wrap_callback(self, callback): + async def wrapper(messages: list[Message], user: str, instance: ChatInterface): + prev = self._explorations.value + self._idle.clear() try: - self._idle.clear() - with self._coordinator.param.update(memory=local_memory): - plan = await callback(contents, user, instance) - self._explorations.value['view'].plan = plan + plan = await callback(messages, prev["view"].context, user, instance) + if any("pipeline" in step[0].output_schema.__annotations__ for step in plan): + exploration = await self._add_exploration(plan) + new_exploration = True + watcher = plan.param.watch(partial(self._add_views, exploration), "views") + else: + exploration = self._explorations.value['view'] + new_exploration = False + if exploration.plan is not None: + plan = exploration.plan.merge(plan) + watcher = None + with plan.param.update(interface=self.interface): + new, out_context = await plan.execute() + if watcher: + plan.param.unwatch(watcher) + if "__error__" in out_context: + # On error we have to sync the conversation, unwatch the plan, + # and remove the exploration if it was newly created + prev['view'].conversation = exploration.conversation + del out_context['__error__'] + if new_exploration: + with hold(): + self._explorations.param.update( + items=self._explorations.items[:-1], + value=prev + ) + await self._update_conversation() + else: + exploration.context = out_context + exploration.view.loading = False + plan.param.watch(partial(self._update_views, exploration), "views") finally: self._explorations.value['view'].conversation = self.interface.objects self._idle.set() - local_memory.remove_on_change('plan', render_plan) - local_memory.remove_on_change('__error__', remove_output) - if not new_exploration: - prev_memory.update(local_memory) return wrapper diff --git a/lumen/ai/utils.py b/lumen/ai/utils.py index d73d42293..51bf8ef6e 100644 --- a/lumen/ai/utils.py +++ b/lumen/ai/utils.py @@ -22,6 +22,7 @@ import pandas as pd import param +import yaml from jinja2 import ( ChoiceLoader, DictLoader, Environment, FileSystemLoader, StrictUndefined, @@ -517,12 +518,7 @@ def describe_data_sync(df): size = df.size shape = df.shape if size < 250: - # Use the first column as index to save tokens - if len(df.columns) > 1: - df = df.set_index(df.columns[0]) - else: - df = df.to_dict("records") - return df + return df.to_markdown(index=False) is_sampled = False if shape[0] > 5000: @@ -598,7 +594,7 @@ def describe_data_sync(df): head_sample = df.head(2).to_dict('records') tail_sample = df.tail(2).to_dict('records') - return { + result = { "summary": { "n_cells": size, "shape": shape, @@ -610,6 +606,8 @@ def describe_data_sync(df): "tail": tail_sample[0] if tail_sample else {}, } + return yaml.dump(result, default_flow_style=False, allow_unicode=True, sort_keys=False) + return await asyncio.to_thread(describe_data_sync, df) @@ -838,20 +836,15 @@ def truncate_iterable(iterable, max_length=150) -> tuple[list, list, bool]: iterable_list = list(iterable) if len(iterable_list) > max_length: half = max_length // 2 - first_half_indices = list(range(half)) - second_half_indices = list(range(len(iterable_list) - half, len(iterable_list))) - first_half_items = iterable_list[:half] second_half_items = iterable_list[-half:] cols_to_show = first_half_items + second_half_items - original_indices = first_half_indices + second_half_indices show_ellipsis = True else: cols_to_show = iterable_list - original_indices = list(range(len(iterable_list))) show_ellipsis = False - return cols_to_show, original_indices, show_ellipsis + return cols_to_show, show_ellipsis async def with_timeout(coro, timeout_seconds=10, default_value=None, error_message=None): diff --git a/lumen/ai/vector_store.py b/lumen/ai/vector_store.py index b133cc57b..47c73017c 100644 --- a/lumen/ai/vector_store.py +++ b/lumen/ai/vector_store.py @@ -1039,6 +1039,8 @@ class DuckDBVectorStore(VectorStore): uri = param.String(default=":memory:", doc="The URI of the DuckDB database") + read_only = param.Boolean(default=False, doc="Whether to open the database in read-only mode") + embeddings = param.ClassSelector( class_=Embeddings, default=None, @@ -1064,14 +1066,16 @@ def __init__(self, **params): return uri_exists = Path(self.uri).exists() try: - connection.execute(f"ATTACH DATABASE '{self.uri}' AS embedded;") + attach_mode = "READ_ONLY" if self.read_only else "READ_WRITE" + connection.execute(f"ATTACH DATABASE '{self.uri}' AS embedded ({attach_mode});") except duckdb.CatalogException: # handle "Failure while replaying WAL file" # remove .wal uri on corruption wal_path = Path(str(self.uri) + ".wal") if wal_path.exists(): wal_path.unlink() - connection.execute(f"ATTACH DATABASE '{self.uri}' AS embedded;") + attach_mode = "READ_ONLY" if self.read_only else "READ_WRITE" + connection.execute(f"ATTACH DATABASE '{self.uri}' AS embedded ({attach_mode});") connection.execute("USE embedded;") self.connection = connection has_documents = ( @@ -1125,6 +1129,8 @@ def _setup_database(self, embedding_dim: int) -> None: "params": {} } for param_name in self.embeddings.param: + if param_name == "api_key": + continue if param_name not in ['name']: value = getattr(self.embeddings, param_name) if isinstance(value, (str, int, float, bool, list, dict)) or value is None: @@ -1252,30 +1258,29 @@ async def _add_items( text_ids = [] - for i in range(len(texts)): - vector = np.array(embeddings[i], dtype=np.float32) - - # Prepare parameters and query - if force_ids is not None: - query = """ - INSERT INTO documents (id, text, metadata, embedding) - VALUES (?, ?, ?::JSON, ?) RETURNING id; - """ - params = [ - force_ids[i], - texts[i], - json.dumps(metadata[i]), - vector.tolist(), - ] - else: - query = """ - INSERT INTO documents (text, metadata, embedding) - VALUES (?, ?::JSON, ?) RETURNING id; - """ - params = [texts[i], json.dumps(metadata[i]), vector.tolist()] - - # Run the potentially blocking DB operation in a thread - async with self._add_items_lock: + # Acquire the lock once for the entire batch operation + async with self._add_items_lock: + for i in range(len(texts)): + vector = np.array(embeddings[i], dtype=np.float32) + + if force_ids is not None: + query = """ + INSERT INTO documents (id, text, metadata, embedding) + VALUES (?, ?, ?::JSON, ?) RETURNING id; + """ + params = [ + force_ids[i], + texts[i], + json.dumps(metadata[i]), + vector.tolist(), + ] + else: + query = """ + INSERT INTO documents (text, metadata, embedding) + VALUES (?, ?::JSON, ?) RETURNING id; + """ + params = [texts[i], json.dumps(metadata[i]), vector.tolist()] + result = await asyncio.to_thread(self._execute_query, query, params) text_ids.append(result) diff --git a/lumen/ai/views.py b/lumen/ai/views.py index 7d07a03c4..72a941650 100644 --- a/lumen/ai/views.py +++ b/lumen/ai/views.py @@ -2,6 +2,7 @@ import traceback from copy import deepcopy +from typing import Any import panel as pn import param @@ -11,8 +12,7 @@ from jsonschema import Draft7Validator, ValidationError from panel.config import config from panel.layout import Column, Row -from panel.pane import Markdown -from panel.param import ParamRef +from panel.param import ParamMethod from panel.viewable import Viewer from panel.widgets import CodeEditor from panel_material_ui import ( @@ -39,7 +39,7 @@ class LumenOutput(Viewer): loading = param.Boolean() - footer = param.List() + footer = param.List(default=[]) spec = param.String(allow_None=True) @@ -48,32 +48,33 @@ class LumenOutput(Viewer): language = "yaml" def __init__(self, **params): - if 'spec' not in params and 'component' in params and params['component'] is not None: + if "spec" in params: + spec_dict = params["component"].to_spec() + else: try: - component_spec = params['component'].to_spec() - params['spec'] = yaml.dump(component_spec, Dumper=NumpyDumper) + params["spec"], spec_dict = self._serialize_component(params["component"]) except Exception: - params['spec'] = None + spec_dict = {} + params["spec"] = None super().__init__(**params) - code_editor = CodeEditor( - value=self.param.spec.rx.or_(""), + self._spec_dict = spec_dict + self._editor = CodeEditor( + value=self.param.spec.rx.or_(f'{self.title} output could not be serialized and may therefore not be edited.'), language=self.language, theme="github_dark" if config.theme == "dark" else "github_light_default", - sizing_mode="stretch_width", + sizing_mode="stretch_both", soft_tabs=True, on_keyup=False, indent=2, + disabled=self.param.spec.rx.is_(None) ) - code_editor.link(self, bidirectional=True, value='spec') - placeholder = Markdown( - f'{self.title} output could not be serialized and may therefore not be edited.' - ) + self._editor.link(self, bidirectional=True, value='spec') copy_icon = IconButton( icon="content_copy", active_icon="check", margin=(5, 0), toggle_duration=1000, description="Copy YAML to clipboard", size="small", color="default", icon_size="0.8em" ) copy_icon.js_on_click( - args={"code_editor": code_editor}, + args={"code_editor": self._editor}, code="navigator.clipboard.writeText(code_editor.code);", ) download_icon = IconButton( @@ -81,7 +82,7 @@ def __init__(self, **params): description="Download YAML to file", size="small", color="default", icon_size="0.9em" ) download_icon.js_on_click( - args={"code_editor": code_editor}, + args={"code_editor": self._editor}, code=""" var text = code_editor.code; var blob = new Blob([text], {type: 'text/plain'}); @@ -94,21 +95,39 @@ def __init__(self, **params): a.parentNode.removeChild(a); //afterwards we remove the element again """, ) - icons = Row( - copy_icon, download_icon, *self.footer, + self._icons = Row( + *([copy_icon, download_icon] + self.footer), margin=(0, 0, 5, 5) ) - code_col = Column( - code_editor, - icons, + self.editor = Column( + self._editor, + self._icons, sizing_mode="stretch_both" ) - no_spec = self.param.spec.rx.is_(None) - self._main = ParamRef(no_spec.rx.where(placeholder, code_col), min_height=no_spec.rx.where(None, 300)) + self._main = ParamMethod(self.render, inplace=True, sizing_mode='stretch_width') self._main.loading = self.param.loading self._rendered = False self._last_output = {} + @classmethod + def _serialize_component(cls, component: Component) -> str: + component_spec = component.to_spec() + return yaml.dump(component_spec, Dumper=NumpyDumper) + + @classmethod + def _deserialize_component(cls, component: Component, yaml_spec: str, spec_dict: dict[str, Any]) -> str: + spec = load_yaml(yaml_spec) + cls._validate_spec(spec) + return type(component).from_spec(spec) + + @classmethod + def _validate_spec(cls, spec): + return spec + + @param.depends("footer", watch=True) + def _update_footer(self): + self._icons[:] = list(self._icons[:2]) + self.footer + async def _render_pipeline(self, pipeline): table = Table( pipeline=pipeline, pagination='remote', @@ -168,9 +187,7 @@ async def render(self): try: if self._rendered: - yaml_spec = load_yaml(self.spec) - self._validate_spec(yaml_spec) - self.component = type(self.component).from_spec(yaml_spec) + self.component = self._deserialize_component(self.component, self.spec, self._spec_dict) if isinstance(self.component, Pipeline): output = await self._render_pipeline(self.component) else: @@ -186,10 +203,6 @@ async def render(self): alert_type="danger", ) - @classmethod - def _validate_spec(cls, spec): - return spec - def __panel__(self): return self._main @@ -242,6 +255,19 @@ def process_error(err): process_error(error) return "\n".join(errors.values()) + @classmethod + def _serialize_component(cls, component: Component) -> tuple[str, dict[str, Any]]: + component_spec = component.to_spec() + vega_spec = component_spec['spec'] + return yaml.dump(vega_spec, Dumper=NumpyDumper), component_spec + + @classmethod + def _deserialize_component(cls, component: Component, yaml_spec: str, spec_dict: dict[str, Any]) -> str: + spec = load_yaml(yaml_spec) + cls._validate_spec(spec) + spec_dict = dict(spec_dict, spec=spec) + return type(component).from_spec(spec_dict) + @classmethod def _validate_spec(cls, spec): if "spec" in spec: @@ -317,8 +343,31 @@ async def _rerun(self, event): view = await asyncio.to_thread(self.analysis, self.pipeline) self.component = view self._rendered = False - spec = view.to_spec() + spec, self._spec_dict = self._serialize_component(view) self.param.update( - spec=yaml.dump(spec), + spec=spec, active=0 ) + + +class SQLOutput(LumenOutput): + + language = "sql" + + @classmethod + def _serialize_component(cls, component: Component) -> tuple[str, dict[str, Any]]: + component_spec = component.to_spec() + sql_spec = component_spec["source"]["tables"][component.table] + return sql_spec, component_spec + + @classmethod + def _deserialize_component(cls, component: Component, sql_spec: str, spec_dict: dict[str, Any]) -> str: + spec_dict = deepcopy(spec_dict) + spec_dict["source"]["tables"][component.table] = sql_spec + return type(component).from_spec(spec_dict) + + def __panel__(self): + return self._main + + def __str__(self): + return f"{self.__class__.__name__}:\n```sql\n{self.spec}\n```" diff --git a/lumen/sources/base.py b/lumen/sources/base.py index 08d831df3..122c595c5 100644 --- a/lumen/sources/base.py +++ b/lumen/sources/base.py @@ -950,6 +950,17 @@ class BaseSQLSource(Source): - Table name only: 'TABLE' - Wildcards: 'SCHEMA.*'""") + table_params = param.Dict( + default={}, + doc=""" + Dictionary mapping table names to SQL parameters. + Parameters can be: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters for :name, %(name)s, etc. syntax + Each table maps to either a list or dict of parameters. + Example: {'my_table': [2024, 'active'], 'other_table': {'year': 2024}}""", + ) + load_schema = param.Boolean(default=True, doc="Whether to load the schema") # Declare this source supports SQL transforms @@ -1015,17 +1026,32 @@ def get_sql_expr(self, table: str | dict): else: table = self.normalize_table(table) - sql_expr = SQLSelectFrom(sql_expr=self.sql_expr).apply(table) + sql_expr = SQLSelectFrom(sql_expr=self.sql_expr, read=self.dialect).apply(table) return sql_expr - def create_sql_expr_source(self, tables: dict[str, str], **kwargs): + def create_sql_expr_source(self, tables: dict[str, str], params: dict[str, list | dict] | None = None, **kwargs): """ Creates a new SQL Source given a set of table names and corresponding SQL expressions. + + Arguments + --------- + tables: dict[str, str] + Mapping from table name to SQL expression. + params: dict[str, list | dict] | None + Optional mapping from table name to parameters: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters for :name, %(name)s, etc. syntax + kwargs: any + Additional keyword arguments. + + Returns + ------- + source: BaseSQLSource subclass """ raise NotImplementedError - def execute(self, sql_query: str, *args, **kwargs) -> pd.DataFrame: + def execute(self, sql_query: str, params: list | dict | None = None, *args, **kwargs) -> pd.DataFrame: """ Executes a SQL query and returns the result as a DataFrame. @@ -1033,8 +1059,13 @@ def execute(self, sql_query: str, *args, **kwargs) -> pd.DataFrame: --------- sql_query : str The SQL Query to execute + params : list | dict | None + Parameters to use in the SQL query: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters for :name, %(name)s, etc. syntax + - None: No parameters *args : list - Positional arguments to pass to the SQL query + Additional positional arguments to pass to the SQL query **kwargs : dict Keyword arguments to pass to the SQL query @@ -1045,7 +1076,7 @@ def execute(self, sql_query: str, *args, **kwargs) -> pd.DataFrame: """ raise NotImplementedError - async def execute_async(self, sql_query: str, *args, **kwargs) -> pd.DataFrame: + async def execute_async(self, sql_query: str, params: list | dict | None = None, *args, **kwargs) -> pd.DataFrame: """ Executes a SQL query asynchronously and returns the result as a DataFrame. @@ -1057,8 +1088,13 @@ async def execute_async(self, sql_query: str, *args, **kwargs) -> pd.DataFrame: --------- sql_query : str The SQL Query to execute + params : list | dict | None + Parameters to use in the SQL query: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters for :name, %(name)s, etc. syntax + - None: No parameters *args : list - Positional arguments to pass to the SQL query + Additional positional arguments to pass to the SQL query **kwargs : dict Keyword arguments to pass to the SQL query @@ -1067,7 +1103,7 @@ async def execute_async(self, sql_query: str, *args, **kwargs) -> pd.DataFrame: pd.DataFrame The result as a pandas DataFrame """ - return await asyncio.to_thread(self.execute, sql_query, *args, **kwargs) + return await asyncio.to_thread(self.execute, sql_query, params, *args, **kwargs) async def get_async(self, table: str, **query) -> DataFrame: """ @@ -1113,12 +1149,12 @@ def get_schema( data_sql_expr = sql_expr for sql_transform in sql_transforms: data_sql_expr = sql_transform.apply(data_sql_expr) - data = self.execute(data_sql_expr) + data = self.execute(data_sql_expr, self.table_params.get(entry, [])) schemas[entry] = schema = get_dataframe_schema(data)['items']['properties'] count_expr = SQLCount(read=self.dialect).apply(sql_expr) count_expr = ' '.join(count_expr.splitlines()) - count_data = self.execute(count_expr) + count_data = self.execute(count_expr, self.table_params.get(entry, [])) count_col = 'count' if 'count' in count_data else 'COUNT' count = int(count_data[count_col].iloc[0]) if limit: @@ -1136,7 +1172,7 @@ def get_schema( for col in enums: distinct_expr = SQLDistinct(columns=[col], read=self.dialect).apply(sql_expr) distinct_expr = ' '.join(distinct_expr.splitlines()) - distinct = self.execute(distinct_expr) + distinct = self.execute(distinct_expr, self.table_params.get(entry, [])) schema[col]['enum'] = distinct[col].tolist() schema['__len__'] = count @@ -1145,7 +1181,7 @@ def get_schema( minmax_expr = SQLMinMax(columns=min_maxes, read=self.dialect).apply(sql_expr) minmax_expr = ' '.join(minmax_expr.splitlines()) - minmax_data = self.execute(minmax_expr) + minmax_data = self.execute(minmax_expr, self.table_params.get(entry, [])) for col in min_maxes: kind = data[col].dtype.kind if kind in 'iu': diff --git a/lumen/sources/bigquery.py b/lumen/sources/bigquery.py index 07d3e9545..8309b7b02 100644 --- a/lumen/sources/bigquery.py +++ b/lumen/sources/bigquery.py @@ -1,4 +1,5 @@ import asyncio +import datetime import json import threading @@ -36,6 +37,18 @@ class BigQuerySource(BaseSQLSource): dialect = "bigquery" + # Mapping from Python types to BigQuery parameter types + _BQ_TYPE_MAPPING = { + bool: "BOOL", + int: "INT64", + float: "FLOAT64", + str: "STRING", + bytes: "BYTES", + datetime.date: "DATE", + datetime.datetime: "TIMESTAMP", + pd.Timestamp: "TIMESTAMP", + } + def __init__(self, **params) -> None: self._metadata__client: Client | None = None self._sql__client: Client | None = None @@ -117,24 +130,118 @@ def _authorize(self) -> None: if self.project_id == "": self.project_id = project_id - def execute(self, sql_query: str) -> pd.DataFrame: - return self._sql_client.query_and_wait(sql_query).to_dataframe() + def _get_bq_param_type(self, value: Any) -> str: + """ + Determine the BigQuery parameter type from a Python value. + + Parameters + ---------- + value : Any + The parameter value - async def execute_async(self, sql_query: str) -> pd.DataFrame: + Returns + ------- + str + The BigQuery type string (e.g., 'STRING', 'INT64', 'FLOAT64', etc.) """ - Execute a BigQuery SQL query asynchronously and return the result as a DataFrame. + # Direct type lookup (handles most cases) + value_type = type(value) + if value_type in self._BQ_TYPE_MAPPING: + return self._BQ_TYPE_MAPPING[value_type] + + # Check MRO (Method Resolution Order) for subclasses + # This handles subclasses of datetime, custom date types, etc. + for base_type in value_type.__mro__[1:]: + if base_type in self._BQ_TYPE_MAPPING: + return self._BQ_TYPE_MAPPING[base_type] + + # Default to STRING for unknown types + return "STRING" + + def _build_query_config(self, params: list | dict) -> bigquery.QueryJobConfig: + """ + Build a BigQuery QueryJobConfig with the appropriate query parameters. Parameters ---------- + params : list | dict + Either a list of positional parameters or dict of named parameters + + Returns + ------- + bigquery.QueryJobConfig + Configuration object with query parameters set + """ + if isinstance(params, list): + # Positional parameters (? placeholders) + query_parameters = [ + bigquery.ScalarQueryParameter(None, self._get_bq_param_type(param), param) + for param in params + ] + else: + # Named parameters (@name placeholders) + query_parameters = [ + bigquery.ScalarQueryParameter(name, self._get_bq_param_type(value), value) + for name, value in params.items() + ] + + return bigquery.QueryJobConfig(query_parameters=query_parameters) + + def execute(self, sql_query: str, params: list | dict | None = None, *args, **kwargs) -> pd.DataFrame: + """ + Executes a SQL query and returns the result as a DataFrame. + + Arguments + --------- + sql_query : str + The SQL Query to execute + params : list | dict | None + Parameters to use in the SQL query: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters for :name, %(name)s, etc. syntax + - None: No parameters + *args : list + Additional positional arguments to pass to the SQL query + **kwargs : dict + Keyword arguments to pass to the SQL query + + Returns + ------- + pd.DataFrame + The result as a pandas DataFrame + """ + if params: + kwargs['job_config'] = self._build_query_config(params) + return self._sql_client.query_and_wait(sql_query, *args, **kwargs).to_dataframe() + + async def execute_async(self, sql_query: str, params: list | dict | None = None, *args, **kwargs) -> pd.DataFrame: + """ + Executes a SQL query asynchronously and returns the result as a DataFrame. + + This implementation runs queries asynchronously using BigQuery's job API. + + Arguments + --------- sql_query : str - The SQL query to execute + The SQL Query to execute + params : list | dict | None + Parameters to use in the SQL query: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters for :name, %(name)s, etc. syntax + - None: No parameters + *args : list + Additional positional arguments to pass to the SQL query + **kwargs : dict + Keyword arguments to pass to the SQL query Returns ------- pd.DataFrame - The query result as a pandas DataFrame + The result as a pandas DataFrame """ - job = self._sql_client.query(sql_query) + if params: + kwargs['job_config'] = self._build_query_config(params) + job = self._sql_client.query(sql_query, *args, **kwargs) while not job.done(): await asyncio.sleep(0.1) @@ -181,12 +288,34 @@ def get_sql_expr(self, table: str) -> str: return self.tables[table] return f"SELECT * FROM {table}" - def create_sql_expr_source(self, tables: dict[str, str], **kwargs): - params = dict(self.param.values(), **kwargs) - params.pop("name", None) - params["tables"] = tables - source = type(self)(**params) - return source + def create_sql_expr_source(self, tables: dict[str, str], params: dict[str, list | dict] | None = None, **kwargs): + """ + Creates a new SQL Source given a set of table names and + corresponding SQL expressions. + + Arguments + --------- + tables: dict[str, str] + Mapping from table name to SQL expression. + params: dict[str, list | dict] | None + Optional mapping from table name to parameters: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters for :name, %(name)s, etc. syntax + kwargs: any + Additional keyword arguments. + + Returns + ------- + source: BigQuerySource + """ + if params is None: + params = {} + source_params = dict(self.param.values(), **kwargs) + source_params.pop("name", None) + source_params["tables"] = tables + if params: + source_params["table_params"] = params + return type(self)(**source_params) def _get_dataset_metadata(self, datasets: list[str] | None = None) -> dict: """Get metadata for all available datasets in the project. @@ -398,7 +527,7 @@ def get(self, table, **query): sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms for st in sql_transforms: sql_expr = st.apply(sql_expr) - return self.execute(sql_expr) + return self.execute(sql_expr, self.table_params.get(table, [])) async def get_async(self, table, **query): """ @@ -425,7 +554,7 @@ async def get_async(self, table, **query): sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms for st in sql_transforms: sql_expr = st.apply(sql_expr) - return await self.execute_async(sql_expr) + return await self.execute_async(sql_expr, self.table_params.get(table, [])) def close(self): """ diff --git a/lumen/sources/duckdb.py b/lumen/sources/duckdb.py index 89cca1085..b92c1f709 100644 --- a/lumen/sources/duckdb.py +++ b/lumen/sources/duckdb.py @@ -67,6 +67,10 @@ class DuckDBSource(BaseSQLSource): tables = param.ClassSelector(class_=(list, dict), doc=""" List or dictionary of tables.""") + table_params = param.Dict(default={}, doc=""" + Dictionary mapping table names to lists of SQL parameters. + Parameters are used with placeholders (?) in SQL expressions.""") + uri = param.String(doc="The URI of the DuckDB database") source_type = 'duckdb' @@ -325,7 +329,7 @@ def from_spec(cls, spec: dict[str, Any] | str) -> Source: return source def create_sql_expr_source( - self, tables: dict[str, str], materialize: bool = True, **kwargs + self, tables: dict[str, str], materialize: bool = True, params: dict[str, list | dict] | None = None, **kwargs ): """ Creates a new SQL Source given a set of table names and @@ -337,6 +341,10 @@ def create_sql_expr_source( Mapping from table name to SQL expression. materialize: bool Whether to materialize new tables + params: dict[str, list | dict] | None + Optional mapping from table name to parameters: + - list: Positional parameters for ? placeholders + - dict: Named parameters for $param_name placeholders kwargs: any Additional keyword arguments. @@ -344,13 +352,19 @@ def create_sql_expr_source( ------- source: DuckDBSource """ - params = dict(self.param.values(), **kwargs) - params['tables'] = tables + if params is None: + params = {} + + source_params = dict(self.param.values(), **kwargs) + source_params['tables'] = tables + if params: + source_params['table_params'] = params # Reuse connection unless it has changed if 'uri' not in kwargs and 'initializers' not in kwargs: - params['_connection'] = self._connection - params.pop('name', None) - source = type(self)(**params) + source_params['_connection'] = self._connection + source_params.pop('name', None) + source = type(self)(**source_params) + if not materialize: return source @@ -360,7 +374,11 @@ def create_sql_expr_source( table_expr = f'CREATE OR REPLACE TEMP TABLE "{table}" AS ({sql_expr})' cursor = self._connection.cursor() try: - cursor.execute(table_expr) + # Execute with parameters if provided for this table + if table in params: + cursor.execute(table_expr, params[table]) + else: + cursor.execute(table_expr) except duckdb.CatalogException as e: original_e = e pattern = r"Table with name\s(\S+)" @@ -383,8 +401,10 @@ def create_sql_expr_source( source._file_based_tables.update(self._file_based_tables) return source - def execute(self, sql_query: str, *args, **kwargs): + def execute(self, sql_query: str, params: list | dict | None = None, *args, **kwargs): with self._connection.cursor() as cursor: + if params: + return cursor.execute(sql_query, params, *args, **kwargs).fetch_df() return cursor.execute(sql_query, *args, **kwargs).fetch_df() def get_tables(self): @@ -416,8 +436,13 @@ def get(self, table, **query): sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms for st in sql_transforms: sql_expr = st.apply(sql_expr) + + # Apply stored SQL parameters if available for this table with self._connection.cursor() as cursor: - rel = cursor.execute(sql_expr) + if table in self.table_params: + rel = cursor.execute(sql_expr, self.table_params[table]) + else: + rel = cursor.execute(sql_expr) has_geom = any(d[0] == 'geometry' and d[1] == 'BINARY' for d in rel.description) df = rel.fetch_df(date_as_object=True) if has_geom: diff --git a/lumen/sources/snowflake.py b/lumen/sources/snowflake.py index e4f8f0fbe..886659b96 100644 --- a/lumen/sources/snowflake.py +++ b/lumen/sources/snowflake.py @@ -63,6 +63,10 @@ class SnowflakeSource(BaseSQLSource): user = param.String(default=None, doc=""" The user to authenticate as.""") + paramstyle = param.Selector(default='qmark', objects=[ + 'qmark', 'numeric', 'format', 'pyformat'], doc=""" + The paramstyle to use for SQL queries.""") + password = param.String(default=None, doc=""" The password to authenticate with (if authenticator is set to "snowflake").""") @@ -113,6 +117,8 @@ def __init__(self, **params): conn_kwargs['host'] = self.host if self.token is not None: conn_kwargs['token'] = self.token + if self.paramstyle is not None: + conn_kwargs['paramstyle'] = self.paramstyle if self.password is not None: conn_kwargs['password'] = self.password if self.private_key is not None: @@ -219,16 +225,35 @@ def resolve_private_key(self): encryption_algorithm=NoEncryption(), ) - def create_sql_expr_source(self, tables: dict[str, str], **kwargs): + def create_sql_expr_source(self, tables: dict[str, str], params: dict[str, list | dict] | None = None, **kwargs): """ Creates a new SQL Source given a set of table names and corresponding SQL expressions. + + Arguments + --------- + tables: dict[str, str] + Mapping from table name to SQL expression. + params: dict[str, list | dict] | None + Optional mapping from table name to parameters: + - list: Positional parameters for placeholder (?) syntax + - dict: Named parameters (for pyformat/format paramstyle) + kwargs: any + Additional keyword arguments. + + Returns + ------- + source: SnowflakeSource """ - params = dict(self.param.values(), **kwargs) - params.pop("name", None) - params['tables'] = tables - params['conn'] = self._conn - return SnowflakeSource(**params) + if params is None: + params = {} + source_params = dict(self.param.values(), **kwargs) + source_params.pop("name", None) + source_params['tables'] = tables + if params: + source_params['table_params'] = params + source_params['conn'] = self._conn + return SnowflakeSource(**source_params) def _cast_to_supported_dtypes(self, df: pd.DataFrame, sample: int = 100) -> pd.DataFrame: """ @@ -256,12 +281,15 @@ def _cast_to_supported_dtypes(self, df: pd.DataFrame, sample: int = 100) -> pd.D df[col] = df[col].astype(str) return df - def execute(self, sql_query: str, *args, **kwargs): + def execute(self, sql_query: str, params: list | dict | None = None, *args, **kwargs): # TODO: remove cast in future, but keep until bokeh has a solution - df = self._cursor.execute(sql_query, *args, **kwargs).fetch_pandas_all() + if params: + df = self._cursor.execute(sql_query, params, *args, **kwargs).fetch_pandas_all() + else: + df = self._cursor.execute(sql_query, *args, **kwargs).fetch_pandas_all() return self._cast_to_supported_dtypes(df) - async def execute_async(self, sql_query: str, *args, **kwargs): + async def execute_async(self, sql_query: str, params: list | dict | None = None, *args, **kwargs): """ Execute a Snowflake SQL query asynchronously and return the result as a DataFrame. @@ -269,8 +297,13 @@ async def execute_async(self, sql_query: str, *args, **kwargs): ---------- sql_query : str The SQL query to execute + params : list | dict | None + Parameters to use in the SQL query: + - list: Positional parameters for placeholder (?) syntax (qmark paramstyle) + - dict: Named parameters (for pyformat/format paramstyle) + - None: No parameters *args : tuple - Positional arguments to pass to the query + Additional positional arguments to pass to the query **kwargs : dict Keyword arguments to pass to the query @@ -279,7 +312,10 @@ async def execute_async(self, sql_query: str, *args, **kwargs): pd.DataFrame The query result as a pandas DataFrame with supported dtypes """ - self._cursor.execute_async(sql_query, *args, **kwargs) + if params: + self._cursor.execute_async(sql_query, params, *args, **kwargs) + else: + self._cursor.execute_async(sql_query, *args, **kwargs) query_id = self._cursor.sfqid while True: @@ -319,7 +355,7 @@ def get(self, table, **query): sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms for st in sql_transforms: sql_expr = st.apply(sql_expr) - return self.execute(sql_expr) + return self.execute(sql_expr, self.table_params.get(table, [])) async def get_async(self, table, **query): """ @@ -345,7 +381,7 @@ async def get_async(self, table, **query): sql_transforms = [SQLFilter(conditions=conditions)] + sql_transforms for st in sql_transforms: sql_expr = st.apply(sql_expr) - return await self.execute_async(sql_expr) + return await self.execute_async(sql_expr, self.table_params.get(table, [])) @contextlib.contextmanager def _timeout_context(self, seconds=None): diff --git a/lumen/tests/ai/test_context.py b/lumen/tests/ai/test_context.py new file mode 100644 index 000000000..8a6283461 --- /dev/null +++ b/lumen/tests/ai/test_context.py @@ -0,0 +1,229 @@ +from typing import Annotated, NotRequired, TypedDict + +from lumen.ai.context import ( + AccumulateSpec, ContextError, ContextModel, ValidationIssue, _dedupe, + _parse_accumulate_meta, isinstance_like, merge_contexts, + render_issues_tree, schema_fields, types_compatible, validate_task_inputs, + validate_taskgroup_exclusions, +) + +# ----------------------- +# Test schemas & tasks +# ----------------------- + +class Source(TypedDict): + id: str + title: str + +class SQLMetaset(TypedDict): + tables: list[str] + +class SQLInputs(ContextModel): + # single source + source: Source + # accumulate into sources from 'source' + sources: Annotated[list[Source], ("accumulate", "source")] + # another required input + sql_metaset: SQLMetaset + +class SummOutputs(TypedDict, total=False): + # produced optional summary + summary: str + +class ReportOutputs(TypedDict): + report: str + +class DummyTask: + def __init__(self, name: str, inputs: type[TypedDict] | None, outputs: type[TypedDict] | None, not_with=None): + self.name = name + self.inputs = inputs + self.outputs = outputs + self.not_with = not_with + +# ----------------------- +# Tests +# ----------------------- + +def test_parse_accumulate_meta_tuple_form(): + ann = Annotated[list[Source], ("accumulate", "source")] + spec = _parse_accumulate_meta(ann) + assert isinstance(spec, AccumulateSpec) + assert spec.from_key == "source" + assert spec.func is list + +def test_dedupe_value_and_callable(): + assert _dedupe([1, 2, 2, 3], "value") == [1, 2, 3] + items = [{"id": 1}, {"id": 2}, {"id": 1}] + out = _dedupe(items, key=lambda x: x["id"]) + assert out == [{"id": 1}, {"id": 2}] + +def test_merge_contexts_accumulates_from_alt_key_and_direct_target(): + class S(TypedDict): + id: int + class Inputs(TypedDict, total=False): + sources: Annotated[list[S], ("accumulate", "source")] + c1 = {"source": {"id": 1}} + c2 = {"sources": [{"id": 2}], "source": {"id": 3}} + merged = merge_contexts(Inputs, [c1, c2]) + assert merged["sources"] == [{"id": 1}, {"id": 2}, {"id": 3}] + +def test_schema_fields_required_optional_and_meta(): + fields = schema_fields(SQLInputs) + assert fields["source"]["required"] is True + assert fields["sources"]["required"] is True + assert fields["sql_metaset"]["required"] is True + assert ("accumulate", "source") in fields["sources"]["meta"] or _accumulate_alt_key(fields["sources"]["meta"]) == "source" + +def test_isinstance_like_typed_dict_and_list(): + src = {"id": "1", "title": "Doc"} + assert isinstance_like(src, Source) + assert isinstance_like([src], list[Source]) + assert not isinstance_like([{"id": 2}], list[Source]) + +def test_types_compatible_list_elem_and_typed_dict(): + assert types_compatible(list[Source], list[Source]) + assert types_compatible(Source, Source) + assert not types_compatible(list[str], list[int]) + +def test_validate_task_inputs_with_concrete_values_ok(): + task = DummyTask("Summarize", SQLInputs, SummOutputs, None) + ctx = { + "source": {"id": "1", "title": "Doc"}, + "sources": [{"id": "2", "title": "Doc2"}], + "sql_metaset": {"tables": ["t1"]} + } + issues = validate_task_inputs(task, ctx, available_types={}, path=("Group", "TaskA")) + assert issues == [] + +def test_validate_task_inputs_uses_upstream_types(): + task = DummyTask("Summarize", SQLInputs, SummOutputs, None) + available_types = { + "source": Source, + "sql_metaset": SQLMetaset, + } + issues = validate_task_inputs(task, {}, available_types, path=("Group", "TaskA")) + assert [i for i in issues if i.key == "sources"] == [] + assert issues == [] + +def test_validate_task_inputs_flags_incompatible_type(): + task = DummyTask("Summarize", SQLInputs, SummOutputs, None) + class Wrong(TypedDict): + x: int + issues = validate_task_inputs(task, {}, available_types={"sql_metaset": Wrong}, path=("G", "Task")) + assert any("Incompatible upstream type" in i.message and i.key == "sql_metaset" for i in issues) + +def test_validate_task_inputs_accumulator_value_single_elem_allowed(): + task = DummyTask("Summarize", SQLInputs, SummOutputs, None) + ctx = {"source": {"id": "1", "title": "Doc"}, "sql_metaset": {"tables": []}} + issues = validate_task_inputs(task, ctx, {}, path=("G", "Task")) + assert [i for i in issues if i.key == "sources"] == [] + assert issues == [] + +def test_render_issues_tree_and_context_error(): + issues = [ + ValidationIssue(path=("Group", "Step[0]", "Summarize"), key="summary", expected=str, actual=int, message="Type mismatch"), + ValidationIssue(path=("Group", "Step[0]", "Summarize"), key="sources", expected=list, actual=None, message="Missing required key 'sources'"), + ValidationIssue(path=("Group", "Step[1]", "Report"), key="report", expected=str, actual=None, message="Missing required key 'report'"), + ] + msg = render_issues_tree(issues, title="Context validation failed") + assert "Context validation failed" in msg + assert "Group" in msg + assert "Summarize" in msg and "Report" in msg + assert "summary [error]" in msg + err = ContextError(issues) + s = str(err) + assert "Context validation failed" in s + assert "Report" in s + assert "Missing required key 'report'" in s + assert "issue(s)" in repr(err) + +def test_validate_taskgroup_exclusions_pairwise_and_nested(): + class A: pass + class B: pass + + t1 = DummyTask("TaskA", None, None, not_with=["B"]) + t2 = DummyTask("TaskB", None, None, not_with=[A]) + t1.__class__ = A + t2.__class__ = B + + issues = validate_taskgroup_exclusions([t1, t2], path="Root") + assert len(issues) == 1 + assert "Mutually exclusive tasks selected" in issues[0].message + assert "A × B" in issues[0].key + +def test_merge_contexts_empty_bucket_absent_by_default(): + class Inputs(TypedDict, total=False): + sources: Annotated[list[dict], ("accumulate", "source")] + merged = merge_contexts(Inputs, [{}]) + assert "sources" not in merged + +def test_types_compatible_typed_dict_key_compatibility(): + class A(TypedDict): + a: int + b: str + class B(TypedDict): + a: int + b: str + c: float + assert types_compatible(A, B) + +def test_isinstance_like_nested_collections(): + class Node(TypedDict): + id: int + tags: list[str] + value = {"id": 1, "tags": ["x", "y"]} + assert isinstance_like([value], list[Node]) + assert not isinstance_like([{"id": "1", "tags": [1]}], list[Node]) + +def test_typed_dict_notrequired_presence_and_type(): + class User(TypedDict): + name: str + age: NotRequired[int] + + # Missing optional key is ok + assert isinstance_like({"name": "Alice"}, User) + # Present with correct type is ok + assert isinstance_like({"name": "Alice", "age": 30}, User) + # Present with wrong type should fail + assert not isinstance_like({"name": "Alice", "age": "30"}, User) + +def test_schema_fields_notrequired_required_and_type_unwrapped(): + class Item(TypedDict): + id: int + note: NotRequired[str] + + fields = schema_fields(Item) + assert fields["id"]["required"] is True + assert fields["note"]["required"] is False + + # When present, optional field type should still be enforced as str + # (accept correct) + assert isinstance_like("hello", fields["note"]["type"]) + # (reject incorrect) + assert not isinstance_like(123, fields["note"]["type"]) + +def test_validate_task_inputs_notrequired_type_mismatch(): + class Inp(ContextModel): + name: str + notes: NotRequired[str] + + task = DummyTask("T", Inp, None, None) + + # Wrong type for optional field when present should be flagged + issues = validate_task_inputs(task, {"name": "ok", "notes": 123}, {}, path=("G", "T")) + assert any(i.key == "notes" and "Type mismatch" in i.message for i in issues) + + # Optional field omitted: no issues + issues2 = validate_task_inputs(task, {"name": "ok"}, {}, path=("G", "T")) + assert issues2 == [] + +def test_isinstance_like_notrequired_vs_optional_none(): + class TD(TypedDict): + a: NotRequired[int] + b: int | None + + assert isinstance_like({"b": None}, TD) + + assert not isinstance_like({"a": None, "b": 1}, TD) + + assert isinstance_like({"a": 5, "b": None}, TD) diff --git a/lumen/tests/ai/test_vector_store.py b/lumen/tests/ai/test_vector_store.py index 000f4f895..5b2d1d5ed 100644 --- a/lumen/tests/ai/test_vector_store.py +++ b/lumen/tests/ai/test_vector_store.py @@ -7,7 +7,7 @@ except ModuleNotFoundError: pytest.skip("lumen.ai could not be imported, skipping tests.", allow_module_level=True) -from lumen.ai.embeddings import Embeddings, NumpyEmbeddings +from lumen.ai.embeddings import Embeddings, NumpyEmbeddings, OpenAIEmbeddings from lumen.ai.vector_store import DuckDBVectorStore, NumpyVectorStore @@ -584,3 +584,24 @@ async def test_check_embeddings_consistency(self, tmp_path): with pytest.raises(ValueError, match="Provided embeddings class"): DuckDBVectorStore(uri=db_path, embeddings=Embeddings()) + + @pytest.mark.asyncio + async def test_api_key_not_stored_in_metadata(self, tmp_path): + """Verifies that api_key parameter is not included in stored embeddings metadata.""" + import json + + db_path = str(tmp_path / "test_duckdb.db") + + embeddings = OpenAIEmbeddings(api_key="sk-test-secret-key-12345") + store = DuckDBVectorStore(uri=db_path, embeddings=embeddings) + store._setup_database(1) + + metadata_result = store.connection.execute( + "SELECT value FROM vector_store_metadata WHERE key = 'embeddings';" + ).fetchone() + assert metadata_result is not None, "Embeddings metadata should be stored" + + metadata = json.loads(metadata_result[0]) + params = metadata.get("params", {}) + assert "api_key" not in params, "api_key should not be stored in metadata" + store.close() diff --git a/lumen/tests/sources/test_duckdb.py b/lumen/tests/sources/test_duckdb.py index d2cd872a0..a4d312e8c 100644 --- a/lumen/tests/sources/test_duckdb.py +++ b/lumen/tests/sources/test_duckdb.py @@ -545,6 +545,170 @@ def test_mirrors_not_in_tables_from_spec(): assert source.mirrors == {} +def test_create_sql_expr_source_with_params(sample_csv_files): + """Test that create_sql_expr_source accepts and uses params correctly.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + # Create initial source with CSV files + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv' + } + ) + + # Create a new source with parameterized SQL expressions + new_tables = { + 'filtered_customers': 'SELECT * FROM customers WHERE id > ?', + 'high_value_orders': 'SELECT * FROM orders WHERE total > ? AND customer_id = ?' + } + + # Define parameters for each table + params = { + 'filtered_customers': [1], # id > 1 + 'high_value_orders': [200, 1] # total > 200 AND customer_id = 1 + } + + new_source = source.create_sql_expr_source(new_tables, params=params) + + # Check that the new source has the tables + tables = new_source.get_tables() + assert 'filtered_customers' in tables + assert 'high_value_orders' in tables + + # Verify the parameterized queries work correctly + filtered_result = new_source.get('filtered_customers') + assert len(filtered_result) == 2 # Only customers with id > 1 (Bob and Charlie) + assert all(filtered_result['id'] > 1) + assert set(filtered_result['name']) == {'Bob', 'Charlie'} + + high_value_result = new_source.get('high_value_orders') + assert len(high_value_result) == 2 # Two orders match: customer_id=1 and total > 200 + assert all(high_value_result['customer_id'] == 1) + assert all(high_value_result['total'] > 200) + # Check specific values + assert set(high_value_result['total']) == {250.5, 300.0} + + finally: + os.chdir(original_cwd) + + +def test_table_params_parameter(sample_csv_files): + """Test using table_params as a public parameter directly in constructor.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + # Create source with table_params directly + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'filtered_customers': 'SELECT * FROM customers WHERE city = ?' + }, + table_params={ + 'filtered_customers': ['NYC'] + } + ) + + # Query the parameterized table + result = source.get('filtered_customers') + assert len(result) == 1 + assert result.iloc[0]['name'] == 'Alice' + assert result.iloc[0]['city'] == 'NYC' + + # Regular table should still work + all_customers = source.get('customers') + assert len(all_customers) == 3 + + finally: + os.chdir(original_cwd) + + +def test_create_sql_expr_source_without_params(sample_csv_files): + """Test that create_sql_expr_source works without params (backward compatibility).""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv' + } + ) + + # Create a new source WITHOUT params - should work as before + new_tables = { + 'all_customers': 'SELECT * FROM customers', + 'all_orders': 'SELECT * FROM orders' + } + + new_source = source.create_sql_expr_source(new_tables) + + # Verify tables work without params + customers_result = new_source.get('all_customers') + assert len(customers_result) == 3 + + orders_result = new_source.get('all_orders') + assert len(orders_result) == 3 + + finally: + os.chdir(original_cwd) + + +def test_create_sql_expr_source_mixed_params(sample_csv_files): + """Test create_sql_expr_source with some tables having params and others not.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv' + } + ) + + # Mix of parameterized and non-parameterized queries + new_tables = { + 'filtered_customers': 'SELECT * FROM customers WHERE id = ?', # Has params + 'all_orders': 'SELECT * FROM orders' # No params + } + + params = { + 'filtered_customers': [2] # Only customer with id=2 (Bob) + } + + new_source = source.create_sql_expr_source(new_tables, params=params) + + # Verify parameterized query + filtered_result = new_source.get('filtered_customers') + assert len(filtered_result) == 1 + assert filtered_result.iloc[0]['id'] == 2 + assert filtered_result.iloc[0]['name'] == 'Bob' + + # Verify non-parameterized query + orders_result = new_source.get('all_orders') + assert len(orders_result) == 3 + + finally: + os.chdir(original_cwd) + + def test_detour_roundtrip(sample_csv_files): """ Test that creating a SQL expression source from an existing source @@ -566,3 +730,229 @@ def test_detour_roundtrip(sample_csv_files): assert read_df.iloc[[0]].equals(df.iloc[[0]]) assert read_source.tables["limited_customers"] == 'SELECT * FROM customers LIMIT 1' assert "customers" in read_source.tables + + +def test_table_params_basic(sample_csv_files): + """Test table_params with single and multiple parameters.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv', + 'by_name': 'SELECT * FROM customers WHERE name = ?', + 'by_customer_total': 'SELECT * FROM orders WHERE customer_id = ? AND total > ?' + }, + table_params={ + 'by_name': ['Alice'], + 'by_customer_total': [1, 200] + } + ) + + # Single parameter + result = source.get('by_name') + assert len(result) == 1 + assert result.iloc[0]['name'] == 'Alice' + + # Multiple parameters + result = source.get('by_customer_total') + assert len(result) == 2 + assert all(result['customer_id'] == 1) + assert all(result['total'] > 200) + finally: + os.chdir(original_cwd) + + +def test_table_params_data_types(sample_csv_files): + """Test table_params with various data types and SQL patterns.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv', + 'by_int': 'SELECT * FROM orders WHERE customer_id = ?', + 'by_float': 'SELECT * FROM orders WHERE total > ?', + 'by_like': 'SELECT * FROM customers WHERE name LIKE ?', + 'by_in': 'SELECT * FROM customers WHERE city IN (?, ?)' + }, + table_params={ + 'by_int': [2], + 'by_float': [200.5], + 'by_like': ['%li%'], + 'by_in': ['NYC', 'LA'] + } + ) + + assert source.get('by_int').iloc[0]['customer_id'] == 2 + assert len(source.get('by_float')) == 2 + assert set(source.get('by_like')['name']) == {'Alice', 'Charlie'} + assert len(source.get('by_in')) == 2 + finally: + os.chdir(original_cwd) + + +def test_table_params_complex_queries(sample_csv_files): + """Test table_params with JOINs, aggregations, CTEs, and subqueries.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + source = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'orders': 'orders.csv', + 'join_query': 'SELECT c.name, o.total FROM customers c JOIN orders o ON c.id = o.customer_id WHERE c.name = ?', + 'agg_query': 'SELECT customer_id, COUNT(*) as cnt, SUM(total) as sum FROM orders WHERE customer_id = ? GROUP BY customer_id', + 'cte_query': 'WITH totals AS (SELECT customer_id, SUM(total) as spent FROM orders GROUP BY customer_id) SELECT c.name, t.spent FROM customers c JOIN totals t ON c.id = t.customer_id WHERE t.spent > ?', + 'subquery': 'SELECT * FROM customers WHERE id IN (SELECT customer_id FROM orders WHERE total > ?)' + }, + table_params={ + 'join_query': ['Alice'], + 'agg_query': [1], + 'cte_query': [200], + 'subquery': [200] + } + ) + + # JOIN + join_result = source.get('join_query') + assert len(join_result) == 2 + assert all(join_result['name'] == 'Alice') + + # Aggregation + agg_result = source.get('agg_query') + assert agg_result.iloc[0]['cnt'] == 2 + assert agg_result.iloc[0]['sum'] == 550.5 + + # CTE + cte_result = source.get('cte_query') + assert cte_result.iloc[0]['name'] == 'Alice' + + # Subquery + subquery_result = source.get('subquery') + assert subquery_result.iloc[0]['name'] == 'Alice' + finally: + os.chdir(original_cwd) + + +def test_table_params_in_create_sql_expr_source(sample_csv_files): + """Test params in create_sql_expr_source and propagation.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + base_source = DuckDBSource( + uri=':memory:', + tables={'customers': 'customers.csv'}, + table_params={'base_filter': [1]} + ) + + # Create new source with params + new_source = base_source.create_sql_expr_source( + tables={'new_filter': 'SELECT * FROM customers WHERE city = ?'}, + params={'new_filter': ['LA']} + ) + + result = new_source.get('new_filter') + assert len(result) == 1 + assert result.iloc[0]['city'] == 'LA' + finally: + os.chdir(original_cwd) + + +def test_table_params_edge_cases(sample_csv_files): + """Test edge cases: empty params, mismatched counts, multiple tables.""" + df = pd.DataFrame({'col1': [1, 2, 3], 'col2': ['a', 'b', 'c']}) + + # Empty params should work fine + source1 = DuckDBSource.from_df({'test': df}, table_params={}) + pd.testing.assert_frame_equal(source1.get('test'), df) + + # Mismatched parameter count - need to create source directly, not from_df + source2 = DuckDBSource( + uri=':memory:', + ephemeral=True, + tables={ + 'test': 'SELECT * FROM test', + 'bad': 'SELECT * FROM test WHERE col1 > ? AND col1 < ?' + }, + table_params={'bad': [1]} # Missing second parameter + ) + # First load the test table + source2._connection.from_df(df).to_view('test') + + # Now try to query with mismatched params + with pytest.raises(Exception): # DuckDB will raise an error about bind parameter count + source2.get('bad') + + # Multiple tables with different params + files = sample_csv_files + original_cwd = os.getcwd() + try: + os.chdir(files['dir']) + source3 = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'by_city': 'SELECT * FROM customers WHERE city = ?', + 'by_id': 'SELECT * FROM customers WHERE id = ?' + }, + table_params={'by_city': ['NYC'], 'by_id': [2]} + ) + assert source3.get('by_city').iloc[0]['city'] == 'NYC' + assert source3.get('by_id').iloc[0]['id'] == 2 + finally: + os.chdir(original_cwd) + + +def test_table_params_serialization(sample_csv_files): + """Test that params survive to_spec/from_spec roundtrip.""" + files = sample_csv_files + original_cwd = os.getcwd() + + try: + os.chdir(files['dir']) + original = DuckDBSource( + uri=':memory:', + tables={ + 'customers': 'customers.csv', + 'filtered': 'SELECT * FROM customers WHERE id = ?' + }, + table_params={'filtered': [2]} + ) + + # Get the original result before serialization + original_result = original.get('filtered') + assert len(original_result) == 1 + assert original_result.iloc[0]['id'] == 2 + + # Serialize to spec + spec = original.to_spec() + + # For ephemeral in-memory sources with CSV files, we need to stay in the same directory + # or use absolute paths for the restored source to find the files + spec['tables'] = { + 'customers': files['customers'], + 'filtered': 'SELECT * FROM customers WHERE id = ?' + } + + restored = DuckDBSource.from_spec(spec) + restored_result = restored.get('filtered') + + pd.testing.assert_frame_equal(original_result, restored_result) + assert len(restored_result) == 1 + assert restored_result.iloc[0]['id'] == 2 + finally: + os.chdir(original_cwd) diff --git a/lumen/transforms/sql.py b/lumen/transforms/sql.py index 0e7b7359a..a802a9f12 100644 --- a/lumen/transforms/sql.py +++ b/lumen/transforms/sql.py @@ -244,6 +244,8 @@ def apply(self, sql_in: str) -> str: else: parameters[k] = v replaced_expression = replace_placeholders(expression, **parameters) + else: + replaced_expression = expression return self.to_sql(replaced_expression,) @@ -252,7 +254,7 @@ class SQLSelectFrom(SQLFormat): sql_expr = param.String( default="SELECT * FROM {table}", doc=""" - The SQL expression to useif the sql_in does NOT + The SQL expression to use if the sql_in does NOT already contain a SELECT statement.""", ) diff --git a/pyproject.toml b/pyproject.toml index bd721a316..531d181a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ tests = ['pytest', 'pytest-rerunfailures', 'pytest-asyncio'] sql = ['duckdb', 'intake-sql', 'sqlalchemy'] ai = [ 'griffe', 'nbformat', 'duckdb >= 1.2.0', 'pyarrow', 'instructor >=1.6.4', 'pydantic >=2.8.0', 'pydantic-extra-types', 'panel-graphic-walker[kernel] >=0.6.4', - 'markitdown', 'semchunk', 'tiktoken', 'chardet', "panel-material-ui >=0.4.0" + 'markitdown', 'semchunk', 'tiktoken', 'chardet', "panel-material-ui >=0.4.0", "tabulate" ] ai-local = ['lumen[ai]', 'huggingface_hub', 'hf_xet'] ai-openai = ['lumen[ai]', 'openai']