import pandas as pd
import chromadb
import os
import erniebot
erniebot.api_type = "aistudio"
erniebot.access_token = "ff1531c8c0f429f92adbc2eaed2e23bfb5349e0f"
chroma_client = chromadb.PersistentClient(path="data/chroma")

vectordb = chroma_client.get_collection(name="collection")

question = "什么是克隆羊"
query = [question,]
response = erniebot.Embedding.create(
    input=query,
    model='ernie-text-embedding',
)
query_embeddings = response.data[0]['embedding']
print(query_embeddings)

results = vectordb.query(
    query_embeddings = query_embeddings,
    n_results=1 #top_k
)

search_result = {'metadatas': results['metadatas'][0], 'documents': results['documents'][0]}
prompt = f"你是一个擅长给小朋友讲故事的小助手，请根据自己的知识和检索到的内容，对小朋友的疑问进行解答。注意，由于小朋友的年龄较小，问题解答需要用故事的形式。尽量使故事生动、有趣、富有哲理。 \n 用户问题: {question} \n\n 搜索结果:\n {search_result}"
print(prompt)

messages = [
    {"role": "user", "content": prompt}
]
response = erniebot.ChatCompletion.create(
    model="ernie-4.0",
    messages=messages
)
story = response.get_result()
print(story)