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)存中的文檔存儲(chǔ)區(qū)
document_store = InMemoryDocumentStore()

#將產(chǎn)品數(shù)據(jù)轉(zhuǎn)換為Haystack文檔對(duì)象
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)建一個(gè)用于編制文檔索引的管道
indexing_pipeline = Pipeline()

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

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

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

#運(yùn)行索引管道來(lái)處理和存儲(chǔ)文檔
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ù)給定的查詢(xún)來(lái)標(biāo)識(shí)產(chǎn)品，并檢索每個(gè)已標(biāo)識(shí)的產(chǎn)品的相關(guān)詳細(xì)信息。

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

返回值:
dict: 一個(gè)字典，其中鍵是產(chǎn)品名稱(chēng)，值是每個(gè)產(chǎn)品的詳細(xì)信息。如果沒(méi)有找到產(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):
"""
為每一類(lèi)產(chǎn)品找到最經(jīng)濟(jì)友好的選擇。

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

返回結(jié)果:
dict: 一個(gè)字典，其中鍵是產(chǎn)品類(lèi)別，值是每個(gè)類(lèi)別的最經(jīng)濟(jì)友好的產(chǎn)品詳細(xì)信息。
"""
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|>
'''

##測(cè)試代理
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)

# 定期對(duì)話(huà)
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()