Merge pull request 'feat: Select individual tools to disable / enable' (#398) from disable-tool into main

Reviewed-on: https://git.biggo.com/Funmula/dive-mcp-host/pulls/398

Author: LouisFunmula

.github/workflows/ci.yml

@@ -60,6 +60,7 @@ jobs:
         AZURE_OPENAI_ENDPOINT: ${{ secrets.AZURE_OPENAI_ENDPOINT }}
         AZURE_OPENAI_DEPLOYMENT_NAME: ${{ secrets.AZURE_OPENAI_DEPLOYMENT_NAME }}
         AZURE_OPENAI_API_VERSION: ${{ secrets.AZURE_OPENAI_API_VERSION }}
+        OAP_TOKEN: ${{ vars.OAP_TOKEN }}
       run: |
         uv run --extra dev --frozen pytest
 

.gitignore

@@ -44,3 +44,4 @@ db.sqlite-wal
 cache/
 upload/
 cli_config.json
+logs/

dive_mcp_host/cli/cli.py

@@ -3,7 +3,8 @@
 import argparse
 from pathlib import Path
 
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.output_parsers import StrOutputParser
 
 from dive_mcp_host.cli.cli_types import CLIArgs
 from dive_mcp_host.host.conf import HostConfig
@@ -67,6 +68,7 @@ async def run() -> None:
         with Path(args.prompt_file).open("r") as f:
             system_prompt = f.read()
 
+    output_parser = StrOutputParser()
     async with DiveMcpHost(config) as mcp_host:
         print("Waiting for tools to initialize...")
         await mcp_host.tools_initialized_event.wait()
@@ -75,7 +77,15 @@ async def run() -> None:
         current_chat_id = chat.chat_id
         async with chat:
             async for response in chat.query(query, stream_mode="messages"):
-                print(response[0].content, end="")  # type: ignore
+                assert isinstance(response, tuple)
+                msg = response[0]
+                if isinstance(msg, AIMessage):
+                    content = output_parser.invoke(msg)
+                    print(content, end="")
+                    continue
+                print(f"\n\n==== Start Of {type(msg)} ===")
+                print(msg)
+                print(f"==== End Of {type(msg)} ===\n")
 
     print()
     print(f"Chat ID: {current_chat_id}")

dive_mcp_host/env.py

@@ -0,0 +1,5 @@
+from os import getenv
+from pathlib import Path
+
+RESOURCE_DIR = Path(getenv("RESOURCE_DIR", Path.cwd()))
+DIVE_CONFIG_DIR = Path(getenv("DIVE_CONFIG_DIR", Path.cwd()))

dive_mcp_host/host/agents/agent_factory.py

@@ -1,3 +1,4 @@
+from enum import StrEnum
 from typing import Literal, Protocol
 
 from langchain_core.messages import AnyMessage, BaseMessage, HumanMessage
@@ -11,6 +12,16 @@
 from dive_mcp_host.host.prompt import PromptType
 
 
+class ConfigurableKey(StrEnum):
+    """Enum for RunnableConfig.configurable keys."""
+
+    # Thread id is also known as chat_id
+    THREAD_ID = "thread_id"
+    USER_ID = "user_id"
+    MAX_INPUT_TOKENS = "max_input_tokens"
+    OVERSIZE_POLICY = "oversize_policy"
+
+
 # XXX is there any better way to do this?
 class AgentFactory[T: MessagesState](Protocol):
     """A factory for creating agents.
@@ -62,10 +73,10 @@ def create_config(
         """
         return {
             "configurable": {
-                "thread_id": thread_id,
-                "user_id": user_id,
-                "max_input_tokens": max_input_tokens,
-                "oversize_policy": oversize_policy,
+                ConfigurableKey.THREAD_ID: thread_id,
+                ConfigurableKey.USER_ID: user_id,
+                ConfigurableKey.MAX_INPUT_TOKENS: max_input_tokens,
+                ConfigurableKey.OVERSIZE_POLICY: oversize_policy,
             },
             "recursion_limit": 102,
         }

dive_mcp_host/host/agents/chat_agent.py

@@ -13,12 +13,13 @@
     HumanMessage,
     RemoveMessage,
     SystemMessage,
+    ToolCall,
     ToolMessage,
 )
 from langchain_core.messages.utils import count_tokens_approximately, trim_messages
 from langchain_core.prompt_values import ChatPromptValue
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import Runnable, RunnableConfig
+from langchain_core.runnables import Runnable, RunnableConfig, RunnablePassthrough
 from langchain_core.tools import BaseTool
 from langgraph.checkpoint.base import BaseCheckpointSaver, V
 from langgraph.graph import END, StateGraph
@@ -31,16 +32,18 @@
 from pydantic import BaseModel
 
 from dive_mcp_host.host.agents.agent_factory import AgentFactory, initial_messages
