122 lines
2.7 KiB
Python
122 lines
2.7 KiB
Python
import json
|
||
from langchain.tools import BaseTool
|
||
import proto.documents_pb2
|
||
|
||
import init
|
||
import doc_client
|
||
import openai
|
||
|
||
# class Eva
|
||
|
||
|
||
# 连续对话
|
||
messages = [
|
||
{
|
||
"role": "system",
|
||
"content": """
|
||
回答问题使用文档，并以提问的语言和Markdown回答，并告诉来源。
|
||
你得用“你”的身份指代用户。请辨别文档中的内容，有一些是不相干的。
|
||
"""
|
||
},
|
||
# {
|
||
# "role": "system",
|
||
# "content": f"""
|
||
# Context: {plain_text}
|
||
# """
|
||
# },
|
||
# {
|
||
# "role": "user",
|
||
# "content": f"""
|
||
# {question}
|
||
# """
|
||
# }
|
||
]
|
||
|
||
|
||
def ask_question(question):
|
||
messages.append({
|
||
"role": "user",
|
||
"content": f"""
|
||
{question}
|
||
"""
|
||
})
|
||
|
||
question_vec = init.text_to_vector(question)
|
||
search_param = {
|
||
"data": [question_vec],
|
||
"anns_field": "vector",
|
||
"param": {"metric_type": "L2"},
|
||
"limit": 5,
|
||
"expr": "user_id == 2",
|
||
"output_fields": ["todo_id", "title", "source", "todo_description", "language", "text", "user_id"],
|
||
}
|
||
res = init.collection.search(**search_param)
|
||
|
||
plain_text = ""
|
||
|
||
for i in range(len(res[0])):
|
||
_chunk_id = res[0][i].id
|
||
|
||
if _chunk_id == "0":
|
||
continue
|
||
|
||
print("正在获取分块 " + str(_chunk_id) + " 的内容...")
|
||
|
||
try:
|
||
_chunk_content = doc_client.stub.GetDocumentChunk(proto.documents_pb2.GetDocumentChunkByIdRequest(
|
||
id=_chunk_id
|
||
))
|
||
|
||
print(_chunk_content)
|
||
|
||
plain_text += ("=== \n" + f"文档 ID:{_chunk_content.document.id}\n"
|
||
+ f"文档内容: {_chunk_content.content}" + "===\n")
|
||
|
||
except Exception as e:
|
||
print(e)
|
||
|
||
messages.append({
|
||
"role": "system",
|
||
"content": f"""
|
||
文档: {plain_text}
|
||
"""
|
||
})
|
||
|
||
print("正在调用 LLM...")
|
||
result = openai.ChatCompletion.create(
|
||
messages=messages, model="gpt-3.5-turbo", temperature=0
|
||
)
|
||
res = result["choices"][0]["message"].to_dict_recursive()
|
||
|
||
# add to
|
||
messages.append({
|
||
"role": "assistant",
|
||
"content": res["content"]
|
||
})
|
||
|
||
|
||
# 删除多余的 system 消息
|
||
for i in range(len(messages)):
|
||
if messages[i]["role"] == "system":
|
||
|
||
if i == 0:
|
||
continue
|
||
|
||
messages.pop(i)
|
||
break
|
||
|
||
return res["content"]
|
||
|
||
|
||
while True:
|
||
|
||
print(messages)
|
||
|
||
if len(messages) > 10:
|
||
messages = messages[-10:]
|
||
print("很抱歉，我只能记住最近 10 条上下文数据。让我们重新开始吧。")
|
||
|
||
question = input("请输入问题：")
|
||
resp = ask_question(question)
|
||
print(resp)
|