|
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 |
|
|
|
|
|
|
|
class Controller: |
|
def __init__(self): |
|
self.agent_config = AgentConfig() |
|
|
|
image = [] |
|
def handle_submission(user_message, selected_tools, url_endpoint, document, image, context): |
|
|
|
log_response("User input \n {}".format(user_message)) |
|
log_response("selected_tools \n {}".format(selected_tools)) |
|
log_response("url_endpoint \n {}".format(url_endpoint)) |
|
log_response("document \n {}".format(document)) |
|
log_response("image \n {}".format(image)) |
|
log_response("context \n {}".format(context)) |
|
|
|
agent = CustomHfAgent( |
|
url_endpoint=url_endpoint, |
|
token=os.environ['HF_token'], |
|
additional_tools=selected_tools, |
|
input_params={"max_new_tokens": 192}, |
|
) |
|
|
|
response = agent.chat(user_message,document=document,image=image, context = context) |
|
|
|
log_response("Agent Response\n {}".format(response)) |
|
|
|
return response |
|
|
|
def cut_text_after_keyword(text, keyword): |
|
index = text.find(keyword) |
|
if index != -1: |
|
return text[:index].strip() |
|
return text |
|
|
|
|
|
def handle_submission_chat(user_message, response): |
|
|
|
agent_chat_bot = ConversationChainSingleton().get_conversation_chain() |
|
|
|
if response is not None: |
|
text = agent_chat_bot.predict(input=user_message + response) |
|
else: |
|
text = agent_chat_bot.predict(input=user_message) |
|
|
|
result = cut_text_after_keyword(text, "Human:") |
|
|
|
print(result) |
|
|
|
return result |
