Update main.py

main.py

@@ -7,17 +7,19 @@ import time
 from typing import List, Optional, Union, Any
 import httpx
 from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from fastapi.responses import JSONResponse, StreamingResponse
-from pydantic import BaseModel, Field
+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"
 SNAPZION_API_KEY = os.environ.get("SNAP", "")
 
 # --- Dummy Model Definitions ---
+# In a real application, these would be defined properly.
 AVAILABLE_MODELS = [
     {"id": "gpt-4-turbo", "object": "model", "created": int(time.time()), "owned_by": "system"},
     {"id": "gpt-4o", "object": "model", "created": int(time.time()), "owned_by": "system"},
@@ -36,56 +38,35 @@ app = FastAPI(
 
 # --- Helper Function for Random ID Generation ---
 def generate_random_id(prefix: str, length: int = 29) -> str:
+    """
+    Generates a cryptographically secure, random alphanumeric ID.
+    """
     population = string.ascii_letters + string.digits
     random_part = "".join(secrets.choice(population) for _ in range(length))
     return f"{prefix}{random_part}"
 
-# === Tool Call Models ===
-class FunctionCall(BaseModel):
-    name: str
-    arguments: str
-
-class ToolCall(BaseModel):
-    id: str
-    type: str = "function"
-    function: FunctionCall
-
-# === Message Models ===
-class Message(BaseModel):
-    role: str
-    content: Optional[str] = None
-    name: Optional[str] = None
-    tool_calls: Optional[List[ToolCall]] = None
-    tool_call_id: Optional[str] = None
-
 # === 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
+
 class ChatRequest(BaseModel):
     messages: List[Message]
     model: str
     stream: Optional[bool] = False
     tools: Optional[Any] = None
 
-def build_tool_prompt(tools: List[Any]) -> str:
-    tool_definitions = "\n".join([
-        f"- {tool['function']['name']}: {tool['function'].get('description', 'No description available')}"
-        for tool in tools
-    ])
-    return f"""You have access to tools. To call a tool, respond with JSON inside <tool_call></tool_call> tags.
-Available Tools:
-{tool_definitions}
-
-Response Format:
-<tool_call>
-{{"name": "tool_name", "parameters": {{"arg1": "value1"}}}}
-</tool_call>"""
-
 @app.post("/v1/chat/completions")
 async def chat_completion(request: ChatRequest):
+    """
+    Handles chat completion requests, supporting both streaming and non-streaming responses.
+    """
     model_id = MODEL_ALIASES.get(request.model, request.model)
     chat_id = generate_random_id("chatcmpl-")
     headers = {
@@ -95,189 +76,191 @@ async def chat_completion(request: ChatRequest):
         'referer': 'https://www.chatwithmono.xyz/',
         'user-agent': 'Mozilla/5.0',
     }
-
-    # Handle tool definitions
     if request.tools:
-        tool_prompt = build_tool_prompt(request.tools)
-        if request.messages and request.messages[0].role == "system":
+        # 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, Message(role="system", content=tool_prompt))
+            request.messages.insert(0, {"role": "system", "content": tool_prompt})
+    request_data = request.model_dump(exclude_unset=True)
 
     payload = {
-        "messages": [msg.model_dump(exclude_none=True) for msg in request.messages],
+        "messages": request_data["messages"],
         "model": model_id
     }
-
-    # Streaming response
     if request.stream:
         async def event_stream():
             created = int(time.time())
-            chunks_buffer = []
-            max_initial_chunks = 4
+            is_first_chunk = True
+            usage_info = None
             is_tool_call = False
-            tool_call_content = ""
-            
+            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:
+                    async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
                         response.raise_for_status()
                         async for line in response.aiter_lines():
-                            if not line: 
-                                continue
+                            if not line: continue
                             if line.startswith("0:"):
                                 try:
                                     content_piece = json.loads(line[2:])
-                                    # Buffer initial chunks
+                                    print(content_piece)
+                                    # Buffer the first few chunks
                                     if len(chunks_buffer) < max_initial_chunks:
                                         chunks_buffer.append(content_piece)
                                         continue
-                                    
-                                    # Check for tool call pattern
-                                    if not is_tool_call:
-                                        full_buffer = ''.join(chunks_buffer + [content_piece])
+                                    # 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
-                                            tool_call_content = full_buffer
-                                            chunks_buffer = []
+
+                                    # 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:
-                                            # Send buffered chunks
-                                            if chunks_buffer:
-                                                delta = {"content": "".join(chunks_buffer)}
-                                                if not tool_call_content:  # Only add role in first chunk
-                                                    delta["role"] = "assistant"
-                                                yield create_chunk(chat_id, created, model_id, delta)
-                                                chunks_buffer = []
-                                            
-                                            # Send current chunk
-                                            delta = {"content": content_piece}
-                                            yield create_chunk(chat_id, created, model_id, delta)
+                                            continue
                                     else:
-                                        # Accumulate tool call content
-                                        tool_call_content += content_piece
-                                        if "</tool_call>" in tool_call_content:
-                                            tool_call_str = tool_call_content.split("<tool_call>")[1].split("</tool_call>")[0].strip()
-                                            try:
-                                                tool_call_data = json.loads(tool_call_str)
-                                                tool_call = ToolCall(
-                                                    id=generate_random_id("call_"),
-                                                    function=FunctionCall(
-                                                        name=tool_call_data["name"],
-                                                        arguments=json.dumps(tool_call_data.get("parameters", {}))
-                                                )
-                                                delta = {
-                                                    "content": None,
-                                                    "tool_calls": [tool_call.model_dump()]
-                                                }
-                                                yield create_chunk(chat_id, created, model_id, delta)
-                                                is_tool_call = False
-                                                tool_call_content = ""
-                                            except (json.JSONDecodeError, KeyError) as e:
-                                                print(f"Tool call parsing error: {e}")
-                                except json.JSONDecodeError:
-                                    continue
-                            
+
+                                        # Regular content
+                                        if is_first_chunk:
+                                            delta = {"content": "".join(chunks_buffer), "tool_calls": None}
+                                            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"
+
+                                        delta = {"content": content_piece, "tool_calls": None}
+
+                                        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"
             except httpx.HTTPStatusError as e:
                 error_content = {
                     "error": {
-                        "message": f"Upstream error: {e.response.status_code}",
-                        "type": "upstream_error",
-                        "code": str(e.response.status_code)
+                        "message": f"Upstream API error: {e.response.status_code}. Details: {e.response.text}",
+                        "type": "upstream_error", "code": str(e.response.status_code)
                     }
                 }
                 yield f"data: {json.dumps(error_content)}\n\n"
             finally:
-                # Finish signal
-                done_chunk = {
-                    "id": chat_id, 
-                    "object": "chat.completion.chunk", 
-                    "created": created, 
-                    "model": model_id,
-                    "choices": [{
-                        "index": 0,
-                        "delta": {},
-                        "finish_reason": "stop"
-                    }]
-                }
-                yield f"data: {json.dumps(done_chunk)}\n\n"
                 yield "data: [DONE]\n\n"
-        
         return StreamingResponse(event_stream(), media_type="text/event-stream")
-    
-    # Non-streaming response
-    else:
+    else: # Non-streaming
+        assistant_response, usage_info = "", {}
+        tool_call_json = None
         try:
             async with httpx.AsyncClient(timeout=120) as client:
-                response = await client.post(
-                    "https://www.chatwithmono.xyz/api/chat",
-                    headers=headers,
-                    json=payload
-                )
-                response.raise_for_status()
-                content = ""
-                for line in response.text.splitlines():
-                    if line.startswith("0:"):
-                        try:
-                            content += json.loads(line[2:])
-                        except json.JSONDecodeError:
-                            continue
+                async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
+                    response.raise_for_status()
+                    async for chunk in response.aiter_lines():
+                        if chunk.startswith("0:"):
+                            try: assistant_response += json.loads(chunk[2:])
+                            except: continue
+                        elif chunk.startswith(("e:", "d:")):
+                            try: usage_info = json.loads(chunk[2:]).get("usage", {})
+                            except: continue
 
-                tool_calls = None
-                if "<tool_call>" in content and "</tool_call>" in content:
-                    try:
-                        tool_call_str = content.split("<tool_call>")[1].split("</tool_call>")[0].strip()
-                        tool_call_data = json.loads(tool_call_str)
-                        tool_call = ToolCall(
-                            id=generate_random_id("call_"),
-                            function=FunctionCall(
-                                name=tool_call_data["name"],
-                                arguments=json.dumps(tool_call_data.get("parameters", {}))
-                        )
-                        tool_calls = [tool_call.model_dump()]
-                        content = None  # Clear content for tool call
-                    except (json.JSONDecodeError, KeyError) as e:
-                        print(f"Tool call parsing error: {e}")
+            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.loads(tool_call_str.strip())
+                tool_call_json = [{"id": generate_random_id("call_"),"function": {"name": tool_call["name"], "arguments": json.dumps(tool_call["parameters"])}}]
 
-                return JSONResponse(content={
-                    "id": chat_id,
-                    "object": "chat.completion",
-                    "created": int(time.time()),
-                    "model": model_id,
-                    "choices": [{
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": content,
-                            "tool_calls": tool_calls
-                        },
-                        "finish_reason": "tool_calls" if tool_calls else "stop"
-                    }],
-                    "usage": {
-                        "prompt_tokens": 0,
-                        "completion_tokens": 0,
-                        "total_tokens": 0
-                    }
-                })
-        
+
+                
+            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 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),
+                    "total_tokens": usage_info.get("promptTokens", 0) + usage_info.get("completionTokens", 0),
+                }
+            })
         except httpx.HTTPStatusError as e:
