AiDA_Python/app/service/chat_robot/script/main.py

import json
import logging

from langchain.agents import Tool
from langchain.callbacks import FileCallbackHandler
from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder
from langchain.schema import SystemMessage, AIMessage
from langchain.utilities import SerpAPIWrapper
from langchain_community.chat_models import ChatTongyi
from loguru import logger

from app.core.config import *
from app.service.chat_robot.script.agents import CustomAgentExecutor, ConversationalFunctionsAgent
from app.service.chat_robot.script.database import CustomDatabase
from app.service.chat_robot.script.memory import UserConversationBufferWindowMemory
from app.service.chat_robot.script.prompt import FASHION_CHAT_BOT_PREFIX, TOOLS_FUNCTIONS_SUFFIX
from app.service.chat_robot.script.service import CallQWen
from app.service.chat_robot.script.tools import (QuerySQLDataBaseTool, InfoSQLDatabaseTool, QuerySQLCheckerTool, ListSQLDatabaseTool)

# os.environ["http_proxy"] = "http://127.0.0.1:7890"
# os.environ["https_proxy"] = "http://127.0.0.1:7890"
# log callbacks
logfile = "logs/chat_debug.log"
logger.add(logfile, colorize=True, enqueue=True)
log_handler = FileCallbackHandler(logfile)

# Initiate our LLM 'gpt-3.5-turbo'
# llm = ChatOpenAI(temperature=0.1,
#                  openai_api_key=OPENAI_API_KEY,
#                  # callbacks=[OpenAICallbackHandler()]
#                  )

llm = ChatTongyi(api_key=QWEN_API_KEY)

search = SerpAPIWrapper()
db = CustomDatabase.from_uri(f'mysql+pymysql://{DB_USERNAME}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/attribute_retrieval_V3',
                             include_tables=['female_top', 'female_skirt', 'female_pants', 'female_dress',
                                             'female_outwear', 'male_bottom', 'male_top', 'male_outwear'],
                             engine_args={"pool_recycle": 7200})
tools = [
    Tool(
        name="internet_search",
        description="Can be used to perform Internet searches",
        func=search.run
    ),
    QuerySQLDataBaseTool(db=db, return_direct=False),
    InfoSQLDatabaseTool(db=db),
    ListSQLDatabaseTool(db=db),
    # QuerySQLCheckerTool(db=db, llm=OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)),
    QuerySQLCheckerTool(db=db, llm=ChatTongyi(temperature=0, api_key=QWEN_API_KEY)),
    # Tool(
    #     name="tutorial_tool",
    #     description="Utilize this tool to retrieve specific statements related to user guidance tutorials."
    #                 "Input is an empty string",
    #     func=CustomTutorialTool(),
    #     return_direct=True
    # )
]

messages = [
    SystemMessage(content=FASHION_CHAT_BOT_PREFIX),
    MessagesPlaceholder(variable_name="history"),
    HumanMessagePromptTemplate.from_template(
        "{input} "
        "Question from a {gender}."
    ),
    AIMessage(content=TOOLS_FUNCTIONS_SUFFIX),
    MessagesPlaceholder(variable_name="agent_scratchpad"),
]

prompt = ChatPromptTemplate(input_variables=["input", "gender", "agent_scratchpad", "history"], messages=messages)
agent = ConversationalFunctionsAgent(
    llm=llm,
    tools=tools,
    prompt=prompt
)

memory = UserConversationBufferWindowMemory.from_redis(
    return_messages=True, k=2, input_key='input', output_key='output'
)
agent_executor = CustomAgentExecutor.from_agent_and_tools(
    agent=agent,
    tools=tools,
    verbose=True,
    memory=memory,
)


def chat(post_data):
    user_id = post_data.user_id
    session_id = post_data.session_id
    input_message = post_data.message

    # final_outputs = agent_executor(
    #     {"input": input_message, "gender": gender},
    #     callbacks=[OpenAITokenRecordCallbackHandler(), log_handler],
    #     session_key=f"buffer:{user_id}:{session_id}",
    # )

    final_outputs = CallQWen.call_with_messages(input_message)
    # api_response = {
    #     'user_id': user_id,
    #     'session_id': session_id,
    #     # 'message_id': message_id,
    #     # 'create_time': created_time,
    #     'input': final_outputs['input'],
    #     # 'conversion': messages,
    #     'output': final_outputs['output'],
    #     # 'gpt_response_time': gpt_response_time,
    #     'total_tokens': final_outputs['total_tokens'],
    #     'total_cost': final_outputs['total_cost'],
    #     'prompt_tokens': final_outputs['prompt_tokens'],
    #     'completion_tokens': final_outputs['completion_tokens'],
    #     'response_type': final_outputs['response_type']
    # }
    # if final_outputs["output"].startswith("["):
    #     final_str = final_outputs["output"].replace("\\", "")
    # else:
    #     final_str = final_outputs["output"]

    api_response = {
        'user_id': user_id,
        'session_id': session_id,
        # 'message_id': message_id,
        # 'create_time': created_time,
        'input': input_message,
        # 'conversion': messages,
        'output': final_outputs["output"],
        # 'gpt_response_time': gpt_response_time,
        'total_tokens': final_outputs['total_tokens'],
        'total_cost': final_outputs['total_cost'],
        'prompt_tokens': final_outputs['prompt_tokens'],
        'completion_tokens': final_outputs['completion_tokens'],
        'response_type': final_outputs["response_type"]
    }
    logging.info(json.dumps(api_response))
    return api_response