import autogen
from autogen import AssistantAgent
from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
import mysql.connector
import random
import requests
import os
from qdrant_client import QdrantClient
import gradio as gr
from groq import Groq

# MySQLに接続
conn = mysql.connector.connect(
        host="www.ryhintl.com",
        user="smairuser",
        password="smairuser",
        port=36000,
        database="smair"
    )

cursor = conn.cursor(dictionary=True)

def get_api_keys():   
    select_one_data_query = "SELECT * FROM agentic_apis_copy where api = 'GROQ_KEYS'"
    cursor.execute(select_one_data_query)
    result = cursor.fetchall()
    keys = [item['key'] for item in result]
    rtnkey = ",".join(map(str, keys))
    
    return rtnkey

# List of Groq tokens
mytokens = get_api_keys()

tokens = eval("["+mytokens+"]")

def get_next_token():
    token = tokens[random.randint(0, len(tokens) - 1)]
    return token

token = get_next_token()

os.environ["GROQ_API_KEY"] = token

# Configure Groq
config_list = [{
    "model": "llama-3.3-70b-versatile",
    "api_key": os.environ["GROQ_API_KEY"],
    "api_type": "groq"
}]

assistant = AssistantAgent(
    name="assistant",
    system_message="You are a helpful assistant.",
    llm_config={"config_list": config_list}
)

client = QdrantClient(url="https://02cbe366-829e-43a6-adf5-3b712a886c21.us-west-1-0.aws.cloud.qdrant.io", api_key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.hWu5_qKaYHUhuMAjUScqw1R_1kkXiidv337wuGKcu9o")

ragproxyagent = RetrieveUserProxyAgent(
    name="ragproxyagent",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=2,
    retrieve_config = {
        "task": "qa",
        "docs_path": [
            "https://www.ryhintl.com/hunting.md"
        ],
        "chunk_token_size": 2000,
        "model": config_list[0]["model"],
        "vector_db": "qdrant",  # Qdrant Cloud database
        "collection_name": "scouter_base",
        "db_config": {"client": client},
        "get_or_create": True,
        "overwrite": True,
        },
    code_execution_config=False,
)

assistant.reset()

def gradio_interface(problem):
    response = ragproxyagent.initiate_chat(assistant, message=ragproxyagent.message_generator, problem=problem)
    groq_assistant_contents = [entry['content'] for entry in response.chat_history if entry['role'] == 'user' and entry['name'] == 'assistant']
    
    if (response.summary == ""):
        # Create the Groq client
        client = Groq(api_key=os.environ["GROQ_API_KEY"])

        # Set the system prompt
        system_prompt = {
            "role": "system",
            "content": "あなたは便利なアシスタントです。質問には簡潔に答えてください。必ず、日本語で答えてください。"
        }

        # Set the user prompt
        user_input = problem
        user_prompt = {
            "role": "user", "content": user_input
        }

        # Initialize the chat history
        chat_history = [system_prompt, user_prompt]

        response = client.chat.completions.create(model="llama3-70b-8192",
            messages=chat_history,
            max_tokens=100,
            temperature=0.1)

        return response.choices[0].message.content
    else:
        return response.summary

# HTML for the title with an image
title_html = """
<div style="display: flex">
    <img src="https://www.ryhintl.com/images/ryhlogo/ryhlogo.png" alt="ハイブリッド・エージェント・アシスタント" style="width:35px;height:35px;">
    &nbsp;&nbsp;&nbsp;ハイブリッド・エージェント・アシスタント
</div>
"""

# Create Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(label="🔎 エージェント・アシスタントの質問", value="トータルソリューションズのAIプロジェクトの立ち上げに必要な専門知識を持っている人材が不足して悩んでいます。AIに関する専門知識を持っている専門家とPYTHONやAGENTの経験を持っている技術者をリスティングして、アドバイスしてください。"),
    outputs=gr.Textbox(label="🗞 エージェント・アシスタントの応答"),
    #title="ハイブリッド・エージェント・アシスタント",
    title=title_html,
    description="質問をするとエージェント・アシスタントから回答が得られます。RAGエージェントから回答が見つからない場合、LLMからの回答を取得します。",
    show_progress=True,
    submit_btn="実行",
    clear_btn="クリア",
    flagging_mode="never"
)

# Launch the Gradio interface
iface.launch()