From 2e1a18503b9da4c5cee86c188483ce39e6701811 Mon Sep 17 00:00:00 2001
From: balalo <balalo.fernandez@gmail.com>
Date: Tue, 18 Mar 2025 10:41:43 +0100
Subject: [PATCH] Set tool choice from context aggregator

---
 src/pipecat/frames/frames.py                   |  6 ++++++
 .../processors/aggregators/llm_response.py     | 18 +++++++++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/pipecat/frames/frames.py b/src/pipecat/frames/frames.py
index 74dd2accb..aec433033 100644
--- a/src/pipecat/frames/frames.py
+++ b/src/pipecat/frames/frames.py
@@ -362,6 +362,12 @@ class LLMSetToolsFrame(DataFrame):
 
     tools: List[dict]
 
+@dataclass
+class LLMSetToolChoiceFrame(DataFrame):
+    """A frame containing a tool choice for an LLM to use for function calling."""
+
+    tool_choice: Literal["none", "auto", "required"]
+
 
 @dataclass
 class LLMEnablePromptCachingFrame(DataFrame):
diff --git a/src/pipecat/processors/aggregators/llm_response.py b/src/pipecat/processors/aggregators/llm_response.py
index d8582e32f..620e34b2c 100644
--- a/src/pipecat/processors/aggregators/llm_response.py
+++ b/src/pipecat/processors/aggregators/llm_response.py
@@ -7,7 +7,7 @@
 import asyncio
 import time
 from abc import abstractmethod
-from typing import List
+from typing import List, Literal
 
 from pipecat.frames.frames import (
     CancelFrame,
@@ -22,6 +22,7 @@ from pipecat.frames.frames import (
     LLMMessagesFrame,
     LLMMessagesUpdateFrame,
     LLMSetToolsFrame,
+    LLMSetToolChoiceFrame,
     LLMTextFrame,
     StartFrame,
     StartInterruptionFrame,
@@ -132,6 +133,11 @@ class BaseLLMResponseAggregator(FrameProcessor):
         """Set LLM tools to be used in the current conversation."""
         pass
 
+    @abstractmethod
+    def set_tool_choice(self, tool_choice):
+        """Set the tool choice. This should modify the LLM context."""
+        pass
+
     @abstractmethod
     def reset(self):
         """Reset the internals of this aggregator. This should not modify the
@@ -185,6 +191,9 @@ class LLMResponseAggregator(BaseLLMResponseAggregator):
     def set_tools(self, tools):
         pass
 
+    def set_tool_choice(self, tool_choice):
+        pass
+
     def reset(self):
         self._aggregation = ""
 
@@ -244,6 +253,9 @@ class LLMContextResponseAggregator(BaseLLMResponseAggregator):
     def set_tools(self, tools: List):
         self._context.set_tools(tools)
 
+    def set_tool_choice(self, tool_choice: Literal["none", "auto", "required"]):
+        self._context.set_tool_choice(tool_choice)
+
     def reset(self):
         self._aggregation = ""
 
@@ -328,6 +340,8 @@ class LLMUserContextAggregator(LLMContextResponseAggregator):
             self.set_messages(frame.messages)
         elif isinstance(frame, LLMSetToolsFrame):
             self.set_tools(frame.tools)
+        elif isinstance(frame, LLMSetToolChoiceFrame):
+            self.set_tool_choice(frame.tool_choice)
         else:
             await self.push_frame(frame, direction)
 
@@ -448,6 +462,8 @@ class LLMAssistantContextAggregator(LLMContextResponseAggregator):
             self.set_messages(frame.messages)
         elif isinstance(frame, LLMSetToolsFrame):
             self.set_tools(frame.tools)
+        elif isinstance(frame, LLMSetToolChoiceFrame):
+            self.set_tool_choice(frame.tool_choice)
         else:
             await self.push_frame(frame, direction)