perplexity model, shell decode

frdel committed Jul 22, 2024 at 14:02 UTC a3ad9655454d0d8b1a9923db6579f92c62771d06
2 files changed +17 -12
python/helpers/perplexity_search.py
+1 -1
@@ -94,7 +94,7 @@ def PerplexitySearchLLM(api_key,model_name="sonar-medium-online",base_url="https
94 return call_model
95
96
97 -call_llm = PerplexitySearchLLM(api_key=api_key_from_env,model_name="sonar-medium-online")
97 +call_llm = PerplexitySearchLLM(api_key=api_key_from_env,model_name="llama-3-sonar-large-32k-online")
98
99 def perplexity_search(search_query: str):
100 return call_llm(search_query)
\ No newline at end of file
python/helpers/shell_ssh.py
+16 -11
@@ -17,7 +17,7 @@ class SSHInteractiveSession:
17 self.client = paramiko.SSHClient()
18 self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
19 self.shell = None
20 - self.full_output = ''
20 + self.full_output = b''
21
22 def connect(self):
23 # try 3 times with wait and then except
@@ -49,29 +49,34 @@ class SSHInteractiveSession:
49 def send_command(self, command: str):
50 if not self.shell:
51 raise Exception("Shell not connected")
52 - self.full_output = ""
52 + self.full_output = b""
53 self.shell.send((command + " \\\n" +SSHInteractiveSession.end_comment + "\n").encode())
54
55 def read_output(self) -> Tuple[str, str]:
56 if not self.shell:
57 raise Exception("Shell not connected")
58
59 - partial_output = ''
59 + partial_output = b''
60 while self.shell.recv_ready():
61 - data = self.shell.recv(1024).decode('utf-8')
62 - data = self.clean_string(data)
61 + data = self.shell.recv(1024)
62 partial_output += data
63 self.full_output += data
64 time.sleep(0.1) # Prevent busy waiting
65
67 - self.full_output = self.clean_string(self.full_output)
66 + # Decode once at the end
67 + decoded_partial_output = partial_output.decode('utf-8', errors='replace')
68 + decoded_full_output = self.full_output.decode('utf-8', errors='replace')
69 +
70 + decoded_partial_output = self.clean_string(decoded_partial_output)
71 + decoded_full_output = self.clean_string(decoded_full_output)
72
69 - # split output at end_comment
70 - if SSHInteractiveSession.end_comment in self.full_output:
71 - self.full_output = self.full_output.split(SSHInteractiveSession.end_comment)[-1].lstrip("\r\n")
72 - partial_output = partial_output.split(SSHInteractiveSession.end_comment)[-1].lstrip("\r\n")
73 + # Split output at end_comment
74 + if SSHInteractiveSession.end_comment in decoded_full_output:
75 + decoded_full_output = decoded_full_output.split(SSHInteractiveSession.end_comment)[-1].lstrip("\r\n")
76 + decoded_partial_output = decoded_partial_output.split(SSHInteractiveSession.end_comment)[-1].lstrip("\r\n")
77
74 - return self.full_output, partial_output
78 + return decoded_full_output, decoded_partial_output
79 +
80
81 def clean_string(self, input_string):
82 # Remove ANSI escape codes