dao-ai 0.0.20__py3-none-any.whl → 0.0.22__py3-none-any.whl

dao_ai/tools/genie.py

@@ -1,20 +1,284 @@
+import bisect
+import json
+import logging
 import os
+import time
+from dataclasses import asdict, dataclass
+from datetime import datetime
 from textwrap import dedent
-from typing import Any, Callable, Optional
+from typing import Annotated, Any, Callable, Optional, Union
 
-from databricks_ai_bridge.genie import GenieResponse
-from databricks_langchain.genie import Genie
-from langchain_core.tools import StructuredTool
+import mlflow
+import pandas as pd
+from databricks.sdk import WorkspaceClient
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import InjectedToolCallId, tool
+from langgraph.prebuilt import InjectedState
+from langgraph.types import Command
+from loguru import logger
+from pydantic import BaseModel, Field
 
-from dao_ai.config import (
-    GenieRoomModel,
-)
+from dao_ai.config import AnyVariable, CompositeVariableModel, GenieRoomModel, value_of
+
+MAX_TOKENS_OF_DATA: int = 20000
+MAX_ITERATIONS: int = 50
+DEFAULT_POLLING_INTERVAL_SECS: int = 2
+
+
+def _count_tokens(text):
+    import tiktoken
+
+    encoding = tiktoken.encoding_for_model("gpt-4o")
+    return len(encoding.encode(text))
+
+
+@dataclass
+class GenieResponse:
+    conversation_id: str
+    result: Union[str, pd.DataFrame]
+    query: Optional[str] = ""
+    description: Optional[str] = ""
+
+    def to_json(self):
+        return json.dumps(asdict(self))
+
+
+class GenieToolInput(BaseModel):
+    """Input schema for the Genie tool."""
+
+    question: str = Field(
+        description="The question 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."
+    )
+
+
+def _truncate_result(dataframe: pd.DataFrame) -> str:
+    query_result = dataframe.to_markdown()
+    tokens_used = _count_tokens(query_result)
+
+    # If the full result fits, return it
+    if tokens_used <= MAX_TOKENS_OF_DATA:
+        return query_result.strip()
+
+    def is_too_big(n):
+        return _count_tokens(dataframe.iloc[:n].to_markdown()) > MAX_TOKENS_OF_DATA
+
+    # Use bisect_left to find the cutoff point of rows within the max token data limit in a O(log n) complexity
+    # Passing True, as this is the target value we are looking for when _is_too_big returns
+    cutoff = bisect.bisect_left(range(len(dataframe) + 1), True, key=is_too_big)
+
+    # Slice to the found limit
+    truncated_df = dataframe.iloc[:cutoff]
+
+    # Edge case: Cannot return any rows because of tokens so return an empty string
+    if len(truncated_df) == 0:
+        return ""
+
+    truncated_result = truncated_df.to_markdown()
+
+    # Double-check edge case if we overshot by one
+    if _count_tokens(truncated_result) > MAX_TOKENS_OF_DATA:
+        truncated_result = truncated_df.iloc[:-1].to_markdown()
+    return truncated_result
+
+
+@mlflow.trace(span_type="PARSER")
+def _parse_query_result(resp, truncate_results) -> Union[str, pd.DataFrame]:
+    output = resp["result"]
+    if not output:
+        return "EMPTY"
+
+    columns = resp["manifest"]["schema"]["columns"]
+    header = [str(col["name"]) for col in columns]
+    rows = []
+
+    for item in output["data_array"]:
+        row = []
+        for column, value in zip(columns, item):
+            type_name = column["type_name"]
+            if value is None:
+                row.append(None)
+                continue
+
+            if type_name in ["INT", "LONG", "SHORT", "BYTE"]:
+                row.append(int(value))
+            elif type_name in ["FLOAT", "DOUBLE", "DECIMAL"]:
+                row.append(float(value))
+            elif type_name == "BOOLEAN":
+                row.append(value.lower() == "true")
+            elif type_name == "DATE" or type_name == "TIMESTAMP":
+                row.append(datetime.strptime(value[:10], "%Y-%m-%d").date())
+            elif type_name == "BINARY":
+                row.append(bytes(value, "utf-8"))
+            else:
+                row.append(value)
+
+        rows.append(row)
+
+    dataframe = pd.DataFrame(rows, columns=header)
+
+    if truncate_results:
+        query_result = _truncate_result(dataframe)
+    else:
+        query_result = dataframe.to_markdown()
+
+    return query_result.strip()
+
+
+class Genie:
+    def __init__(
+        self,
+        space_id,
+        client: WorkspaceClient | None = None,
+        truncate_results: bool = False,
+        polling_interval: int = DEFAULT_POLLING_INTERVAL_SECS,
+    ):
+        self.space_id = space_id
+        workspace_client = client or WorkspaceClient()
+        self.genie = workspace_client.genie
+        self.description = self.genie.get_space(space_id).description
+        self.headers = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+        self.truncate_results = truncate_results
+        if polling_interval < 1 or polling_interval > 30:
+            raise ValueError("poll_interval must be between 1 and 30 seconds")
+        self.poll_interval = polling_interval
+
+    @mlflow.trace()
+    def start_conversation(self, content):
+        resp = self.genie._api.do(
+            "POST",
+            f"/api/2.0/genie/spaces/{self.space_id}/start-conversation",
+            body={"content": content},
+            headers=self.headers,
+        )
+        return resp
+
+    @mlflow.trace()
+    def create_message(self, conversation_id, content):
+        resp = self.genie._api.do(
+            "POST",
+            f"/api/2.0/genie/spaces/{self.space_id}/conversations/{conversation_id}/messages",
+            body={"content": content},
+            headers=self.headers,
+        )
+        return resp
+
+    @mlflow.trace()
+    def poll_for_result(self, conversation_id, message_id):
+        @mlflow.trace()
+        def poll_query_results(attachment_id, query_str, description):
+            iteration_count = 0
+            while iteration_count < MAX_ITERATIONS:
+                iteration_count += 1
+                resp = self.genie._api.do(
+                    "GET",
+                    f"/api/2.0/genie/spaces/{self.space_id}/conversations/{conversation_id}/messages/{message_id}/attachments/{attachment_id}/query-result",
+                    headers=self.headers,
+                )["statement_response"]
+                state = resp["status"]["state"]
+                if state == "SUCCEEDED":
+                    result = _parse_query_result(resp, self.truncate_results)
+                    return GenieResponse(
+                        conversation_id, result, query_str, description
+                    )
+                elif state in ["RUNNING", "PENDING"]:
+                    logging.debug("Waiting for query result...")
+                    time.sleep(self.poll_interval)
+                else:
+                    return GenieResponse(
+                        conversation_id,
+                        f"No query result: {resp['state']}",
+                        query_str,
+                        description,
+                    )
+            return GenieResponse(
+                conversation_id,
+                f"Genie query for result timed out after {MAX_ITERATIONS} iterations of {self.poll_interval} seconds",
+                query_str,
+                description,
+            )
+
+        @mlflow.trace()
+        def poll_result():
+            iteration_count = 0
+            while iteration_count < MAX_ITERATIONS:
+                iteration_count += 1
+                resp = self.genie._api.do(
+                    "GET",
+                    f"/api/2.0/genie/spaces/{self.space_id}/conversations/{conversation_id}/messages/{message_id}",
+                    headers=self.headers,
+                )
+                if resp["status"] == "COMPLETED":
+                    # Check if attachments key exists in response
+                    attachments = resp.get("attachments", [])
+                    if not attachments:
+                        # Handle case where response has no attachments
+                        return GenieResponse(
+                            conversation_id,
+                            result=f"Genie query completed but no attachments found. Response: {resp}",
+                        )
+
+                    attachment = next((r for r in attachments if "query" in r), None)
+                    if attachment:
+                        query_obj = attachment["query"]
+                        description = query_obj.get("description", "")
+                        query_str = query_obj.get("query", "")
+                        attachment_id = attachment["attachment_id"]
+                        return poll_query_results(attachment_id, query_str, description)
+                    if resp["status"] == "COMPLETED":
+                        text_content = next(
+                            (r for r in attachments if "text" in r), None
+                        )
+                        if text_content:
+                            return GenieResponse(
+                                conversation_id, result=text_content["text"]["content"]
+                            )
+                        return GenieResponse(
+                            conversation_id,
+                            result="Genie query completed but no text content found in attachments.",
+                        )
+                elif resp["status"] in {"CANCELLED", "QUERY_RESULT_EXPIRED"}:
+                    return GenieResponse(
+                        conversation_id, result=f"Genie query {resp['status'].lower()}."
+                    )
+                elif resp["status"] == "FAILED":
+                    return GenieResponse(
+                        conversation_id,
+                        result=f"Genie query failed with error: {resp.get('error', 'Unknown error')}",
+                    )
+                # includes EXECUTING_QUERY, Genie can retry after this status
+                else:
+                    logging.debug(f"Waiting...: {resp['status']}")
+                    time.sleep(self.poll_interval)
+            return GenieResponse(
+                conversation_id,
+                f"Genie query timed out after {MAX_ITERATIONS} iterations of {self.poll_interval} seconds",
+            )
+
+        return poll_result()
+
+    @mlflow.trace()
+    def ask_question(self, question: str, conversation_id: str | None = None):
+        logger.debug(
+            f"ask_question called with question: {question}, conversation_id: {conversation_id}"
+        )
+        if conversation_id:
+            resp = self.create_message(conversation_id, question)
+        else:
+            resp = self.start_conversation(question)
+        logger.debug(f"ask_question response: {resp}")
+        return self.poll_for_result(resp["conversation_id"], resp["message_id"])
 
 
 def create_genie_tool(
     genie_room: GenieRoomModel | dict[str, Any],
     name: Optional[str] = None,
     description: Optional[str] = None,
+    persist_conversation: bool = False,
+    truncate_results: bool = False,
+    poll_interval: int = DEFAULT_POLLING_INTERVAL_SECS,
 ) -> Callable[[str], GenieResponse]:
     """
     Create a tool for interacting with Databricks Genie for natural language queries to databases.
@@ -24,22 +288,33 @@ def create_genie_tool(
     answering questions about inventory, sales, and other structured retail data.
 
     Args:
-        space_id: Databricks workspace ID where Genie is configured. If None, tries to
-                get it from DATABRICKS_GENIE_SPACE_ID environment variable.
+        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
 
     Returns:
-        A callable tool function that processes natural language queries through Genie
+        A LangGraph tool that processes natural language queries through Genie
     """
 
     if isinstance(genie_room, dict):
         genie_room = GenieRoomModel(**genie_room)
 
-    space_id: str = genie_room.space_id or os.environ.get("DATABRICKS_GENIE_SPACE_ID")
-
-    genie: Genie = Genie(
-        space_id=space_id,
-        client=genie_room.workspace_client,
+    space_id: AnyVariable = genie_room.space_id or os.environ.get(
+        "DATABRICKS_GENIE_SPACE_ID"
+    )
+    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)
+
+    # genie: Genie = Genie(
+    #     space_id=space_id,
+    #     client=genie_room.workspace_client,
+    #     truncate_results=truncate_results,
+    #     polling_interval=poll_interval,
+    # )
 
     default_description: str = dedent("""
     This tool lets you have a conversation and chat with tabular data about <topic>. You should ask
@@ -49,29 +324,66 @@ def create_genie_tool(
     Prefer to call this tool multiple times rather than asking a complex question.
     """)
 
-    if description is None:
-        description = default_description
+    tool_description: str = (
+        description if description is not None else default_description
+    )
+    tool_name: str = name if name is not None else "genie_tool"
 
-    doc_signature: str = dedent("""
-    Args:
-        question (str): The question to ask to ask Genie
+    function_docs = """
 
-    Returns:
-        response (GenieResponse): An object containing the Genie response
-    """)
+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.
 
-    doc: str = description + "\n" + doc_signature
+Returns:
+GenieResponse: A response object containing the conversation ID and result from Genie."""
+    tool_description = tool_description + function_docs
 
-    async def genie_tool(question: str) -> GenieResponse:
-        # Use sync API for now since Genie doesn't support async yet
-        # Can be easily updated to await when Genie gets async support
-        response: GenieResponse = genie.ask_question(question)
-        return response
+    @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:
+        genie: Genie = Genie(
+            space_id=space_id,
+            client=genie_room.workspace_client,
+            truncate_results=truncate_results,
+            polling_interval=poll_interval,
+        )
 
-    name: str = name if name else genie_tool.__name__
+        """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}"
+        )
 
-    structured_tool: StructuredTool = StructuredTool.from_function(
-        coroutine=genie_tool, name=name, description=doc, parse_docstring=False
-    )
+        response: GenieResponse = genie.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(), 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
+
+        logger.debug(f"State update: {update}")
+
+        return Command(update=update)
 
-    return structured_tool
+    return genie_tool

dao_ai/utils.py

@@ -112,3 +112,7 @@ def load_function(function_name: str) -> Callable[..., Any]:
     except (ImportError, AttributeError, TypeError) as e:
         # Provide a detailed error message that includes the original exception
         raise ImportError(f"Failed to import {function_name}: {e}")
+
+
+def is_in_model_serving() -> bool:
+    return os.environ.get("IS_IN_DB_MODEL_SERVING_ENV", "false").lower() == "true"

{dao_ai-0.0.20.dist-info → dao_ai-0.0.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.20
+Version: 0.0.22
 Summary: DAO AI: A modular, multi-agent orchestration framework for complex AI workflows. Supports agent handoff, tool integration, and dynamic configuration via YAML.
 Project-URL: Homepage, https://github.com/natefleming/dao-ai
 Project-URL: Documentation, https://natefleming.github.io/dao-ai
@@ -24,9 +24,9 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
 Requires-Python: >=3.12
-Requires-Dist: databricks-agents>=1.6.0
+Requires-Dist: databricks-agents>=1.6.1
 Requires-Dist: databricks-langchain>=0.8.0
-Requires-Dist: databricks-sdk[openai]>=0.66.0
+Requires-Dist: databricks-sdk[openai]>=0.67.0
 Requires-Dist: duckduckgo-search>=8.0.2
 Requires-Dist: grandalf>=0.8
 Requires-Dist: langchain-mcp-adapters>=0.1.10

{dao_ai-0.0.20.dist-info → dao_ai-0.0.22.dist-info}/RECORD

@@ -3,38 +3,38 @@ dao_ai/agent_as_code.py,sha256=kPSeDz2-1jRaed1TMs4LA3VECoyqe9_Ed2beRLB9gXQ,472
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/chat_models.py,sha256=uhwwOTeLyHWqoTTgHrs4n5iSyTwe4EQcLKnh3jRxPWI,8626
 dao_ai/cli.py,sha256=Aez2TQW3Q8Ho1IaIkRggt0NevDxAAVPjXkePC5GPJF0,20429
-dao_ai/config.py,sha256=ZO5ei45gnhqg1BtD0R9aekJz4ClmiTw2GHhOk4Idil4,51958
-dao_ai/graph.py,sha256=gmD9mxODfXuvn9xWeBfewm1FiuVAWMLEdnZz7DNmSH0,7859
+dao_ai/config.py,sha256=Kh0oJwWr2dhqrOriOEgjMs2CGhOdMYh1hTGO7sAPdw8,55268
+dao_ai/graph.py,sha256=APYc2y3cig4P52X4sOHSFSZNK8j5EtEPJLFwWeJ3KQQ,7956
 dao_ai/guardrails.py,sha256=4TKArDONRy8RwHzOT1plZ1rhy3x9GF_aeGpPCRl6wYA,4016
 dao_ai/messages.py,sha256=xl_3-WcFqZKCFCiov8sZOPljTdM3gX3fCHhxq-xFg2U,7005
-dao_ai/models.py,sha256=Xb23U-lhDG8KyNRIijcJ4InluadlaGNy4rrYx7Cjgfg,26939
+dao_ai/models.py,sha256=8r8GIG3EGxtVyWsRNI56lVaBjiNrPkzh4HdwMZRq8iw,31689
 dao_ai/nodes.py,sha256=SSuFNTXOdFaKg_aX-yUkQO7fM9wvNGu14lPXKDapU1U,8461
 dao_ai/prompts.py,sha256=vpmIbWs_szXUgNNDs5Gh2LcxKZti5pHDKSfoClUcgX0,1289
-dao_ai/state.py,sha256=GwbMbd1TWZx1T5iQrEOX6_rpxOitlmyeJ8dMr2o_pag,1031
+dao_ai/state.py,sha256=_lF9krAYYjvFDMUwZzVKOn0ZnXKcOrbjWKdre0C5B54,1137
 dao_ai/types.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dao_ai/utils.py,sha256=dkZTXNN6q0xwkrvSWdNq8937W2xGuLCRWRb6hRQM6kA,4217
+dao_ai/utils.py,sha256=A7BR28Rv7tH9B_lLVKlgBnWtcPeQ-Bafm2f57oTrUlQ,4337
 dao_ai/vector_search.py,sha256=jlaFS_iizJ55wblgzZmswMM3UOL-qOp2BGJc0JqXYSg,2839
 dao_ai/hooks/__init__.py,sha256=LlHGIuiZt6vGW8K5AQo1XJEkBP5vDVtMhq0IdjcLrD4,417
 dao_ai/hooks/core.py,sha256=ZShHctUSoauhBgdf1cecy9-D7J6-sGn-pKjuRMumW5U,6663
 dao_ai/memory/__init__.py,sha256=1kHx_p9abKYFQ6EYD05nuc1GS5HXVEpufmjBGw_7Uho,260
 dao_ai/memory/base.py,sha256=99nfr2UZJ4jmfTL_KrqUlRSCoRxzkZyWyx5WqeUoMdQ,338
-dao_ai/memory/core.py,sha256=g7chjBgVgx3iKjR2hghl0QL1j3802uIM_e7mgszur9M,4151
-dao_ai/memory/postgres.py,sha256=pxxMjGotgqjrKhx0lVR3EAjSZTQgBpiPZOB0-cyjprc,12505
+dao_ai/memory/core.py,sha256=DnEjQO3S7hXr3CDDd7C2eE7fQUmcCS_8q9BXEgjPH3U,4271
+dao_ai/memory/postgres.py,sha256=vvI3osjx1EoU5GBA6SCUstTBKillcmLl12hVgDMjfJY,15346
 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=fZ8mGotfA3W3t5yUej2xGmGHSybjBFYr895mOctT418,28203
+dao_ai/providers/databricks.py,sha256=CFZ2RojcTjiJ1aGwNI3_0qCGf339w2o5h9CRDKNesLs,39825
 dao_ai/tools/__init__.py,sha256=ye6MHaJY7tUnJ8336YJiLxuZr55zDPNdOw6gm7j5jlc,1103
 dao_ai/tools/agent.py,sha256=WbQnyziiT12TLMrA7xK0VuOU029tdmUBXbUl-R1VZ0Q,1886
 dao_ai/tools/core.py,sha256=Kei33S8vrmvPOAyrFNekaWmV2jqZ-IPS1QDSvU7RZF0,1984
-dao_ai/tools/genie.py,sha256=GzV5lfDYKmzW_lSLxAsPaTwnzX6GxQOB1UcLaTDqpfY,2787
+dao_ai/tools/genie.py,sha256=8HSOCzSg6PlBzBYXMmNfUnl-LO03p3Ki3fxLPm_dhPg,15051
 dao_ai/tools/human_in_the_loop.py,sha256=yk35MO9eNETnYFH-sqlgR-G24TrEgXpJlnZUustsLkI,3681
 dao_ai/tools/mcp.py,sha256=auEt_dwv4J26fr5AgLmwmnAsI894-cyuvkvjItzAUxs,4419
 dao_ai/tools/python.py,sha256=XcQiTMshZyLUTVR5peB3vqsoUoAAy8gol9_pcrhddfI,1831
 dao_ai/tools/time.py,sha256=Y-23qdnNHzwjvnfkWvYsE7PoWS1hfeKy44tA7sCnNac,8759
 dao_ai/tools/unity_catalog.py,sha256=uX_h52BuBAr4c9UeqSMI7DNz3BPRLeai5tBVW4sJqRI,13113
 dao_ai/tools/vector_search.py,sha256=EDYQs51zIPaAP0ma1D81wJT77GQ-v-cjb2XrFVWfWdg,2621
-dao_ai-0.0.20.dist-info/METADATA,sha256=gWNRLhswz5sCe1vxbBQ6dGlgiObI9nI829Q5DQRqRRY,41380
-dao_ai-0.0.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dao_ai-0.0.20.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.20.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.20.dist-info/RECORD,,
+dao_ai-0.0.22.dist-info/METADATA,sha256=kqyr-YBFC_fs-PHknnvm4Ahhad8Pfac0gTb8vKydHMw,41380
+dao_ai-0.0.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dao_ai-0.0.22.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.22.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.22.dist-info/RECORD,,