hadadrjt commited on
Commit
f4fd6dc
·
1 Parent(s): 612075b

ai: Optimize connection to the server.

Browse files
Files changed (1) hide show
  1. jarvis.py +13 -6
jarvis.py CHANGED
@@ -16,6 +16,7 @@ import pandas as pd
16
  import pptx
17
  import fitz
18
  import io
 
19
  from pathlib import Path
20
  from PIL import Image
21
  from pptx import Presentation
@@ -37,10 +38,14 @@ META_TAGS = os.getenv("META_TAGS")
37
 
38
  ALLOWED_EXTENSIONS = json.loads(os.getenv("ALLOWED_EXTENSIONS"))
39
 
 
 
 
 
 
 
40
  def create_session():
41
- s = requests.Session()
42
- s.headers.update({"Connection": "keep-alive"})
43
- return s
44
 
45
  def get_model_key(display_name):
46
  return next((k for k, v in MODEL_MAPPING.items() if v == display_name), MODEL_CHOICES[0])
@@ -90,16 +95,18 @@ def extract_file_content(file_path):
90
  def chat_with_model(history, user_input, selected_model_display, sess):
91
  if not LINUX_SERVER_PROVIDER_KEYS or not LINUX_SERVER_HOSTS:
92
  return RESPONSES["RESPONSE_3"]
 
 
93
  selected_model = get_model_key(selected_model_display)
94
  model_config = MODEL_CONFIG.get(selected_model, DEFAULT_CONFIG)
95
  messages = [{"role": "user", "content": user} for user, _ in history]
96
  messages += [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
97
  messages.append({"role": "user", "content": user_input})
98
- data = {"model": selected_model, "messages": messages, **model_config}
99
  random.shuffle(LINUX_SERVER_PROVIDER_KEYS)
100
  random.shuffle(LINUX_SERVER_HOSTS)
101
- for api_key in LINUX_SERVER_PROVIDER_KEYS[:2]:
102
- for host in LINUX_SERVER_HOSTS[:2]:
103
  try:
104
  response = sess.post(host, json=data, headers={"Authorization": f"Bearer {api_key}"}, timeout=5)
105
  if response.status_code < 400:
 
16
  import pptx
17
  import fitz
18
  import io
19
+ import uuid
20
  from pathlib import Path
21
  from PIL import Image
22
  from pptx import Presentation
 
38
 
39
  ALLOWED_EXTENSIONS = json.loads(os.getenv("ALLOWED_EXTENSIONS"))
40
 
41
+ class SessionWithID(requests.Session):
42
+ def __init__(self):
43
+ super().__init__()
44
+ self.headers.update({"Connection": "keep-alive"})
45
+ self.session_id = str(uuid.uuid4())
46
+
47
  def create_session():
48
+ return SessionWithID()
 
 
49
 
50
  def get_model_key(display_name):
51
  return next((k for k, v in MODEL_MAPPING.items() if v == display_name), MODEL_CHOICES[0])
 
95
  def chat_with_model(history, user_input, selected_model_display, sess):
96
  if not LINUX_SERVER_PROVIDER_KEYS or not LINUX_SERVER_HOSTS:
97
  return RESPONSES["RESPONSE_3"]
98
+ if not hasattr(sess, "session_id"):
99
+ sess.session_id = str(uuid.uuid4())
100
  selected_model = get_model_key(selected_model_display)
101
  model_config = MODEL_CONFIG.get(selected_model, DEFAULT_CONFIG)
102
  messages = [{"role": "user", "content": user} for user, _ in history]
103
  messages += [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
104
  messages.append({"role": "user", "content": user_input})
105
+ data = {"model": selected_model, "messages": messages, **model_config, "session_id": sess.session_id}
106
  random.shuffle(LINUX_SERVER_PROVIDER_KEYS)
107
  random.shuffle(LINUX_SERVER_HOSTS)
108
+ for api_key in LINUX_SERVER_PROVIDER_KEYS:
109
+ for host in LINUX_SERVER_HOSTS:
110
  try:
111
  response = sess.post(host, json=data, headers={"Authorization": f"Bearer {api_key}"}, timeout=5)
112
  if response.status_code < 400: