|
import os |
|
from app_agent_config import AgentConfig |
|
from utils.logger import log_response |
|
from model.custom_agent import CustomHfAgent |
|
from model.conversation_chain_singleton import ConversationChainSingleton |
|
|
|
def cut_text_after_keyword(text, keyword): |
|
index = text.find(keyword) |
|
if index != -1: |
|
return text[:index].strip() |
|
return text |
|
|
|
|
|
class Controller: |
|
def __init__(self): |
|
self.agent_config = AgentConfig() |
|
|
|
|
|
image = [] |
|
def handle_submission(self, user_message ): |
|
log_response("User input \n {}".format(user_message)) |
|
log_response("selected_tools \n {}".format(self.agent_config.s_tool_checkboxes)) |
|
log_response("url_endpoint \n {}".format(self.agent_config.url_endpoint)) |
|
log_response("document \n {}".format(self.agent_config.document)) |
|
log_response("image \n {}".format(self.agent_config.image)) |
|
log_response("context \n {}".format(self.agent_config.context)) |
|
agent = CustomHfAgent( |
|
url_endpoint=self.agent_config.url_endpoint, |
|
token=os.environ['HF_token'], |
|
additional_tools=self.agent_config.tool_checkboxes, |
|
input_params={"max_new_tokens": 192}, |
|
) |
|
|
|
selected_tools = [self.agent_config.tool_loader.tools[idx] for idx, checkbox in enumerate(self.agent_config.s_tool_checkboxe) if checkbox] |
|
|
|
angent_respone = agent.chat(user_message,document=self.agent_config.document,image=self.agent_config.image, context = self.agent_config.context) |
|
|
|
log_response("Agent Response\n {}".format(angent_respone)) |
|
|
|
return angent_respone |
|
|
|
def handle_submission_chat(self, user_message, angent_respone): |
|
|
|
agent_chat_bot = ConversationChainSingleton().get_conversation_chain() |
|
if angent_respone is not None: |
|
text = agent_chat_bot.predict(input=user_message + angent_respone) |
|
else: |
|
text = agent_chat_bot.predict(input=user_message) |
|
|
|
result = cut_text_after_keyword(text, "Human:") |
|
print(result) |
|
return result |
