Feat: Add tool call caching

api/src/main/models/base.py

@@ -7,6 +7,8 @@
 from ..backend import BackendType
 from ..utils import FunctionCalling, FunctionCallResult
 
+import json
+
 
 @dataclass
 class Tags:
@@ -60,12 +62,13 @@ def clean_up(self):
         self.__class__.__instance = None
 
     def parse_tool_calling(
-        self,
-        outputs,
-        chat_history: ChatHistory,
-        tools: List[Dict[str, str]],
-        stream: bool = True,
-        print_output: bool = False
+            self,
+            outputs,
+            chat_history: ChatHistory,
+            tools: List[Dict[str, str]],
+            stream: bool = True,
+            print_output: bool = False,
+            tool_call_caches: Optional[dict[str, str]] = None
     ) -> Union[Generator[str, None, None], str]:
         """ Parse tool calling from the model's output """
         result_obj = FunctionCallResult()
@@ -78,12 +81,13 @@ def parse_tool_calling(
                 if self.special_tags.TOOLCALL in word:  # Start of a tool call
                     started = True
                     if buffer:
-                        buffer = 0
+                        buffer = ""
                 elif self.special_tags.TOOLCALL_END in word:  # End of a tool call
                     if buffer:
                         result_obj.stage(
                             buffer,
-                            (self.special_tags.TOOLCALL, self.special_tags.TOOLCALL_END)
+                            (self.special_tags.TOOLCALL, self.special_tags.TOOLCALL_END),
+                            tool_call_caches
                         )
                         state = result_obj.state
                         if state is not None:
@@ -119,8 +123,10 @@ def parse_tool_calling(
             final_result = result_obj.finalize(
                 chat_history,
                 (self.special_tags.TOOLCALL, self.special_tags.TOOLCALL_END),
-                print_output=print_output
+                print_output=print_output,
+                tool_call_caches=tool_call_caches
             )
+
             if queued > 0 and stat % 2 == 0:
                 print(f"\r{spinner[(stat//2) % len(spinner)]} Waiting for tool calls to finish...", end="", flush=True)
             if final_result is False:
@@ -136,21 +142,22 @@ def parse_tool_calling(
                 return outputs + final_result
 
     def chat(
-        self,
-        chat_history: ChatHistory,
-        user_prompt: str,
-        system_prompt: str = "",
-        tools: Optional[List[Dict[str, str]]] = None,
-        temperature: float = 0.2,
-        top_p: float = 0.95,
-        top_k: int = 40,
-        min_p: float = 0.05,
-        typical_p: float = 1.0,
-        stream: bool = True,
-        max_new_tokens: int = 1024,
-        repeat_penalty: float = 1.0,
-        print_output: bool = False,
-        **kwargs
+            self,
+            chat_history: ChatHistory,
+            user_prompt: str,
+            system_prompt: str = "",
+            tools: Optional[List[Dict[str, str]]] = None,
+            temperature: float = 0.2,
+            top_p: float = 0.95,
+            top_k: int = 40,
+            min_p: float = 0.05,
+            typical_p: float = 1.0,
+            stream: bool = True,
+            max_new_tokens: int = 1024,
+            repeat_penalty: float = 1.0,
+            print_output: bool = False,
+            tool_call_caches: Optional[dict[str, str]] = None,
+            **kwargs
     ) -> Union[Generator[str, None, None], str]:
         """ Process a chat request
 
@@ -229,7 +236,8 @@ def adaptive_special_tag_buffering(outs, wait_tokens_for=6):
                 chat_history=chat_history,
                 tools=tools,
                 stream=stream,
-                print_output=print_output
+                print_output=print_output,
+                tool_call_caches= tool_call_caches
             )
 
             if stream:

api/src/main/models/llama3/model.py

@@ -47,21 +47,22 @@ def _get_runtime(self, backend: BackendType | None = None):
         )
 
     def chat(
-        self,
-        chat_history: ChatHistory,
-        user_prompt: str,
-        system_prompt: str = system_prompt,
-        tools: Optional[List[Dict[str, str]]] = None,
-        temperature: float = 0.2,
-        top_p: float = 0.95,
-        top_k: int = 40,
-        min_p: float = 0.05,
-        typical_p: float = 1.0,
-        stream: bool = True,
-        max_new_tokens: int = 0,
-        repeat_penalty: float = 1.0,
-        print_output: bool = False,
-        **kwargs
+            self,
+            chat_history: ChatHistory,
+            user_prompt: str,
+            system_prompt: str = system_prompt,
+            tools: Optional[List[Dict[str, str]]] = None,
+            temperature: float = 0.2,
+            top_p: float = 0.95,
+            top_k: int = 40,
+            min_p: float = 0.05,
+            typical_p: float = 1.0,
+            stream: bool = True,
+            max_new_tokens: int = 0,
+            repeat_penalty: float = 1.0,
+            print_output: bool = False,
+            tool_call_caches: Optional[dict[str, str]] = None,
+            **kwargs
     ) -> Union[Generator[str, None, None], str]:
         return super().chat(
             chat_history=chat_history,
@@ -77,5 +78,6 @@ def chat(
             max_new_tokens=max_new_tokens,
             repeat_penalty=repeat_penalty,
             print_output=print_output,
+            tool_call_caches=tool_call_caches,
             **kwargs
         )

api/src/main/models/qwen3/model.py

@@ -37,6 +37,14 @@
 - When interpreting relative time expressions, always use calendar week boundaries (Monday-Sunday), not rolling periods from today
 - Always prioritize accuracy over speed
 
+**[Cache-First Principle]**
+The primary goal of tool use is to maximize efficiency by actively utilizing cached data. Before invoking a new tool, your thought process **MUST** follow the steps below.
+1.  **Relevance Check:** First, check if there is a previous tool call in the conversation history that could be helpful in answering the user's current query. (e.g., same tool, similar parameters)
+2.  **Retrieve Cache Data:** If a relevant record exists, this is the **mandatory first step**. You must immediately call `get_cache_data()` with the `cache_id` to retrieve the full data.
+3.  **Data Analysis:** Analyze the cached data you have retrieved. Is it **sufficient to fully answer** the user's current question?
+4.  **Utilize Cache:** If the data is sufficient, use the cached results to generate your response. **Calling the original tool again is absolutely forbidden.**
+5.  **New Tool Call:** You should only proceed with a new tool call if there is no relevant cache, or if the cached content is irrelevant or insufficient for the question.
+
 Remember: Your role is to be a reliable, knowledgeable professional assistant who thinks carefully before responding and actively seeks current information when needed."""
 print("INFO:     Use default system prompt -", system_prompt)
 
@@ -88,6 +96,7 @@ def chat(
         max_new_tokens: int = 0,
         repeat_penalty: float = 1.0,
         print_output: bool = False,
+        tool_call_caches: Optional[dict[str, str]] = None,
         **kwargs
     ) -> Union[Generator[str, None, None], str]:
         return super().chat(
@@ -108,5 +117,6 @@ def chat(
             max_new_tokens=max_new_tokens,
             repeat_penalty=repeat_penalty,
             print_output=print_output,
+            tool_call_caches=tool_call_caches,
             **kwargs
         )

api/src/main/server.py

@@ -94,7 +94,11 @@ async def chat(request: Request, user_prompt: str, history: Optional[List[Messag
         if not session_id:
             raise HTTPException(status_code=400, detail="Session ID is required.")
 
-        model = Session(session_id=session_id).model
+        session = Session(session_id=session_id)
+
+        model = session.model
+        tool_call_caches = session.tool_call_caches
+
     except ValueError:
         traceback.print_exc()
         raise HTTPException(status_code=404, detail="The session is not found.")
@@ -103,7 +107,7 @@ async def chat(request: Request, user_prompt: str, history: Optional[List[Messag
     if history:
         chat_history.extend([h.model_dump() for h in history])
 
-    response = model.chat(chat_history, user_prompt, stream=False, print_output=True)
+    response = model.chat(chat_history, user_prompt, tool_call_caches=tool_call_caches, stream=False, print_output=True)
     del model
     return response
 
@@ -115,7 +119,13 @@ async def chat_with_streaming(websocket: WebSocket):
 
     try:
         session_id = json.loads(await websocket.receive_text()).get("session_id")
-        model = Session(session_id=session_id).model
+        session = Session(session_id=session_id)
+        model = session.model
+
+        tool_call_caches = session.tool_call_caches
+
+        print(f"Cache ID in WebSocket: {id(session.tool_call_caches)}")
+
     except Exception:
         traceback.print_exc()
         await websocket.close(code=1008, reason="Invalid session ID or model not found.")
@@ -125,7 +135,7 @@ async def chat_with_streaming(websocket: WebSocket):
     chat_history.extend(json.loads(await websocket.receive_text()))
     user_prompt = await websocket.receive_text()
 
-    for token in model.chat(chat_history, user_prompt, print_output=True):
+    for token in model.chat(chat_history, user_prompt, tool_call_caches=tool_call_caches, print_output=True):
         await websocket.send_text(token)
         await asyncio.sleep(0.0001)  # 0.1ms delay between tokens
 

api/src/main/settings.py

@@ -52,6 +52,7 @@ def __init__(self, model_id: str = None, session_id: str = None):
             self.__sessions[self.session_id] = self
             print("INFO:     Current sessions:", list(self.__sessions))
             self._model = None
+            self.tool_call_caches = {}
 
     @property
     def model(self):

api/src/main/utils/__init__.py

@@ -12,6 +12,7 @@
 from . import currency
 from . import calculator
 from . import web_search
+from . import cache
 #from . import embedding
 
 
@@ -97,11 +98,12 @@ def state(self):
                 return "\n" + "\n".join(queue)
 
     def finalize(
-        self,
-        history_list: list,
-        tag: tuple[str, str] = ("<tool_call>", "</tool_call>"),
-        print_output: bool = False
-    ) -> Union[str, False]:
+            self,
+            history_list: list,
+            tag: tuple[str, str] = ("<tool_call>", "</tool_call>"),
+            print_output: bool = False,
+            tool_call_caches: dict[str, dict] = None
+    ) -> Union[str, bool]:
         # Check if there are any pending tool calls
         with self.__queue_mutex:
             if len(self.job_list) != self.__completed_jobs:
@@ -117,12 +119,18 @@ def finalize(
 
                 # Dump client-side tool call history
                 state = ""
+                caches : dict[str, str] = {}
                 if self.__message_queue:
                     state = "\n" + "\n".join(self.__message_queue)
                 for data in self[1:]:
                     if 'tool_call_id' in data:
+                        caches[data['tool_call_id']] = data['content']
                         data['content'] = f"<cached_result:{data['tool_call_id']}>"
-                result = state + "\n" + tag[0] + "\n" + dumps(dict(history=self, ensure_ascii=False)) + "\n" + tag[1]
+
+                if tool_call_caches is not None and caches:
+                    tool_call_caches.update(caches)
+
+                result = state + "\n" + tag[0] + "\n" + dumps(dict(history=self), ensure_ascii=False) + "\n" + tag[1]
 
                 # Clear the job list and message queue
                 self.job_list = []
@@ -131,7 +139,7 @@ def finalize(
                 self.__message_queue = []
                 return result
 
-    def stage(self, calling: str, tag: tuple[str, str] = ("<tool_call>", "</tool_call>")):
+    def stage(self, calling: str, tag: tuple[str, str] = ("<tool_call>", "</tool_call>"), tool_call_caches: dict[str, str] = None):
         with self.__queue_mutex:
             job_id = datetime.now().strftime("call_%Y%m%d%H%M%S")
             try:
@@ -148,20 +156,26 @@ def stage(self, calling: str, tag: tuple[str, str] = ("<tool_call>", "</tool_cal
                 name=name, arguments=deepcopy(arguments)
             ))), ensure_ascii=False) + "\n" + tag[1])
 
-            self.__thread_pool.submit(self.do, job_id, name, arguments, tag)
+            self.__thread_pool.submit(self.do, job_id, name, arguments, tag, tool_call_caches)
 
     def do(
-        self,
-        job_id: str,
-        name: str,
-        arguments: dict,
-        tag: tuple[str, str] = ("<tool_call>", "</tool_call>")
+            self,
+            job_id: str,
+            name: str,
+            arguments: dict,
+            tag: tuple[str, str] = ("<tool_call>", "</tool_call>"),
+            tool_call_caches: dict[str, str] = None
     ):
         # Execute the function
         try:
             if name not in self.implementations:
                 raise ValueError(f"Function '{name}' is not registered.")
-            result = self.implementations[name](**arguments)
+
+            if name in ["get_cache_data"]:
+                result = self.implementations[name](**arguments, tool_call_caches=tool_call_caches)
+            else:
+                result = self.implementations[name](**arguments)
+
         except Exception as e:
             result = str(e)
 
@@ -375,6 +389,20 @@ def do(
                 },
                 "required": ["url"]
             }
+        ),
+        FunctionSchema(
+            name="get_cache_data",
+            description="Get results for previous tool_call",
+            parameters={
+                "type": "object",
+                "properties": {
+                    "tool_call_cache_id": {
+                        "type": "string",
+                        "description": "The id of the previous tool_call to retrieve from cache"
+                    }
+                },
+                "required": ["tool_call_cache_id"]
+            }
         )
     ],
 
@@ -387,6 +415,7 @@ def do(
         calculate=calculator.calculate,
         search_web=web_search.search_web,
         search_website=web_search.search_website,
-        fetch_webpage=web_search.fetch_webpage
+        fetch_webpage=web_search.fetch_webpage,
+        get_cache_data=cache.get_cache_data
     )
 )

api/src/main/utils/cache.py

@@ -0,0 +1,12 @@
+import json
+import threading
+import datetime
+from typing import Union
+
+
+def get_cache_data(tool_call_cache_id: str, tool_call_caches: dict[str, str]) -> str:  # 함수명 오타 수정
+    """ Get a specific cache data by name """
+    if tool_call_caches and tool_call_cache_id in tool_call_caches:
+        return json.dumps(tool_call_caches[tool_call_cache_id], default=str, ensure_ascii=False)
+    else:
+        return f"Cache '{tool_call_cache_id}' not found"