-            return JSONResponse(
-                status_code=502,
-                content={"error": {"message": f"Upstream error: {e.response.text}", "type": "upstream_error"}}
-            )
+            return JSONResponse(status_code=e.response.status_code, content={"error": {"message": f"Upstream API error. Details: {e.response.text}", "type": "upstream_error"}})
 
-def create_chunk(chat_id: str, created: int, model: str, delta: dict) -> str:
-    chunk = {
-        "id": chat_id,
-        "object": "chat.completion.chunk",
-        "created": created,
-        "model": model,
-        "choices": [{"index": 0, "delta": delta}]
-    }
-    return f"data: {json.dumps(chunk)}\n\n"
 
 # === Image Generation ===
 class ImageGenerationRequest(BaseModel):
@@ -289,65 +272,40 @@ class ImageGenerationRequest(BaseModel):
 
 @app.post("/v1/images/generations")
 async def generate_images(request: ImageGenerationRequest):
+    """Handles image generation requests."""
     results = []
     try:
         async with httpx.AsyncClient(timeout=120) as client:
             for _ in range(request.n):
-                if request.model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
-                    # Mono image generation
-                    resp = await client.post(
-                        "https://www.chatwithmono.xyz/api/image",
-                        json={"prompt": request.prompt, "model": request.model},
-                        headers={'Content-Type': 'application/json'}
-                    )
+                model = request.model or "default"
+                if model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
+                    headers = {'Content-Type': 'application/json', 'User-Agent': 'Mozilla/5.0', 'Referer': 'https://www.chatwithmono.xyz/'}
+                    payload = {"prompt": request.prompt, "model": model}
+                    resp = await client.post("https://www.chatwithmono.xyz/api/image", headers=headers, json=payload)
                     resp.raise_for_status()
                     data = resp.json()
                     b64_image = data.get("image")
