dao-ai 0.0.33__py3-none-any.whl → 0.0.34__py3-none-any.whl

dao_ai/config.py

@@ -236,9 +236,21 @@ class Privilege(str, Enum):
 
 class PermissionModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    principals: list[str] = Field(default_factory=list)
+    principals: list[ServicePrincipalModel | str] = Field(default_factory=list)
     privileges: list[Privilege]
 
+    @model_validator(mode="after")
+    def resolve_principals(self) -> Self:
+        """Resolve ServicePrincipalModel objects to their client_id."""
+        resolved: list[str] = []
+        for principal in self.principals:
+            if isinstance(principal, ServicePrincipalModel):
+                resolved.append(value_of(principal.client_id))
+            else:
+                resolved.append(principal)
+        self.principals = resolved
+        return self
+
 
 class SchemaModel(BaseModel, HasFullName):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
@@ -1581,9 +1593,21 @@ class Entitlement(str, Enum):
 
 class AppPermissionModel(BaseModel):
     model_config = ConfigDict(use_enum_values=True, extra="forbid")
-    principals: list[str] = Field(default_factory=list)
+    principals: list[ServicePrincipalModel | str] = Field(default_factory=list)
     entitlements: list[Entitlement]
 
+    @model_validator(mode="after")
+    def resolve_principals(self) -> Self:
+        """Resolve ServicePrincipalModel objects to their client_id."""
+        resolved: list[str] = []
+        for principal in self.principals:
+            if isinstance(principal, ServicePrincipalModel):
+                resolved.append(value_of(principal.client_id))
+            else:
+                resolved.append(principal)
+        self.principals = resolved
+        return self
+
 
 class LogLevel(str, Enum):
     TRACE = "TRACE"

dao_ai/models.py

@@ -331,13 +331,23 @@ class LanggraphResponsesAgent(ResponsesAgent):
         context: Context = self._convert_request_to_context(request)
         custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
+        # Build the graph input state, including genie_conversation_ids if provided
+        graph_input: dict[str, Any] = {"messages": messages}
+        if request.custom_inputs and "genie_conversation_ids" in request.custom_inputs:
+            graph_input["genie_conversation_ids"] = request.custom_inputs[
+                "genie_conversation_ids"
+            ]
+            logger.debug(
+                f"Including genie_conversation_ids in graph input: {graph_input['genie_conversation_ids']}"
+            )
+
         # Use async ainvoke internally for parallel execution
         import asyncio
 
         async def _async_invoke():
             try:
                 return await self.graph.ainvoke(
-                    {"messages": messages}, context=context, config=custom_inputs
+                    graph_input, context=context, config=custom_inputs
                 )
             except Exception as e:
                 logger.error(f"Error in graph.ainvoke: {e}")
@@ -399,6 +409,16 @@ class LanggraphResponsesAgent(ResponsesAgent):
         context: Context = self._convert_request_to_context(request)
         custom_inputs: dict[str, Any] = {"configurable": context.model_dump()}
 
+        # Build the graph input state, including genie_conversation_ids if provided
+        graph_input: dict[str, Any] = {"messages": messages}
+        if request.custom_inputs and "genie_conversation_ids" in request.custom_inputs:
+            graph_input["genie_conversation_ids"] = request.custom_inputs[
+                "genie_conversation_ids"
+            ]
+            logger.debug(
+                f"Including genie_conversation_ids in graph input: {graph_input['genie_conversation_ids']}"
+            )
+
         # Use async astream internally for parallel execution
         import asyncio
 
