shelljun/LlamaindexRAG

1import streamlit as st2from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings3from llama_index.embeddings.huggingface import HuggingFaceEmbedding4from llama_index.legacy.callbacks import CallbackManager5from llama_index.llms.openai_like import OpenAILike6 7# Create an instance of CallbackManager8callback_manager = CallbackManager()9 10api_base_url =  "https://internlm-chat.intern-ai.org.cn/puyu/api/v1/"11model = "internlm2.5-latest"12api_key = "eyJ0eXBlIjoiSldUIiwiYWxnIjoiSFM1MTIifQ.eyJqdGkiOiIxNzAwODcxOSIsInJvbCI6IlJPTEVfUkVHSVNURVIiLCJpc3MiOiJPcGVuWExhYiIsImlhdCI6MTczMzI4MTc1MiwiY2xpZW50SWQiOiJlYm1ydm9kNnlvMG5semFlazF5cCIsInBob25lIjoiMTc2MjE5NTM5MjkiLCJ1dWlkIjoiOWMyNzhkZDktM2I1My00YjdjLWI5MmQtMzVmNjQxZmRhZDk0IiwiZW1haWwiOiIiLCJleHAiOjE3NDg4MzM3NTJ9.7ZH8A7DLLj5djjYEAF23AMebrc2q7gIhvBbfzeGB3wZnZ2BwJLbBTnLrN-eehcBPiIQneBV8GcKLzJwEYJ1uDA"13 14# api_base_url =  "https://api.siliconflow.cn/v1"15# model = "internlm/internlm2_5-7b-chat"16# api_key = "请填写 API Key"17 18llm =OpenAILike(model=model, api_base=api_base_url, api_key=api_key, is_chat_model=True,callback_manager=callback_manager)19 20 21 22st.set_page_config(page_title="llama_index_demo", page_icon="🦜🔗")23st.title("llama_index_demo")24 25# 初始化模型26@st.cache_resource27def init_models():28    embed_model = HuggingFaceEmbedding(29        model_name="model"30    )31    Settings.embed_model = embed_model32 33    #用初始化llm34    Settings.llm = llm35 36    documents = SimpleDirectoryReader("data").load_data()37    index = VectorStoreIndex.from_documents(documents)38    query_engine = index.as_query_engine()39 40    return query_engine41 42# 检查是否需要初始化模型43if 'query_engine' not in st.session_state:44    st.session_state['query_engine'] = init_models()45 46def greet2(question):47    response = st.session_state['query_engine'].query(question)48    return response49 50      51# Store LLM generated responses52if "messages" not in st.session_state.keys():53    st.session_state.messages = [{"role": "assistant", "content": "你好，我是你的助手，有什么我可以帮助你的吗？"}]    54 55    # Display or clear chat messages56for message in st.session_state.messages:57    with st.chat_message(message["role"]):58        st.write(message["content"])59 60def clear_chat_history():61    st.session_state.messages = [{"role": "assistant", "content": "你好，我是你的助手，有什么我可以帮助你的吗？"}]62 63st.sidebar.button('Clear Chat History', on_click=clear_chat_history)64 65# Function for generating LLaMA2 response66def generate_llama_index_response(prompt_input):67    return greet2(prompt_input)68 69# User-provided prompt70if prompt := st.chat_input():71    st.session_state.messages.append({"role": "user", "content": prompt})72    with st.chat_message("user"):73        st.write(prompt)74 75# Gegenerate_llama_index_response last message is not from assistant76if st.session_state.messages[-1]["role"] != "assistant":77    with st.chat_message("assistant"):78        with st.spinner("Thinking..."):79            response = generate_llama_index_response(prompt)80            placeholder = st.empty()81            placeholder.markdown(response)82    message = {"role": "assistant", "content": response}83    st.session_state.messages.append(message)84