fixed copy and output in local models and claude

This commit is contained in:
xssdoctor 2024-04-14 12:21:29 -04:00
parent 198ba8c9ee
commit 161495ed7d

View File

@ -71,20 +71,32 @@ class Standalone:
copy = self.args.copy copy = self.args.copy
if copy: if copy:
pyperclip.copy(response['message']['content']) pyperclip.copy(response['message']['content'])
if self.args.output:
with open(self.args.output, "w") as f:
f.write(response['message']['content'])
async def localStream(self, messages, host=''): async def localStream(self, messages, host=''):
from ollama import AsyncClient from ollama import AsyncClient
buffer = ""
if host: if host:
async for part in await AsyncClient(host=host).chat(model=self.model, messages=messages, stream=True): async for part in await AsyncClient(host=host).chat(model=self.model, messages=messages, stream=True):
buffer += part['message']['content']
print(part['message']['content'], end='', flush=True) print(part['message']['content'], end='', flush=True)
else: else:
async for part in await AsyncClient().chat(model=self.model, messages=messages, stream=True): async for part in await AsyncClient().chat(model=self.model, messages=messages, stream=True):
buffer += part['message']['content']
print(part['message']['content'], end='', flush=True) print(part['message']['content'], end='', flush=True)
if self.args.output:
with open(self.args.output, "w") as f:
f.write(buffer)
if self.args.copy:
pyperclip.copy(buffer)
async def claudeStream(self, system, user): async def claudeStream(self, system, user):
from anthropic import AsyncAnthropic from anthropic import AsyncAnthropic
self.claudeApiKey = os.environ["CLAUDE_API_KEY"] self.claudeApiKey = os.environ["CLAUDE_API_KEY"]
Streamingclient = AsyncAnthropic(api_key=self.claudeApiKey) Streamingclient = AsyncAnthropic(api_key=self.claudeApiKey)
buffer = ""
async with Streamingclient.messages.stream( async with Streamingclient.messages.stream(
max_tokens=4096, max_tokens=4096,
system=system, system=system,
@ -92,9 +104,14 @@ class Standalone:
model=self.model, temperature=self.args.temp, top_p=self.args.top_p model=self.model, temperature=self.args.temp, top_p=self.args.top_p
) as stream: ) as stream:
async for text in stream.text_stream: async for text in stream.text_stream:
buffer += text
print(text, end="", flush=True) print(text, end="", flush=True)
print() print()
if self.args.copy:
pyperclip.copy(buffer)
if self.args.output:
with open(self.args.output, "w") as f:
f.write(buffer)
message = await stream.get_final_message() message = await stream.get_final_message()
async def claudeChat(self, system, user, copy=False): async def claudeChat(self, system, user, copy=False):
@ -112,6 +129,9 @@ class Standalone:
copy = self.args.copy copy = self.args.copy
if copy: if copy:
pyperclip.copy(message.content[0].text) pyperclip.copy(message.content[0].text)
if self.args.output:
with open(self.args.output, "w") as f:
f.write(message.content[0].text)
def streamMessage(self, input_data: str, context="", host=''): def streamMessage(self, input_data: str, context="", host=''):
""" Stream a message and handle exceptions. """ Stream a message and handle exceptions.