from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.storage import InMemoryStore
from langchain.retrievers.multi_vector import MultiVectorRetriever
from langchain_core.pydantic_v1 import BaseModel
from operator import itemgetter
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain.schema import StrOutputParser
elon_system_template = """
# Role: ElonMuskGPT
## Profile
- Author: mtl
- Version: 0.1
- Language: 中文
- Description: 你是埃隆·马斯克，一位著名的企业家、工程师、发明家和未来主义者。你创立了SpaceX, Tesla, Inc., Neuralink, 和 The Boring Company。你的愿景是实现人类生活的可持续性，包括促进人类成为多行星物种。
### Skill-1
1. 深入了解太空探索、可再生能源和人工智能。
2. 具备前瞻性思维，能够预见科技趋势和社会需求。
### Skill-2
1. 强大的问题解决能力，能够提出创新解决方案。
2. 出色的领导力和团队管理能力。
## Rules
1. 不要打破角色设定。
2. 提供富有启发性和建设性的建议。
## Workflow
1. 深呼吸，一步步解决问题。
2. 首先，根据用户提出的问题或情境，进行分析。
3. 然后，提供基于你的经验和视角的建议或解决方案。
4. 最后，鼓励用户采取行动，追求创新和卓越。
## Initialization
"""
elon_human_template = """
## 指令:
根据<专业知识库数据>中的已知信息，简洁和专业地回答问题。如果根据已知信息无法回答问题，则直接基于问题本身提供答案。
## 专业知识库数据:
{context}
## 问题:
{question}
## 推理过程:
首先，根据<专业知识库数据>和<相关历史记录>进行分析。如果这些信息足以回答问题，则结合它们提供答案。如果这些信息不足以回答问题，或与问题不相关，则直接基于问题本身提供答案。
## 回答要求:
使用对话的风格，自然地回答问题。如果包含Markdown内容，需按Markdown格式返回。
"""
embedding = OpenAIEmbeddings(
openai_api_key='empty',
openai_api_base='api_base',
openai_api_type='OpenAI',
model="text-embedding-ada-002"
)
loader = WebBaseLoader(
"https://baike.baidu.com/item/%E5%9F%83%E9%9A%86%C2%B7%E9%A9%AC%E6%96%AF%E5%85%8B/3776526?fr=ge_ala",
)
loader.requests_kwargs = {'verify': False}
docs = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000)
docs = text_splitter.split_documents(docs)
vectorstore = Chroma(
collection_name="elon",
embedding_function=embedding
)
store = InMemoryStore()
id_key = "doc_id"
retriever = MultiVectorRetriever(
vectorstore=vectorstore,
docstore=store,
id_key=id_key,
search_kwargs={'k': 1}
)
import uuid
doc_ids = [str(uuid.uuid4()) for _ in docs]
sub_docs = []
for i, doc in enumerate(docs):
_id = doc_ids[i]
_sub_docs = RecursiveCharacterTextSplitter(chunk_size=400).split_documents([doc])
c = 0
for _doc in _sub_docs:
_doc.metadata[id_key] = _id
sub_docs.extend(_sub_docs)
retriever.vectorstore.add_documents(sub_docs)
retriever.docstore.mset(list(zip(doc_ids, docs)))
_prompt = ChatPromptTemplate.from_messages(
[
(
"system",
elon_system_template,
),
("human", elon_human_template),
]
)
_model = ChatOpenAI(
openai_api_key='api_key',
openai_api_base='api_base',
model_name='Qwen-72B-Chat-Int4',
temperature=0.0,
max_tokens=None
)
class Question(BaseModel):
question: str
chain = (
itemgetter("question")
| RunnableParallel({"context": retriever, "question": RunnablePassthrough() })
| _prompt
| _model
| StrOutputParser()
)
chain = chain.with_types(input_type=Question)
if __name__ == "__main__":
print(chain.invoke({
"question": "5月30日发生了什么？"
}))
