From 5d73db53a0110303f8e226523e4f0198efb603a1 Mon Sep 17 00:00:00 2001
From: Chad Bailey <chadbailey@gmail.com>
Date: Sun, 21 Jul 2024 01:15:44 +0000
Subject: [PATCH] initial pseudo function calling

---
 src/pipecat/processors/frameworks/rtvi.py | 97 ++++++++++++++++++++++-
 1 file changed, 96 insertions(+), 1 deletion(-)

diff --git a/src/pipecat/processors/frameworks/rtvi.py b/src/pipecat/processors/frameworks/rtvi.py
index 376a4f034..663359ad5 100644
--- a/src/pipecat/processors/frameworks/rtvi.py
+++ b/src/pipecat/processors/frameworks/rtvi.py
@@ -14,6 +14,8 @@ from pipecat.frames.frames import (
     BotInterruptionFrame,
     Frame,
     InterimTranscriptionFrame,
+    LLMFullResponseEndFrame,
+    LLMFullResponseStartFrame,
     LLMMessagesAppendFrame,
     LLMMessagesUpdateFrame,
     LLMModelUpdateFrame,
@@ -21,6 +23,7 @@ from pipecat.frames.frames import (
     SystemFrame,
     TTSSpeakFrame,
     TTSVoiceUpdateFrame,
+    TextFrame,
     TranscriptionFrame,
     TransportMessageFrame,
     UserStartedSpeakingFrame,
@@ -31,7 +34,7 @@ from pipecat.processors.aggregators.llm_response import (
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
 from pipecat.services.ai_services import AIService
 from pipecat.services.cartesia import CartesiaTTSService
-from pipecat.services.openai import OpenAILLMService
+from pipecat.services.openai import OpenAILLMService, OpenAILLMContext, OpenAILLMContextFrame
 from pipecat.transports.base_transport import BaseTransport
 
 DEFAULT_MESSAGES = [
@@ -152,6 +155,92 @@ class RTVIUserStoppedSpeakingMessage(BaseModel):
     type: Literal["user-stopped-speaking"] = "user-stopped-speaking"
 
 
+class RTVIJSONCompletion(BaseModel):
+    label: Literal["realtime-ai"] = "realtime-ai"
+    type: Literal["json-completion"] = "json-completion"
+    data: str
+
+
+class FunctionCaller(FrameProcessor):
+    def __init__(self, context):
+        super().__init__()
+        self._checking = False
+        self._aggregating = False
+        self._emitted_start = False
+        self._aggregation = ""
+        self._context = context
+
+        self._callbacks = {}
+        self._start_callbacks = {}
+
+    def register_function(self, function_name: str, callback, start_callback=None):
+        self._callbacks[function_name] = callback
+        if start_callback:
+            self._start_callbacks[function_name] = start_callback
+
+    def unregister_function(self, function_name: str):
+        del self._callbacks[function_name]
+        if self._start_callbacks[function_name]:
+            del self._start_callbacks[function_name]
+
+    def has_function(self, function_name: str):
+        return function_name in self._callbacks.keys()
+
+    async def call_function(self, function_name: str, args):
+        if function_name in self._callbacks.keys():
+            return await self._callbacks[function_name](self, args)
+        return None
+
+    async def call_start_function(self, function_name: str):
+        if function_name in self._start_callbacks.keys():
+            await self._start_callbacks[function_name](self)
+
+    async def process_frame(self, frame: Frame, direction: FrameDirection):
+        if isinstance(frame, LLMFullResponseStartFrame):
+            self._checking = True
+            await self.push_frame(frame, direction)
+        elif isinstance(frame, TextFrame) and self._checking:
+            # TODO-CB: should we expand this to any non-text character to start the completion?
+            if frame.text.strip().startswith("{") or frame.text.strip().startswith("```"):
+                self._emitted_start = False
+                self._checking = False
+                self._aggregation = frame.text
+                self._aggregating = True
+            else:
+                self._checking = False
+                self._aggregating = False
+                self._aggregation = ""
+                self._emitted_start = False
+                await self.push_frame(frame, direction)
+        elif isinstance(frame, TextFrame) and self._aggregating:
+            self._aggregation += frame.text
+            # TODO-CB: We can probably ignore function start I think
+            # if not self._emitted_start:
+            #     fn = re.search(r'{"function_name":\s*"(.*)",', self._aggregation)
+            #     if fn and fn.group(1):
+            #         await self.call_start_function(fn.group(1))
+            #         self._emitted_start = True
+        elif isinstance(frame, LLMFullResponseEndFrame) and self._aggregating:
+            try:
+                self._aggregation = self._aggregation.replace("```json", "").replace("```", "")
+                self._context.add_message({"role": "assistant", "content": self._aggregation})
+                message = RTVIJSONCompletion(data=self._aggregation)
+                msg = message.model_dump(exclude_none=True)
+                await self.push_frame(TransportMessageFrame(message=msg))
+
+            except Exception as e:
+                print(f"Error parsing function call json: {e}")
+                print(f"aggregation was: {self._aggregation}")
+
+            self._aggregating = False
+            self._aggregation = ""
+            self._emitted_start = False
+        elif isinstance(frame, LLMFullResponseEndFrame):
+            await self.push_frame(frame, direction)
+        else:
+            await self.push_frame(frame, direction)
+
+
 class RTVIProcessor(FrameProcessor):
 
     def __init__(
@@ -311,9 +400,15 @@ class RTVIProcessor(FrameProcessor):
 
         self._tts = self._tts_cls(api_key=self._tts_api_key, voice_id=voice)
 
+        # TODO-CB: Eventually we'll need to switch the context aggregators to use the
+        # OpenAI context frames instead of message frames
+        context = OpenAILLMContext(messages=messages)
+        self._fc = FunctionCaller(context)
+
         pipeline = Pipeline([
             self._tma_in,
             self._llm,
+            self._fc,
             self._tts,
             self._tma_out,
             self._transport.output(),