AIMaster7 commited on
Commit
06dc2f1
·
verified ·
1 Parent(s): c18e9c8

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +284 -85
main.py CHANGED
@@ -5,7 +5,7 @@ import secrets
5
  import string
6
  import time
7
  import tempfile
8
- import ast # <-- NEW IMPORT for safe literal evaluation
9
  from typing import List, Optional, Union, Any
10
 
11
  import httpx
@@ -14,11 +14,13 @@ from fastapi import FastAPI, HTTPException
14
  from fastapi.responses import JSONResponse, StreamingResponse
15
  from pydantic import BaseModel, Field, model_validator
16
 
17
- # New import for OCR
18
  from gradio_client import Client, handle_file
19
 
20
  # --- Configuration ---
21
  load_dotenv()
 
 
22
  IMAGE_API_URL = os.environ.get("IMAGE_API_URL", "https://image.api.example.com")
23
  SNAPZION_UPLOAD_URL = "https://upload.snapzion.com/api/public-upload"
24
  SNAPZION_API_KEY = os.environ.get("SNAP", "")
@@ -41,36 +43,43 @@ MODEL_ALIASES = {}
41
  app = FastAPI(
42
  title="OpenAI Compatible API",
43
  description="An adapter for various services to be compatible with the OpenAI API specification.",
44
- version="1.1.2" # Incremented version for the new fix
45
  )
 
 
46
  try:
47
  ocr_client = Client("multimodalart/Florence-2-l4")
48
  except Exception as e:
49
  print(f"Warning: Could not initialize Gradio client for OCR: {e}")
50
  ocr_client = None
51
 
 
52
  # --- Pydantic Models ---
53
- # (Pydantic models are unchanged and remain the same as before)
54
  class Message(BaseModel):
55
  role: str
56
  content: str
 
57
  class ChatRequest(BaseModel):
58
  messages: List[Message]
59
  model: str
60
  stream: Optional[bool] = False
61
  tools: Optional[Any] = None
 
62
  class ImageGenerationRequest(BaseModel):
63
  prompt: str
64
  aspect_ratio: Optional[str] = "1:1"
65
  n: Optional[int] = 1
66
  user: Optional[str] = None
67
  model: Optional[str] = "default"
 
68
  class ModerationRequest(BaseModel):
69
  input: Union[str, List[str]]
70
  model: Optional[str] = "text-moderation-stable"
 
71
  class OcrRequest(BaseModel):
72
  image_url: Optional[str] = Field(None, description="URL of the image to process.")
73
  image_b64: Optional[str] = Field(None, description="Base64 encoded string of the image to process.")
 
74
  @model_validator(mode='before')
75
  @classmethod
76
  def check_sources(cls, data: Any) -> Any:
@@ -80,116 +89,258 @@ class OcrRequest(BaseModel):
80
  if data.get('image_url') and data.get('image_b64'):
81
  raise ValueError('Provide either image_url or image_b64, not both.')
82
  return data
 
83
  class OcrResponse(BaseModel):
84
  ocr_text: str
85
  raw_response: dict
86
 
 
87
  # --- Helper Function ---
88
  def generate_random_id(prefix: str, length: int = 29) -> str:
 
89
  population = string.ascii_letters + string.digits
90
  random_part = "".join(secrets.choice(population) for _ in range(length))
91
  return f"{prefix}{random_part}"
92
 
 
93
  # === API Endpoints ===
94
 
95
  @app.get("/v1/models", tags=["Models"])
96
  async def list_models():
 
97
  return {"object": "list", "data": AVAILABLE_MODELS}
98
 
99
- # (Chat, Image Generation, and Moderation endpoints are unchanged and remain correct)
100
  @app.post("/v1/chat/completions", tags=["Chat"])
101
  async def chat_completion(request: ChatRequest):
102
- model_id=MODEL_ALIASES.get(request.model,request.model);chat_id=generate_random_id("chatcmpl-");headers={'accept':'text/event-stream','content-type':'application/json','origin':'https://www.chatwithmono.xyz','referer':'https://www.chatwithmono.xyz/','user-agent':'Mozilla/5.0'}
 
 
 
 
 
 
 
 
 
 
103
  if request.tools:
104
- 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 tags. Respond in the format {{"name": tool name, "parameters": dictionary of argument name and its value}}. Do not use variables.
105
  Tools: {";".join(f"<tool>{tool}</tool>" for tool in request.tools)}
106
  Response Format for tool call:
107
  <tool_call>
108
  {{"name": <function-name>, "arguments": <args-json-object>}}
109
  </tool_call>"""
110
- if request.messages[0].role=="system":request.messages[0].content+="\n\n"+tool_prompt
111
- else:request.messages.insert(0,Message(role="system",content=tool_prompt))
112
- payload={"messages":[msg.model_dump()for msg in request.messages],"model":model_id}
 
 
 
 
113
  if request.stream:
114
  async def event_stream():
115
- created=int(time.time());usage_info=None;is_first_chunk=True;tool_call_buffer="";in_tool_call=False
 
 
 
 
 
116
  try:
117
- async with httpx.AsyncClient(timeout=120)as client:
118
- async with client.stream("POST",CHAT_API_URL,headers=headers,json=payload)as response:
119
  response.raise_for_status()
120
  async for line in response.aiter_lines():
121
- if not line:continue
 
 
122
  if line.startswith("0:"):
123
- try:content_piece=json.loads(line[2:])
124
- except json.JSONDecodeError:continue
125
- current_buffer=content_piece
126
- if in_tool_call:current_buffer=tool_call_buffer+content_piece
127
- if"</tool_call>"in current_buffer:
128
- tool_str=current_buffer.split("<tool_call>")[1].split("</tool_call>")[0];tool_json=json.loads(tool_str.strip());delta={"content":None,"tool_calls":[{"index":0,"id":generate_random_id("call_"),"type":"function","function":{"name":tool_json["name"],"arguments":json.dumps(tool_json["parameters"])}}]}
129
- chunk={"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)}\n\n"
130
- in_tool_call=False;tool_call_buffer="";remaining_text=current_buffer.split("</tool_call>",1)[1]
131
- if remaining_text:content_piece=remaining_text
132
- else:continue
133
- if"<tool_call>"in content_piece:
134
- in_tool_call=True;tool_call_buffer+=content_piece.split("<tool_call>",1)[1];text_before=content_piece.split("<tool_call>",1)[0]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
135
  if text_before:
136
- delta={"content":text_before,"tool_calls":None};chunk={"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)}\n\n"
137
- if"</tool_call>"not in tool_call_buffer:continue
 
 
 
 
 
138
  if not in_tool_call:
139
- delta={"content":content_piece}
140
- if is_first_chunk:delta["role"]="assistant";is_first_chunk=False
141
- chunk={"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)}\n\n"
142
- elif line.startswith(("e:","d:")):
143
- try:usage_info=json.loads(line[2:]).get("usage")
144
- except(json.JSONDecodeError,AttributeError):pass
 
 
 
 
 
 
 
145
  break
146
- final_usage=None
147
- if usage_info:final_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)}
148
- done_chunk={"id":chat_id,"object":"chat.completion.chunk","created":created,"model":model_id,"choices":[{"index":0,"delta":{},"finish_reason":"stop"if not in_tool_call else"tool_calls"}],"usage":final_usage};yield f"data: {json.dumps(done_chunk)}\n\n"
149
- except httpx.HTTPStatusError as e:error_content={"error":{"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"
150
- finally:yield"data: [DONE]\n\n"
151
- return StreamingResponse(event_stream(),media_type="text/event-stream")
152
- else:
153
- full_response,usage_info="",{}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
154
  try:
155
- async with httpx.AsyncClient(timeout=120)as client:
156
- async with client.stream("POST",CHAT_API_URL,headers=headers,json=payload)as response:
157
  response.raise_for_status()
158
  async for chunk in response.aiter_lines():
159
  if chunk.startswith("0:"):
160
- try:full_response+=json.loads(chunk[2:])
161
- except:continue
162
- elif chunk.startswith(("e:","d:")):
163
- try:usage_info=json.loads(chunk[2:]).get("usage",{})
164
- except:continue
165
- tool_calls=None;content_response=full_response
166
- if"<tool_call>"in full_response and"</tool_call>"in full_response:
167
- tool_call_str=full_response.split("<tool_call>")[1].split("</tool_call>")[0];tool_call=json.loads(tool_call_str.strip());tool_calls=[{"id":generate_random_id("call_"),"type":"function","function":{"name":tool_call["name"],"arguments":json.dumps(tool_call["parameters"])}}];content_response=None
168
- return JSONResponse(content={"id":chat_id,"object":"chat.completion","created":int(time.time()),"model":model_id,"choices":[{"index":0,"message":{"role":"assistant","content":content_response,"tool_calls":tool_calls},"finish_reason":"stop"if not tool_calls else"tool_calls"}],"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)}})
169
- except httpx.HTTPStatusError as e:return JSONResponse(status_code=e.response.status_code,content={"error":{"message":f"Upstream API error. Details: {e.response.text}","type":"upstream_error"}})
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
170
 
171
  @app.post("/v1/images/generations", tags=["Images"])
172
  async def generate_images(request: ImageGenerationRequest):
173
- results=[]
 
174
  try:
175
- async with httpx.AsyncClient(timeout=120)as client:
176
  for _ in range(request.n):
177
- model=request.model or"default"
178
- if model in["gpt-image-1","dall-e-3","dall-e-2","nextlm-image-1"]:
179
- 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(IMAGE_GEN_API_URL,headers=headers,json=payload);resp.raise_for_status();data=resp.json();b64_image=data.get("image")
180
- if not b64_image:return JSONResponse(status_code=502,content={"error":"Missing base64 image in response"})
181
- image_url=f"data:image/png;base64,{b64_image}"
 
 
 
 
 
 
 
182
  if SNAPZION_API_KEY:
183
- 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)
184
- if upload_resp.status_code==200:image_url=upload_resp.json().get("url",image_url)
185
- results.append({"url":image_url,"b64_json":b64_image,"revised_prompt":data.get("revised_prompt")})
186
- else: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")})
187
- except httpx.HTTPStatusError as e:return JSONResponse(status_code=502,content={"error":f"Image generation failed. Upstream error: {e.response.status_code}","details":e.response.text})
188
- except Exception as e:return JSONResponse(status_code=500,content={"error":"An internal error occurred.","details":str(e)})
189
- return{"created":int(time.time()),"data":results}
 
 
 
 
 
 
 
 
 
 
 
 
190
 
191
 
192
- # === REVISED AND FIXED OCR Endpoint ===
193
  @app.post("/v1/ocr", response_model=OcrResponse, tags=["OCR"])
194
  async def perform_ocr(request: OcrRequest):
195
  """
@@ -218,7 +369,7 @@ async def perform_ocr(request: OcrRequest):
218
  raw_output = prediction[0]
219
  raw_result_dict = {}
220
 
221
- # --- START: ROBUST PARSING LOGIC ---
222
  if isinstance(raw_output, str):
223
  try:
224
  # First, try to parse as standard JSON
@@ -234,17 +385,18 @@ async def perform_ocr(request: OcrRequest):
234
  raw_result_dict = {"result": str(parsed_output)}
235
  except (ValueError, SyntaxError):
236
  # If all parsing fails, assume the string is the direct OCR text.
237
- raw_result_dict = {"ocr_text": raw_output}
238
  elif isinstance(raw_output, dict):
239
  # It's already a dictionary, use it directly
240
  raw_result_dict = raw_output
241
  else:
242
  # Handle other unexpected data types
243
  raise HTTPException(status_code=502, detail=f"Unexpected data type from OCR service: {type(raw_output)}")
244
- # --- END: ROBUST PARSING LOGIC ---
245
 
246
- # Extract text from the dictionary, with fallbacks
247
- ocr_text = raw_result_dict.get("OCR", raw_result_dict.get("ocr_text", str(raw_result_dict)))
 
 
248
 
249
  return OcrResponse(ocr_text=ocr_text, raw_response=raw_result_dict)
