databao-agent 0.1.4.dev5__tar.gz → 0.1.4.dev8__tar.gz

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databao-agent
-Version: 0.1.4.dev5
+Version: 0.1.4.dev8
 Summary: databao-agent: NL queries for data
 Project-URL: Homepage, https://databao.app/
 Project-URL: Source, https://github.com/JetBrains/databao-agent
@@ -9,6 +9,7 @@ License-File: LICENSE.md
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: <3.15,>=3.11
+Requires-Dist: claude-agent-sdk>=0.1.48
 Requires-Dist: databao-context-engine[mysql,postgresql,snowflake]~=0.6.0
 Requires-Dist: dbt-core~=1.9.0
 Requires-Dist: dbt-duckdb>=1.10.0

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/databao/agent/api.py

@@ -8,10 +8,13 @@ from databao.agent.configs.agent import DEFAULT_AGENT_CONFIG, AgentConfig
 from databao.agent.configs.llm import LLMConfig, LLMConfigDirectory
 from databao.agent.core import Agent, Cache, Executor, Visualizer
 from databao.agent.core.domain import Domain, _DCEProjectDomain, _InMemoryDomain
-from databao.agent.executors import ReactDuckDBExecutor
+from databao.agent.executors import (
+    ClaudeCodeExecutor,
+    DbtProjectExecutor,
+    LighthouseExecutor,
+    ReactDuckDBExecutor,
+)
 from databao.agent.executors.dbt.config import DbtConfig
-from databao.agent.executors.dbt.executor import DbtProjectExecutor
-from databao.agent.executors.lighthouse.executor import LighthouseExecutor
 from databao.agent.visualizers.vega_chat import VegaChatVisualizer
 
 
@@ -49,6 +52,8 @@ def agent(
                 data_executor = DbtProjectExecutor(dbt_config=dbt_config, writer=writer)
             case "react_duckdb":
                 data_executor = ReactDuckDBExecutor(writer=writer)
+            case "claude":
+                data_executor = ClaudeCodeExecutor(writer=writer)
             case _:
                 raise ValueError(f"Invalid executor type: {executor_type}")
 

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/databao/agent/core/agent.py

@@ -55,7 +55,7 @@ class Agent:
         self.__llm_config = llm
         self.__agent_config = agent_config
 
-        self.__executor = data_executor
+        self.__executor: Executor = data_executor
         self.__visualizer = visualizer
         self.__cache = cache
         self.__mcp: McpManager = McpManager()

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/databao/agent/core/executor.py

@@ -4,7 +4,6 @@ import re
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, TextIO
 
-from langchain_core.tools import BaseTool
 from pandas import DataFrame
 from pydantic import BaseModel, ConfigDict
 
@@ -149,8 +148,8 @@ class Executor(ABC):
     """
 
     @abstractmethod
-    def register_tools(self, tools: list[BaseTool]) -> None:
-        """Register additional LangChain tools to be available during execution."""
+    def register_tools(self, tools: list[Any]) -> None:
+        """Register additional tools to be available during execution."""
 
     @abstractmethod
     def drop_last_opa_group(self, cache: "Cache", n: int = 1) -> None:
@@ -184,7 +183,8 @@ class Executor(ABC):
         """
         pass
 
-    def prepare_for_execution(self, domain: "Domain") -> None:
+    @staticmethod
+    def prepare_for_execution(domain: "Domain") -> None:
         if domain.supports_context and not domain.is_context_built():
             logger.warning(
                 "Context has not been built yet. Building it now — this may take a while. "

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/databao/agent/databases/snowflake_adapter.py

@@ -11,6 +11,7 @@ from databao_context_engine import (
     SnowflakeSSOAuth,
 )
 from databao_context_engine.pluginlib.build_plugin import AbstractConfigFile
+from snowflake.connector.network import SNOWFLAKE_HOST_SUFFIX
 from sqlalchemy import Connection, Engine, make_url
 
 from databao.agent.databases.database_adapter import DatabaseAdapter
@@ -92,8 +93,13 @@ class SnowflakeAdapter(DatabaseAdapter):
         if "dbname" in content:
             content[DATABASE_KEY] = content.pop("dbname")
 
+        host: str | None = content.pop("host", None)
+        account: str = content.get(ACCOUNT_KEY, "")
+        if host and host.endswith(SNOWFLAKE_HOST_SUFFIX):
+            account = host[: -len(SNOWFLAKE_HOST_SUFFIX)]
+
         return SnowflakeConnectionProperties(
-            account=content.get(ACCOUNT_KEY),
+            account=account,
             warehouse=content.get(WAREHOUSE_KEY),
             database=content.get(DATABASE_KEY),
             user=content.get(USER_KEY),

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/databao/agent/duckdb/react_tools.py

@@ -12,7 +12,7 @@ from pydantic import BaseModel
 
 from databao.agent.core import Domain
 from databao.agent.duckdb.schema_inspection import inspect_duckdb_schema, summarize_duckdb_schema
-from databao.agent.executors.tools import make_search_context_tool
+from databao.agent.executors.langchain_tools import make_search_context_tool
 
 _LOGGER = logging.getLogger(__name__)
 

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/databao/agent/executors/__init__.py

@@ -1,5 +1,6 @@
+from databao.agent.executors.claude_code.executor import ClaudeCodeExecutor
 from databao.agent.executors.dbt.executor import DbtProjectExecutor
 from databao.agent.executors.lighthouse.executor import LighthouseExecutor
 from databao.agent.executors.react_duckdb.executor import ReactDuckDBExecutor
 
-__all__ = ["DbtProjectExecutor", "LighthouseExecutor", "ReactDuckDBExecutor"]
+__all__ = ["ClaudeCodeExecutor", "DbtProjectExecutor", "LighthouseExecutor", "ReactDuckDBExecutor"]

{databao_agent-0.1.4.dev5 → databao_agent-0.1.4.dev8}/databao/agent/executors/base.py

@@ -1,9 +1,9 @@
+import logging
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from typing import Any, TextIO
 
 import duckdb
-from databao_context_engine import DuckDBConnectionConfig
 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
@@ -20,11 +20,13 @@ from databao.agent.databases import register_db_in_duckdb
 from databao.agent.executors.frontend.text_frontend import TextStreamFrontend
 from databao.agent.executors.history_cleaning import clean_tool_history
 
+logger = logging.getLogger(__name__)
 
-class GraphExecutor(Executor, ABC):
+
+class DuckDBExecutor(Executor, ABC):
     """
-    Base class for LangGraph executors that execute with a DuckDB connection and LLM configuration.
-    Provides common functionality for graph caching, message handling, and OPA processing.
+    Base class for executors that execute with a DuckDB connection and LLM configuration.
+    Provides common functionality for message handling and OPA processing.
     """
 
     def __init__(self, writer: TextIO | None = None) -> None:
@@ -34,16 +36,11 @@ class GraphExecutor(Executor, ABC):
             writer: Optional TextIO for streaming output. If provided, streaming
                     output will be written to this writer instead of stdout.
         """
-        self._graph_recursion_limit = 50
         self._writer = writer
         self._duckdb_connection: duckdb.DuckDBPyConnection = duckdb.connect(":memory:")
         self._registered_dbs: dict[str, DBDataSource] = {}
         self._registered_dfs: dict[str, DFDataSource] = {}
         self._registered_dbts: dict[str, DBTDataSource] = {}
-        self._extra_tools: dict[str, BaseTool] = {}
-        self._compiled_graph: CompiledStateGraph[Any] | None = None
-        self._compiled_tools_version: int = 0
-        self._compiled_at_version: int = -1
 
     def _init_sources_from_domain(self, domain: Domain, *, register_in_duckdb: bool = True) -> None:
         """Sync sources from the domain into the executor's registered dictionaries.
@@ -61,7 +58,11 @@ class GraphExecutor(Executor, ABC):
         for name, db_source in sources.dbs.items():
             if name not in self._registered_dbs:
                 if register_in_duckdb:
-                    register_db_in_duckdb(self._duckdb_connection, db_source.config, name)
+                    try:
+                        register_db_in_duckdb(self._duckdb_connection, db_source.config, name)
+                    except Exception as exc:
+                        logger.warning("Datasource '%s' is not available and will be skipped: %s", name, exc)
+                        continue
                 self._registered_dbs[name] = db_source
 
         for name, df_source in sources.dfs.items():
@@ -74,25 +75,26 @@ class GraphExecutor(Executor, ABC):
             if name not in self._registered_dbts:
                 self._registered_dbts[name] = dbt_source
 
-    def prepare_for_execution(self, domain: "Domain") -> None:
-        if not domain.supports_context or domain.is_context_built():
-            return
 
-        # DuckDB does not allow two connections to hold a file open simultaneously.
-        # Temporarily detach file-based DuckDB sources so the context engine can attach them.
-        duckdb_file_sources = {
-            name: source
-            for name, source in self._registered_dbs.items()
-            if isinstance(source.config, DuckDBConnectionConfig)
-        }
-        for name in duckdb_file_sources:
-            self._duckdb_connection.execute(f'DETACH "{name}"')
-
-        try:
-            super().prepare_for_execution(domain)
-        finally:
-            for name, source in duckdb_file_sources.items():
-                register_db_in_duckdb(self._duckdb_connection, source.config, name)
+class GraphExecutor(DuckDBExecutor, ABC):
+    """
+    Base class for LangGraph executors that execute with a DuckDB connection and LLM configuration.
+    Provides common functionality for graph caching, message handling, and OPA processing.
+    """
+
+    def __init__(self, writer: TextIO | None = None) -> None:
+        """Initialize agent with graph caching infrastructure.
+
+        Args:
+            writer: Optional TextIO for streaming output. If provided, streaming
+                    output will be written to this writer instead of stdout.
+        """
+        super().__init__(writer)
+        self._extra_tools: dict[str, BaseTool] = {}
+        self._graph_recursion_limit = 50
+        self._compiled_graph: CompiledStateGraph[Any] | None = None
+        self._compiled_tools_version: int = 0
+        self._compiled_at_version: int = -1
 
     def register_tools(self, tools: list[BaseTool]) -> None:
         """Register additional LangChain tools and invalidate the cached compiled graph."""
@@ -100,6 +102,18 @@ class GraphExecutor(Executor, ABC):
             self._extra_tools[t.name] = t
         self._compiled_tools_version += 1
 
+    def drop_last_opa_group(self, cache: Cache, n: int = 1) -> None:
+        """Drop last n groups of operations from the message history."""
+        messages = cache.get("state", default={}).get("messages", [])
+        human_messages = [m for m in messages if isinstance(m, HumanMessage)]
+        if len(human_messages) < n:
+            raise ValueError(f"Cannot drop last {n} operations - only {len(human_messages)} operations found.")
+        c = 0
+        while c < n:
+            m = messages.pop()
+            if isinstance(m, HumanMessage):
+                c += 1
+
     @abstractmethod
     def _compile_graph(
         self,
@@ -120,18 +134,6 @@ class GraphExecutor(Executor, ABC):
             self._compiled_at_version = self._compiled_tools_version
         return self._compiled_graph
 
-    def drop_last_opa_group(self, cache: Cache, n: int = 1) -> None:
-        """Drop last n groups of operations from the message history."""
-        messages = cache.get("state", default={}).get("messages", [])
-        human_messages = [m for m in messages if isinstance(m, HumanMessage)]
-        if len(human_messages) < n:
-            raise ValueError(f"Cannot drop last {n} operations - only {len(human_messages)} operations found.")
-        c = 0
-        while c < n:
-            m = messages.pop()
-            if isinstance(m, HumanMessage):
-                c += 1
-
     def _process_opas(self, opas: list[Opa], cache: Cache) -> list[Any]:
         """
         Process a single opa and convert it to a message, appending to message history.

databao_agent-0.1.4.dev8/databao/agent/executors/claude_code/claude_model_wrapper.py

@@ -0,0 +1,313 @@
+import asyncio
+import json
+import logging
+import queue
+import threading
+from collections.abc import Generator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, TextIO
+
+import pandas as pd
+from _duckdb import DuckDBPyConnection
+from claude_agent_sdk import (
+    ClaudeAgentOptions,
+    ClaudeSDKClient,
+    SdkMcpTool,
+    create_sdk_mcp_server,
+    tool,
+)
+from claude_agent_sdk.types import McpSdkServerConfig, ResultMessage, SystemPromptPreset
+from claude_agent_sdk.types import Message as ClaudeMessage
+from claude_agent_sdk.types import SystemMessage as ClaudeSystemMessage
+from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
+from mcp.types import ToolAnnotations
+
+from databao.agent.configs.llm import LLMConfig
+from databao.agent.core.executor import ExecutionResult
+from databao.agent.executors.claude_code.utils import cast_claude_message_to_langchain_message
+from databao.agent.executors.frontend.messages import get_tool_call
+from databao.agent.executors.frontend.text_frontend import TextStreamFrontend
+from databao.agent.executors.lighthouse.graph import RUN_SQL_QUERY_TOOL_DESCRIPTION
+from databao.agent.executors.utils import run_sql_query
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class QueryResult:
+    sql: str
+    df: pd.DataFrame
+
+
+class ClaudeModelWrapper:
+    DISPLAY_ROW_LIMIT = 12
+    """Max number of rows to return in SQL tool calls."""
+
+    DISPLAY_CELL_CHAR_LIMIT = 1024
+    """Max number of characters a dataframe cell can have before it is trimmed."""
+
+    __runtime_mcp_server: McpSdkServerConfig | None = None
+
+    def __init__(
+        self,
+        *,
+        config: LLMConfig,
+        connection: DuckDBPyConnection,
+        system_prompt: str,
+        append_system_prompt: bool = False,
+        session_id: str | None = None,
+        limit_max_rows: int | None = None,
+        max_turns: int | None = 100,
+    ):
+        self._duckdb_connection = connection
+        self._limit_max_rows = limit_max_rows
+        self.config = config
+        self.sdk_mcp_tools = self._build_tools()
+        self._tool_server_name = Path(__file__).stem + "_mcp_server"
+        self.mcp_tool_names = [self._get_full_tool_name(t.name) for t in self.sdk_mcp_tools]
+
+        self.options = ClaudeAgentOptions(
+            max_turns=max_turns,
+            cwd=".",
+            allowed_tools=self.mcp_tool_names,
+            model=self.config.name,
+            mcp_servers={self._tool_server_name: self._build_tool_server()},
+            permission_mode="acceptEdits",
+            resume=session_id,
+            system_prompt=system_prompt
+            if not append_system_prompt
+            else SystemPromptPreset(
+                type="preset",  # Append to Claude's internal system prompt
+                preset="claude_code",
+                append=system_prompt,
+            ),
+        )
+        self.client = ClaudeSDKClient(options=self.options)
+        self._query_cache: dict[int, QueryResult] = {}
+        self._ready_event: threading.Event
+        self._exit_event: asyncio.Event
+        self._visualization_prompt: str | None = None
+
+    def __enter__(self) -> "ClaudeModelWrapper":
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(target=self._loop.run_forever, daemon=True, name=f"{self._tool_server_name}")
+        self._thread.start()
+
+        self._ready_event = threading.Event()
+
+        async def _lifecycle() -> None:
+            self._exit_event = asyncio.Event()
+            async with self.client:
+                self._ready_event.set()
+                await self._exit_event.wait()
+
+        self._lifecycle_task = asyncio.run_coroutine_threadsafe(_lifecycle(), self._loop)
+        self._ready_event.wait()
+        return self
+
+    def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
+        self._loop.call_soon_threadsafe(self._exit_event.set)
+        self._lifecycle_task.result()
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join()
+
+    def _get_full_tool_name(self, tool_name: str) -> str:
+        return f"mcp__{self._tool_server_name}__{tool_name}"
+
+    def _build_tools(self) -> list[SdkMcpTool[Any]]:
+        # Set read only hints to enable parallel tool execution
+        # (see https://platform.claude.com/docs/en/agent-sdk/agent-loop#parallel-tool-execution)
+
+        tools = []
+
+        @tool(
+            "run_sql_query",
+            RUN_SQL_QUERY_TOOL_DESCRIPTION,
+            {"sql": str},
+            annotations=ToolAnnotations(readOnlyHint=True),
+        )
+        async def _run_sql_query(args: dict[str, Any]) -> dict[str, Any]:
+            result = run_sql_query(
+                args.get("sql", ""),
+                con=self._duckdb_connection,
+                sql_row_limit=self._limit_max_rows,
+                display_row_limit=self.DISPLAY_ROW_LIMIT,
+                display_cell_char_limit=self.DISPLAY_CELL_CHAR_LIMIT,
+            )
+            if "error" in result:
+                return {"content": [{"type": "text", "text": json.dumps(result, default=str)}]}
+
+            result_for_llm: dict[str, Any] = {"csv": result.get("csv", "")}
+
+            if (sql := result.get("sql")) and (df := result.get("df")) is not None:
+                query_id = len(self._query_cache) + 1
+                self._query_cache[query_id] = QueryResult(sql=sql, df=df)
+                result_for_llm["query_id"] = query_id
+
+            return {"content": [{"type": "text", "text": json.dumps(result_for_llm, default=str)}]}
+
+        tools.append(_run_sql_query)
+
+        @tool(
+            "submit_query_id",
+            """\
+This tool call must be the last tool to be called by the model.
+It will provide to the user the generated sql and the output thereof resulting from the query with
+the respective query id. You will find the query ids of the error-free queries in the outputs of
+the run_sql_query tool in the `query_id` key. The `query_id` itself need not be the one of the last
+generated query, it rather needs to reference the query which most closely matches the
+user's question.
+
+Args:
+query_id: The ID of the query to submit.""",
+            {"query_id": int, "visualization_prompt": str},
+            annotations=ToolAnnotations(readOnlyHint=True),
+        )
+        async def submit_query_id(args: dict[str, Any]) -> dict[str, Any]:
+            query_id: int | None = args.get("query_id")
+            self._visualization_prompt = args.get("visualization_prompt")
+
+            if query_id not in self._query_cache:
+                return {"content": [{"type": "text", "text": json.dumps({"error": f"Query id {query_id} not found"})}]}
+            return {"content": [{"type": "text", "text": json.dumps({"query_id": query_id})}]}
+
+        tools.append(submit_query_id)
+
+        return tools
+
+    def _build_tool_server(self) -> McpSdkServerConfig:
+        tools = self._build_tools()
+        if self.__runtime_mcp_server is None:
+            self.__runtime_mcp_server = create_sdk_mcp_server(
+                name=self._tool_server_name,
+                version="1.0.0",
+                tools=tools,
+            )
+        return self.__runtime_mcp_server
+
+    def _check_mcp_tool_availability(self, first_message: ClaudeMessage) -> None:
+        """
+        Each conversation begins with an initial init system message. This SystemMessage
+        carries the information about the tools available to claude. To prevent
+        the system from running with the mcp tools being silently not available, we
+        explicitly look for them and raise and error if any of them is missing.
+        """
+        if not isinstance(first_message, ClaudeSystemMessage):
+            raise TypeError(
+                f"The first message should be a system message, got {type(first_message)}. "
+                "Check if you are actually calling this function on the first message of the conversation."
+            )
+
+        if missing_tools := set(self.mcp_tool_names).difference(first_message.data["tools"]):
+            raise ValueError(
+                f"The following mcp tools are not available: {missing_tools}. "
+                "Check the connection to the mcp servers by running /mcp in the claude console."
+            )
+
+    def _get_tool_query_id_results(self, message: ToolMessage) -> QueryResult | None:
+        try:
+            payload = json.loads(message.text)
+        except json.JSONDecodeError as e:
+            _LOGGER.warning("Failed to parse tool call payload: %s", message.text, exc_info=e)
+            payload = {}
+        query_id = payload.get("query_id")
+        if query_id is not None:
+            return self._query_cache.get(query_id)
+        return None
+
+    def solve(self, prompt: str) -> Generator[ClaudeMessage, None, None]:
+        _LOGGER.info(f"Querying {prompt}")
+
+        _sentinel = object()
+        q: queue.Queue[Any] = queue.Queue()
+
+        async def _produce() -> None:
+            await self.client.query(prompt=prompt)
+            messages = self.client.receive_response()
+            async for message in messages:
+                q.put(message)
+            q.put(_sentinel)
+
+        asyncio.run_coroutine_threadsafe(_produce(), self._loop)
+
+        first_message = q.get()
+        self._check_mcp_tool_availability(first_message)
+        yield first_message
+        _LOGGER.info(first_message)
+
+        n_messages = 1
+        while (message := q.get()) is not _sentinel:
+            _LOGGER.info(message)
+            n_messages += 1
+            yield message
+
+        _LOGGER.info(f"End of conversation. Got {n_messages} messages.\n\n")
+
+    def ask(
+        self,
+        prompt: str,
+        *,
+        stream: bool = False,
+        writer: TextIO | None = None,
+    ) -> tuple[ExecutionResult, str | None]:
+        """
+        Iterate through the messages from claude, cast them into BaseMessage
+        object so that they are compatible with the Experiment class and pack
+        them into a SolverResult object.
+        """
+        session_id: str | None = None
+        max_init_query_id = max(self._query_cache) if self._query_cache else 0
+        message_log: list[BaseMessage] = []
+        submitted_query_result: QueryResult | None = None
+        frontend = TextStreamFrontend({"messages": message_log}, writer=writer)
+        for message in self.solve(prompt):
+            if isinstance(message, ClaudeSystemMessage) and session_id is None:
+                # Child subagents have their own system messages, but we want the parent one only
+                session_id = message.data.get("session_id", "default")
+
+            # Skip the final text-only ResultMessage, as the previous AssistantMessage already contains the text
+            # of this message.
+            if isinstance(message, ResultMessage):
+                continue
+
+            lc_message = cast_claude_message_to_langchain_message(message)
+
+            if isinstance(lc_message, ToolMessage):
+                tool_call = get_tool_call(message_log, lc_message)
+                if tool_call is not None:
+                    if tool_call["name"] == self._get_full_tool_name("run_sql_query"):  # noqa: SIM102
+                        if query_result := self._get_tool_query_id_results(lc_message):
+                            lc_message.artifact = {
+                                "sql": query_result.sql,
+                                "df": query_result.df,
+                            }  # To show when streaming
+                    if tool_call["name"] == self._get_full_tool_name("submit_query_id"):  # noqa: SIM102
+                        if query_result := self._get_tool_query_id_results(lc_message):
+                            submitted_query_result = query_result
+
+            message_log.append(lc_message)
+
+            if stream:
+                if isinstance(lc_message, AIMessage):
+                    frontend.write_full_ai_message(lc_message)
+                frontend.write_stream_chunk("values", {"messages": message_log})
+
+        if stream:
+            frontend.end()
+
+        if submitted_query_result is None:
+            # Fallback to the last executed query if no query was submitted
+            max_query_id = max(self._query_cache) if self._query_cache else 0
+            if max_query_id > max_init_query_id:
+                submitted_query_result = self._query_cache[max_query_id]
+
+        return ExecutionResult(
+            text=message_log[-1].text if message_log else "",
+            meta={
+                "visualization_prompt": self._visualization_prompt,
+                ExecutionResult.META_MESSAGES_KEY: message_log,
+            },
+            code=submitted_query_result.sql if submitted_query_result else "",
+            df=submitted_query_result.df if submitted_query_result else None,
+        ), session_id