dao-ai 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

dao_ai/genie/cache/lru.py

@@ -124,9 +124,7 @@ class LRUCacheService(GenieServiceBase):
         if self._cache:
             oldest_key: str = next(iter(self._cache))
             del self._cache[oldest_key]
-            logger.trace(
-                "Evicted cache entry", layer=self.name, key_prefix=oldest_key[:50]
-            )
+            logger.trace("Evicted cache entry", layer=self.name, key=oldest_key[:50])
 
     def _get(self, key: str) -> SQLCacheEntry | None:
         """Get from cache, returning None if not found or expired."""
@@ -137,7 +135,7 @@ class LRUCacheService(GenieServiceBase):
 
         if self._is_expired(entry):
             del self._cache[key]
-            logger.trace("Expired cache entry", layer=self.name, key_prefix=key[:50])
+            logger.trace("Expired cache entry", layer=self.name, key=key[:50])
             return None
 
         self._cache.move_to_end(key)
@@ -157,11 +155,11 @@ class LRUCacheService(GenieServiceBase):
             conversation_id=response.conversation_id,
             created_at=datetime.now(),
         )
-        logger.info(
+        logger.debug(
             "Stored cache entry",
             layer=self.name,
-            key_prefix=key[:50],
-            sql_prefix=response.query[:50] if response.query else None,
+            key=key[:50],
+            sql=response.query[:50] if response.query else None,
             cache_size=len(self._cache),
             capacity=self.capacity,
         )
@@ -180,7 +178,7 @@ class LRUCacheService(GenieServiceBase):
         w: WorkspaceClient = self.warehouse.workspace_client
         warehouse_id: str = str(self.warehouse.warehouse_id)
 
-        logger.trace("Executing cached SQL", layer=self.name, sql_prefix=sql[:100])
+        logger.trace("Executing cached SQL", layer=self.name, sql=sql[:100])
 
         statement_response: StatementResponse = w.statement_execution.execute_statement(
             statement=sql,
@@ -258,13 +256,17 @@ class LRUCacheService(GenieServiceBase):
             cached: SQLCacheEntry | None = self._get(key)
 
         if cached is not None:
+            cache_age_seconds = (datetime.now() - cached.created_at).total_seconds()
             logger.info(
                 "Cache HIT",
                 layer=self.name,
-                question_prefix=question[:50],
+                question=question[:80],
                 conversation_id=conversation_id,
+                cached_sql=cached.query[:80] if cached.query else None,
+                cache_age_seconds=round(cache_age_seconds, 1),
                 cache_size=self.size,
                 capacity=self.capacity,
+                ttl_seconds=self.parameters.time_to_live_seconds,
             )
 
             # Re-execute the cached SQL to get fresh data
@@ -286,17 +288,19 @@ class LRUCacheService(GenieServiceBase):
         logger.info(
             "Cache MISS",
             layer=self.name,
-            question_prefix=question[:50],
+            question=question[:80],
             conversation_id=conversation_id,
             cache_size=self.size,
             capacity=self.capacity,
+            ttl_seconds=self.parameters.time_to_live_seconds,
             delegating_to=type(self.impl).__name__,
         )
 
         result: CacheResult = self.impl.ask_question(question, conversation_id)
         with self._lock:
             self._put(key, result.response)
-        return CacheResult(response=result.response, cache_hit=False, served_by=None)
+        # Propagate the inner cache's result - if it was a hit there, preserve that info
+        return result
 
     @property
     def space_id(self) -> str:

dao_ai/genie/cache/semantic.py

@@ -497,6 +497,7 @@ class SemanticCacheService(GenieServiceBase):
         conversation_context: str,
         question_embedding: list[float],
         context_embedding: list[float],
+        conversation_id: str | None = None,
     ) -> tuple[SQLCacheEntry, float] | None:
         """
         Find a semantically similar cached entry using dual embedding matching.
@@ -509,6 +510,7 @@ class SemanticCacheService(GenieServiceBase):
             conversation_context: The conversation context string
             question_embedding: The embedding vector of just the question
             context_embedding: The embedding vector of the conversation context
+            conversation_id: Optional conversation ID (for logging)
 
         Returns:
             Tuple of (SQLCacheEntry, combined_similarity_score) if found, None otherwise
@@ -576,8 +578,9 @@ class SemanticCacheService(GenieServiceBase):
                     logger.info(
                         "Cache MISS (no entries)",
                         layer=self.name,
-                        question_prefix=question[:50],
+                        question=question[:50],
                         space=self.space_id,
+                        delegating_to=type(self.impl).__name__,
                     )
                     return None
 
@@ -602,8 +605,8 @@ class SemanticCacheService(GenieServiceBase):
                     context_sim=f"{context_similarity:.4f}",
                     combined_sim=f"{combined_similarity:.4f}",
                     is_valid=is_valid,
-                    cached_question_prefix=cached_question[:50],
-                    cached_context_prefix=cached_context[:80],
+                    cached_question=cached_question[:50],
+                    cached_context=cached_context[:80],
                 )
 
                 # Check BOTH similarity thresholds (dual embedding precision check)
@@ -613,6 +616,7 @@ class SemanticCacheService(GenieServiceBase):
                         layer=self.name,
                         question_sim=f"{question_similarity:.4f}",
                         threshold=self.parameters.similarity_threshold,
+                        delegating_to=type(self.impl).__name__,
                     )
                     return None
 
@@ -622,6 +626,7 @@ class SemanticCacheService(GenieServiceBase):
                         layer=self.name,
                         context_sim=f"{context_similarity:.4f}",
                         threshold=self.parameters.context_similarity_threshold,
+                        delegating_to=type(self.impl).__name__,
                     )
                     return None
 
@@ -635,17 +640,32 @@ class SemanticCacheService(GenieServiceBase):
                         layer=self.name,
                         combined_sim=f"{combined_similarity:.4f}",
                         ttl_seconds=ttl_seconds,
-                        cached_question_prefix=cached_question[:50],
+                        cached_question=cached_question[:50],
+                        delegating_to=type(self.impl).__name__,
                     )
                     return None
 
+                from datetime import datetime as dt
+
+                cache_age_seconds = (
+                    (dt.now(created_at.tzinfo) - created_at).total_seconds()
+                    if created_at
+                    else None
+                )
                 logger.info(
                     "Cache HIT",
                     layer=self.name,
-                    question_sim=f"{question_similarity:.4f}",
-                    context_sim=f"{context_similarity:.4f}",
-                    combined_sim=f"{combined_similarity:.4f}",
-                    cached_question_prefix=cached_question[:50],
+                    question=question[:80],
+                    conversation_id=conversation_id,
+                    matched_question=cached_question[:80],
+                    cache_age_seconds=round(cache_age_seconds, 1)
+                    if cache_age_seconds
+                    else None,
+                    question_similarity=f"{question_similarity:.4f}",
+                    context_similarity=f"{context_similarity:.4f}",
+                    combined_similarity=f"{combined_similarity:.4f}",
+                    cached_sql=sql_query[:80] if sql_query else None,
+                    ttl_seconds=self.parameters.time_to_live_seconds,
                 )
 
                 entry = SQLCacheEntry(
@@ -696,12 +716,12 @@ class SemanticCacheService(GenieServiceBase):
                         response.conversation_id,
                     ),
                 )
-                logger.info(
+                logger.debug(
                     "Stored cache entry",
                     layer=self.name,
-                    question_prefix=question[:50],
-                    context_prefix=conversation_context[:80],
-                    sql_prefix=response.query[:50] if response.query else None,
+                    question=question[:50],
+                    context=conversation_context[:80],
+                    sql=response.query[:50] if response.query else None,
                     space=self.space_id,
                     table=self.table_name,
                 )
@@ -796,7 +816,11 @@ class SemanticCacheService(GenieServiceBase):
 
         # Check cache using dual embedding similarity
         cache_result: tuple[SQLCacheEntry, float] | None = self._find_similar(
-            question, conversation_context, question_embedding, context_embedding
+            question,
+            conversation_context,
+            question_embedding,
+            context_embedding,
+            conversation_id,
         )
 
         if cache_result is not None:
@@ -805,7 +829,8 @@ class SemanticCacheService(GenieServiceBase):
                 "Semantic cache hit",
                 layer=self.name,
                 combined_similarity=f"{combined_similarity:.3f}",
-                question_prefix=question[:50],
+                question=question[:50],
+                conversation_id=conversation_id,
             )
 
             # Re-execute the cached SQL to get fresh data
@@ -825,16 +850,25 @@ class SemanticCacheService(GenieServiceBase):
             return CacheResult(response=response, cache_hit=True, served_by=self.name)
 
         # Cache miss - delegate to wrapped service
-        logger.trace("Cache miss", layer=self.name, question_prefix=question[:50])
+        logger.info(
+            "Cache MISS",
+            layer=self.name,
+            question=question[:80],
+            conversation_id=conversation_id,
+            space_id=self.space_id,
+            similarity_threshold=self.similarity_threshold,
+            delegating_to=type(self.impl).__name__,
+        )
 
         result: CacheResult = self.impl.ask_question(question, conversation_id)
 
         # Store in cache if we got a SQL query
         if result.response.query:
-            logger.info(
+            logger.debug(
                 "Storing new cache entry",
                 layer=self.name,
-                question_prefix=question[:50],
+                question=question[:50],
+                conversation_id=conversation_id,
                 space=self.space_id,
             )
             self._store_entry(
@@ -848,7 +882,7 @@ class SemanticCacheService(GenieServiceBase):
             logger.warning(
                 "Not caching: response has no SQL query",
                 layer=self.name,
-                question_prefix=question[:50],
+                question=question[:50],
             )
 
         return CacheResult(response=result.response, cache_hit=False, served_by=None)

dao_ai/memory/postgres.py

@@ -3,6 +3,7 @@ import atexit
 import threading
 from typing import Any, Optional
 
+from databricks_ai_bridge.lakebase import AsyncLakebasePool, LakebasePool
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.postgres import ShallowPostgresSaver
 from langgraph.checkpoint.postgres.aio import AsyncShallowPostgresSaver
@@ -86,13 +87,22 @@ async def _create_async_pool(
 
 
 class AsyncPostgresPoolManager:
+    """
+    Asynchronous PostgreSQL connection pool manager that shares pools
+    based on database configuration.
+
+    For Lakebase connections (when instance_name is provided), uses AsyncLakebasePool
+    from databricks_ai_bridge which handles automatic token rotation and host resolution.
+    For standard PostgreSQL connections, uses psycopg_pool.AsyncConnectionPool.
+    """
+
     _pools: dict[str, AsyncConnectionPool] = {}
+    _lakebase_pools: dict[str, AsyncLakebasePool] = {}
     _lock: asyncio.Lock = asyncio.Lock()
 
     @classmethod
     async def get_pool(cls, database: DatabaseModel) -> AsyncConnectionPool:
         connection_key: str = database.name
-        connection_params: dict[str, Any] = database.connection_params
 
         async with cls._lock:
             if connection_key in cls._pools:
@@ -103,19 +113,43 @@ class AsyncPostgresPoolManager:
 
             logger.debug("Creating new async PostgreSQL pool", database=database.name)
 
-            kwargs: dict[str, Any] = {
-                "row_factory": dict_row,
-                "autocommit": True,
-            } | database.connection_kwargs or {}
-
-            # Create connection pool
-            pool: AsyncConnectionPool = await _create_async_pool(
-                connection_params=connection_params,
-                database_name=database.name,
-                max_pool_size=database.max_pool_size,
-                timeout_seconds=database.timeout_seconds,
-                kwargs=kwargs,
-            )
+            if database.is_lakebase:
+                # Use AsyncLakebasePool for Lakebase connections
+                # AsyncLakebasePool handles automatic token rotation and host resolution
+                lakebase_pool = AsyncLakebasePool(
+                    instance_name=database.instance_name,
+                    workspace_client=database.workspace_client,
+                    min_size=1,
+                    max_size=database.max_pool_size,
+                    timeout=float(database.timeout_seconds),
+                )
+                # Open the async pool
+                await lakebase_pool.open()
+                # Store the AsyncLakebasePool for proper cleanup
+                cls._lakebase_pools[connection_key] = lakebase_pool
+                # Get the underlying AsyncConnectionPool
+                pool = lakebase_pool.pool
+                logger.success(
+                    "Async Lakebase connection pool created",
+                    database=database.name,
+                    instance_name=database.instance_name,
+                    pool_size=database.max_pool_size,
+                )
+            else:
+                # Use standard async PostgreSQL pool for non-Lakebase connections
+                connection_params: dict[str, Any] = database.connection_params
+                kwargs: dict[str, Any] = {
+                    "row_factory": dict_row,
+                    "autocommit": True,
+                } | database.connection_kwargs or {}
+
+                pool = await _create_async_pool(
+                    connection_params=connection_params,
+                    database_name=database.name,
+                    max_pool_size=database.max_pool_size,
+                    timeout_seconds=database.timeout_seconds,
+                    kwargs=kwargs,
+                )
 
             cls._pools[connection_key] = pool
             return pool
@@ -125,7 +159,13 @@ class AsyncPostgresPoolManager:
         connection_key: str = database.name
 
         async with cls._lock:
-            if connection_key in cls._pools:
+            # Close AsyncLakebasePool if it exists (handles underlying pool cleanup)
+            if connection_key in cls._lakebase_pools:
+                lakebase_pool = cls._lakebase_pools.pop(connection_key)
+                await lakebase_pool.close()
+                cls._pools.pop(connection_key, None)
+                logger.debug("Async Lakebase pool closed", database=database.name)
+            elif connection_key in cls._pools:
                 pool = cls._pools.pop(connection_key)
                 await pool.close()
                 logger.debug("Async PostgreSQL pool closed", database=database.name)
@@ -133,9 +173,32 @@ class AsyncPostgresPoolManager:
     @classmethod
     async def close_all_pools(cls):
         async with cls._lock:
+            # Close all AsyncLakebasePool instances first
+            for connection_key, lakebase_pool in cls._lakebase_pools.items():
+                try:
+                    await asyncio.wait_for(lakebase_pool.close(), timeout=2.0)
+                    logger.debug("Async Lakebase pool closed", pool=connection_key)
+                except asyncio.TimeoutError:
+                    logger.warning(
+                        "Timeout closing async Lakebase pool, forcing closure",
+                        pool=connection_key,
+                    )
+                except asyncio.CancelledError:
+                    logger.warning(
+                        "Async Lakebase pool closure cancelled (shutdown in progress)",
+                        pool=connection_key,
+                    )
+                except Exception as e:
+                    logger.error(
+                        "Error closing async Lakebase pool",
+                        pool=connection_key,
+                        error=str(e),
+                    )
+            cls._lakebase_pools.clear()
+
+            # Close any remaining standard async PostgreSQL pools
             for connection_key, pool in cls._pools.items():
                 try:
-                    # Use a short timeout to avoid blocking on pool closure
                     await asyncio.wait_for(pool.close(), timeout=2.0)
                     logger.debug("Async PostgreSQL pool closed", pool=connection_key)
                 except asyncio.TimeoutError:
@@ -309,15 +372,19 @@ class PostgresPoolManager:
     """
     Synchronous PostgreSQL connection pool manager that shares pools
     based on database configuration.
+
+    For Lakebase connections (when instance_name is provided), uses LakebasePool
+    from databricks_ai_bridge which handles automatic token rotation and host resolution.
+    For standard PostgreSQL connections, uses psycopg_pool.ConnectionPool.
     """
 
     _pools: dict[str, ConnectionPool] = {}
+    _lakebase_pools: dict[str, LakebasePool] = {}
     _lock: threading.Lock = threading.Lock()
 
     @classmethod
     def get_pool(cls, database: DatabaseModel) -> ConnectionPool:
         connection_key: str = str(database.name)
-        connection_params: dict[str, Any] = database.connection_params
 
         with cls._lock:
             if connection_key in cls._pools:
@@ -326,19 +393,41 @@ class PostgresPoolManager:
 
             logger.debug("Creating new PostgreSQL pool", database=database.name)
 
-            kwargs: dict[str, Any] = {
-                "row_factory": dict_row,
-                "autocommit": True,
-            } | database.connection_kwargs or {}
-
-            # Create connection pool
-            pool: ConnectionPool = _create_pool(
-                connection_params=connection_params,
-                database_name=database.name,
-                max_pool_size=database.max_pool_size,
-                timeout_seconds=database.timeout_seconds,
-                kwargs=kwargs,
-            )
+            if database.is_lakebase:
+                # Use LakebasePool for Lakebase connections
+                # LakebasePool handles automatic token rotation and host resolution
+                lakebase_pool = LakebasePool(
+                    instance_name=database.instance_name,
+                    workspace_client=database.workspace_client,
+                    min_size=1,
+                    max_size=database.max_pool_size,
+                    timeout=float(database.timeout_seconds),
+                )
+                # Store the LakebasePool for proper cleanup
+                cls._lakebase_pools[connection_key] = lakebase_pool
+                # Get the underlying ConnectionPool
+                pool = lakebase_pool.pool
+                logger.success(
+                    "Lakebase connection pool created",
+                    database=database.name,
+                    instance_name=database.instance_name,
+                    pool_size=database.max_pool_size,
+                )
+            else:
+                # Use standard PostgreSQL pool for non-Lakebase connections
+                connection_params: dict[str, Any] = database.connection_params
+                kwargs: dict[str, Any] = {
+                    "row_factory": dict_row,
+                    "autocommit": True,
+                } | database.connection_kwargs or {}
+
+                pool = _create_pool(
+                    connection_params=connection_params,
+                    database_name=database.name,
+                    max_pool_size=database.max_pool_size,
+                    timeout_seconds=database.timeout_seconds,
+                    kwargs=kwargs,
+                )
 
             cls._pools[connection_key] = pool
             return pool
@@ -348,7 +437,13 @@ class PostgresPoolManager:
         connection_key: str = database.name
 
         with cls._lock:
-            if connection_key in cls._pools:
+            # Close LakebasePool if it exists (handles underlying pool cleanup)
+            if connection_key in cls._lakebase_pools:
+                lakebase_pool = cls._lakebase_pools.pop(connection_key)
+                lakebase_pool.close()
+                cls._pools.pop(connection_key, None)
+                logger.debug("Lakebase pool closed", database=database.name)
+            elif connection_key in cls._pools:
                 pool = cls._pools.pop(connection_key)
                 pool.close()
                 logger.debug("PostgreSQL pool closed", database=database.name)
@@ -356,16 +451,32 @@ class PostgresPoolManager:
     @classmethod
     def close_all_pools(cls):
         with cls._lock:
-            for connection_key, pool in cls._pools.items():
+            # Close all LakebasePool instances first
+            for connection_key, lakebase_pool in cls._lakebase_pools.items():
                 try:
-                    pool.close()
-                    logger.debug("PostgreSQL pool closed", pool=connection_key)
+                    lakebase_pool.close()
+                    logger.debug("Lakebase pool closed", pool=connection_key)
                 except Exception as e:
                     logger.error(
-                        "Error closing PostgreSQL pool",
+                        "Error closing Lakebase pool",
                         pool=connection_key,
                         error=str(e),
                     )
+            cls._lakebase_pools.clear()
+
+            # Close any remaining standard PostgreSQL pools
+            for connection_key, pool in cls._pools.items():
+                # Skip if already closed via LakebasePool
+                if connection_key not in cls._lakebase_pools:
+                    try:
+                        pool.close()
+                        logger.debug("PostgreSQL pool closed", pool=connection_key)
+                    except Exception as e:
+                        logger.error(
+                            "Error closing PostgreSQL pool",
+                            pool=connection_key,
+                            error=str(e),
+                        )
             cls._pools.clear()
 
 

dao_ai/orchestration/core.py

@@ -9,7 +9,7 @@ This module provides the foundational utilities for multi-agent orchestration:
 - Main orchestration graph factory
 """
 
-from typing import Awaitable, Callable, Literal
+from typing import Any, Awaitable, Callable, Literal
 
 from langchain.tools import ToolRuntime, tool
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
@@ -179,8 +179,16 @@ def create_agent_node_handler(
             "messages": filtered_messages,
         }
 
-        # Invoke the agent
-        result: AgentState = await agent.ainvoke(agent_state, context=runtime.context)
+        # Build config with configurable from context for langmem compatibility
+        # langmem tools expect user_id to be in config.configurable
+        config: dict[str, Any] = {}
+        if runtime.context:
+            config = {"configurable": runtime.context.model_dump()}
+
+        # Invoke the agent with both context and config
+        result: AgentState = await agent.ainvoke(
+            agent_state, context=runtime.context, config=config
+        )
 
         # Extract agent response based on output mode
         result_messages = result.get("messages", [])
@@ -227,15 +235,31 @@ def create_handoff_tool(
         tool_call_id: str = runtime.tool_call_id
         logger.debug("Handoff to agent", target_agent=target_agent_name)
 
+        # Get the AIMessage that triggered this handoff (required for tool_use/tool_result pairing)
+        # LLMs expect tool calls to be paired with their responses, so we must include both
+        # the AIMessage containing the tool call and the ToolMessage acknowledging it.
+        messages: list[BaseMessage] = runtime.state.get("messages", [])
+        last_ai_message: AIMessage | None = None
+        for msg in reversed(messages):
+            if isinstance(msg, AIMessage) and msg.tool_calls:
+                last_ai_message = msg
+                break
+
+        # Build message list with proper pairing
+        update_messages: list[BaseMessage] = []
+        if last_ai_message:
+            update_messages.append(last_ai_message)
+        update_messages.append(
+            ToolMessage(
+                content=f"Transferred to {target_agent_name}",
+                tool_call_id=tool_call_id,
+            )
+        )
+
         return Command(
             update={
                 "active_agent": target_agent_name,
-                "messages": [
-                    ToolMessage(
-                        content=f"Transferred to {target_agent_name}",
-                        tool_call_id=tool_call_id,
-                    )
-                ],
+                "messages": update_messages,
             },
             goto=target_agent_name,
             graph=Command.PARENT,

dao_ai/orchestration/supervisor.py

@@ -13,7 +13,7 @@ from langchain.agents import create_agent
 from langchain.agents.middleware import AgentMiddleware as LangchainAgentMiddleware
 from langchain.tools import ToolRuntime, tool
 from langchain_core.language_models import LanguageModelLike
-from langchain_core.messages import ToolMessage
+from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
 from langchain_core.tools import BaseTool
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph import StateGraph
@@ -75,15 +75,30 @@ def _create_handoff_back_to_supervisor_tool() -> BaseTool:
         tool_call_id: str = runtime.tool_call_id
         logger.debug("Agent handing back to supervisor", summary_preview=summary[:100])
 
+        # Get the AIMessage that triggered this handoff (required for tool_use/tool_result pairing)
+        # LLMs expect tool calls to be paired with their responses, so we must include both
+        # the AIMessage containing the tool call and the ToolMessage acknowledging it.
+        messages: list[BaseMessage] = runtime.state.get("messages", [])
+        last_ai_message: AIMessage | None = None
+        for msg in reversed(messages):
+            if isinstance(msg, AIMessage) and msg.tool_calls:
+                last_ai_message = msg
+                break
+
+        # Build message list with proper pairing
+        update_messages: list[BaseMessage] = []
+        if last_ai_message:
+            update_messages.append(last_ai_message)
+        update_messages.append(
+            ToolMessage(
+                content=f"Task completed: {summary}",
+                tool_call_id=tool_call_id,
+            )
+        )
+
         return Command(
             update={
-                "active_agent": None,
-                "messages": [
-                    ToolMessage(
-                        content=f"Task completed: {summary}",
-                        tool_call_id=tool_call_id,
-                    )
-                ],
+                "messages": update_messages,
             },
             goto=SUPERVISOR_NODE,
             graph=Command.PARENT,

dao_ai/{prompts.py → prompts/__init__.py}

@@ -2,9 +2,11 @@
 Prompt utilities for DAO AI agents.
 
 This module provides utilities for creating dynamic prompts using
-LangChain v1's @dynamic_prompt middleware decorator pattern.
+LangChain v1's @dynamic_prompt middleware decorator pattern, as well as
+paths to prompt template files.
 """
 
+from pathlib import Path
 from typing import Any, Optional
 
 from langchain.agents.middleware import (
@@ -18,6 +20,13 @@ from loguru import logger
 from dao_ai.config import PromptModel
 from dao_ai.state import Context
 
+PROMPTS_DIR = Path(__file__).parent
+
+
+def get_prompt_path(name: str) -> Path:
+    """Get the path to a prompt template file."""
+    return PROMPTS_DIR / name
+
 
 def make_prompt(
     base_system_prompt: Optional[str | PromptModel],