dao-ai 0.0.36__py3-none-any.whl → 0.1.1__py3-none-any.whl

dao_ai-0.1.1.dist-info/RECORD

@@ -0,0 +1,62 @@
+dao_ai/__init__.py,sha256=18P98ExEgUaJ1Byw440Ct1ty59v6nxyWtc5S6Uq2m9Q,1062
+dao_ai/agent_as_code.py,sha256=sviZQV7ZPxE5zkZ9jAbfegI681nra5i8yYxw05e3X7U,552
+dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
+dao_ai/cli.py,sha256=AoGw4erjF1C8_2XEppj9yU8agQKSRnJr-x-I8Y-oAx8,31121
+dao_ai/config.py,sha256=VjFFoBNer7ylE6u3lKgqGVUN4VDlIyeUlTcXuMkNHVo,101027
+dao_ai/graph.py,sha256=1-uQlo7iXZQTT3uU8aYu0N5rnhw5_g_2YLwVsAs6M-U,1119
+dao_ai/messages.py,sha256=4ZBzO4iFdktGSLrmhHzFjzMIt2tpaL-aQLHOQJysGnY,6959
+dao_ai/models.py,sha256=hOUEPyWZNOL0SmA0V9e0fyaO2ZtuAd-rr59M_3CoggQ,77734
+dao_ai/nodes.py,sha256=2DeiR6WmLlfiWMIaT8Gj9FOC2Ewk5bammMKXmVb1VTc,7299
+dao_ai/optimization.py,sha256=KwdKR9njYXe1aIlwGXJ8f5dE7G83aV1CxNqhvjBfGis,22780
+dao_ai/prompts.py,sha256=5Mh-8OYMASMHid58iFAj-rr6K1B41gmNN91Ia8Ciq_I,4846
+dao_ai/state.py,sha256=fNRZ8_M2VBqO8lXS3Egh53PH6LsD3bB8npOG44XWxGA,5641
+dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dao_ai/utils.py,sha256=elkkhNe0L976oTdaEjxHsmW5Fc9VLLS4IjS7q5HgRdA,11401
+dao_ai/vector_search.py,sha256=jlaFS_iizJ55wblgzZmswMM3UOL-qOp2BGJc0JqXYSg,2839
+dao_ai/genie/__init__.py,sha256=vdEyGhrt6L8GlK75SyYvTnl8QpHKDCJC5hJKLg4DesQ,1063
+dao_ai/genie/core.py,sha256=HPKbocvhnnw_PkQwfoq5bpgQmL9lZyyS6_goTJL8yiY,1073
+dao_ai/genie/cache/__init__.py,sha256=JfgCJl1NYQ1aZvZ4kly4T6uQK6ZCJ6PX_htuq7nJF50,1203
+dao_ai/genie/cache/base.py,sha256=_MhHqYrHejVGrJjSLX26TdHwvQZb-HgiantRYSB8fJY,1961
+dao_ai/genie/cache/core.py,sha256=Oe7MSreefcaOri4eSpJJFQBBfNP0MoW7qQdbA0vq89U,2439
+dao_ai/genie/cache/lru.py,sha256=4oXINJnTb-5FvKVd-U3J6gUI9o9jUEoBLOHsxlkvW34,11465
+dao_ai/genie/cache/semantic.py,sha256=RZAzrniuW_VmeWZoHWfbPr2eP3m9gj4fi4vEDM_fvYA,36892
+dao_ai/hooks/__init__.py,sha256=uA4DQdP9gDf4SyNjNx9mWPoI8UZOcTyFsCXV0NraFvQ,463
+dao_ai/hooks/core.py,sha256=upTAI11RncWZ5uu2wBA6nJMaqRMVHJFuzci-iVL8MFw,1700
+dao_ai/memory/__init__.py,sha256=Us3wFehvug_h83m-UJ7OXdq2qZ0e9nHBQE7m5RwoAd8,559
+dao_ai/memory/base.py,sha256=99nfr2UZJ4jmfTL_KrqUlRSCoRxzkZyWyx5WqeUoMdQ,338
+dao_ai/memory/core.py,sha256=vhN-cDZxhtykwXFnRK5HWBbMvlbYyafRoIgew-7ufqI,5233
+dao_ai/memory/databricks.py,sha256=PJSNKD1C48-jKXNg3jPYEANKbP4bH-zWLpMZG5SHo0A,14136
+dao_ai/memory/postgres.py,sha256=5T9qTSijxUl2zCEiW3REqgsdsPjgdxfQTGTf6gBY2M8,15362
+dao_ai/middleware/__init__.py,sha256=epSCtCtttIogl21nVK768Ln35L0mOShVczyURtR6Ln8,3609
+dao_ai/middleware/assertions.py,sha256=EgFVqxtD_UWBTQkXHmoduVYRuBYZtBbfcDCieMj3PXM,26651
+dao_ai/middleware/base.py,sha256=uG2tpdnjL5xY5jCKvb_m3UTBtl4ZC6fJQUkDsQvV8S4,1279
+dao_ai/middleware/core.py,sha256=Umq4wgxlsDHN_-LYMjlrQur2R6j9LBYozP7uCEFC3uI,1852
+dao_ai/middleware/guardrails.py,sha256=U2U5bzWTJ5UdceXbEwSZS4brhpFZhFhRmx9RmbcikJE,13967
+dao_ai/middleware/human_in_the_loop.py,sha256=hivU5CugD4D-9j_d5AHqu6atIc6s6GHel5VxHTbgcys,7435
+dao_ai/middleware/message_validation.py,sha256=zBAAlI_th1SdxX06p4_4hCKKkM4v0Dkh8TEWGHeCDYc,18471
+dao_ai/middleware/summarization.py,sha256=VmpFXtUmFA5UybGvrODzzerZ1aoVQEFvgTTU90clo4M,7152
+dao_ai/orchestration/__init__.py,sha256=i85CLfRR335NcCFhaXABcMkn6WZfXnJ8cHH4YZsZN0s,1622
+dao_ai/orchestration/core.py,sha256=4EdorNTs-NV3KkX_K3RtwRsKH3-K4sT_ORO-lsMYpYo,9418
+dao_ai/orchestration/supervisor.py,sha256=tjZwwcwXzKHjwf1x7x0MT0QqywAWP6ODvQhrK9jJ3Jk,9433
+dao_ai/orchestration/swarm.py,sha256=d7lq5KFFBX3qxSoTAIEKVZJtDJiEjSLviGUtas3-4pY,7603
+dao_ai/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dao_ai/providers/base.py,sha256=-fjKypCOk28h6vioPfMj9YZSw_3Kcbi2nMuAyY7vX9k,1383
+dao_ai/providers/databricks.py,sha256=rzFNVLmZYO2Lm3U3kGGQsEdn4W8mm04NUdnSoiPD12o,54714
+dao_ai/tools/__init__.py,sha256=WLb4mgC7WUbaSDOVpjlf9tJkm5Dr4vEmxG5hnyCgeSc,1568
+dao_ai/tools/agent.py,sha256=WbQnyziiT12TLMrA7xK0VuOU029tdmUBXbUl-R1VZ0Q,1886
+dao_ai/tools/core.py,sha256=f6unLEtYfLZ9eDhoqgaREFylrb9JZ6zlZBYV_45Albo,3785
+dao_ai/tools/email.py,sha256=p7FHlA9gq0LR9-YLuJL1Glq3n465wFu2tqDPX-UW0aY,9662
+dao_ai/tools/genie.py,sha256=zfHAosWMLl-us2vG1c7ziEe8oQehBrxZISWqmC7GZ5s,8942
+dao_ai/tools/mcp.py,sha256=C1NAUKkgl3QUkMTkCZKO1uIMU0gViWOFSjiq-t_tNGk,8454
+dao_ai/tools/memory.py,sha256=lwObKimAand22Nq3Y63tsv-AXQ5SXUigN9PqRjoWKes,1836
+dao_ai/tools/python.py,sha256=D021e4kPlH-CkrQGpzM3n7DZJtTCv2LCW0VPVQGb5Sk,1734
+dao_ai/tools/search.py,sha256=AdQf8UouvJjw5UdCmD1-jrX_PJ-cNZ1oY_jh01dHGxY,433
+dao_ai/tools/slack.py,sha256=NWtzxpUJOxXbngThpb2KYfVvqSOi312crLCCA0xtv-w,4775
+dao_ai/tools/time.py,sha256=Y-23qdnNHzwjvnfkWvYsE7PoWS1hfeKy44tA7sCnNac,8759
+dao_ai/tools/unity_catalog.py,sha256=5Wf47t810HOLMue7qhixuy041B8ykRYwzovbLTL13No,14579
+dao_ai/tools/vector_search.py,sha256=tiump9N6o558Y2tIsSox7-R2nQYmdJh0XTyxgRWDmlY,12881
+dao_ai-0.1.1.dist-info/METADATA,sha256=PK7dUXOuoWkCAYg-KnAgzoFVQMt6uzGHvoc494oMKfI,79539
+dao_ai-0.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.1.1.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.1.1.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.1.1.dist-info/RECORD,,

dao_ai/chat_models.py

@@ -1,204 +0,0 @@
-import json
-from typing import Any, Iterator, Optional, Sequence
-
-from databricks_langchain import ChatDatabricks
-from langchain_core.callbacks import CallbackManagerForLLMRun
-from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
-from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
-from loguru import logger
-
-
-class ChatDatabricksFiltered(ChatDatabricks):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def _preprocess_messages(
-        self, messages: Sequence[BaseMessage]
-    ) -> Sequence[BaseMessage]:
-        logger.debug(f"Preprocessing {len(messages)} messages for filtering")
-
-        logger.trace(
-            f"Original messages:\n{json.dumps([msg.model_dump() for msg in messages], indent=2)}"
-        )
-
-        # Diagnostic logging to understand what types of messages we're getting
-        message_types = {}
-        remove_message_count = 0
-        empty_content_count = 0
-
-        for msg in messages:
-            msg_type = msg.__class__.__name__
-            message_types[msg_type] = message_types.get(msg_type, 0) + 1
-
-            if msg_type == "RemoveMessage":
-                remove_message_count += 1
-            elif hasattr(msg, "content") and (msg.content == "" or msg.content is None):
-                empty_content_count += 1
-
-        logger.debug(f"Message type breakdown: {message_types}")
-        logger.debug(
-            f"RemoveMessage count: {remove_message_count}, Empty content count: {empty_content_count}"
-        )
-
-        filtered_messages = []
-        for i, msg in enumerate(messages):
-            # First, filter out RemoveMessage objects completely - they're LangGraph-specific
-            # and should never be sent to an LLM
-            if hasattr(msg, "__class__") and msg.__class__.__name__ == "RemoveMessage":
-                logger.debug(f"Filtering out RemoveMessage at index {i}")
-                continue
-
-            # Be very conservative with filtering - only filter out messages that are:
-            # 1. Have empty or None content AND
-            # 2. Are not tool-related messages AND
-            # 3. Don't break tool_use/tool_result pairing
-            # 4. Are not the only remaining message (to avoid filtering everything)
-            has_empty_content = hasattr(msg, "content") and (
-                msg.content == "" or msg.content is None
-            )
-
-            # Check if this message has tool calls (non-empty list)
-            has_tool_calls = (
-                hasattr(msg, "tool_calls")
-                and msg.tool_calls
-                and len(msg.tool_calls) > 0
-            )
-
-            # Check if this is a tool result message
-            is_tool_result = hasattr(msg, "tool_call_id") or isinstance(
-                msg, ToolMessage
-            )
-
-            # Check if the previous message had tool calls (this message might be a tool result)
-            prev_had_tool_calls = False
-            if i > 0:
-                prev_msg = messages[i - 1]
-                prev_had_tool_calls = (
-                    hasattr(prev_msg, "tool_calls")
-                    and prev_msg.tool_calls
-                    and len(prev_msg.tool_calls) > 0
-                )
-
-            # Check if the next message is a tool result (this message might be a tool use)
-            next_is_tool_result = False
-            if i < len(messages) - 1:
-                next_msg = messages[i + 1]
-                next_is_tool_result = hasattr(next_msg, "tool_call_id") or isinstance(
-                    next_msg, ToolMessage
-                )
-
-            # Special handling for empty AIMessages - they might be placeholders or incomplete responses
-            # Don't filter them if they're the only AI response or seem important to the conversation flow
-            is_empty_ai_message = has_empty_content and isinstance(msg, AIMessage)
-
-            # Only filter out messages with empty content that are definitely not needed
-            should_filter = (
-                has_empty_content
-                and not has_tool_calls
-                and not is_tool_result
-                and not prev_had_tool_calls  # Don't filter if previous message had tool calls
-                and not next_is_tool_result  # Don't filter if next message is a tool result
-                and not (
-                    is_empty_ai_message and len(messages) <= 2
-                )  # Don't filter empty AI messages in short conversations
-            )
-
-            if should_filter:
-                logger.debug(f"Filtering out message at index {i}: {msg.model_dump()}")
-                continue
-            else:
-                filtered_messages.append(msg)
-
-        logger.debug(
-            f"Filtered {len(messages)} messages down to {len(filtered_messages)} messages"
-        )
-
-        # Log diagnostic information if all messages were filtered out
-        if len(filtered_messages) == 0:
-            logger.warning(
-                f"All {len(messages)} messages were filtered out! This indicates a problem with the conversation state."
-            )
-            logger.debug(f"Original message types: {message_types}")
-
-            if remove_message_count == len(messages):
-                logger.warning(
-                    "All messages were RemoveMessage objects - this suggests a bug in summarization logic"
-                )
-            elif empty_content_count > 0:
-                logger.debug(f"{empty_content_count} messages had empty content")
-
-        return filtered_messages
-
-    def _postprocess_message(self, message: BaseMessage) -> BaseMessage:
-        return message
-
-    def _generate(
-        self,
-        messages: Sequence[BaseMessage],
-        stop: Optional[Sequence[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        """Override _generate to apply message preprocessing and postprocessing."""
-        # Apply message preprocessing
-        processed_messages: Sequence[BaseMessage] = self._preprocess_messages(messages)
-
-        if len(processed_messages) == 0:
-            logger.error(
-                "All messages were filtered out during preprocessing. This indicates a serious issue with the conversation state."
-            )
-            empty_generation = ChatGeneration(
-                message=AIMessage(content="", id="empty-response")
-            )
-            return ChatResult(generations=[empty_generation])
-
-        logger.trace(
-            f"Processed messages:\n{json.dumps([msg.model_dump() for msg in processed_messages], indent=2)}"
-        )
-
-        result: ChatResult = super()._generate(
-            processed_messages, stop, run_manager, **kwargs
-        )
-
-        if result.generations:
-            for generation in result.generations:
-                if isinstance(generation, ChatGeneration) and generation.message:
-                    generation.message = self._postprocess_message(generation.message)
-
-        return result
-
-    def _stream(
-        self,
-        messages: Sequence[BaseMessage],
-        stop: Optional[Sequence[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> Iterator[ChatGeneration]:
-        """Override _stream to apply message preprocessing and postprocessing."""
-        # Apply message preprocessing
-        processed_messages: Sequence[BaseMessage] = self._preprocess_messages(messages)
-
-        # Handle the edge case where all messages were filtered out
-        if len(processed_messages) == 0:
-            logger.error(
-                "All messages were filtered out during preprocessing. This indicates a serious issue with the conversation state."
-            )
-            # Return an empty streaming result without calling the underlying API
-            # This prevents API errors while making the issue visible through an empty response
-            empty_chunk = ChatGenerationChunk(
-                message=AIMessage(content="", id="empty-response")
-            )
-            yield empty_chunk
-            return
-
-        logger.trace(
-            f"Processed messages:\n{json.dumps([msg.model_dump() for msg in processed_messages], indent=2)}"
-        )
-
-        # Call the parent ChatDatabricks implementation
-        for chunk in super()._stream(processed_messages, stop, run_manager, **kwargs):
-            chunk: ChatGenerationChunk
-            # Apply message postprocessing to each chunk
-            if isinstance(chunk, ChatGeneration) and chunk.message:
-                chunk.message = self._postprocess_message(chunk.message)
-            yield chunk

dao_ai/guardrails.py

@@ -1,112 +0,0 @@
-from typing import Any, Literal, Optional, Type
-
-from langchain_core.language_models import LanguageModelLike
-from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
-from langchain_core.runnables import RunnableConfig
-from langchain_core.runnables.base import RunnableLike
-from langgraph.graph import END, START, MessagesState, StateGraph
-from langgraph.graph.state import CompiledStateGraph
-from langgraph.managed import RemainingSteps
-from loguru import logger
-from openevals.llm import create_llm_as_judge
-
-from dao_ai.config import GuardrailModel
-from dao_ai.messages import last_ai_message, last_human_message
-from dao_ai.state import SharedState
-
-
-class MessagesWithSteps(MessagesState):
-    guardrails_remaining_steps: RemainingSteps
-
-
-def end_or_reflect(state: MessagesWithSteps) -> Literal[END, "graph"]:
-    if state["guardrails_remaining_steps"] < 2:
-        return END
-    if len(state["messages"]) == 0:
-        return END
-    last_message = state["messages"][-1]
-    if isinstance(last_message, HumanMessage):
-        return "graph"
-    else:
-        return END
-
-
-def create_reflection_graph(
-    graph: CompiledStateGraph,
-    reflection: CompiledStateGraph,
-    state_schema: Optional[Type[Any]] = None,
-    config_schema: Optional[Type[Any]] = None,
-) -> StateGraph:
-    logger.debug("Creating reflection graph")
-    _state_schema = state_schema or graph.builder.schema
-
-    if "guardrails_remaining_steps" in _state_schema.__annotations__:
-        raise ValueError(
-            "Has key 'guardrails_remaining_steps' in state_schema, this shadows a built in key"
-        )
-
-    if "messages" not in _state_schema.__annotations__:
-        raise ValueError("Missing required key 'messages' in state_schema")
-
-    class StateSchema(_state_schema):
-        guardrails_remaining_steps: RemainingSteps
-
-    rgraph = StateGraph(StateSchema, config_schema=config_schema)
-    rgraph.add_node("graph", graph)
-    rgraph.add_node("reflection", reflection)
-    rgraph.add_edge(START, "graph")
-    rgraph.add_edge("graph", "reflection")
-    rgraph.add_conditional_edges("reflection", end_or_reflect)
-    return rgraph
-
-
-def with_guardrails(
-    graph: CompiledStateGraph, guardrail: CompiledStateGraph
-) -> CompiledStateGraph:
-    logger.debug("Creating graph with guardrails")
-    return create_reflection_graph(
-        graph, guardrail, state_schema=SharedState, config_schema=RunnableConfig
-    ).compile()
-
-
-def judge_node(guardrails: GuardrailModel) -> RunnableLike:
-    def judge(state: SharedState, config: RunnableConfig) -> dict[str, BaseMessage]:
-        llm: LanguageModelLike = guardrails.model.as_chat_model()
-
-        evaluator = create_llm_as_judge(
-            prompt=guardrails.prompt,
-            judge=llm,
-        )
-
-        ai_message: AIMessage = last_ai_message(state["messages"])
-        human_message: HumanMessage = last_human_message(state["messages"])
-
-        logger.debug(f"Evaluating response: {ai_message.content}")
-        eval_result = evaluator(
-            inputs=human_message.content, outputs=ai_message.content
-        )
-
-        if eval_result["score"]:
-            logger.debug("Response approved by judge")
-            logger.debug(f"Judge's comment: {eval_result['comment']}")
-            return
-        else:
-            # Otherwise, return the judge's critique as a new user message
-            logger.warning("Judge requested improvements")
-            comment: str = eval_result["comment"]
-            logger.warning(f"Judge's critique: {comment}")
-            content: str = "\n".join([human_message.content, comment])
-            return {"messages": [HumanMessage(content=content)]}
-
-    return judge
-
-
-def reflection_guardrail(guardrails: GuardrailModel) -> CompiledStateGraph:
-    judge: CompiledStateGraph = (
-        StateGraph(SharedState, config_schema=RunnableConfig)
-        .add_node("judge", judge_node(guardrails=guardrails))
-        .add_edge(START, "judge")
-        .add_edge("judge", END)
-        .compile()
-    )
-    return judge

dao_ai/tools/genie/__init__.py

@@ -1,236 +0,0 @@
-"""
-Genie tools for natural language queries to databases.
-
-This package provides tools for interacting with Databricks Genie to translate
-natural language questions into SQL queries.
-
-Main exports:
-- create_genie_tool: Factory function to create a Genie tool with optional caching
-
-Cache implementations are available in the genie cache package:
-- dao_ai.genie.cache.lru: LRU (Least Recently Used) cache
-- dao_ai.genie.cache.semantic: Semantic similarity cache using pg_vector
-"""
-
-import json
-import os
-from textwrap import dedent
-from typing import Annotated, Any, Callable
-
-import pandas as pd
-from databricks_ai_bridge.genie import Genie, GenieResponse
-from langchain.tools import tool
-from langchain_core.messages import ToolMessage
-from langchain_core.tools import InjectedToolCallId
-from langgraph.prebuilt import InjectedState
-from langgraph.types import Command
-from loguru import logger
-from pydantic import BaseModel
-
-from dao_ai.config import (
-    AnyVariable,
-    CompositeVariableModel,
-    GenieLRUCacheParametersModel,
-    GenieRoomModel,
-    GenieSemanticCacheParametersModel,
-    value_of,
-)
-from dao_ai.genie import GenieService
-from dao_ai.genie.cache import (
-    CacheResult,
-    GenieServiceBase,
-    LRUCacheService,
-    SemanticCacheService,
-    SQLCacheEntry,
-)
-
-
-class GenieToolInput(BaseModel):
-    """Input schema for Genie tool - only includes user-facing parameters."""
-
-    question: str
-
-
-def _response_to_json(response: GenieResponse) -> str:
-    """Convert GenieResponse to JSON string, handling DataFrame results."""
-    # Convert result to string if it's a DataFrame
-    result: str | pd.DataFrame = response.result
-    if isinstance(result, pd.DataFrame):
-        result = result.to_markdown()
-
-    data: dict[str, Any] = {
-        "result": result,
-        "query": response.query,
-        "description": response.description,
-        "conversation_id": response.conversation_id,
-    }
-    return json.dumps(data)
-
-
-def create_genie_tool(
-    genie_room: GenieRoomModel | dict[str, Any],
-    name: str | None = None,
-    description: str | None = None,
-    persist_conversation: bool = True,
-    truncate_results: bool = False,
-    lru_cache_parameters: GenieLRUCacheParametersModel | dict[str, Any] | None = None,
-    semantic_cache_parameters: GenieSemanticCacheParametersModel
-    | dict[str, Any]
-    | None = None,
-) -> Callable[..., Command]:
-    """
-    Create a tool for interacting with Databricks Genie for natural language queries to databases.
-
-    This factory function generates a tool that leverages Databricks Genie to translate natural
-    language questions into SQL queries and execute them against retail databases. This enables
-    answering questions about inventory, sales, and other structured retail data.
-
-    Args:
-        genie_room: GenieRoomModel or dict containing Genie configuration
-        name: Optional custom name for the tool. If None, uses default "genie_tool"
-        description: Optional custom description for the tool. If None, uses default description
-        persist_conversation: Whether to persist conversation IDs across tool calls for
-            multi-turn conversations within the same Genie space
-        truncate_results: Whether to truncate large query results to fit token limits
-        lru_cache_parameters: Optional LRU cache configuration for SQL query caching
-        semantic_cache_parameters: Optional semantic cache configuration using pg_vector
-            for similarity-based query matching
-
-    Returns:
-        A LangGraph tool that processes natural language queries through Genie
-    """
-    logger.debug("create_genie_tool")
-    logger.debug(f"genie_room type: {type(genie_room)}")
-    logger.debug(f"genie_room: {genie_room}")
-    logger.debug(f"persist_conversation: {persist_conversation}")
-    logger.debug(f"truncate_results: {truncate_results}")
-    logger.debug(f"name: {name}")
-    logger.debug(f"description: {description}")
-    logger.debug(f"genie_room: {genie_room}")
-    logger.debug(f"persist_conversation: {persist_conversation}")
-    logger.debug(f"truncate_results: {truncate_results}")
-    logger.debug(f"lru_cache_parameters: {lru_cache_parameters}")
-    logger.debug(f"semantic_cache_parameters: {semantic_cache_parameters}")
-
-    if isinstance(genie_room, dict):
-        genie_room = GenieRoomModel(**genie_room)
-
-    if isinstance(lru_cache_parameters, dict):
-        lru_cache_parameters = GenieLRUCacheParametersModel(**lru_cache_parameters)
-
-    if isinstance(semantic_cache_parameters, dict):
-        semantic_cache_parameters = GenieSemanticCacheParametersModel(
-            **semantic_cache_parameters
-        )
-
-    space_id: AnyVariable = genie_room.space_id or os.environ.get(
-        "DATABRICKS_GENIE_SPACE_ID"
-    )
-    if isinstance(space_id, dict):
-        space_id = CompositeVariableModel(**space_id)
-    space_id = value_of(space_id)
-
-    default_description: str = dedent("""
-    This tool lets you have a conversation and chat with tabular data about <topic>. You should ask
-    questions about the data and the tool will try to answer them.
-    Please ask simple clear questions that can be answer by sql queries. If you need to do statistics or other forms of testing defer to using another tool.
-    Try to ask for aggregations on the data and ask very simple questions.
-    Prefer to call this tool multiple times rather than asking a complex question.
-    """)
-
-    tool_description: str = (
-        description if description is not None else default_description
-    )
-    tool_name: str = name if name is not None else "genie_tool"
-
-    function_docs = """
-
-Args:
-question (str): The question to ask to ask Genie about your data. Ask simple, clear questions about your tabular data. For complex analysis, ask multiple simple questions rather than one complex question.
-
-Returns:
-GenieResponse: A response object containing the conversation ID and result from Genie."""
-    tool_description = tool_description + function_docs
-
-    genie: Genie = Genie(
-        space_id=space_id,
-        client=genie_room.workspace_client,
-        truncate_results=truncate_results,
-    )
-
-    genie_service: GenieServiceBase = GenieService(genie)
-
-    # Wrap with semantic cache first (checked second due to decorator pattern)
-    if semantic_cache_parameters is not None:
-        genie_service = SemanticCacheService(
-            impl=genie_service,
-            parameters=semantic_cache_parameters,
-            genie_space_id=space_id,
-        )
-
-    # Wrap with LRU cache last (checked first - fast O(1) exact match)
-    if lru_cache_parameters is not None:
-        genie_service = LRUCacheService(
-            impl=genie_service,
-            parameters=lru_cache_parameters,
-        )
-
-    @tool(
-        name_or_callable=tool_name,
-        description=tool_description,
-    )
-    def genie_tool(
-        question: Annotated[str, "The question to ask Genie about your data"],
-        state: Annotated[dict, InjectedState],
-        tool_call_id: Annotated[str, InjectedToolCallId],
-    ) -> Command:
-        """Process a natural language question through Databricks Genie."""
-        # Get existing conversation mapping and retrieve conversation ID for this space
-        conversation_ids: dict[str, str] = state.get("genie_conversation_ids", {})
-        existing_conversation_id: str | None = conversation_ids.get(space_id)
-        logger.debug(
-            f"Existing conversation ID for space {space_id}: {existing_conversation_id}"
-        )
-
-        response: GenieResponse = genie_service.ask_question(
-            question, conversation_id=existing_conversation_id
-        )
-
-        current_conversation_id: str = response.conversation_id
-        logger.debug(
-            f"Current conversation ID for space {space_id}: {current_conversation_id}"
-        )
-
-        # Update the conversation mapping with the new conversation ID for this space
-
-        update: dict[str, Any] = {
-            "messages": [
-                ToolMessage(_response_to_json(response), tool_call_id=tool_call_id)
-            ],
-        }
-
-        if persist_conversation:
-            updated_conversation_ids: dict[str, str] = conversation_ids.copy()
-            updated_conversation_ids[space_id] = current_conversation_id
-            update["genie_conversation_ids"] = updated_conversation_ids
-
-        return Command(update=update)
-
-    return genie_tool
-
-
-# Re-export cache types for convenience
-__all__ = [
-    # Main tool
-    "create_genie_tool",
-    # Input types
-    "GenieToolInput",
-    # Service base classes
-    "GenieService",
-    "GenieServiceBase",
-    # Cache types (from cache subpackage)
-    "CacheResult",
-    "LRUCacheService",
-    "SemanticCacheService",
-    "SQLCacheEntry",
-]