why10w/answerQuestions.py

import pandas as pd
import chromadb
import os
import erniebot

# 需要设置token
erniebot.api_type = "aistudio"
erniebot.access_token = ""

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


if __name__ == "__main__":
    question = "什么是克隆羊"
    print(generateStory(question))