from haystack import Pipeline, Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.writers import DocumentWriter
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
from haystack.components.generators import OpenAIGenerator
from haystack.utils import Secret
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.builders import PromptBuilder
from haystack.components.embedders import SentenceTransformersTextEmbedder
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.dataclasses import ChatMessage
import pandas as pd

#從CSV加載產(chǎn)品數(shù)據(jù)
df = pd.read_csv("product_sample.csv")

#初始化內(nèi)存中的文檔存儲區(qū)
document_store = InMemoryDocumentStore()

#將產(chǎn)品數(shù)據(jù)轉(zhuǎn)換為Haystack文檔對象
documents = [
Document(
content=item.product_name, 
meta={
"id": item.uniq_id, 
"price": item.selling_price, 
"url": item.product_url
}
) for item in df.itertuples()
]

#創(chuàng)建一個用于編制文檔索引的管道
indexing_pipeline = Pipeline()

#使用句子轉(zhuǎn)換器模型向管道中添加一個文檔嵌入器
indexing_pipeline.add_component(
instance=SentenceTransformersDocumentEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"), name="doc_embedder"
)

# 向管道中添加文檔寫入器，以便在文檔存儲區(qū)中存儲文檔
indexing_pipeline.add_component(instance=DocumentWriter(document_store=document_store), name="doc_writer")

#將嵌入器的輸出連接到寫入器的輸入
indexing_pipeline.connect("doc_embedder.documents", "doc_writer.documents")

#運行索引管道來處理和存儲文檔
indexing_pipeline.run({"doc_embedder": {"documents": documents}})

template = """
Understand the user query and list of products the user is interested in and return product names as list.
You should always return a Python list. Do not return any explanation.

Examples:
Question: I am interested in camping boots, charcoal and disposable rain jacket.
Answer: ["camping_boots","charcoal","disposable_rain_jacket"]

Question: Need a laptop, wireless mouse, and noise-cancelling headphones for work.
Answer: ["laptop","wireless_mouse","noise_cancelling_headphones"]

Question: {{ question }}
Answer:
"""

product_identifier = Pipeline()

product_identifier.add_component("prompt_builder", PromptBuilder(template=template))
product_identifier.add_component("llm", generator())

product_identifier.connect("prompt_builder", "llm")

template = """
Return product name, price, and url as a python dictionary. 
You should always return a Python dictionary with keys price, name and url for single product.
You should always return a Python list of dictionaries with keys price, name and url for multiple products.
Do not return any explanation.

Legitimate Response Schema:
{"price": "float", "name": "string", "url": "string"}
Legitimate Response Schema for multiple products:
[{"price": "float", "name": "string", "url": "string"},{"price": "float", "name": "string", "url": "string"}]

Context:
{% for document in documents %}
product_price: {{ document.meta['price'] }}
product_url: {{ document.meta['url'] }}
product_id: {{ document.meta['id'] }}
product_name: {{ document.content }}
{% endfor %}
Question: {{ question }}
Answer:
"""

rag_pipe = Pipeline()
rag_pipe.add_component("embedder", SentenceTransformersTextEmbedder(model="sentence-transformers/all-MiniLM-L6-v2"))
rag_pipe.add_component("retriever", InMemoryEmbeddingRetriever(document_store=document_store, top_k=5))
rag_pipe.add_component("prompt_builder", PromptBuilder(template=template))
rag_pipe.add_component("llm", generator())

rag_pipe.connect("embedder.embedding", "retriever.query_embedding")
rag_pipe.connect("retriever", "prompt_builder.documents")
rag_pipe.connect("prompt_builder", "llm")

def product_identifier_func(query: str):
"""
根據(jù)給定的查詢來標識產(chǎn)品，并檢索每個已標識的產(chǎn)品的相關(guān)詳細信息。

參數(shù):
query (str): 用于標識產(chǎn)品的查詢字符串。

返回值:
dict: 一個字典，其中鍵是產(chǎn)品名稱，值是每個產(chǎn)品的詳細信息。如果沒有找到產(chǎn)品，則返回“No product found”。
"""
product_understanding = product_identifier.run({"prompt_builder": {"question": query}})

