209 lines
6.6 KiB
Python
209 lines
6.6 KiB
Python
|
import contextlib
|
||
|
|
import json
|
||
|
|
|
||
|
|
import aiohttp
|
||
|
|
|
||
|
|
from jarvis import CFG
|
||
|
|
from jarvis.ai_agent.agent_utils import must_not_be_valid_json, get_thoughts, get_function, execute_function
|
||
|
|
from jarvis.ai_agent.base_agent import BaseAgent
|
||
|
|
from jarvis.functional_modules.functional_module import CallerContext, moduleRegistry
|
||
|
|
from jarvis.json_utils.json_fix_llm import fix_json_using_multiple_techniques
|
||
|
|
from jarvis.json_utils.utilities import validate_json
|
||
|
|
from jarvis.logger import logger
|
||
|
|
|
||
|
|
|
||
|
|
def _generate_system_prompt():
|
||
|
|
return """Since now, every your response should satisfy the following JSON format, a 'function' must be chosen:
|
||
|
|
```
|
||
|
|
{
|
||
|
|
"thoughts": {
|
||
|
|
"text": "<Your thought>",
|
||
|
|
"reasoning": "<Your reasoning, think step by step>",
|
||
|
|
"speak": "<what you want to say to me>"
|
||
|
|
},
|
||
|
|
"function": {
|
||
|
|
"name": "<mandatory, one of listed functions>",
|
||
|
|
"args": {
|
||
|
|
"arg name": "<value>"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
```
|
||
|
|
I will ask you questions or ask you to do something. You should:
|
||
|
|
First, you should determine if you know the answer of the question or you can accomplish the task directly.
|
||
|
|
If so, you should response directly.
|
||
|
|
If not, you should try to complete the task by calling the functions below.
|
||
|
|
If you can't accomplish the task by yourself and no function is able to accomplish the task, say "Dear master, sorry, I'm not able to do that."
|
||
|
|
|
||
|
|
```
|
||
|
|
Available functions:
|
||
|
|
```
|
||
|
|
""" + moduleRegistry.to_prompt() + """
|
||
|
|
```
|
||
|
|
Your setup:
|
||
|
|
```
|
||
|
|
{
|
||
|
|
"author": "OpenDAN",
|
||
|
|
"name": "Jarvis",
|
||
|
|
}
|
||
|
|
Example:
|
||
|
|
```
|
||
|
|
Tom: generate a picture of me.
|
||
|
|
Jarvis: {
|
||
|
|
"function": {
|
||
|
|
"name": "stable_diffusion",
|
||
|
|
"args": {
|
||
|
|
"prompt": "me"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
```
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
def _generate_request(prompt: str):
|
||
|
|
return {
|
||
|
|
'prompt': prompt,
|
||
|
|
'max_new_tokens': 1000,
|
||
|
|
'do_sample': True,
|
||
|
|
'temperature': 0.5,
|
||
|
|
'top_p': 0.5,
|
||
|
|
'typical_p': 1,
|
||
|
|
'repetition_penalty': 1.18,
|
||
|
|
'top_k': 40,
|
||
|
|
'min_length': 0,
|
||
|
|
'no_repeat_ngram_size': 0,
|
||
|
|
'num_beams': 1,
|
||
|
|
'penalty_alpha': 0,
|
||
|
|
'length_penalty': 1,
|
||
|
|
'early_stopping': False,
|
||
|
|
'seed': -1,
|
||
|
|
'add_bos_token': True,
|
||
|
|
'truncation_length': 2048,
|
||
|
|
'ban_eos_token': False,
|
||
|
|
'skip_special_tokens': True,
|
||
|
|
'stopping_strings': ["Tom: "]
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _convert_role(role: str):
|
||
|
|
if role == 'user':
|
||
|
|
return 'Tom'
|
||
|
|
if role == 'assistant':
|
||
|
|
return 'Jarvis'
|
||
|
|
return role
|
||
|
|
|
||
|
|
|
||
|
|
async def _completion(prompt):
|
||
|
|
async with aiohttp.ClientSession() as session:
|
||
|
|
# body = json.dumps(_generate_request(prompt))
|
||
|
|
async with session.post(CFG.private_ai_address, json=_generate_request(prompt)) as response:
|
||
|
|
if response.status == 200:
|
||
|
|
resp_obj = await response.json()
|
||
|
|
logger.debug(f"Completion result: {json.dumps(resp_obj, indent=2)}")
|
||
|
|
result = resp_obj["results"][0]['text']
|
||
|
|
return result
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
class WebuiAgent(BaseAgent):
|
||
|
|
_system_prompt: str
|
||
|
|
_history = []
|
||
|
|
|
||
|
|
def __init__(self, context: CallerContext):
|
||
|
|
super().__init__(context)
|
||
|
|
self._system_prompt = _generate_system_prompt()
|
||
|
|
|
||
|
|
async def feed_prompt(self, prompt):
|
||
|
|
prompt = f'Tom: {prompt}'
|
||
|
|
self._history.append(prompt)
|
||
|
|
final_prompt = self._system_prompt + '\n' + '\n'.join(self._history)
|
||
|
|
logger.debug(f"Final prompt: {final_prompt}")
|
||
|
|
reply = await self._feed_prompt_to_get_respones(final_prompt)
|
||
|
|
await self._handle_reply(reply)
|
||
|
|
|
||
|
|
async def _feed_prompt_to_get_respones(self, prompt):
|
||
|
|
assistant_reply = await _completion(prompt)
|
||
|
|
|
||
|
|
reply = {
|
||
|
|
"thoughts": None,
|
||
|
|
"reasoning": None,
|
||
|
|
"speak": None,
|
||
|
|
"function": None,
|
||
|
|
"arguments": None,
|
||
|
|
}
|
||
|
|
|
||
|
|
if must_not_be_valid_json(assistant_reply):
|
||
|
|
raise Exception(f"AI replied an invalid response: {assistant_reply}!")
|
||
|
|
else:
|
||
|
|
assistant_reply_json = await fix_json_using_multiple_techniques(assistant_reply)
|
||
|
|
|
||
|
|
# Print Assistant thoughts
|
||
|
|
if assistant_reply_json != {}:
|
||
|
|
validate_json(assistant_reply_json, "llm_response_format_1")
|
||
|
|
try:
|
||
|
|
get_thoughts(reply, assistant_reply_json)
|
||
|
|
get_function(reply, assistant_reply_json)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"AI replied an invalid response: {assistant_reply}. Error: {str(e)}")
|
||
|
|
raise e
|
||
|
|
else:
|
||
|
|
raise Exception(f"AI replied an invalid response: {assistant_reply}!")
|
||
|
|
|
||
|
|
function_name = reply["function"]
|
||
|
|
if function_name is None or function_name == '':
|
||
|
|
raise Exception(f"Missing a function")
|
||
|
|
arguments = reply["arguments"]
|
||
|
|
|
||
|
|
if not isinstance(arguments, dict):
|
||
|
|
raise Exception(f"Invalid arguments, it MUST be a dict")
|
||
|
|
return reply
|
||
|
|
|
||
|
|
async def _handle_reply(self, reply):
|
||
|
|
# TODO: It's not reliable now, thus do nothing now.
|
||
|
|
return
|
||
|
|
if reply is None:
|
||
|
|
await self._caller_context.reply_text("Sorry, but I don't understand what you want me to do.")
|
||
|
|
return
|
||
|
|
|
||
|
|
# Execute function
|
||
|
|
function_name: str = reply["function"]
|
||
|
|
arguments: dict = reply["arguments"]
|
||
|
|
|
||
|
|
await self._caller_context.reply_text(reply["speak"])
|
||
|
|
execute_error = None
|
||
|
|
try:
|
||
|
|
function_result = await execute_function(self._caller_context, function_name, **arguments)
|
||
|
|
except Exception as e:
|
||
|
|
function_result = "Failed"
|
||
|
|
execute_error = e
|
||
|
|
result = f"Function {function_name} returned: " f"{function_result}"
|
||
|
|
|
||
|
|
if function_name is not None:
|
||
|
|
if result is not None:
|
||
|
|
self._caller_context.append_history_message("system", result)
|
||
|
|
logger.debug(f"SYSTEM: {result}")
|
||
|
|
else:
|
||
|
|
self._caller_context.append_history_message("system", "Unable to execute function")
|
||
|
|
logger.debug("SYSTEM: Unable to execute function")
|
||
|
|
|
||
|
|
if execute_error is not None:
|
||
|
|
raise execute_error
|
||
|
|
|
||
|
|
def append_history_message(self, role: str, content: str):
|
||
|
|
self._history.append({'role': role, 'content': content})
|
||
|
|
|
||
|
|
def clear_history_messages(self):
|
||
|
|
self._history.clear()
|
||
|
|
|
||
|
|
def save_history(self, to_where):
|
||
|
|
with open(to_where, "w") as f:
|
||
|
|
s = json.dumps(self._history)
|
||
|
|
f.write(s)
|
||
|
|
|
||
|
|
def load_history(self, from_where):
|
||
|
|
with contextlib.suppress(Exception):
|
||
|
|
with open(from_where, "r") as f:
|
||
|
|
self._history = json.loads(f.read())
|