+from dive_mcp_host.host.agents.file_in_additional_kwargs import FileMsgConverter
 from dive_mcp_host.host.agents.message_order import tool_call_order
 from dive_mcp_host.host.agents.tools_in_prompt import (
     convert_messages,
     extract_tool_calls,
 )
 from dive_mcp_host.host.helpers import today_datetime
 from dive_mcp_host.host.prompt import PromptType, tools_prompt
+from dive_mcp_host.host.store.base import StoreManagerProtocol
 
-StructuredResponse = dict | BaseModel
-StructuredResponseSchema = dict | type[BaseModel]
+type StructuredResponse = dict | BaseModel
+type StructuredResponseSchema = dict | type[BaseModel]
 
 
 class AgentState(MessagesState):
@@ -92,6 +95,41 @@ def _func(state: AgentState | ChatPromptValue) -> list[BaseMessage]:
     return prompt_runnable
 
 
+class HackedToolNode(ToolNode):
+    """hacked tool node to inject tool_call_id into the config.
+
+    This is a hack. If langgraph support tool_call_id, we will remove this class.
+    """
+
+    async def _arun_one(
+        self,
+        call: ToolCall,
+        input_type: Literal["list", "dict", "tool_calls"],
+        config: RunnableConfig,
+    ) -> ToolMessage:
+        if "metadata" in config:
+            config["metadata"]["tool_call_id"] = call["id"]
+        else:
+            config["metadata"] = {
+                "tool_call_id": call["id"],
+            }
+        return await super()._arun_one(call, input_type, config)
+
+    def _run_one(
+        self,
+        call: ToolCall,
+        input_type: Literal["list", "dict", "tool_calls"],
+        config: RunnableConfig,
+    ) -> ToolMessage:
+        if "metadata" in config:
+            config["metadata"]["tool_call_id"] = call["id"]
+        else:
+            config["metadata"] = {
+                "tool_call_id": call["id"],
+            }
+        return super()._run_one(call, input_type, config)
+
+
 class ChatAgentFactory(AgentFactory[AgentState]):
     """A factory for ChatAgents."""
 
@@ -100,6 +138,7 @@ def __init__(
         model: BaseChatModel,
         tools: Sequence[BaseTool] | ToolNode,
         tools_in_prompt: bool = False,
+        store: StoreManagerProtocol | None = None,
     ) -> None:
         """Initialize the chat agent factory."""
         self._model = model
@@ -110,6 +149,12 @@ def __init__(
             StructuredResponseSchema | tuple[str, StructuredResponseSchema] | None
         ) = None
 
+        self._file_msg_converter = (
+            FileMsgConverter(model_provider=self._model_class, store=store).runnable
+            if store
+            else RunnablePassthrough()
+        )
+
         # changed when self._build_graph is called
         self._tool_classes: list[BaseTool] = []
         self._should_return_direct: set[str] = set()
@@ -160,22 +205,28 @@ def _check_more_steps_needed(
             )
         )
 
-    def _call_model(self, state: AgentState, config: RunnableConfig) -> AgentState:
+    async def _call_model(
+        self, state: AgentState, config: RunnableConfig
+    ) -> AgentState:
+        # TODO: _validate_chat_history
         if not self._tools_in_prompt:
             model = self._model
             if self._tool_classes:
                 model = self._model.bind_tools(self._tool_classes)
-            model_runnable = self._prompt | drop_empty_messages | model
+            model_runnable = (
+                self._prompt | self._file_msg_converter | drop_empty_messages | model
+            )
         else:
             model_runnable = (
                 self._prompt
                 | self._tool_prompt
                 | convert_messages
+                | self._file_msg_converter
                 | drop_empty_messages
                 | self._model
             )
 
-        response = model_runnable.invoke(state, config)
+        response = await model_runnable.ainvoke(state, config)
         if isinstance(response, AIMessage):
             response = extract_tool_calls(response)
         if self._check_more_steps_needed(state, response):
@@ -254,7 +305,9 @@ def _build_graph(self) -> None:
         graph.add_edge("before_agent", "agent")
 
         tool_node = (
-            self._tools if isinstance(self._tools, ToolNode) else ToolNode(self._tools)
+            self._tools
+            if isinstance(self._tools, ToolNode)
+            else HackedToolNode(self._tools)
         )
         self._tool_classes = list(tool_node.tools_by_name.values())
         graph.add_node("tools", tool_node)
@@ -323,9 +376,15 @@ def get_chat_agent_factory(
     model: BaseChatModel,
     tools: Sequence[BaseTool] | ToolNode,
     tools_in_prompt: bool = False,
+    store: StoreManagerProtocol | None = None,
 ) -> ChatAgentFactory:
     """Get an agent factory."""
-    return ChatAgentFactory(model, tools, tools_in_prompt)
+    return ChatAgentFactory(
+        model=model,
+        tools=tools,
+        tools_in_prompt=tools_in_prompt,
+        store=store,
+    )
 
 
 @RunnableCallable