try:
product_list = literal_eval(product_understanding["llm"]["replies"][0])
except:
return "No product found"

results = {}

for product in product_list:
response = rag_pipe.run({"embedder": {"text": product}, "prompt_builder": {"question": product}})
try:
results[product] = literal_eval(response["llm"]["replies"][0])
except:
results[product] = {}

return results

query = "I want crossbow and woodstock puzzle"
#執(zhí)行函數(shù)
product_identifier_func(query)

# {'crossbow': {'name': 'DB Longboards CoreFlex Crossbow 41" Bamboo Fiberglass '
#                        'Longboard Complete',
#                'price': 237.68,
#                'url': 'https://www.amazon.com/DB-Longboards-CoreFlex-Fiberglass-Longboard/dp/B07KMVJJK7'},
#  'woodstock_puzzle': {'name': 'Woodstock- Collage 500 pc Puzzle',
#                       'price': 17.49,
#                       'url': 'https://www.amazon.com/Woodstock-Collage-500-pc-Puzzle/dp/B07MX21WWX'}}

{
"product_key": {
"name": "string",
"price": "float",
"url": "string"
}
}

def find_budget_friendly_option(selected_product_details):
"""
為每一類產(chǎn)品找到最經(jīng)濟友好的選擇。

參數(shù):
selected_product_details (dict): 一個字典，其中的鍵是產(chǎn)品類別和值是列表的產(chǎn)品細節(jié)。每個產(chǎn)品的細節(jié)都應(yīng)該是一個包含一個“price”鍵的字典。

返回結(jié)果:
dict: 一個字典，其中鍵是產(chǎn)品類別，值是每個類別的最經(jīng)濟友好的產(chǎn)品詳細信息。
"""
budget_friendly_options = {}

for category, items in selected_product_details.items():
if isinstance(items, list):
lowest_price_item = min(items, key=lambda x: x['price'])
else:
lowest_price_item = items

budget_friendly_options[category] = lowest_price_item

return budget_friendly_options

chat_template = '''<|start_header_id|>system<|end_header_id|>

You are a function calling AI model. You are provided with function signatures within <tools></tools> XML tags. You may call one or more functions to assist with the user query. Don't make assumptions about what values to plug into functions. For each function call return a json object with function name and arguments within <tool_call></tool_call> XML tags as follows:
<tool_call>
{"name": <function-name>,"arguments": <args-dict>}
</tool_call>

Here are the available tools:
<tools>
{
"name": "product_identifier_func",
"description": "To understand user interested products and its details",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The query to use in the search. Infer this from the user's message. It should be a question or a statement"
}
},
"required": ["query"]
}
},
{
"name": "find_budget_friendly_option",
"description": "Get the most cost-friendly option. If selected_product_details has morethan one key this should return most cost-friendly options",
"parameters": {
"type": "object",
"properties": {
"selected_product_details": {
"type": "dict",
"description": "Input data is a dictionary where each key is a category name, and its value is either a single dictionary with 'price', 'name', and 'url' keys or a list of such dictionaries; example: {'category1': [{'price': 10.5, 'name': 'item1', 'url': 'http://example.com/item1'}, {'price': 8.99, 'name': 'item2', 'url': 'http://example.com/item2'}], 'category2': {'price': 15.0, 'name': 'item3', 'url': 'http://example.com/item3'}}"
}
},
"required": ["selected_product_details"]
}
}
</tools><|eot_id|><|start_header_id|>user<|end_header_id|>

I need to buy a crossbow<|eot_id|><|start_header_id|>assistant<|end_header_id|>

<tool_call>
{"id":"call_deok","name":"product_identifier_func","arguments":{"query":"I need to buy a crossbow"}}
</tool_call><|eot_id|><|start_header_id|>tool<|end_header_id|>

<tool_response>
{"id":"call_deok","result":{'crossbow': {'price': 237.68,'name': 'crossbow','url': 'https://www.amazon.com/crossbow/dp/B07KMVJJK7'}}}
</tool_response><|eot_id|><|start_header_id|>assistant<|end_header_id|>
'''