-                    
-                    if not b64_image:
-                        raise HTTPException(502, "Missing image in response")
-                    
-                    # Upload to Snapzion if API key available
+                    if not b64_image: return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})
                     if SNAPZION_API_KEY:
-                        files = {'file': ('image.png', base64.b64decode(b64_image), 'image/png')}
-                        upload_resp = await client.post(
-                            SNAPZION_UPLOAD_URL,
-                            files=files,
-                            headers={"Authorization": SNAPZION_API_KEY}
-                        )
+                        upload_headers = {"Authorization": SNAPZION_API_KEY}
+                        upload_files = {'file': ('image.png', base64.b64decode(b64_image), 'image/png')}
+                        upload_resp = await client.post(SNAPZION_UPLOAD_URL, headers=upload_headers, files=upload_files)
                         upload_resp.raise_for_status()
-                        image_url = upload_resp.json().get("url")
+                        upload_data = upload_resp.json()
+                        image_url = upload_data.get("url")
                     else:
                         image_url = f"data:image/png;base64,{b64_image}"
-                    
-                    results.append({
-                        "url": image_url,
-                        "b64_json": b64_image,
-                        "revised_prompt": data.get("revised_prompt", request.prompt)
-                    })
+                    results.append({"url": image_url, "b64_json": b64_image, "revised_prompt": data.get("revised_prompt")})
                 else:
-                    # Default image generation
-                    params = {"prompt": request.prompt, "aspect_ratio": request.aspect_ratio}
+                    params = {"prompt": request.prompt, "aspect_ratio": request.aspect_ratio, "link": "typegpt.net"}
                     resp = await client.get(IMAGE_API_URL, params=params)
                     resp.raise_for_status()
                     data = resp.json()