@@ -408,7 +428,7 @@ class LanggraphResponsesAgent(ResponsesAgent):
 
             try:
                 async for nodes, stream_mode, messages_batch in self.graph.astream(
-                    {"messages": messages},
+                    graph_input,
                     context=context,
                     config=custom_inputs,
                     stream_mode=["messages", "custom"],

dao_ai/providers/databricks.py

@@ -1151,7 +1151,7 @@ class DatabricksProvider(ServiceProvider):
         If an explicit version or alias is specified in the prompt_model, uses that directly.
         Otherwise, tries to load prompts in this order:
         1. champion alias
-        2. latest version (max version number from search_prompt_versions)
+        2. latest alias
         3. default alias
         4. Register default_template if provided
 
@@ -1166,7 +1166,6 @@ class DatabricksProvider(ServiceProvider):
         """
 
         prompt_name: str = prompt_model.full_name
-        mlflow_client: MlflowClient = MlflowClient()
 
         # If explicit version or alias is specified, use it directly
         if prompt_model.version or prompt_model.alias:
@@ -1197,19 +1196,13 @@ class DatabricksProvider(ServiceProvider):
         except Exception as e:
             logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
 
-        # 2. Try to get latest version by finding the max version number
+        # 2. Try latest alias
         try:
-            versions = mlflow_client.search_prompt_versions(
-                prompt_name, max_results=100
-            )
-            if versions:
-                latest = max(versions, key=lambda v: int(v.version))
-                logger.info(
-                    f"Loaded prompt '{prompt_name}' version {latest.version} (latest by max version)"
-                )
-                return latest
+            prompt_version = load_prompt(f"prompts:/{prompt_name}@latest")
+            logger.info(f"Loaded prompt '{prompt_name}' from latest alias")
+            return prompt_version
         except Exception as e:
-            logger.debug(f"Failed to find latest version for '{prompt_name}': {e}")
+            logger.debug(f"Latest alias not found for '{prompt_name}': {e}")
 
         # 3. Try default alias
         try:
@@ -1225,7 +1218,7 @@ class DatabricksProvider(ServiceProvider):
                 f"No existing prompt found for '{prompt_name}', "
                 "attempting to register default_template"
             )
-            return self._sync_default_template_to_registry(
+            return self._register_default_template(
                 prompt_name, prompt_model.default_template, prompt_model.description
             )
 
@@ -1235,49 +1228,17 @@ class DatabricksProvider(ServiceProvider):
             "and no default_template provided"
         )
 
-    def _sync_default_template_to_registry(
+    def _register_default_template(
         self, prompt_name: str, default_template: str, description: str | None = None
     ) -> PromptVersion:
-        """Get the best available prompt version, or register default_template if possible.
-
-        Tries to load prompts in order: champion → latest (max version) → default.
-        If none found and we have write permissions, registers the default_template.
-        If registration fails (e.g., in Model Serving), logs the error and raises.
-        """
-        mlflow_client: MlflowClient = MlflowClient()
-
-        # Try to find an existing prompt version in priority order
-        # 1. Try champion alias
-        try:
-            champion = mlflow.genai.load_prompt(f"prompts:/{prompt_name}@champion")
-            logger.info(f"Loaded prompt '{prompt_name}' from champion alias")
-            return champion
-        except Exception as e:
-            logger.debug(f"Champion alias not found for '{prompt_name}': {e}")
+        """Register default_template as a new prompt version.
 
-        # 2. Try to get the latest version by finding the max version number
-        try:
-            versions = mlflow_client.search_prompt_versions(
-                prompt_name, max_results=100
-            )
-            if versions:
-                latest = max(versions, key=lambda v: int(v.version))
-                logger.info(
-                    f"Loaded prompt '{prompt_name}' version {latest.version} (latest by max version)"
-                )
-                return latest
-        except Exception as e:
-            logger.debug(f"Failed to search versions for '{prompt_name}': {e}")
-
-        # 3. Try default alias
-        try:
-            default = mlflow.genai.load_prompt(f"prompts:/{prompt_name}@default")
-            logger.info(f"Loaded prompt '{prompt_name}' from default alias")
-            return default
-        except Exception as e:
-            logger.debug(f"Default alias not found for '{prompt_name}': {e}")
+        Called when no existing prompt version is found (champion, latest, default all failed).
+        Registers the template and sets both 'default' and 'champion' aliases.
 
-        # No existing prompt found - try to register if we have a template
+        If registration fails (e.g., in Model Serving with restricted permissions),
+        logs the error and raises.
+        """
         logger.info(
             f"No existing prompt found for '{prompt_name}', attempting to register default_template"
         )

dao_ai/tools/genie.py

@@ -5,8 +5,9 @@ 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, tool
+from langchain_core.tools import InjectedToolCallId
 from langgraph.prebuilt import InjectedState
 from langgraph.types import Command
 from loguru import logger
@@ -43,7 +44,7 @@ def create_genie_tool(
     genie_room: GenieRoomModel | dict[str, Any],
     name: str | None = None,
     description: str | None = None,
-    persist_conversation: bool = False,
+    persist_conversation: bool = True,
     truncate_results: bool = False,
 ) -> Callable[..., Command]:
     """
@@ -64,6 +65,16 @@ def create_genie_tool(
     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}")
 
     if isinstance(genie_room, dict):
         genie_room = GenieRoomModel(**genie_room)
@@ -106,14 +117,13 @@ GenieResponse: A response object containing the conversation ID and result from
         state: Annotated[dict, InjectedState],
         tool_call_id: Annotated[str, InjectedToolCallId],
     ) -> Command:
-        """Process a natural language question through Databricks Genie."""
-        # Create Genie instance using databricks_langchain implementation
         genie: Genie = Genie(
             space_id=space_id,
             client=genie_room.workspace_client,
             truncate_results=truncate_results,
         )
 