##測試代理
messages = [
ChatMessage.from_system(
chat_template
),
ChatMessage.from_user("I need to buy a crossbow for my child and Pokémon for myself."),
]

chat_generator = get_chat_generator()
response = chat_generator.run(messages=messages)
pprint(response)

## 響應(yīng)結(jié)果
{'replies': [ChatMessage(content='<tool_call>\n'
'{"id": 0, "name": "product_identifier_func", '
'"arguments": {"query": "I need to buy a '
'crossbow for my child"}}\n'
'</tool_call>\n'
'<tool_call>\n'
'{"id": 1, "name": "product_identifier_func", '
'"arguments": {"query": "I need to buy a '
'Pokemon for myself"}}\n'
'</tool_call>',
role=<ChatRole.ASSISTANT: 'assistant'>,
name=None,
meta={'finish_reason': 'stop',
'index': 0,
'model': 'llama3-groq-70b-8192-tool-use-preview',
'usage': {'completion_time': 0.217823967,
'completion_tokens': 70,
'prompt_time': 0.041348261,
'prompt_tokens': 561,
'total_time': 0.259172228,
'total_tokens': 631}})]}

def extract_tool_calls(tool_calls_str):
json_objects = re.findall(r'<tool_call>(.*?)</tool_call>', tool_calls_str, re.DOTALL)

result_list = [json.loads(obj) for obj in json_objects]

return result_list

available_functions = {
"product_identifier_func": product_identifier_func, 
"find_budget_friendly_option": find_budget_friendly_option
}

messages.append(ChatMessage.from_user(message))
response = chat_generator.run(messages=messages)

if response and "<tool_call>" in response["replies"][0].content:
function_calls = extract_tool_calls(response["replies"][0].content)
for function_call in function_calls:
# 解析函數(shù)調(diào)用信息
function_name = function_call["name"]
function_args = function_call["arguments"]

#找到相應(yīng)的函數(shù)并用給定的參數(shù)調(diào)用它
function_to_call = available_functions[function_name]
function_response = function_to_call(**function_args)

# 使用`ChatMessage.from_function`在消息列表中附加函數(shù)響應(yīng)
messages.append(ChatMessage.from_function(content=json.dumps(function_response), name=function_name))
response = chat_generator.run(messages=messages)

import gradio as gr

messages = [ChatMessage.from_system(chat_template)]
chat_generator = get_chat_generator()

def chatbot_with_fc(message, messages):
messages.append(ChatMessage.from_user(message))
response = chat_generator.run(messages=messages)

while True:
if response and "<tool_call>" in response["replies"][0].content:
function_calls = extract_tool_calls(response["replies"][0].content)
for function_call in function_calls:
#解析函數(shù)調(diào)用信息
function_name = function_call["name"]
function_args = function_call["arguments"]

#找到相應(yīng)的函數(shù)并用給定的參數(shù)調(diào)用它
function_to_call = available_functions[function_name]
function_response = function_to_call(**function_args)

# 使用`ChatMessage.from_function`在消息列表中附加函數(shù)響應(yīng)
messages.append(ChatMessage.from_function(content=json.dumps(function_response), name=function_name))
response = chat_generator.run(messages=messages)

# 定期對話
else:
messages.append(response["replies"][0])
break
return response["replies"][0].content


def chatbot_interface(user_input, state):
response_content = chatbot_with_fc(user_input, state)
return response_content, state

with gr.Blocks() as demo:
gr.Markdown("# AI Purchase Assistant")
gr.Markdown("Ask me about products you want to buy!")

state = gr.State(value=messages)

with gr.Row():
user_input = gr.Textbox(label="Your message:")
response_output = gr.Markdown(label="Response:")

user_input.submit(chatbot_interface, [user_input, state], [response_output, state])
gr.Button("Send").click(chatbot_interface, [user_input, state], [response_output, state])


demo.launch()