Update api/utils.py

api/utils.py

@@ -147,6 +147,7 @@ def create_chat_completion_data(
     prompt_tokens: int = 0,
     completion_tokens: int = 0,
     finish_reason: Optional[str] = None,
+    function_call: Optional[Dict] = None,
 ) -> Dict[str, Any]:
     usage = None
     if finish_reason == "stop":
@@ -161,15 +162,28 @@ def create_chat_completion_data(
         "created": timestamp,
         "model": model,
         "system_fingerprint": system_fingerprint,
-        "choices": [{"index": 0, "delta": {"content": content, "role": "assistant"}, "finish_reason": finish_reason}],
+        "choices": [{
+            "index": 0,
+            "delta": {
+                "content": content if not function_call else None,
+                "role": "assistant",
+                "function_call": function_call
+            },
+            "finish_reason": finish_reason
+        }],
         "usage": usage,
     }
 
-def message_to_dict(message, model_prefix: Optional[str] = None):
+async def handle_function_call(request: ChatRequest, function_call_details: Dict) -> Dict[str, Any]:
+    # Placeholder for function calling logic
+    logger.info(f"Handling function call for model: {request.model}")
+    logger.info(f"Function call details: {function_call_details}")
+    return {"message": "Function call handled successfully", "details": function_call_details}
+
+def message_to_dict(message, model_prefix: Optional[str] = None, tools: Optional[List[Dict]] = None) -> Dict[str, Any]:
     """
     Convert a ChatRequest message to a dict for the request payload.
-    Supports up to three images with type-based structure and sends multiple formats.
-    Prepends model_prefix to text content if specified.
+    Supports function calling, images, and model prefixes.
     """
     content = ""
     images_data = []
@@ -242,13 +256,12 @@ async def process_streaming_response(request: ChatRequest):
         logger.error("No h-value for validation.")
         raise HTTPException(status_code=500, detail="Missing h-value.")
 
-    messages = [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages]
+    messages = [message_to_dict(msg, model_prefix=model_prefix, tools=request.tools) for msg in request.messages]
 
     json_data = {
         "agentMode": agent_mode,
         "clickedAnswer2": False,
         "clickedAnswer3": False,
-        "reasoningMode": False,
         "clickedForceWebSearch": False,
         "codeInterpreterMode": False,
         "codeModelMode": True,
@@ -275,15 +288,16 @@ async def process_streaming_response(request: ChatRequest):
         "visitFromDelta": False,
         "webSearchModePrompt": False,
         "vscodeClient": False,
-        "designerMode": False,
-        "workspaceId": "",
-        "beastMode": False,
         "customProfile": {"name": "", "occupation": "", "traits": [], "additionalInfo": "", "enableNewChats": False},
         "webSearchModeOption": {"autoMode": False, "webMode": False, "offlineMode": True},
         "session": {
-            "user": {"name": random_name, "email": "[email protected]", "image": "https://lh3.googleusercontent.com/a/...=s96-c", "subscriptionStatus": "PREMIUM"},
+            "user": {"name": random_name, "email": random_email, "image": "https://lh3.googleusercontent.com/a/...=s96-c", "subscriptionStatus": "PREMIUM"},
             "expires": datetime.now(timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
             "subscriptionCache": {"customerId": random_customer_id, "status": "PREMIUM", "isTrialSubscription": "False", "expiryTimestamp": 1744652408, "lastChecked": int(time.time() * 1000)},
+            "beastMode": False,
+            "reasoningMode": False,
+            "designerMode": False,
+            "workspaceId": "",
         },
     }
 
@@ -320,7 +334,27 @@ async def process_streaming_response(request: ChatRequest):
                         final_snapzion_links.extend(snapzion_urls)
                     cleaned_content = strip_model_prefix(chunk, model_prefix)
                     completion_tokens += calculate_tokens(cleaned_content, request.model)
-                    yield "data: " + json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp, request_id, system_fingerprint, prompt_tokens, completion_tokens)) + "\n\n"
+                    
+                    # Handle function call responses
+                    function_call = None
+                    if cleaned_content and cleaned_content.startswith("{"):
+                        try:
+                            function_call = json.loads(cleaned_content)
+                            cleaned_content = None  # Content must be null for function calls
+                        except json.JSONDecodeError:
+                            pass
+                    
+                    yield "data: " + json.dumps(create_chat_completion_data(
+                        cleaned_content, 
+                        request.model,
+                        timestamp,
+                        request_id,
+                        system_fingerprint,
+                        prompt_tokens,
+                        completion_tokens,
+                        finish_reason=None,
+                        function_call=function_call
+                    )) + "\n\n"
                 yield "data: " + json.dumps(create_chat_completion_data("", request.model, timestamp, request_id, system_fingerprint, prompt_tokens, completion_tokens, "stop")) + "\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e:
@@ -374,13 +408,12 @@ async def process_non_streaming_response(request: ChatRequest):
         logger.error("Failed to retrieve h-value.")
         raise HTTPException(status_code=500, detail="Missing h-value.")
 
-    messages = [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages]
+    messages = [message_to_dict(msg, model_prefix=model_prefix, tools=request.tools) for msg in request.messages]
 
     json_data = {
         "agentMode": agent_mode,
         "clickedAnswer2": False,
         "clickedAnswer3": False,
-        "reasoningMode": False,
         "clickedForceWebSearch": False,
         "codeInterpreterMode": False,
         "codeModelMode": True,
@@ -407,15 +440,16 @@ async def process_non_streaming_response(request: ChatRequest):
         "visitFromDelta": False,
         "webSearchModePrompt": False,
         "vscodeClient": False,
-        "designerMode": False,
-        "workspaceId": "",
-        "beastMode": False,
         "customProfile": {"name": "", "occupation": "", "traits": [], "additionalInfo": "", "enableNewChats": False},
         "webSearchModeOption": {"autoMode": False, "webMode": False, "offlineMode": True},
         "session": {
-            "user": {"name": random_name, "email": "[email protected]", "image": "https://lh3.googleusercontent.com/a/...=s96-c", "subscriptionStatus": "PREMIUM"},
+            "user": {"name": random_name, "email": random_email, "image": "https://lh3.googleusercontent.com/a/...=s96-c", "subscriptionStatus": "PREMIUM"},
             "expires": datetime.now(timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
             "subscriptionCache": {"customerId": random_customer_id, "status": "PREMIUM", "isTrialSubscription": "False", "expiryTimestamp": 1744652408, "lastChecked": int(time.time() * 1000)},
+            "beastMode": False,
+            "reasoningMode": False,
+            "designerMode": False,
+            "workspaceId": "",
         },
     }