AIMaster7 commited on
Commit
3f3c9af
·
verified ·
1 Parent(s): 374cc4b

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +53 -66
main.py CHANGED
@@ -1,11 +1,11 @@
1
  from fastapi import FastAPI, Request
2
  from fastapi.responses import StreamingResponse, JSONResponse
3
  from pydantic import BaseModel
4
- import requests
5
  import time
6
  import json
7
- from typing import List, Optional
8
- from models import AVAILABLE_MODELS
9
 
10
  app = FastAPI()
11
 
@@ -27,6 +27,8 @@ async def list_models():
27
 
28
  @app.post("/v1/chat/completions")
29
  async def chat_completion(request: ChatRequest):
 
 
30
  headers = {
31
  'accept': 'text/event-stream',
32
  'content-type': 'application/json',
@@ -37,96 +39,81 @@ async def chat_completion(request: ChatRequest):
37
 
38
  payload = {
39
  "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
40
- "model": request.model
41
  }
42
 
43
  if request.stream:
44
- def event_stream():
45
  chat_id = f"chatcmpl-{unix_id()}"
46
  created = int(time.time())
47
  sent_done = False
48
 
49
- with requests.post(
50
- "https://www.chatwithmono.xyz/api/chat",
51
- headers=headers,
52
- json=payload,
53
- stream=True,
54
- timeout=120
55
- ) as response:
56
- for line in response.iter_lines(decode_unicode=True):
57
- if line.startswith("0:"):
58
- try:
59
- content_piece = json.loads(line[2:])
60
- chunk_data = {
61
- "id": chat_id,
62
- "object": "chat.completion.chunk",
63
- "created": created,
64
- "model": request.model,
65
- "choices": [
66
- {
67
  "delta": {"content": content_piece},
68
  "index": 0,
69
  "finish_reason": None
70
- }
71
- ]
72
- }
73
- yield f"data: {json.dumps(chunk_data)}\n\n"
74
- except:
75
- continue
76
- elif line.startswith(("e:", "d:")) and not sent_done:
77
- sent_done = True
78
- done_chunk = {
79
- "id": chat_id,
80
- "object": "chat.completion.chunk",
81
- "created": created,
82
- "model": request.model,
83
- "choices": [
84
- {
85
  "delta": {},
86
  "index": 0,
87
  "finish_reason": "stop"
88
- }
89
- ]
90
- }
91
- yield f"data: {json.dumps(done_chunk)}\n\ndata: [DONE]\n\n"
92
-
93
  return StreamingResponse(event_stream(), media_type="text/event-stream")
94
 
95
  else:
96
  assistant_response = ""
97
  usage_info = {}
98
 
99
- with requests.post(
100
- "https://www.chatwithmono.xyz/api/chat",
101
- headers=headers,
102
- json=payload,
103
- stream=True,
104
- timeout=120
105
- ) as response:
106
- for chunk in response.iter_lines(decode_unicode=True):
107
- if chunk.startswith("0:"):
108
- try:
109
- piece = json.loads(chunk[2:])
110
- assistant_response += piece # this is just a string fragment
111
- except:
112
- continue
113
- elif chunk.startswith(("e:", "d:")):
114
- try:
115
- data = json.loads(chunk[2:])
116
- usage_info = data.get("usage", {})
117
- except:
118
- continue
119
 
120
  return JSONResponse(content={
121
  "id": f"chatcmpl-{unix_id()}",
122
  "object": "chat.completion",
123
  "created": int(time.time()),
124
- "model": request.model,
125
  "choices": [{
126
  "index": 0,
127
  "message": {
128
  "role": "assistant",
129
- "content": assistant_response # correctly concatenated string
130
  },
131
  "finish_reason": "stop"
132
  }],
@@ -135,4 +122,4 @@ async def chat_completion(request: ChatRequest):
135
  "completion_tokens": usage_info.get("completionTokens", 0),
136
  "total_tokens": usage_info.get("promptTokens", 0) + usage_info.get("completionTokens", 0),
137
  }
138
- })
 
1
  from fastapi import FastAPI, Request
2
  from fastapi.responses import StreamingResponse, JSONResponse
3
  from pydantic import BaseModel
4
+ from typing import List, Optional
5
  import time
6
  import json
7
+ import httpx
8
+ from models import AVAILABLE_MODELS, MODEL_ALIASES
9
 
10
  app = FastAPI()
11
 
 
27
 
28
  @app.post("/v1/chat/completions")
29
  async def chat_completion(request: ChatRequest):
30
+ model_id = MODEL_ALIASES.get(request.model, request.model)
31
+
32
  headers = {
33
  'accept': 'text/event-stream',
34
  'content-type': 'application/json',
 
39
 
40
  payload = {
41
  "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
42
+ "model": model_id
43
  }
44
 
45
  if request.stream:
46
+ async def event_stream():
47
  chat_id = f"chatcmpl-{unix_id()}"
48
  created = int(time.time())
49
  sent_done = False
50
 
51
+ async with httpx.AsyncClient(timeout=120) as client:
52
+ async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
53
+ async for line in response.aiter_lines():
54
+ if line.startswith("0:"):
55
+ try:
56
+ content_piece = json.loads(line[2:])
57
+ chunk_data = {
58
+ "id": chat_id,
59
+ "object": "chat.completion.chunk",
60
+ "created": created,
61
+ "model": model_id,
62
+ "choices": [{
 
 
 
 
 
 
63
  "delta": {"content": content_piece},
64
  "index": 0,
65
  "finish_reason": None
66
+ }]
67
+ }
68
+ yield f"data: {json.dumps(chunk_data)}\n\n"
69
+ except:
70
+ continue
71
+ elif line.startswith(("e:", "d:")) and not sent_done:
72
+ sent_done = True
73
+ done_chunk = {
74
+ "id": chat_id,
75
+ "object": "chat.completion.chunk",
76
+ "created": created,
77
+ "model": model_id,
78
+ "choices": [{
 
 
79
  "delta": {},
80
  "index": 0,
81
  "finish_reason": "stop"
82
+ }]
83
+ }
84
+ yield f"data: {json.dumps(done_chunk)}\n\ndata: [DONE]\n\n"
 
 
85
  return StreamingResponse(event_stream(), media_type="text/event-stream")
86
 
87
  else:
88
  assistant_response = ""
89
  usage_info = {}
90
 
91
+ async with httpx.AsyncClient(timeout=120) as client:
92
+ async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
93
+ async for chunk in response.aiter_lines():
94
+ if chunk.startswith("0:"):
95
+ try:
96
+ piece = json.loads(chunk[2:])
97
+ assistant_response += piece
98
+ except:
99
+ continue
100
+ elif chunk.startswith(("e:", "d:")):
101
+ try:
102
+ data = json.loads(chunk[2:])
103
+ usage_info = data.get("usage", {})
104
+ except:
105
+ continue
 
 
 
 
 
106
 
107
  return JSONResponse(content={
108
  "id": f"chatcmpl-{unix_id()}",
109
  "object": "chat.completion",
110
  "created": int(time.time()),
111
+ "model": model_id,
112
  "choices": [{
113
  "index": 0,
114
  "message": {
115
  "role": "assistant",
116
+ "content": assistant_response
117
  },
118
  "finish_reason": "stop"
119
  }],
 
122
  "completion_tokens": usage_info.get("completionTokens", 0),
123
  "total_tokens": usage_info.get("promptTokens", 0) + usage_info.get("completionTokens", 0),
124
  }
125
+ })