-                    results.append({
-                        "url": data.get("image_link"),
-                        "b64_json": data.get("base64_output"),
-                        "revised_prompt": request.prompt
-                    })
-                    
+                    results.append({"url": data.get("image_link"), "b64_json": data.get("base64_output")})
     except httpx.HTTPStatusError as e:
-        return JSONResponse(
-            status_code=502,
-            content={"error": f"Image service error: {e.response.status_code}"}
-        )
+        return JSONResponse(status_code=502, content={"error": f"Image generation failed. Upstream error: {e.response.status_code}", "details": e.response.text})
     except Exception as e:
-        return JSONResponse(
-            status_code=500,
-            content={"error": f"Internal error: {str(e)}"}
-        )
-    
+        return JSONResponse(status_code=500, content={"error": "An internal error occurred.", "details": str(e)})
     return {"created": int(time.time()), "data": results}
 
 # === Moderation Endpoint ===
@@ -357,76 +315,67 @@ class ModerationRequest(BaseModel):
 
 @app.post("/v1/moderations")
 async def create_moderation(request: ModerationRequest):
+    """
+    Handles moderation requests, conforming to the OpenAI API specification.
+    Includes a custom 'reason' field in the result if provided by the upstream API.
+    """
     input_texts = [request.input] if isinstance(request.input, str) else request.input
     if not input_texts:
-        return JSONResponse(
-            status_code=400,
-            content={"error": "At least one input string is required"}
-        )
-    
+         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',
+        '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 = []
-    async with httpx.AsyncClient(timeout=30) as client:
-        for text in input_texts:
-            try:
-                resp = await client.post(
-                    "https://www.chatwithmono.xyz/api/moderation",
-                    json={"text": text},
-                    headers=headers
-                )
+    try:
+        async with httpx.AsyncClient(timeout=30) as client:
+            for text_input in input_texts:
+                payload = {"text": text_input}
+                resp = await client.post(moderation_url, headers=headers, json=payload)
                 resp.raise_for_status()
-                data = resp.json()
-                
-                # Transform to OpenAI format
-                flagged = data.get("overall_sentiment") == "flagged"
-                categories = data.get("categories", {})
+                upstream_data = resp.json()
+                # --- Transform upstream response to OpenAI format ---
+                upstream_categories = upstream_data.get("categories", {})
                 openai_categories = {
-                    "hate": categories.get("hate", False),
-                    "hate/threatening": False,
-                    "self-harm": categories.get("self-harm", False),
-                    "sexual": categories.get("sexual", False),
-                    "sexual/minors": False,
-                    "violence": categories.get("violence", False),
-                    "violence/graphic": False,
+                    "hate": upstream_categories.get("hate", False), "hate/threatening": False,
+                    "harassment": False, "harassment/threatening": False,
+                    "self-harm": upstream_categories.get("self-harm", False), "self-harm/intent": False, "self-harm/instructions": False,
+                    "sexual": upstream_categories.get("sexual", False), "sexual/minors": False,
+                    "violence": upstream_categories.get("violence", False), "violence/graphic": False,
                 }
-                
+                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": {k: 1.0 if v else 0.0 for k, v in openai_categories.items()}
+                    "category_scores": category_scores,
                 }
-                
-                # Add reason if available
-                if "reason" in data:
-                    result_item["reason"] = data["reason"]
-                    
+
+                # --- 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)
-                
-            except httpx.HTTPStatusError:
-                results.append({
-                    "flagged": False,
-                    "categories": {k: False for k in [
-                        "hate", "hate/threatening", "self-harm", 
-                        "sexual", "sexual/minors", "violence", "violence/graphic"
-                    ]},
-                    "category_scores": {k: 0.0 for k in [
-                        "hate", "hate/threatening", "self-harm", 
-                        "sexual", "sexual/minors", "violence", "violence/graphic"
-                    ]}
-                })
-    
-    return {
+    except httpx.HTTPStatusError as e:
+        return JSONResponse(
+            status_code=502, # Bad Gateway
+            content={"error": {"message": f"Moderation failed. Upstream error: {e.response.status_code}", "type": "upstream_error", "details": e.response.text}}
+        )
+    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-"),
         "model": request.model,
-        "results": results
+        "results": results,
     }
+    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)