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