ParthSadaria commited on
Commit
d03f8cc
·
verified ·
1 Parent(s): df1c187

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +25 -6
main.py CHANGED
@@ -143,9 +143,10 @@ async def get_models():
143
  async def fetch_models():
144
  return await get_models()
145
 
146
- server_status = False #working
 
147
  @app.post("/chat/completions")
148
- @app.post("/v1/chat/completions")
149
  async def get_completion(payload: Payload, request: Request):
150
  # Check server status
151
  if not server_status:
@@ -162,13 +163,17 @@ async def get_completion(payload: Payload, request: Request):
162
  status_code=400,
163
  detail=f"Model '{model_to_use}' is not available. Check /models for the available model list."
164
  )
 
165
  usage_tracker.record_request(model=model_to_use, endpoint="/chat/completions")
166
- # Proceed with the request handling
 
167
  payload_dict = payload.dict()
168
  payload_dict["model"] = model_to_use
169
 
170
  # Select the appropriate endpoint
171
  endpoint = secret_api_endpoint_2 if model_to_use in alternate_models else secret_api_endpoint
 
 
172
  current_time = (datetime.datetime.utcnow() + datetime.timedelta(hours=5, minutes=30)).strftime("%Y-%m-%d %I:%M:%S %p")
173
  aaip = request.client.host
174
  print(f"Time: {current_time}, {aaip}")
@@ -176,11 +181,25 @@ async def get_completion(payload: Payload, request: Request):
176
 
177
  async def stream_generator(payload_dict):
178
  scraper = cloudscraper.create_scraper() # Create a CloudScraper session
 
 
 
 
 
 
 
 
 
179
  try:
180
- # Send POST request using CloudScraper
181
- response = scraper.post(f"{endpoint}/v1/chat/completions", json=payload_dict, stream=True)
 
 
 
 
 
182
 
183
- # Check response status
184
  if response.status_code == 422:
185
  raise HTTPException(status_code=422, detail="Unprocessable entity. Check your payload.")
186
  elif response.status_code == 400:
 
143
  async def fetch_models():
144
  return await get_models()
145
 
146
+ ENDPOINT_ORIGIN = os.getenv('ENDPOINT_ORIGIN', 'https://default-endpoint.com')
147
+
148
  @app.post("/chat/completions")
149
+ @app.post("api/v1/chat/completions")
150
  async def get_completion(payload: Payload, request: Request):
151
  # Check server status
152
  if not server_status:
 
163
  status_code=400,
164
  detail=f"Model '{model_to_use}' is not available. Check /models for the available model list."
165
  )
166
+
167
  usage_tracker.record_request(model=model_to_use, endpoint="/chat/completions")
168
+
169
+ # Prepare payload
170
  payload_dict = payload.dict()
171
  payload_dict["model"] = model_to_use
172
 
173
  # Select the appropriate endpoint
174
  endpoint = secret_api_endpoint_2 if model_to_use in alternate_models else secret_api_endpoint
175
+
176
+ # Current time and IP logging
177
  current_time = (datetime.datetime.utcnow() + datetime.timedelta(hours=5, minutes=30)).strftime("%Y-%m-%d %I:%M:%S %p")
178
  aaip = request.client.host
179
  print(f"Time: {current_time}, {aaip}")
 
181
 
182
  async def stream_generator(payload_dict):
183
  scraper = cloudscraper.create_scraper() # Create a CloudScraper session
184
+
185
+ # Prepare custom headers
186
+ custom_headers = {
187
+ 'DNT': '1',
188
+ 'Origin': ENDPOINT_ORIGIN,
189
+ 'Priority': 'u=1, i',
190
+ 'Referer': ENDPOINT_ORIGIN
191
+ }
192
+
193
  try:
194
+ # Send POST request using CloudScraper with custom headers
195
+ response = scraper.post(
196
+ f"{endpoint}/v1/chat/completions",
197
+ json=payload_dict,
198
+ headers=custom_headers,
199
+ stream=True
200
+ )
201
 
202
+ # Error handling remains the same as in previous version
203
  if response.status_code == 422:
204
  raise HTTPException(status_code=422, detail="Unprocessable entity. Check your payload.")
205
  elif response.status_code == 400: