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")

def generateStory(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()
    return story



question = "什么是克隆羊"
print(generateStory(question))