File size: 1,856 Bytes
c7e94cf
b8b0b89
 
 
 
55a5dbd
003c5fb
b2b25fd
b8b0b89
 
 
b2b25fd
b8b0b89
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
003c5fb
b8b0b89
b2b25fd
b8b0b89
3e43065
b8b0b89
3e43065
b8b0b89
 
 
 
 
3e43065
 
b8b0b89
 
 
55a5dbd
b8b0b89
 
 
 
55a5dbd
b8b0b89
 
 
6827c6b
b8b0b89
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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):
        # os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HF_token']
        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