250
 
@@ -256,22 +408,69 @@ async def perform_ocr(request: OcrRequest):
256
  if temp_file_path and os.path.exists(temp_file_path):
257
  os.unlink(temp_file_path)
258
 
 
259
  @app.post("/v1/moderations", tags=["Moderation"])
260
  async def create_moderation(request: ModerationRequest):
261
- input_texts=[request.input]if isinstance(request.input,str)else request.input
262
- if not input_texts:return JSONResponse(status_code=400,content={"error":{"message":"Request must have at least one input string."}})
263
- headers={'Content-Type':'application/json','User-Agent':'Mozilla/5.0','Referer':'https://www.chatwithmono.xyz/'};results=[]
 
 
 
 
 
264
  try:
265
- async with httpx.AsyncClient(timeout=30)as client:
266
  for text_input in input_texts:
267
- resp=await client.post(MODERATION_API_URL,headers=headers,json={"text":text_input});resp.raise_for_status();upstream_data=resp.json();upstream_categories=upstream_data.get("categories",{})
268
- openai_categories={"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}
269
- result_item={"flagged":upstream_data.get("overall_sentiment")=="flagged","categories":openai_categories,"category_scores":{k:1.0 if v else 0.0 for k,v in openai_categories.items()}}
270
- if reason:=upstream_data.get("reason"):result_item["reason"]=reason
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
271
  results.append(result_item)
272
- except httpx.HTTPStatusError as e:return JSONResponse(status_code=502,content={"error":{"message":f"Moderation failed. Upstream error: {e.response.status_code}","details":e.response.text}})
273
- except Exception as e:return JSONResponse(status_code=500,content={"error":{"message":"An internal error occurred during moderation.","details":str(e)}})
274
- return JSONResponse(content={"id":generate_random_id("modr-"),"model":request.model,"results":results})
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
275
 
276
 
277
  # --- Main Execution ---
 
5
  import string
6
  import time
7
  import tempfile
8
+ import ast
9
  from typing import List, Optional, Union, Any
10
 
11
  import httpx
 
14
  from fastapi.responses import JSONResponse, StreamingResponse
15
  from pydantic import BaseModel, Field, model_validator
16
 
17
+ # Import for OCR functionality
18
  from gradio_client import Client, handle_file
19
 
20
  # --- Configuration ---
21
  load_dotenv()
22
+
23
+ # Environment variables for external services
24
  IMAGE_API_URL = os.environ.get("IMAGE_API_URL", "https://image.api.example.com")
25
  SNAPZION_UPLOAD_URL = "https://upload.snapzion.com/api/public-upload"
26
  SNAPZION_API_KEY = os.environ.get("SNAP", "")
 
43
  app = FastAPI(
44
  title="OpenAI Compatible API",
45
  description="An adapter for various services to be compatible with the OpenAI API specification.",
46
+ version="1.1.3" # Version reflects final formatting and fixes
47
  )
48
+
49
+ # Initialize Gradio client globally to avoid re-initialization on each request
50
  try:
51
  ocr_client = Client("multimodalart/Florence-2-l4")
52
  except Exception as e:
53
  print(f"Warning: Could not initialize Gradio client for OCR: {e}")
54
  ocr_client = None
55
 
56
+
57
  # --- Pydantic Models ---
 
58
  class Message(BaseModel):
59
  role: str
60
  content: str
61
+
62
  class ChatRequest(BaseModel):
63
  messages: List[Message]
64
  model: str
65
  stream: Optional[bool] = False
66
  tools: Optional[Any] = None
67
+
68
  class ImageGenerationRequest(BaseModel):
69
  prompt: str
70
  aspect_ratio: Optional[str] = "1:1"
71
  n: Optional[int] = 1
72
  user: Optional[str] = None
73
  model: Optional[str] = "default"
74
+
75
  class ModerationRequest(BaseModel):
76
  input: Union[str, List[str]]
77
  model: Optional[str] = "text-moderation-stable"
78
+
79
  class OcrRequest(BaseModel):
80
  image_url: Optional[str] = Field(None, description="URL of the image to process.")
81
  image_b64: Optional[str] = Field(None, description="Base64 encoded string of the image to process.")
82
+
83
  @model_validator(mode='before')
84
  @classmethod
85
  def check_sources(cls, data: Any) -> Any:
 
89
  if data.get('image_url') and data.get('image_b64'):
90
  raise ValueError('Provide either image_url or image_b64, not both.')
91
  return data
92
+
93
  class OcrResponse(BaseModel):
94
  ocr_text: str
95
  raw_response: dict
96
 
97
+
98
  # --- Helper Function ---
99
  def generate_random_id(prefix: str, length: int = 29) -> str:
100
+ """Generates a cryptographically secure, random alphanumeric ID."""
101
  population = string.ascii_letters + string.digits
102
  random_part = "".join(secrets.choice(population) for _ in range(length))
103
  return f"{prefix}{random_part}"
104
 
105
+
106
  # === API Endpoints ===
107
 
108
  @app.get("/v1/models", tags=["Models"])
109
  async def list_models():
110
+ """Lists the available models."""
111
  return {"object": "list", "data": AVAILABLE_MODELS}
112
 
113
+
114
  @app.post("/v1/chat/completions", tags=["Chat"])
115
  async def chat_completion(request: ChatRequest):
116
+ """Handles chat completion requests, supporting streaming and non-streaming."""
117
+ model_id = MODEL_ALIASES.get(request.model, request.model)
118
+ chat_id = generate_random_id("chatcmpl-")
119
+ headers = {
120
+ 'accept': 'text/event-stream',
121
+ 'content-type': 'application/json',
122
+ 'origin': 'https://www.chatwithmono.xyz',
123
+ 'referer': 'https://www.chatwithmono.xyz/',
124
+ 'user-agent': 'Mozilla/5.0',
125
+ }
126
+
127
  if request.tools:
128
+ 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 tags. Respond in the format {{"name": tool name, "parameters": dictionary of argument name and its value}}. Do not use variables.
129
  Tools: {";".join(f"<tool>{tool}</tool>" for tool in request.tools)}
130
  Response Format for tool call:
131
  <tool_call>
132
  {{"name": <function-name>, "arguments": <args-json-object>}}
133
  </tool_call>"""
134
+ if request.messages[0].role == "system":
135
+ request.messages[0].content += "\n\n" + tool_prompt
136
+ else:
137
+ request.messages.insert(0, Message(role="system", content=tool_prompt))
138
+
139
+ payload = {"messages": [msg.model_dump() for msg in request.messages], "model": model_id}
140
+
141
  if request.stream:
142
  async def event_stream():
143
+ created = int(time.time())
144
+ usage_info = None
145
+ is_first_chunk = True
146
+ tool_call_buffer = ""
147
+ in_tool_call = False
148
+
149
  try:
150
+ async with httpx.AsyncClient(timeout=120) as client:
151
+ async with client.stream("POST", CHAT_API_URL, headers=headers, json=payload) as response:
152
  response.raise_for_status()
153
  async for line in response.aiter_lines():
154
+ if not line:
155
+ continue
156
+
157
  if line.startswith("0:"):
158
+ try:
159
+ content_piece = json.loads(line[2:])
160
+ except json.JSONDecodeError:
161
+ continue
162
+
163
+ current_buffer = content_piece
164
+ if in_tool_call:
165
+ current_buffer = tool_call_buffer + content_piece
166
+
167
+ if "</tool_call>" in current_buffer:
168
+ tool_str = current_buffer.split("<tool_call>")[1].split("</tool_call>")[0]
169
+ tool_json = json.loads(tool_str.strip())
170
+ delta = {
171
+ "content": None,
172
+ "tool_calls": [{"index": 0, "id": generate_random_id("call_"), "type": "function",
173
+ "function": {"name": tool_json["name"], "arguments": json.dumps(tool_json["parameters"])}}]
174
+ }
175
+ chunk = {"id": chat_id, "object": "chat.completion.chunk", "created": created, "model": model_id,
176
+ "choices": [{"index": 0, "delta": delta, "finish_reason": None}], "usage": None}
177
+ yield f"data: {json.dumps(chunk)}\n\n"
178
+
179
+ in_tool_call = False
180
+ tool_call_buffer = ""
181
+ remaining_text = current_buffer.split("</tool_call>", 1)[1]
182
+ if remaining_text:
183
+ content_piece = remaining_text
184
+ else:
185
+ continue
186
+
187
+ if "<tool_call>" in content_piece:
188
+ in_tool_call = True
189
+ tool_call_buffer += content_piece.split("<tool_call>", 1)[1]
190
+ text_before = content_piece.split("<tool_call>", 1)[0]
191
  if text_before:
192
+ delta = {"content": text_before, "tool_calls": None}
193
+ chunk = {"id": chat_id, "object": "chat.completion.chunk", "created": created, "model": model_id,
194
+ "choices": [{"index": 0, "delta": delta, "finish_reason": None}], "usage": None}
195
+ yield f"data: {json.dumps(chunk)}\n\n"
196
+ if "</tool_call>" not in tool_call_buffer:
197
+ continue
198
+
199
  if not in_tool_call:
200
+ delta = {"content": content_piece}
201
+ if is_first_chunk:
202
+ delta["role"] = "assistant"
203
+ is_first_chunk = False
204
+ chunk = {"id": chat_id, "object": "chat.completion.chunk", "created": created, "model": model_id,
205
+ "choices": [{"index": 0, "delta": delta, "finish_reason": None}], "usage": None}
206
+ yield f"data: {json.dumps(chunk)}\n\n"
207
+
208
+ elif line.startswith(("e:", "d:")):
209
+ try:
210
+ usage_info = json.loads(line[2:]).get("usage")
211
+ except (json.JSONDecodeError, AttributeError):
212
+ pass
213
  break
214
+
215
+ final_usage = None
216
+ if usage_info:
217
+ prompt_tokens = usage_info.get("promptTokens", 0)
218
+ completion_tokens = usage_info.get("completionTokens", 0)
219
+ final_usage = {
220
+ "prompt_tokens": prompt_tokens,
221
+ "completion_tokens": completion_tokens,
222
+ "total_tokens": prompt_tokens + completion_tokens
223
+ }
224
+
225
+ finish_reason = "tool_calls" if in_tool_call else "stop"
226
+ done_chunk = {"id": chat_id, "object": "chat.completion.chunk", "created": created, "model": model_id,
227
+ "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}], "usage": final_usage}
228
+ yield f"data: {json.dumps(done_chunk)}\n\n"
229
+
230
+ except httpx.HTTPStatusError as e:
231
+ error_content = {"error": {"message": f"Upstream API error: {e.response.status_code}. Details: {e.response.text}", "type": "upstream_error", "code": str(e.response.status_code)}}
232
+ yield f"data: {json.dumps(error_content)}\n\n"
233
+ finally:
234
+ yield "data: [DONE]\n\n"
235
+
236
+ return StreamingResponse(event_stream(), media_type="text/event-stream")
237
+
238
+ else: # Non-streaming response
239
+ full_response, usage_info = "", {}
240
  try:
241
+ async with httpx.AsyncClient(timeout=120) as client:
242
+ async with client.stream("POST", CHAT_API_URL, headers=headers, json=payload) as response:
243
  response.raise_for_status()
244
  async for chunk in response.aiter_lines():
245
  if chunk.startswith("0:"):
246
+ try:
247
+ full_response += json.loads(chunk[2:])
248
+ except:
249
+ continue
250
+ elif chunk.startswith(("e:", "d:")):
251
+ try:
252
+ usage_info = json.loads(chunk[2:]).get("usage", {})
253
+ except:
254
+ continue
255
+
256
+ tool_calls = None
257
+ content_response = full_response
258
+ finish_reason = "stop"
259
+ if "<tool_call>" in full_response and "</tool_call>" in full_response:
260
+ tool_call_str = full_response.split("<tool_call>")[1].split("</tool_call>")[0]
261
+ tool_call = json.loads(tool_call_str.strip())
262
+ tool_calls = [{
263
+ "id": generate_random_id("call_"),
264
+ "type": "function",
265
+ "function": {
266
+ "name": tool_call["name"],
267
+ "arguments": json.dumps(tool_call["parameters"])
268
+ }
269
+ }]
270
+ content_response = None
271
+ finish_reason = "tool_calls"
272
+
273
+ prompt_tokens = usage_info.get("promptTokens", 0)
274
+ completion_tokens = usage_info.get("completionTokens", 0)
275
+
276
+ return JSONResponse(content={
277
+ "id": chat_id,
278
+ "object": "chat.completion",
279
+ "created": int(time.time()),
280
+ "model": model_id,
281
+ "choices": [{
282
+ "index": 0,
283
+ "message": {
284
+ "role": "assistant",
285
+ "content": content_response,
286
+ "tool_calls": tool_calls
287
+ },
288
+ "finish_reason": finish_reason
289
+ }],
290
+ "usage": {
291
+ "prompt_tokens": prompt_tokens,
292
+ "completion_tokens": completion_tokens,
293
+ "total_tokens": prompt_tokens + completion_tokens
294
+ }
295
+ })
296
+ except httpx.HTTPStatusError as e:
297
+ return JSONResponse(
298
+ status_code=e.response.status_code,
299
+ content={"error": {"message": f"Upstream API error. Details: {e.response.text}", "type": "upstream_error"}}
300
+ )
301
+
302
 
303
  @app.post("/v1/images/generations", tags=["Images"])
304
  async def generate_images(request: ImageGenerationRequest):
305
+ """Handles image generation requests."""
306
+ results = []
307
  try:
308
+ async with httpx.AsyncClient(timeout=120) as client:
309
  for _ in range(request.n):
310
+ model = request.model or "default"
311
+ if model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
312
+ headers = {'Content-Type': 'application/json', 'User-Agent': 'Mozilla/5.0', 'Referer': 'https://www.chatwithmono.xyz/'}
313
+ payload = {"prompt": request.prompt, "model": model}
314
+ resp = await client.post(IMAGE_GEN_API_URL, headers=headers, json=payload)
315
+ resp.raise_for_status()
316
+ data = resp.json()
317
+ b64_image = data.get("image")
318
+ if not b64_image:
319
+ return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})
320
+
321
+ image_url = f"data:image/png;base64,{b64_image}"
322
  if SNAPZION_API_KEY:
323
+ upload_headers = {"Authorization": SNAPZION_API_KEY}
324
+ upload_files = {'file': ('image.png', base64.b64decode(b64_image), 'image/png')}
325
+ upload_resp = await client.post(SNAPZION_UPLOAD_URL, headers=upload_headers, files=upload_files)
326
+ if upload_resp.status_code == 200:
327
+ image_url = upload_resp.json().get("url", image_url)
328
+
329
+ results.append({"url": image_url, "b64_json": b64_image, "revised_prompt": data.get("revised_prompt")})
330
+ else:
331
+ params = {"prompt": request.prompt, "aspect_ratio": request.aspect_ratio, "link": "typegpt.net"}
332
+ resp = await client.get(IMAGE_API_URL, params=params)
333
+ resp.raise_for_status()
334
+ data = resp.json()
335
+ results.append({"url": data.get("image_link"), "b64_json": data.get("base64_output")})
336
+ except httpx.HTTPStatusError as e:
337
+ return JSONResponse(status_code=502, content={"error": f"Image generation failed. Upstream error: {e.response.status_code}", "details": e.response.text})
338
+ except Exception as e:
339
+ return JSONResponse(status_code=500, content={"error": "An internal error occurred.", "details": str(e)})
340
+
341
+ return {"created": int(time.time()), "data": results}
342
 
343
 
 
344
  @app.post("/v1/ocr", response_model=OcrResponse, tags=["OCR"])
345
  async def perform_ocr(request: OcrRequest):
346
  """
 
369
  raw_output = prediction[0]
370
  raw_result_dict = {}
371
 
372
+ # --- Robust Parsing Logic ---
373
  if isinstance(raw_output, str):
374
  try:
375
  # First, try to parse as standard JSON
 
385
  raw_result_dict = {"result": str(parsed_output)}
386
  except (ValueError, SyntaxError):
387
  # If all parsing fails, assume the string is the direct OCR text.
388
+ raw_result_dict = {"ocr_text_from_string": raw_output}
389
  elif isinstance(raw_output, dict):
390
  # It's already a dictionary, use it directly
391
  raw_result_dict = raw_output
392
  else:
393
  # Handle other unexpected data types
394
  raise HTTPException(status_code=502, detail=f"Unexpected data type from OCR service: {type(raw_output)}")
 
395
 
396
+ # Extract text from the dictionary, with multiple fallbacks
397
+ ocr_text = raw_result_dict.get("OCR",
398
+ raw_result_dict.get("ocr_text_from_string",
399
+ str(raw_result_dict)))
400
 
401
  return OcrResponse(ocr_text=ocr_text, raw_response=raw_result_dict)
402
 
 
408
  if temp_file_path and os.path.exists(temp_file_path):
409
  os.unlink(temp_file_path)
410
 
411
+
412
  @app.post("/v1/moderations", tags=["Moderation"])
413
  async def create_moderation(request: ModerationRequest):
414
+ """Handles moderation requests, conforming to the OpenAI API specification."""
415
+ input_texts = [request.input] if isinstance(request.input, str) else request.input
416
+ if not input_texts:
417
+ return JSONResponse(status_code=400, content={"error": {"message": "Request must have at least one input string."}})
418
+
419
+ headers = {'Content-Type': 'application/json', 'User-Agent': 'Mozilla/5.0', 'Referer': 'https://www.chatwithmono.xyz/'}
420
+ results = []
421
+
422
  try:
423
+ async with httpx.AsyncClient(timeout=30) as client:
424
  for text_input in input_texts:
425
+ payload = {"text": text_input}
426
+ resp = await client.post(MODERATION_API_URL, headers=headers, json=payload)
427
+ resp.raise_for_status()
428
+
429
+ upstream_data = resp.json()
430
+ upstream_categories = upstream_data.get("categories", {})
431
+
432
+ openai_categories = {
433
+ "hate": upstream_categories.get("hate", False),
434
+ "hate/threatening": False,
435
+ "harassment": False,
436
+ "harassment/threatening": False,
437
+ "self-harm": upstream_categories.get("self-harm", False),
438
+ "self-harm/intent": False,
439
+ "self-harm/instructions": False,
440
+ "sexual": upstream_categories.get("sexual", False),
441
+ "sexual/minors": False,
442
+ "violence": upstream_categories.get("violence", False),
443
+ "violence/graphic": False,
444
+ }
445
+
446
+ result_item = {
447
+ "flagged": upstream_data.get("overall_sentiment") == "flagged",
448
+ "categories": openai_categories,
449
+ "category_scores": {k: 1.0 if v else 0.0 for k, v in openai_categories.items()},
450
+ }
451
+
452
+ if reason := upstream_data.get("reason"):
453
+ result_item["reason"] = reason
454
+
455
  results.append(result_item)
456
+
457
+ except httpx.HTTPStatusError as e:
458
+ return JSONResponse(
459
+ status_code=502,
460
+ content={"error": {"message": f"Moderation failed. Upstream error: {e.response.status_code}", "details": e.response.text}}
461
+ )
462
+ except Exception as e:
463
+ return JSONResponse(
464
+ status_code=500,
465
+ content={"error": {"message": "An internal error occurred during moderation.", "details": str(e)}}
466
+ )
467
+
468
+ final_response = {
469
+ "id": generate_random_id("modr-"),
470
+ "model": request.model,
471
+ "results": results,
472
+ }
473
+ return JSONResponse(content=final_response)
474
 
475
 
476
  # --- Main Execution ---