+        """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)
@@ -131,6 +141,7 @@ GenieResponse: A response object containing the conversation ID and result from
         )
 
         # 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)

{dao_ai-0.0.33.dist-info → dao_ai-0.0.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dao-ai
-Version: 0.0.33
+Version: 0.0.34
 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

{dao_ai-0.0.33.dist-info → dao_ai-0.0.34.dist-info}/RECORD

@@ -3,11 +3,11 @@ dao_ai/agent_as_code.py,sha256=sviZQV7ZPxE5zkZ9jAbfegI681nra5i8yYxw05e3X7U,552
 dao_ai/catalog.py,sha256=sPZpHTD3lPx4EZUtIWeQV7VQM89WJ6YH__wluk1v2lE,4947
 dao_ai/chat_models.py,sha256=uhwwOTeLyHWqoTTgHrs4n5iSyTwe4EQcLKnh3jRxPWI,8626
 dao_ai/cli.py,sha256=gq-nsapWxDA1M6Jua3vajBvIwf0Oa6YLcB58lEtMKUo,22503
-dao_ai/config.py,sha256=Uj0FgOhjnYp0qEmY44mCnp3Ijafg-381FNXt8R_QuWw,78513
+dao_ai/config.py,sha256=Jzb0ePrt2TM2WuXI_LtmTafbseKBlJ8J8J2ExyBowbM,79491
 dao_ai/graph.py,sha256=9kjJx0oFZKq5J9-Kpri4-0VCJILHYdYyhqQnj0_noxQ,8913
 dao_ai/guardrails.py,sha256=4TKArDONRy8RwHzOT1plZ1rhy3x9GF_aeGpPCRl6wYA,4016
 dao_ai/messages.py,sha256=xl_3-WcFqZKCFCiov8sZOPljTdM3gX3fCHhxq-xFg2U,7005
-dao_ai/models.py,sha256=8r8GIG3EGxtVyWsRNI56lVaBjiNrPkzh4HdwMZRq8iw,31689
+dao_ai/models.py,sha256=hvEZO2N0EC2sQoMgjJ9mbKmDWcdxnnAb2NqzpXh4Wgk,32691
 dao_ai/nodes.py,sha256=iQ_5vL6mt1UcRnhwgz-l1D8Ww4CMQrSMVnP_Lu7fFjU,8781
 dao_ai/prompts.py,sha256=iA2Iaky7yzjwWT5cxg0cUIgwo1z1UVQua__8WPnvV6g,1633
 dao_ai/state.py,sha256=_lF9krAYYjvFDMUwZzVKOn0ZnXKcOrbjWKdre0C5B54,1137
@@ -22,11 +22,11 @@ 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=rPBMdGcJvdGBRK9FZeBxkLfcTpXyxU1cs14YllyZKbY,67857
+dao_ai/providers/databricks.py,sha256=WEigNPGRTlIPVjwp97My8o1zOHn5ftuMsMrpqrBeaLg,66012
 dao_ai/tools/__init__.py,sha256=G5-5Yi6zpQOH53b5IzLdtsC6g0Ep6leI5GxgxOmgw7Q,1203
 dao_ai/tools/agent.py,sha256=WbQnyziiT12TLMrA7xK0VuOU029tdmUBXbUl-R1VZ0Q,1886
 dao_ai/tools/core.py,sha256=kN77fWOzVY7qOs4NiW72cUxCsSTC0DnPp73s6VJEZOQ,1991
-dao_ai/tools/genie.py,sha256=BPM_1Sk5bf7QSCFPPboWWkZKYwBwDwbGhMVp5-QDd10,5956
+dao_ai/tools/genie.py,sha256=hWDLLGUNz1wgwOb69pXnMiLJnMbG_1YmMdfVKt1Qe8o,6426
 dao_ai/tools/human_in_the_loop.py,sha256=yk35MO9eNETnYFH-sqlgR-G24TrEgXpJlnZUustsLkI,3681
 dao_ai/tools/mcp.py,sha256=5aQoRtx2z4xm6zgRslc78rSfEQe-mfhqov2NsiybYfc,8416
 dao_ai/tools/python.py,sha256=XcQiTMshZyLUTVR5peB3vqsoUoAAy8gol9_pcrhddfI,1831
@@ -34,8 +34,8 @@ dao_ai/tools/slack.py,sha256=SCvyVcD9Pv_XXPXePE_fSU1Pd8VLTEkKDLvoGTZWy2Y,4775
 dao_ai/tools/time.py,sha256=Y-23qdnNHzwjvnfkWvYsE7PoWS1hfeKy44tA7sCnNac,8759
 dao_ai/tools/unity_catalog.py,sha256=K9t8M4spsbxbecWmV5yEZy16s_AG7AfaoxT-7IDW43I,14438
 dao_ai/tools/vector_search.py,sha256=3cdiUaFpox25GSRNec7FKceY3DuLp7dLVH8FRA0BgeY,12624
-dao_ai-0.0.33.dist-info/METADATA,sha256=aa4BvkiG1dEvLorpgADosf1LCKRVBg-n8LtReVYJNxc,42761
-dao_ai-0.0.33.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-dao_ai-0.0.33.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
-dao_ai-0.0.33.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
-dao_ai-0.0.33.dist-info/RECORD,,
+dao_ai-0.0.34.dist-info/METADATA,sha256=vq51NEV-pg7WTOD5z56jyOrC5_6Q-nUIL51RI5lL-Hg,42761
+dao_ai-0.0.34.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dao_ai-0.0.34.dist-info/entry_points.txt,sha256=Xa-UFyc6gWGwMqMJOt06ZOog2vAfygV_DSwg1AiP46g,43
+dao_ai-0.0.34.dist-info/licenses/LICENSE,sha256=YZt3W32LtPYruuvHE9lGk2bw6ZPMMJD8yLrjgHybyz4,1069
+dao_ai-0.0.34.dist-info/RECORD,,