intentkit 0.7.5.dev3__py3-none-any.whl → 0.8.34.dev7__py3-none-any.whl

intentkit/core/engine.py

@@ -17,32 +17,33 @@ import textwrap
 import time
 import traceback
 from datetime import datetime
-from typing import Optional
 
 import sqlalchemy
 from epyxid import XID
-from fastapi import HTTPException
+from langchain.agents import create_agent as create_langchain_agent
 from langchain_core.messages import (
     BaseMessage,
     HumanMessage,
 )
 from langchain_core.tools import BaseTool
+from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.errors import GraphRecursionError
 from langgraph.graph.state import CompiledStateGraph
-from langgraph.prebuilt import create_react_agent
 from sqlalchemy import func, update
 from sqlalchemy.exc import SQLAlchemyError
 
 from intentkit.abstracts.graph import AgentContext, AgentError, AgentState
 from intentkit.config.config import config
-from intentkit.core.agent import agent_store
 from intentkit.core.chat import clear_thread_memory
 from intentkit.core.credit import expense_message, expense_skill
-from intentkit.core.node import PreModelNode, post_model_node
-from intentkit.core.prompt import (
-    create_formatted_prompt_function,
-    explain_prompt,
+from intentkit.core.middleware import (
+    CreditCheckMiddleware,
+    DynamicPromptMiddleware,
+    SummarizationMiddleware,
+    ToolBindingMiddleware,
+    TrimMessagesMiddleware,
 )
+from intentkit.core.prompt import explain_prompt
 from intentkit.models.agent import Agent, AgentTable
 from intentkit.models.agent_data import AgentData, AgentQuota
 from intentkit.models.app_setting import AppSetting, SystemMessageType
@@ -53,26 +54,51 @@ from intentkit.models.chat import (
     ChatMessageSkillCall,
 )
 from intentkit.models.credit import CreditAccount, OwnerType
-from intentkit.models.db import get_langgraph_checkpointer, get_session
-from intentkit.models.llm import LLMModelInfo, LLMProvider, create_llm_model
-from intentkit.models.skill import AgentSkillData, ThreadSkillData
+from intentkit.models.db import (
+    get_checkpointer,
+    get_session,
+)
+from intentkit.models.llm import LLMModelInfo, create_llm_model
+from intentkit.models.skill import AgentSkillData, ChatSkillData, Skill
 from intentkit.models.user import User
 from intentkit.utils.error import IntentKitAPIError
 
 logger = logging.getLogger(__name__)
 
-
 # Global variable to cache all agent executors
 _agents: dict[str, CompiledStateGraph] = {}
-_private_agents: dict[str, CompiledStateGraph] = {}
 
 # Global dictionaries to cache agent update times
 _agents_updated: dict[str, datetime] = {}
-_private_agents_updated: dict[str, datetime] = {}
 
 
-async def create_agent(agent: Agent, is_private: bool = False) -> CompiledStateGraph:
-    """Create an AI agent with specified configuration and tools.
+def _extract_text_content(content: object) -> str:
+    if isinstance(content, list):
+        texts: list[str] = []
+        for item in content:
+            if isinstance(item, dict):
+                t = item.get("text")
+                ty = item.get("type")
+                if t is not None and (ty == "text" or ty is None):
+                    texts.append(t)
+            elif isinstance(item, str):
+                texts.append(item)
+        return "".join(texts)
+    if isinstance(content, dict):
+        if content.get("type") == "text" and "text" in content:
+            return content["text"]
+        if "text" in content:
+            return content["text"]
+        return ""
+    if isinstance(content, str):
+        return content
+    return ""
+
+
+async def build_agent(
+    agent: Agent, agent_data: AgentData, custom_skills: list[BaseTool] = []
+) -> CompiledStateGraph:
+    """Build an AI agent with specified configuration and tools.
 
     This function:
     1. Initializes LLM with specified model
@@ -82,13 +108,13 @@ async def create_agent(agent: Agent, is_private: bool = False) -> CompiledStateG
 
     Args:
         agent (Agent): Agent configuration object
-        is_private (bool, optional): Flag indicating whether the agent is private. Defaults to False.
-        has_search (bool, optional): Flag indicating whether to include search tools. Defaults to False.
+        agent_data (AgentData): Agent data object
+        custom_skills (list[BaseTool], optional): Designed for advanced user who directly
+            call this function to inject custom skills into the agent tool node.
 
     Returns:
         CompiledStateGraph: Initialized LangChain agent
     """
-    agent_data = await AgentData.get(agent.id)
 
     # Create the LLM model instance
     llm_model = await create_llm_model(
@@ -98,17 +124,15 @@ async def create_agent(agent: Agent, is_private: bool = False) -> CompiledStateG
         presence_penalty=agent.presence_penalty,
     )
 
-    # Get the LLM instance
-    llm = await llm_model.create_instance(config)
-
-    # Get the token limit from the model info
-    input_token_limit = min(config.input_token_limit, llm_model.info.context_length)
-
     # ==== Store buffered conversation history in memory.
-    memory = get_langgraph_checkpointer()
+    try:
+        checkpointer = get_checkpointer()
+    except RuntimeError:
+        checkpointer = InMemorySaver()
 
     # ==== Load skills
     tools: list[BaseTool | dict] = []
+    private_tools: list[BaseTool | dict] = []
 
     if agent.skills:
         for k, v in agent.skills.items():
@@ -117,55 +141,65 @@ async def create_agent(agent: Agent, is_private: bool = False) -> CompiledStateG
             try:
                 skill_module = importlib.import_module(f"intentkit.skills.{k}")
                 if hasattr(skill_module, "get_skills"):
+                    # all
                     skill_tools = await skill_module.get_skills(
-                        v, is_private, agent_store, agent_id=agent.id
+                        v, False, agent_id=agent.id, agent=agent
                     )
                     if skill_tools and len(skill_tools) > 0:
                         tools.extend(skill_tools)
+                    # private
+                    skill_private_tools = await skill_module.get_skills(
+                        v, True, agent_id=agent.id, agent=agent
+                    )
+                    if skill_private_tools and len(skill_private_tools) > 0:
+                        private_tools.extend(skill_private_tools)
                 else:
                     logger.error(f"Skill {k} does not have get_skills function")
             except ImportError as e:
                 logger.error(f"Could not import skill module: {k} ({e})")
 
+    # add custom skills to private tools
+    if custom_skills and len(custom_skills) > 0:
+        private_tools.extend(custom_skills)
+
     # filter the duplicate tools
     tools = list({tool.name: tool for tool in tools}.values())
+    private_tools = list({tool.name: tool for tool in private_tools}.values())
 
-    # Add search tools if requested
-    if (
-        llm_model.info.provider == LLMProvider.OPENAI
-        and llm_model.info.supports_search
-        and not agent.model.startswith(
-            "gpt-5"
-        )  # tmp disable gpt-5 search since package bugs
-    ):
-        tools.append({"type": "web_search_preview"})
+    for tool in private_tools:
+        logger.info(
+            f"[{agent.id}] loaded tool: {tool.name if isinstance(tool, BaseTool) else tool}"
+        )
 
-    # Create the formatted_prompt function using the refactored prompt module
-    formatted_prompt = create_formatted_prompt_function(agent, agent_data)
+    base_model = await llm_model.create_instance()
 
-    for tool in tools:
-        logger.info(
-            f"[{agent.id}{'-private' if is_private else ''}] loaded tool: {tool.name if isinstance(tool, BaseTool) else tool}"
+    middleware = [
+        ToolBindingMiddleware(llm_model, tools, private_tools),
+        DynamicPromptMiddleware(agent, agent_data),
+    ]
+
+    if agent.short_term_memory_strategy == "trim":
+        middleware.append(TrimMessagesMiddleware(max_summary_tokens=2048))
+    elif agent.short_term_memory_strategy == "summarize":
+        summarize_llm = await create_llm_model(model_name="gpt-5-mini")
+        summarize_model = await summarize_llm.create_instance()
+        middleware.append(
+            SummarizationMiddleware(
+                model=summarize_model,
+                max_tokens_before_summary=llm_model.info.context_length // 2,
+            )
         )
 
-    # Pre model hook
-    pre_model_hook = PreModelNode(
-        model=llm,
-        short_term_memory_strategy=agent.short_term_memory_strategy,
-        max_tokens=input_token_limit // 2,
-        max_summary_tokens=2048,
-    )
+    if config.payment_enabled:
+        middleware.append(CreditCheckMiddleware())
 
-    # Create ReAct Agent using the LLM and CDP Agentkit tools.
-    executor = create_react_agent(
-        model=llm,
-        tools=tools,
-        prompt=formatted_prompt,
-        pre_model_hook=pre_model_hook,
-        post_model_hook=post_model_node if config.payment_enabled else None,
+    executor = create_langchain_agent(
+        model=base_model,
+        tools=private_tools,
+        middleware=middleware,
         state_schema=AgentState,
         context_schema=AgentContext,
-        checkpointer=memory,
+        checkpointer=checkpointer,
         debug=config.debug_checkpoint,
         name=agent.id,
     )
@@ -173,7 +207,23 @@ async def create_agent(agent: Agent, is_private: bool = False) -> CompiledStateG
     return executor
 
 
-async def initialize_agent(aid, is_private=False):
+async def create_agent(agent: Agent) -> CompiledStateGraph:
+    """Create an AI agent with specified configuration and tools.
+
+    This function maintains backward compatibility by calling build_agent internally.
+
+    Args:
+        agent (Agent): Agent configuration object
+        is_private (bool, optional): Flag indicating whether the agent is private. Defaults to False.
+
+    Returns:
+        CompiledStateGraph: Initialized LangChain agent
+    """
+    agent_data = await AgentData.get(agent.id)
+    return await build_agent(agent, agent_data)
+
+
+async def initialize_agent(aid):
     """Initialize an AI agent with specified configuration and tools.
 
     This function:
@@ -192,50 +242,41 @@ async def initialize_agent(aid, is_private=False):
         HTTPException: If agent not found (404) or database error (500)
     """
     # get the agent from the database
-    agent: Optional[Agent] = await Agent.get(aid)
+    agent: Agent | None = await Agent.get(aid)
     if not agent:
-        raise HTTPException(status_code=404, detail="Agent not found")
+        raise IntentKitAPIError(
+            status_code=404, key="AgentNotFound", message="Agent not found"
+        )
 
     # Create the agent using the new create_agent function
-    executor = await create_agent(agent, is_private)
+    executor = await create_agent(agent)
 
     # Cache the agent executor
-    if is_private:
-        _private_agents[aid] = executor
-        _private_agents_updated[aid] = agent.updated_at
-    else:
-        _agents[aid] = executor
-        _agents_updated[aid] = agent.updated_at
+    _agents[aid] = executor
+    _agents_updated[aid] = agent.deployed_at if agent.deployed_at else agent.updated_at
 
 
-async def agent_executor(
-    agent_id: str, is_private: bool
-) -> (CompiledStateGraph, float):
+async def agent_executor(agent_id: str) -> tuple[CompiledStateGraph, float]:
     start = time.perf_counter()
     agent = await Agent.get(agent_id)
     if not agent:
-        raise HTTPException(status_code=404, detail="Agent not found")
-    agents = _private_agents if is_private else _agents
-    agents_updated = _private_agents_updated if is_private else _agents_updated
-
+        raise IntentKitAPIError(
+            status_code=404, key="AgentNotFound", message="Agent not found"
+        )
+    updated_at = agent.deployed_at if agent.deployed_at else agent.updated_at
     # Check if agent needs reinitialization due to updates
     needs_reinit = False
-    if agent_id in agents:
-        if (
-            agent_id not in agents_updated
-            or agent.updated_at != agents_updated[agent_id]
-        ):
+    if agent_id in _agents:
+        if agent_id not in _agents_updated or updated_at != _agents_updated[agent_id]:
             needs_reinit = True
-            logger.info(
-                f"Reinitializing agent {agent_id} due to updates, private mode: {is_private}"
-            )
+            logger.info(f"Reinitializing agent {agent_id} due to updates")
 
     # cold start or needs reinitialization
     cold_start_cost = 0.0
-    if (agent_id not in agents) or needs_reinit:
-        await initialize_agent(agent_id, is_private)
+    if (agent_id not in _agents) or needs_reinit:
+        await initialize_agent(agent_id)
         cold_start_cost = time.perf_counter() - start
-    return agents[agent_id], cold_start_cost
+    return _agents[agent_id], cold_start_cost
 
 
 async def stream_agent(message: ChatMessageCreate):
@@ -254,24 +295,60 @@ async def stream_agent(message: ChatMessageCreate):
     Yields:
         ChatMessage: Individual response messages including timing information
     """
+    agent = await Agent.get(message.agent_id)
+    executor, cold_start_cost = await agent_executor(message.agent_id)
+    message.cold_start_cost = cold_start_cost
+    async for chat_message in stream_agent_raw(message, agent, executor):
+        yield chat_message
+
+
+async def stream_agent_raw(
+    message: ChatMessageCreate, agent: Agent, executor: CompiledStateGraph
+):
     start = time.perf_counter()
     # make sure reply_to is set
     message.reply_to = message.id
 
     # save input message first
-    input = await message.save()
+    user_message = await message.save()
 
-    # agent
-    agent = await Agent.get(input.agent_id)
+    # temporary debug logging for telegram messages
+    if user_message.author_type == AuthorType.TELEGRAM:
+        logger.info(
+            f"[TELEGRAM DEBUG] Agent: {user_message.agent_id} | Chat: {user_message.chat_id} | Message: {user_message.message}"
+        )
+
+    if re.search(
+        r"(@clear|/clear)(?!\w)",
+        user_message.message.strip(),
+        re.IGNORECASE,
+    ):
+        await clear_thread_memory(user_message.agent_id, user_message.chat_id)
+
+        confirmation_message = ChatMessageCreate(
+            id=str(XID()),
+            agent_id=user_message.agent_id,
+            chat_id=user_message.chat_id,
+            user_id=user_message.user_id,
+            author_id=user_message.agent_id,
+            author_type=AuthorType.AGENT,
+            model=agent.model,
+            thread_type=user_message.author_type,
+            reply_to=user_message.id,
+            message="Memory in context has been cleared.",
+            time_cost=time.perf_counter() - start,
+        )
+
+        yield await confirmation_message.save()
+        return
 
-    # model
     model = await LLMModelInfo.get(agent.model)
 
     payment_enabled = config.payment_enabled
 
     # check user balance
     if payment_enabled:
-        if not input.user_id or not agent.owner:
+        if not user_message.user_id or not agent.owner:
             raise IntentKitAPIError(
                 500,
                 "PaymentError",
@@ -282,23 +359,25 @@ async def stream_agent(message: ChatMessageCreate):
             if owner and agent.fee_percentage > 100 + owner.nft_count * 10:
                 error_message_create = await ChatMessageCreate.from_system_message(
                     SystemMessageType.SERVICE_FEE_ERROR,
-                    agent_id=input.agent_id,
-                    chat_id=input.chat_id,
-                    user_id=input.user_id,
-                    author_id=input.agent_id,
-                    thread_type=input.author_type,
-                    reply_to=input.id,
+                    agent_id=user_message.agent_id,
+                    chat_id=user_message.chat_id,
+                    user_id=user_message.user_id,
+                    author_id=user_message.agent_id,
+                    thread_type=user_message.author_type,
+                    reply_to=user_message.id,
                     time_cost=time.perf_counter() - start,
                 )
                 error_message = await error_message_create.save()
                 yield error_message
                 return
         # payer
-        payer = input.user_id
-        if input.author_type in [
+        payer = user_message.user_id
+        if user_message.author_type in [
             AuthorType.TELEGRAM,
+            AuthorType.DISCORD,
             AuthorType.TWITTER,
             AuthorType.API,
+            AuthorType.X402,
         ]:
             payer = agent.owner
         # user account
@@ -318,12 +397,12 @@ async def stream_agent(message: ChatMessageCreate):
             if quota and quota.free_income_daily > 24000:
                 error_message_create = await ChatMessageCreate.from_system_message(
                     SystemMessageType.DAILY_USAGE_LIMIT_EXCEEDED,
-                    agent_id=input.agent_id,
-                    chat_id=input.chat_id,
-                    user_id=input.user_id,
-                    author_id=input.agent_id,
-                    thread_type=input.author_type,
-                    reply_to=input.id,
+                    agent_id=user_message.agent_id,
+                    chat_id=user_message.chat_id,
+                    user_id=user_message.user_id,
+                    author_id=user_message.agent_id,
+                    thread_type=user_message.author_type,
+                    reply_to=user_message.id,
                     time_cost=time.perf_counter() - start,
                 )
                 error_message = await error_message_create.save()
@@ -336,12 +415,12 @@ async def stream_agent(message: ChatMessageCreate):
         if not user_account.has_sufficient_credits(avg_count):
             error_message_create = await ChatMessageCreate.from_system_message(
                 SystemMessageType.INSUFFICIENT_BALANCE,
-                agent_id=input.agent_id,
-                chat_id=input.chat_id,
-                user_id=input.user_id,
-                author_id=input.agent_id,
-                thread_type=input.author_type,
-                reply_to=input.id,
+                agent_id=user_message.agent_id,
+                chat_id=user_message.chat_id,
+                user_id=user_message.user_id,
+                author_id=user_message.agent_id,
+                thread_type=user_message.author_type,
+                reply_to=user_message.id,
                 time_cost=time.perf_counter() - start,
             )
             error_message = await error_message_create.save()
@@ -349,56 +428,41 @@ async def stream_agent(message: ChatMessageCreate):
             return
 
     is_private = False
-    if input.user_id == agent.owner:
+    if user_message.user_id == agent.owner:
         is_private = True
 
-    executor, cold_start_cost = await agent_executor(input.agent_id, is_private)
-    last = start + cold_start_cost
+    last = start
 
     # Extract images from attachments
     image_urls = []
-    if input.attachments:
+    if user_message.attachments:
         image_urls = [
             att["url"]
-            for att in input.attachments
+            for att in user_message.attachments
             if "type" in att and att["type"] == "image" and "url" in att
         ]
 
     # Process input message to handle @skill patterns
-    if config.admin_llm_skill_control:
-        input_message = await explain_prompt(input.message)
-    else:
-        input_message = input.message
+    input_message = await explain_prompt(user_message.message)
 
     # super mode
     recursion_limit = 30
-    if re.search(r"\b@super\b", input_message):
+    if (
+        re.search(r"@super\b", input_message)
+        or user_message.super_mode
+        or agent.has_super()
+    ):
         recursion_limit = 300
-        # Remove @super from the message
-        input_message = re.sub(r"\b@super\b", "", input_message).strip()
+        input_message = re.sub(r"@super\b", "", input_message).strip()
 
     # llm native search
-    if re.search(r"\b@search\b", input_message) or re.search(
-        r"\b@web\b", input_message
-    ):
-        if model.supports_search:
-            input_message = re.sub(
-                r"\b@search\b",
-                "(You have native search tool, you can use it to get more recent information)",
-                input_message,
-            ).strip()
-            input_message = re.sub(
-                r"\b@web\b",
-                "(You have native search tool, you can use it to get more recent information)",
-                input_message,
-            ).strip()
-        else:
-            input_message = re.sub(r"\b@search\b", "", input_message).strip()
-            input_message = re.sub(r"\b@web\b", "", input_message).strip()
+    search = user_message.search_mode if user_message.search_mode is not None else False
+    if re.search(r"@search\b", input_message) or re.search(r"@web\b", input_message):
+        search = True
 
     # content to llm
-    content = [
-        {"type": "text", "text": input_message},
+    messages = [
+        HumanMessage(content=input_message),
     ]
     # if the model doesn't natively support image parsing, add the image URLs to the message
     if image_urls:
@@ -406,25 +470,24 @@ async def stream_agent(message: ChatMessageCreate):
             agent.has_image_parser_skill(is_private=is_private)
             and not model.supports_image_input
         ):
-            input_message += f"\n\nImages:\n{'\n'.join(image_urls)}"
-            content = [
-                {"type": "text", "text": input_message},
+            image_urls_text = "\n".join(image_urls)
+            input_message += f"\n\nImages:\n{image_urls_text}"
+            messages = [
+                HumanMessage(content=input_message),
             ]
         else:
             # anyway, pass it directly to LLM
-            content.extend(
+            messages.extend(
                 [
-                    {"type": "image_url", "image_url": {"url": image_url}}
+                    HumanMessage(
+                        content={"type": "image_url", "image_url": {"url": image_url}}
+                    )
                     for image_url in image_urls
                 ]
             )
 
-    messages = [
-        HumanMessage(content=content),
-    ]
-
     # stream config
-    thread_id = f"{input.agent_id}-{input.chat_id}"
+    thread_id = f"{user_message.agent_id}-{user_message.chat_id}"
     stream_config = {
         "configurable": {
             "thread_id": thread_id,
@@ -432,13 +495,18 @@ async def stream_agent(message: ChatMessageCreate):
         "recursion_limit": recursion_limit,
     }
 
+    def get_agent() -> Agent:
+        return agent
+
     context = AgentContext(
-        agent_id=input.agent_id,
-        chat_id=input.chat_id,
-        user_id=input.user_id,
-        app_id=input.app_id,
-        entrypoint=input.author_type,
+        agent_id=user_message.agent_id,
+        get_agent=get_agent,
+        chat_id=user_message.chat_id,
+        user_id=user_message.user_id,
+        app_id=user_message.app_id,
+        entrypoint=user_message.author_type,
         is_private=is_private,
+        search=search,
         payer=payer if payment_enabled else None,
     )
 
@@ -446,42 +514,84 @@ async def stream_agent(message: ChatMessageCreate):
     cached_tool_step = None
     try:
         async for chunk in executor.astream(
-            {"messages": messages}, context=context, config=stream_config
+            {"messages": messages},
+            context=context,
+            config=stream_config,
+            stream_mode=["updates", "custom"],
         ):
             this_time = time.perf_counter()
             logger.debug(f"stream chunk: {chunk}", extra={"thread_id": thread_id})
-            if "agent" in chunk and "messages" in chunk["agent"]:
-                if len(chunk["agent"]["messages"]) != 1:
+
+            if isinstance(chunk, tuple) and len(chunk) == 2:
+                event_kind, payload = chunk
+                chunk = payload
+
+            if isinstance(chunk, dict) and "credit_check" in chunk:
+                credit_payload = chunk.get("credit_check", {})
+                content = credit_payload.get("message")
+                if content:
+                    credit_message_create = ChatMessageCreate(
+                        id=str(XID()),
+                        agent_id=user_message.agent_id,
+                        chat_id=user_message.chat_id,
+                        user_id=user_message.user_id,
+                        author_id=user_message.agent_id,
+                        author_type=AuthorType.AGENT,
+                        model=agent.model,
+                        thread_type=user_message.author_type,
+                        reply_to=user_message.id,
+                        message=content,
+                        input_tokens=0,
+                        output_tokens=0,
+                        time_cost=this_time - last,
+                    )
+                    last = this_time
+                    credit_message = await credit_message_create.save()
+                    yield credit_message
+
+                    error_message_create = await ChatMessageCreate.from_system_message(
+                        SystemMessageType.INSUFFICIENT_BALANCE,
+                        agent_id=user_message.agent_id,
+                        chat_id=user_message.chat_id,
+                        user_id=user_message.user_id,
+                        author_id=user_message.agent_id,
+                        thread_type=user_message.author_type,
+                        reply_to=user_message.id,
+                        time_cost=0,
+                    )
+                    error_message = await error_message_create.save()
+                    yield error_message
+                return
+
+            if not isinstance(chunk, dict):
+                continue
+
+            if "model" in chunk and "messages" in chunk["model"]:
+                if len(chunk["model"]["messages"]) != 1:
                     logger.error(
-                        "unexpected agent message: " + str(chunk["agent"]["messages"]),
+                        "unexpected model message: " + str(chunk["model"]["messages"]),
                         extra={"thread_id": thread_id},
                     )
-                msg = chunk["agent"]["messages"][0]
-                if hasattr(msg, "tool_calls") and msg.tool_calls:
-                    # tool calls, save for later use, if it is deleted by post_model_hook, will not be used.
+                msg = chunk["model"]["messages"][0]
+                has_tools = hasattr(msg, "tool_calls") and bool(msg.tool_calls)
+                if has_tools:
                     cached_tool_step = msg
-                if hasattr(msg, "content") and msg.content:
-                    content = msg.content
-                    if isinstance(msg.content, list):
-                        # in new version, content item maybe a list
-                        content = msg.content[0]
-                    if isinstance(content, dict):
-                        if "text" in content:
-                            content = content["text"]
-                        else:
-                            content = str(content)
-                            logger.error(f"unexpected content type: {content}")
-                    # agent message
+                content = (
+                    _extract_text_content(msg.content)
+                    if hasattr(msg, "content")
+                    else ""
+                )
+                if content and not has_tools:
                     chat_message_create = ChatMessageCreate(
                         id=str(XID()),
-                        agent_id=input.agent_id,
-                        chat_id=input.chat_id,
-                        user_id=input.user_id,
-                        author_id=input.agent_id,
+                        agent_id=user_message.agent_id,
+                        chat_id=user_message.chat_id,
+                        user_id=user_message.user_id,
+                        author_id=user_message.agent_id,
                         author_type=AuthorType.AGENT,
                         model=agent.model,
-                        thread_type=input.author_type,
-                        reply_to=input.id,
+                        thread_type=user_message.author_type,
+                        reply_to=user_message.id,
                         message=content,
                         input_tokens=(
                             msg.usage_metadata.get("input_tokens", 0)
@@ -496,19 +606,13 @@ async def stream_agent(message: ChatMessageCreate):
                         time_cost=this_time - last,
                     )
                     last = this_time
-                    if cold_start_cost > 0:
-                        chat_message_create.cold_start_cost = cold_start_cost
-                        cold_start_cost = 0
-                    # handle message and payment in one transaction
                     async with get_session() as session:
-                        # payment
                         if payment_enabled:
                             amount = await model.calculate_cost(
                                 chat_message_create.input_tokens,
                                 chat_message_create.output_tokens,
                             )
 
-                            # Check for web_search_call in additional_kwargs
                             if (
                                 hasattr(msg, "additional_kwargs")
                                 and msg.additional_kwargs
@@ -519,7 +623,7 @@ async def stream_agent(message: ChatMessageCreate):
                                 for tool_output in tool_outputs:
                                     if tool_output.get("type") == "web_search_call":
                                         logger.info(
-                                            f"[{input.agent_id}] Found web_search_call in additional_kwargs"
+                                            f"[{user_message.agent_id}] Found web_search_call in additional_kwargs"
                                         )
                                         amount += 35
                                         break
@@ -527,11 +631,13 @@ async def stream_agent(message: ChatMessageCreate):
                                 session,
                                 payer,
                                 chat_message_create.id,
-                                input.id,
+                                user_message.id,
                                 amount,
                                 agent,
                             )
-                            logger.info(f"[{input.agent_id}] expense message: {amount}")
+                            logger.info(
+                                f"[{user_message.agent_id}] expense message: {amount}"
+                            )
                             chat_message_create.credit_event_id = credit_event.id
                             chat_message_create.credit_cost = credit_event.total_amount
                         chat_message = await chat_message_create.save_in_session(
@@ -566,7 +672,8 @@ async def stream_agent(message: ChatMessageCreate):
                                 "parameters": call["args"],
                                 "success": True,
                             }
-                            if msg.status == "error":
+                            status = getattr(msg, "status", None)
+                            if status == "error":
                                 skill_call["success"] = False
                                 skill_call["error_message"] = str(msg.content)
                             else:
@@ -576,20 +683,21 @@ async def stream_agent(message: ChatMessageCreate):
                                     skill_call["response"] = textwrap.shorten(
                                         str(msg.content), width=1000, placeholder="..."
                                     )
-                                if msg.artifact:
-                                    cached_attachments.extend(msg.artifact)
+                                artifact = getattr(msg, "artifact", None)
+                                if artifact:
+                                    cached_attachments.extend(artifact)
                             skill_calls.append(skill_call)
                             break
                 skill_message_create = ChatMessageCreate(
                     id=str(XID()),
-                    agent_id=input.agent_id,
-                    chat_id=input.chat_id,
-                    user_id=input.user_id,
-                    author_id=input.agent_id,
+                    agent_id=user_message.agent_id,
+                    chat_id=user_message.chat_id,
+                    user_id=user_message.user_id,
+                    author_id=user_message.agent_id,
                     author_type=AuthorType.SKILL,
                     model=agent.model,
-                    thread_type=input.author_type,
-                    reply_to=input.id,
+                    thread_type=user_message.author_type,
+                    reply_to=user_message.id,
                     message="",
                     skill_calls=skill_calls,
                     attachments=cached_attachments,
@@ -610,13 +718,8 @@ async def stream_agent(message: ChatMessageCreate):
                     time_cost=this_time - last,
                 )
                 last = this_time
-                if cold_start_cost > 0:
-                    skill_message_create.cold_start_cost = cold_start_cost
-                    cold_start_cost = 0
-                # save message and credit in one transaction
                 async with get_session() as session:
                     if payment_enabled:
-                        # message payment, only first call in a group has message bill
                         if have_first_call_in_cache:
                             message_amount = await model.calculate_cost(
                                 skill_message_create.input_tokens,
@@ -626,7 +729,7 @@ async def stream_agent(message: ChatMessageCreate):
                                 session,
                                 payer,
                                 skill_message_create.id,
-                                input.id,
+                                user_message.id,
                                 message_amount,
                                 agent,
                             )
@@ -636,15 +739,17 @@ async def stream_agent(message: ChatMessageCreate):
                             skill_message_create.credit_cost = (
                                 message_payment_event.total_amount
                             )
-                        # skill payment
                         for skill_call in skill_calls:
                             if not skill_call["success"]:
                                 continue
+                            skill = await Skill.get(skill_call["name"])
+                            if not skill:
+                                continue
                             payment_event = await expense_skill(
                                 session,
                                 payer,
                                 skill_message_create.id,
-                                input.id,
+                                user_message.id,
                                 skill_call["id"],
                                 skill_call["name"],
                                 agent,
@@ -652,73 +757,61 @@ async def stream_agent(message: ChatMessageCreate):
                             skill_call["credit_event_id"] = payment_event.id
                             skill_call["credit_cost"] = payment_event.total_amount
                             logger.info(
-                                f"[{input.agent_id}] skill payment: {skill_call}"
+                                f"[{user_message.agent_id}] skill payment: {skill_call}"
                             )
                     skill_message_create.skill_calls = skill_calls
                     skill_message = await skill_message_create.save_in_session(session)
                     await session.commit()
                     yield skill_message
-            elif "pre_model_hook" in chunk:
-                pass
-            elif "post_model_hook" in chunk:
-                logger.debug(
-                    f"post_model_hook: {chunk}",
-                    extra={"thread_id": thread_id},
-                )
-                if chunk["post_model_hook"] and "error" in chunk["post_model_hook"]:
+            else:
+                for node_name, update in chunk.items():
                     if (
-                        chunk["post_model_hook"]["error"]
-                        == AgentError.INSUFFICIENT_CREDITS
+                        node_name.endswith("CreditCheckMiddleware.after_model")
+                        and isinstance(update, dict)
+                        and update.get("error") == AgentError.INSUFFICIENT_CREDITS
                     ):
-                        if "messages" in chunk["post_model_hook"]:
-                            msg = chunk["post_model_hook"]["messages"][-1]
-                            content = msg.content
-                            if isinstance(msg.content, list):
-                                # in new version, content item maybe a list
-                                content = msg.content[0]
-                            post_model_message_create = ChatMessageCreate(
-                                id=str(XID()),
-                                agent_id=input.agent_id,
-                                chat_id=input.chat_id,
-                                user_id=input.user_id,
-                                author_id=input.agent_id,
-                                author_type=AuthorType.AGENT,
-                                model=agent.model,
-                                thread_type=input.author_type,
-                                reply_to=input.id,
-                                message=content,
-                                input_tokens=0,
-                                output_tokens=0,
-                                time_cost=this_time - last,
-                            )
-                            last = this_time
-                            if cold_start_cost > 0:
-                                post_model_message_create.cold_start_cost = (
-                                    cold_start_cost
-                                )
-                                cold_start_cost = 0
-                            post_model_message = await post_model_message_create.save()
-                            yield post_model_message
+                        ai_messages = [
+                            message
+                            for message in update.get("messages", [])
+                            if isinstance(message, BaseMessage)
+                        ]
+                        content = ""
+                        if ai_messages:
+                            content = _extract_text_content(ai_messages[-1].content)
+                        post_model_message_create = ChatMessageCreate(
+                            id=str(XID()),
+                            agent_id=user_message.agent_id,
+                            chat_id=user_message.chat_id,
+                            user_id=user_message.user_id,
+                            author_id=user_message.agent_id,
+                            author_type=AuthorType.AGENT,
+                            model=agent.model,
+                            thread_type=user_message.author_type,
+                            reply_to=user_message.id,
+                            message=content,
+                            input_tokens=0,
+                            output_tokens=0,
+                            time_cost=this_time - last,
+                        )
+                        last = this_time
+                        post_model_message = await post_model_message_create.save()
+                        yield post_model_message
+
                         error_message_create = (
                             await ChatMessageCreate.from_system_message(
                                 SystemMessageType.INSUFFICIENT_BALANCE,
-                                agent_id=input.agent_id,
-                                chat_id=input.chat_id,
-                                user_id=input.user_id,
-                                author_id=input.agent_id,
-                                thread_type=input.author_type,
-                                reply_to=input.id,
+                                agent_id=user_message.agent_id,
+                                chat_id=user_message.chat_id,
+                                user_id=user_message.user_id,
+                                author_id=user_message.agent_id,
+                                thread_type=user_message.author_type,
+                                reply_to=user_message.id,
                                 time_cost=0,
                             )
                         )
                         error_message = await error_message_create.save()
                         yield error_message
-            else:
-                error_traceback = traceback.format_exc()
-                logger.error(
-                    f"unexpected message type: {str(chunk)}\n{error_traceback}",
-                    extra={"thread_id": thread_id},
-                )
+                        return
     except SQLAlchemyError as e:
         error_traceback = traceback.format_exc()
         logger.error(
@@ -727,12 +820,12 @@ async def stream_agent(message: ChatMessageCreate):
         )
         error_message_create = await ChatMessageCreate.from_system_message(
             SystemMessageType.AGENT_INTERNAL_ERROR,
-            agent_id=input.agent_id,
-            chat_id=input.chat_id,
-            user_id=input.user_id,
-            author_id=input.agent_id,
-            thread_type=input.author_type,
-            reply_to=input.id,
+            agent_id=user_message.agent_id,
+            chat_id=user_message.chat_id,
+            user_id=user_message.user_id,
+            author_id=user_message.agent_id,
+            thread_type=user_message.author_type,
+            reply_to=user_message.id,
             time_cost=time.perf_counter() - start,
         )
         error_message = await error_message_create.save()
@@ -742,16 +835,16 @@ async def stream_agent(message: ChatMessageCreate):
         error_traceback = traceback.format_exc()
         logger.error(
             f"reached recursion limit: {str(e)}\n{error_traceback}",
-            extra={"thread_id": thread_id, "agent_id": input.agent_id},
+            extra={"thread_id": thread_id, "agent_id": user_message.agent_id},
         )
         error_message_create = await ChatMessageCreate.from_system_message(
             SystemMessageType.STEP_LIMIT_EXCEEDED,
-            agent_id=input.agent_id,
-            chat_id=input.chat_id,
-            user_id=input.user_id,
-            author_id=input.agent_id,
-            thread_type=input.author_type,
-            reply_to=input.id,
+            agent_id=user_message.agent_id,
+            chat_id=user_message.chat_id,
+            user_id=user_message.user_id,
+            author_id=user_message.agent_id,
+            thread_type=user_message.author_type,
+            reply_to=user_message.id,
             time_cost=time.perf_counter() - start,
         )
         error_message = await error_message_create.save()
@@ -761,21 +854,21 @@ async def stream_agent(message: ChatMessageCreate):
         error_traceback = traceback.format_exc()
         logger.error(
             f"failed to execute agent: {str(e)}\n{error_traceback}",
-            extra={"thread_id": thread_id, "agent_id": input.agent_id},
+            extra={"thread_id": thread_id, "agent_id": user_message.agent_id},
         )
         error_message_create = await ChatMessageCreate.from_system_message(
             SystemMessageType.AGENT_INTERNAL_ERROR,
-            agent_id=input.agent_id,
-            chat_id=input.chat_id,
-            user_id=input.user_id,
-            author_id=input.agent_id,
-            thread_type=input.author_type,
-            reply_to=input.id,
+            agent_id=user_message.agent_id,
+            chat_id=user_message.chat_id,
+            user_id=user_message.user_id,
+            author_id=user_message.agent_id,
+            thread_type=user_message.author_type,
+            reply_to=user_message.id,
             time_cost=time.perf_counter() - start,
         )
         error_message = await error_message_create.save()
         yield error_message
-        await clear_thread_memory(input.agent_id, input.chat_id)
+        await clear_thread_memory(user_message.agent_id, user_message.chat_id)
         return
 
 
@@ -830,14 +923,15 @@ async def clean_agent_memory(
     # get the agent from the database
     try:
         if not clean_skill and not clean_agent:
-            raise HTTPException(
+            raise IntentKitAPIError(
                 status_code=400,
-                detail="at least one of skills data or agent memory should be true.",
+                key="InvalidCleanupParameters",
+                message="at least one of skills data or agent memory should be true",
             )
 
         if clean_skill:
             await AgentSkillData.clean_data(agent_id)
-            await ThreadSkillData.clean_data(agent_id, chat_id)
+            await ChatSkillData.clean_data(agent_id, chat_id)
 
         async with get_session() as db:
             if clean_agent:
@@ -879,7 +973,7 @@ async def clean_agent_memory(
     except SQLAlchemyError as e:
         # Handle other SQLAlchemy-related errors
         logger.error(e)
-        raise HTTPException(status_code=500, detail=str(e))
+        raise IntentKitAPIError(status_code=500, key="DatabaseError", message=str(e))
     except Exception as e:
         logger.error("failed to cleanup the agent memory: " + str(e))
         raise e
@@ -888,10 +982,7 @@ async def clean_agent_memory(
 async def thread_stats(agent_id: str, chat_id: str) -> list[BaseMessage]:
     thread_id = f"{agent_id}-{chat_id}"
     stream_config = {"configurable": {"thread_id": thread_id}}
-    is_private = False
-    if chat_id.startswith("owner") or chat_id.startswith("autonomous"):
-        is_private = True
-    executor, _ = await agent_executor(agent_id, is_private)
+    executor, _ = await agent_executor(agent_id)
     snap = await executor.aget_state(stream_config)
     if snap.values and "messages" in snap.values:
         return snap.values["messages"]