import dataiku
import json
from dataiku import SQLExecutor2
from duckduckgo_search import DDGS

PROJECT = "" # DSS project key goes here
client = dataiku.api_client()
project = client.get_project(PROJECT)
LLM_ID = "" # LLM ID for the LLM Mesh connection + model goes here
llm = project.get_llm(LLM_ID)
chat = llm.new_completion()

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_customer_info",
            "description": "Get customer details from the database given their ID",
            "parameters": {
                "type": "object",
                "properties": {
                    "customer_id": {
                        "type": "string",
                        "description": "The unique identifier for the customer",
                    },
                },
                "required": ["customer_id"],
            },
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_company_info",
            "description": "Get company information from internet search",
            "parameters": {
                "type": "object",
                "properties": {
                    "company_name": {
                        "type": "string",
                        "description": "Name of the company to search for",
                    },
                },
                "required": ["company_name"],
            },
        }
    }
]

chat.settings["tools"] = tools

CONTEXT = '''
  You are a helpful assistant with access to customer and company information.
  You have two tools available:
  - get_customer_info: retrieves customer details from our database
  - get_company_info: searches the internet for company information
  Use these tools to provide comprehensive responses about customers and their companies.
'''

CONTENT = 'Who are you and what is your purpose?'

chat.with_message(CONTEXT, role="system")
chat.with_message(CONTENT, role="user")

response = chat.execute()
response.text

# I am an AI assistant designed to help you gather information about customers
# and companies. My main functions include retrieving customer details from our
# database and searching the internet for company information. If you have
# specific questions or tasks related to customers or companies, feel free to
# ask!

chat.with_message(response.text, role="assistant")

customer_id = "fdouetteau"
CONTENT = f"The customer's id is {customer_id}"
chat.with_message(CONTENT, role="user")
response = chat.execute()

tool_calls = response.tool_calls
tool_calls

# [{'type': 'function',
#   'function': {'name': 'get_customer_info',
#    'arguments': '{"customer_id":"fdouetteau"}'},
#   'id': 'call_cQECgVOCgU7OLLb5mrBOZrg5'}]

chat.with_tool_calls(tool_calls, role="assistant")

def get_customer_details(customer_id):
    dataset = dataiku.Dataset("pro_customers_sql")
    table_name = dataset.get_location_info().get('info', {}).get('table')
    executor = SQLExecutor2(dataset=dataset)
    customer_id = customer_id.replace("'", "\\'")
    query_reader = executor.query_to_iter(
        f"""SELECT name, job, company FROM "{table_name}" WHERE id = '{customer_id}'""")
    for (name, job, company) in query_reader.iter_tuples():
        return f"The customer's name is \"{name}\", holding the position \"{job}\" at the company named {company}"
    return f"No information can be found about the customer {customer_id}"

def search_company_info(company_name):
    results = DDGS().answers(company_name + " (company)")
    result = "Information found about " + company_name + ": " + results[0]["text"] + "\n" \
        if len(results) > 0 and "text" in results[0] \
        else None
    if not result:
        results = DDGS().answers(company_name)
        result = "Information found about " + company_name + ": " + results[0]["text"] + "\n" \
            if len(results) > 0 and "text" in results[0] \
            else "No information can be found about the company " + company_name
    return result

def process_tool_calls(tool_calls):
    tool_name = tool_calls[0]["function"]["name"]
    llm_args = json.loads(tool_calls[0]["function"]["arguments"])
    if tool_name == "get_customer_info":
        return get_customer_details(llm_args["customer_id"])
    elif tool_name == "get_company_info":
        return search_company_info(llm_args["company_name"])

tool_call_result = process_tool_calls(tool_calls)

chat.with_tool_output(tool_call_result, tool_call_id=tool_calls[0]["id"])

# Continue the conversation
CONTENT = "Find more information about the company from a search."

chat.with_message(CONTENT, role="user")
response = chat.execute()

tool_calls = response.tool_calls
tool_calls

# [{'type': 'function',
#   'function': {'name': 'get_company_info',
#    'arguments': '{"company_name":"Dataiku"}'},
#   'id': 'call_4lg3yspLrMdJvISnL2aBtzfn'}]

chat.with_tool_calls(tool_calls, role="assistant")

tool_call_result = process_tool_calls(tool_calls)

chat.with_tool_output(tool_call_result, tool_call_id=tool_calls[0]["id"])

def pretty_print_conversation(messages):
    ROLE_EMOJIS = {
        "system": "💻",
        "user": "👤", 
        "assistant": "🤖",
        "tool": "🛠️"
    }
    
    for message in messages:
        role = message["role"]
        content = message.get("content")
        tool_calls = message.get("toolCalls")
            
        # Print main message content first if it exists
        if content and content != "No content":
            print(f"{ROLE_EMOJIS.get(role, '❓')} : {content}\n")
        
        # Then print any tool calls
        if tool_calls:
            for tc in tool_calls:
                print(f"   🤖 Assistant: Let me look that up")
                print(f"   🔧 Tool call ({tc['id']}): Calling {tc['function']['name']} with {tc['function']['arguments']}")
        elif role == "tool":
            print(f"   🛠️ Tool response:  {str(message['toolOutputs'][0]['output'])}\n\n")

pretty_print_conversation(chat.cq["messages"])

# 💻 : 
#   You are a helpful assistant with access to customer and company information.
#   You have two tools available:
#   - get_customer_info: retrieves customer details from our database
#   - get_company_info: searches the internet for company information
#   Use these tools to provide comprehensive responses about customers and their companies.


# 👤 : Who are you and what is your purpose?

# 🤖 : I am an AI assistant designed to provide information and support regarding customers and companies. My purpose is to assist you by retrieving customer details from our database or by gathering company information from online resources. If you have any specific questions or need information about a customer or a company, feel free to ask!

# 👤 : The customer's id is fdouetteau

#    🤖 Assistant: Let me look that up
#    🔧 Tool call (call_9qCXkF7NAt515MBrjycLuVTz): Calling get_customer_info with {"customer_id":"fdouetteau"}
#    🛠️ Tool response:  The customer's name is "Florian Douetteau", holding the position "CEO" at the company named Dataiku


# 👤 : Find more information about the company from a search.

#    🤖 Assistant: Let me look that up
#    🔧 Tool call (call_rRlHPcIcbce18eH8Zmhak1vc): Calling get_company_info with {"company_name":"Dataiku"}
#    🛠️ Tool response:  Information found about Dataiku: Dataiku is an American artificial intelligence and machine learning company which was founded in 2013. In December 2019, Dataiku announced that CapitalG—the late-stage growth venture capital fund financed by Alphabet Inc.—joined Dataiku as an investor and that it had achieved unicorn status. As of 2021, Dataiku is valued at $4.6 billion. Dataiku currently employs more than 1,000 people worldwide between offices in New York, Denver, Washington DC, Los Angeles, Paris, London, Munich, Frankfurt, Sydney, Singapore, Tokyo, and Dubai.