Update main.py

main.py

@@ -4,8 +4,7 @@ import os
 import secrets
 import string
 import time
-from typing import List, Optional, Union
-
+from typing import List, Optional, Union, Any
 import httpx
 from dotenv import load_dotenv
 from fastapi import FastAPI
@@ -13,9 +12,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 
 # --- Configuration ---
-
 load_dotenv()
-
 # Env variables for external services
 IMAGE_API_URL = os.environ.get("IMAGE_API_URL", "https://image.api.example.com")
 SNAPZION_UPLOAD_URL = "https://upload.snapzion.com/api/public-upload"
@@ -30,18 +27,15 @@ AVAILABLE_MODELS = [
     {"id": "dall-e-3", "object": "model", "created": int(time.time()), "owned_by": "system"},
     {"id": "text-moderation-stable", "object": "model", "created": int(time.time()), "owned_by": "system"},
 ]
-
 MODEL_ALIASES = {}
 
 # --- FastAPI Application ---
-
 app = FastAPI(
     title="OpenAI Compatible API",
     description="An adapter for various services to be compatible with the OpenAI API specification.",
     version="1.0.0"
 )
 
-
 # --- Helper Function for Random ID Generation ---
 def generate_random_id(prefix: str, length: int = 29) -> str:
     """
@@ -51,17 +45,13 @@ def generate_random_id(prefix: str, length: int = 29) -> str:
     random_part = "".join(secrets.choice(population) for _ in range(length))
     return f"{prefix}{random_part}"
 
-
 # === API Endpoints ===
-
 @app.get("/v1/models")
 async def list_models():
     """Lists the available models."""
     return {"object": "list", "data": AVAILABLE_MODELS}
 
-
 # === Chat Completion ===
-
 class Message(BaseModel):
     role: str
     content: str
@@ -70,6 +60,7 @@ class ChatRequest(BaseModel):
     messages: List[Message]
     model: str
     stream: Optional[bool] = False
+    tools: Optional[Any] = None
 
 @app.post("/v1/chat/completions")
 async def chat_completion(request: ChatRequest):
@@ -85,17 +76,44 @@ async def chat_completion(request: ChatRequest):
         'referer': 'https://www.chatwithmono.xyz/',
         'user-agent': 'Mozilla/5.0',
     }
+    if request.tools:
+        # Handle tool by giving in system prompt.
+        # Tool call must be encoded in <tool_call><tool_call> XML tag.
+        tool_prompt = f"""You have access to the following tools . To call a tool, please respond with JSON for a tool call within <tool_call><tool_call> XML tag. Respond in the format {{"name": tool name, "parameters": dictionary of argument name and its value}}. Do not use variables.
+Tools:
+{";".join(f"<tool>{tool}</tool>" for tool in request.tools)}
+
+Response Format for tool call:
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{{"name": <function-name>, "arguments": <args-json-object>}}
+</tool_call>
+
+Example of tool calling:
+<tool_call>
+{{"name": "get_weather", "parameters": {{"city": "New York"}}}}
+</tool_call>
+
+Using tools is recommended.
+        """
+        if request.messages[0].role == "system":
+            request.messages[0].content += "\n\n" + tool_prompt
+        else:
+            request.messages.insert(0, {"role": "system", "content": tool_prompt})
+    request_data = request.model_dump(exclude_unset=True)
+
     payload = {
-        "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
+        "messages": request_data["messages"],
         "model": model_id
     }
-
     if request.stream:
         async def event_stream():
             created = int(time.time())
             is_first_chunk = True
             usage_info = None
-
+            is_tool_call = False
+            chunks_buffer = []
+            max_initial_chunks = 4  # Number of initial chunks to buffer
             try:
                 async with httpx.AsyncClient(timeout=120) as client:
                     async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
@@ -105,41 +123,114 @@ async def chat_completion(request: ChatRequest):
                             if line.startswith("0:"):
                                 try:
                                     content_piece = json.loads(line[2:])
-                                    delta = {"content": content_piece, "function_call": None, "tool_calls": None}
-                                    if is_first_chunk:
-                                        delta["role"] = "assistant"
-                                        is_first_chunk = False
-                                    chunk_data = {
-                                        "id": chat_id, "object": "chat.completion.chunk", "created": created,
-                                        "model": model_id,
-                                        "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
-                                        "usage": None
-                                    }
-                                    yield f"data: {json.dumps(chunk_data)}\n\n"
+                                    # print(content_piece)
+                                    # Buffer the first few chunks
+                                    if len(chunks_buffer) < max_initial_chunks:
+                                        chunks_buffer.append(content_piece)
+                                        continue
+                                    # Process the buffered chunks if we haven't already
+                                    if chunks_buffer and not is_tool_call:
+                                        full_buffer = ''.join(chunks_buffer)
+                                        if "<tool_call>" in full_buffer:
+                                            print("Tool call detected")
+                                            is_tool_call = True
+                                        else:
+                                            # No tool call, send buffered chunks as regular content
+                                            delta = {"content": full_buffer, "tool_calls": None}
+                                            if is_first_chunk:
+                                                delta["role"] = "assistant"
+                                                is_first_chunk = False
+                                            chunk_data = {
+                                                "id": chat_id, "object": "chat.completion.chunk", "created": created,
+                                                "model": model_id,
+                                                "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
+                                                "usage": None
+                                            }
+                                            yield f"data: {json.dumps(chunk_data)}\n\n"
+
+                                    # Process the current chunk
+                                    if is_tool_call:
+                                        chunks_buffer.append(content_piece)
+
+                                        full_buffer = ''.join(chunks_buffer)
+
+                                        if "</tool_call>" in full_buffer:
+                                            print("Tool call End detected")
+                                            # Process tool call in the current chunk
+                                            tool_call_str = full_buffer.split("<tool_call>")[1].split("</tool_call>")[0]
+                                            tool_call_json = json.loads(tool_call_str.strip())
+                                            delta = {
+                                                "content": None,
+                                                "tool_calls": [{
+                                                    "index": 0,
+                                                    "id": generate_random_id("call_"),
+                                                    "type": "function",
+                                                    "function": {
+                                                        "name": tool_call_json["name"],
+                                                        "arguments": json.dumps(tool_call_json["parameters"])
+                                                    }
+                                                }]
+                                            }
+                                            chunk_data = {
+                                                "id": chat_id, "object": "chat.completion.chunk", "created": created,
+                                                "model": model_id,
+                                                "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
+                                                "usage": None
+                                            }
+                                            yield f"data: {json.dumps(chunk_data)}\n\n"
+                                        else:
+                                            continue
+                                    else:
+                                        # Regular content
+                                        delta = {"content": content_piece, "tool_calls": None}
+                                        if is_first_chunk:
+                                            delta["role"] = "assistant"
+                                            is_first_chunk = False
+                                        chunk_data = {
+                                            "id": chat_id, "object": "chat.completion.chunk", "created": created,
+                                            "model": model_id,
+                                            "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
+                                            "usage": None
+                                        }
+                                        yield f"data: {json.dumps(chunk_data)}\n\n"
                                 except json.JSONDecodeError: continue
                             elif line.startswith(("e:", "d:")):
                                 try:
                                     usage_info = json.loads(line[2:]).get("usage")
                                 except (json.JSONDecodeError, AttributeError): pass
                                 break
-                final_usage = None
-                if usage_info:
-                    prompt_tokens = usage_info.get("promptTokens", 0)
-                    completion_tokens = usage_info.get("completionTokens", 0)
-                    final_usage = {
-                        "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens,
-                        "total_tokens": prompt_tokens + completion_tokens,
-                    }
-                done_chunk = {
-                    "id": chat_id, "object": "chat.completion.chunk", "created": created, "model": model_id,
-                    "choices": [{
-                        "index": 0,
-                        "delta": {"role": "assistant", "content": None, "function_call": None, "tool_calls": None},
-                        "finish_reason": "stop"
-                    }],
-                    "usage": final_usage
-                }
-                yield f"data: {json.dumps(done_chunk)}\n\n"
+                        # Handle any remaining buffer content
+                        if chunks_buffer and not is_tool_call:
+                            full_buffer = ''.join(chunks_buffer)
+                            delta = {"content": full_buffer, "tool_calls": None}
+                            if is_first_chunk:
+                                delta["role"] = "assistant"
+                                is_first_chunk = False
+                            chunk_data = {
+                                "id": chat_id, "object": "chat.completion.chunk", "created": created,
+                                "model": model_id,
+                                "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
+                                "usage": None
+                            }
+                            yield f"data: {json.dumps(chunk_data)}\n\n"
+                        final_usage = None
+                        if usage_info:
+                            prompt_tokens = usage_info.get("promptTokens", 0)
+                            completion_tokens = usage_info.get("completionTokens", 0)
+                            final_usage = {
+                                "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens,
+                                "total_tokens": prompt_tokens + completion_tokens,
+                            }
+                        done_chunk = {
+                            "id": chat_id, "object": "chat.completion.chunk", "created": created, "model": model_id,
+                            "choices": [{
+                                "index": 0,
+                                "delta": {"role": "assistant", "content": None, "function_call": None, "tool_calls": None},
+                                "finish_reason": "stop"
+                            }],
+                            "usage": final_usage
+                        }
+                        yield f"data: {json.dumps(done_chunk)}\n\n"
             except httpx.HTTPStatusError as e:
                 error_content = {
                     "error": {
@@ -153,6 +244,7 @@ async def chat_completion(request: ChatRequest):
         return StreamingResponse(event_stream(), media_type="text/event-stream")
     else: # Non-streaming
         assistant_response, usage_info = "", {}
+        tool_call_json = None
         try:
             async with httpx.AsyncClient(timeout=120) as client:
                 async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
@@ -164,9 +256,15 @@ async def chat_completion(request: ChatRequest):
                         elif chunk.startswith(("e:", "d:")):
                             try: usage_info = json.loads(chunk[2:]).get("usage", {})
                             except: continue
+
+            if "<tool_call>" in assistant_response and "</tool_call>" in assistant_response:
+                tool_call_str = assistant_response.split("<tool_call>")[1].split("</tool_call>")[0]
+                tool_call_json = json.loads(tool_call_str.strip())
+
+                
             return JSONResponse(content={
                 "id": chat_id, "object": "chat.completion", "created": int(time.time()), "model": model_id,
-                "choices": [{"index": 0, "message": {"role": "assistant", "content": assistant_response}, "finish_reason": "stop"}],
+                "choices": [{"index": 0, "message": {"role": "assistant", "content": assistant_response if tool_call_json is None else None, "tool_calls": tool_call_json}, "finish_reason": "stop"}],
                 "usage": {
                     "prompt_tokens": usage_info.get("promptTokens", 0),
                     "completion_tokens": usage_info.get("completionTokens", 0),
@@ -178,7 +276,6 @@ async def chat_completion(request: ChatRequest):
 
 
 # === Image Generation ===
-
 class ImageGenerationRequest(BaseModel):
     prompt: str
     aspect_ratio: Optional[str] = "1:1"
@@ -224,9 +321,7 @@ async def generate_images(request: ImageGenerationRequest):
         return JSONResponse(status_code=500, content={"error": "An internal error occurred.", "details": str(e)})
     return {"created": int(time.time()), "data": results}
 
-
 # === Moderation Endpoint ===
-
 class ModerationRequest(BaseModel):
     input: Union[str, List[str]]
     model: Optional[str] = "text-moderation-stable"
@@ -240,14 +335,12 @@ async def create_moderation(request: ModerationRequest):
     input_texts = [request.input] if isinstance(request.input, str) else request.input
     if not input_texts:
          return JSONResponse(status_code=400, content={"error": {"message": "Request must have at least one input string.", "type": "invalid_request_error"}})
-
     moderation_url = "https://www.chatwithmono.xyz/api/moderation"
     headers = {
         'Content-Type': 'application/json',
         'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36',
         'Referer': 'https://www.chatwithmono.xyz/',
     }
-
     results = []
     try:
         async with httpx.AsyncClient(timeout=30) as client:
@@ -256,7 +349,6 @@ async def create_moderation(request: ModerationRequest):
                 resp = await client.post(moderation_url, headers=headers, json=payload)
                 resp.raise_for_status()
                 upstream_data = resp.json()
-
                 # --- Transform upstream response to OpenAI format ---
                 upstream_categories = upstream_data.get("categories", {})
                 openai_categories = {
@@ -268,21 +360,19 @@ async def create_moderation(request: ModerationRequest):
                 }
                 category_scores = {k: 1.0 if v else 0.0 for k, v in openai_categories.items()}
                 flagged = upstream_data.get("overall_sentiment") == "flagged"
-
                 result_item = {
                     "flagged": flagged,
                     "categories": openai_categories,
                     "category_scores": category_scores,
                 }
-                
+
                 # --- NEW: Conditionally add the 'reason' field ---
                 # This is a custom extension to the OpenAI spec to provide more detail.
                 reason = upstream_data.get("reason")
                 if reason:
                     result_item["reason"] = reason
-                
-                results.append(result_item)
 
+                results.append(result_item)
     except httpx.HTTPStatusError as e:
         return JSONResponse(
             status_code=502, # Bad Gateway
@@ -290,7 +380,6 @@ async def create_moderation(request: ModerationRequest):
         )
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": {"message": "An internal error occurred during moderation.", "type": "internal_error", "details": str(e)}})
-
     # Build the final OpenAI-compatible response
     final_response = {
         "id": generate_random_id("modr-"),
@@ -299,9 +388,7 @@ async def create_moderation(request: ModerationRequest):
     }
     return JSONResponse(content=final_response)
 
-
 # --- Main Execution ---
-
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)