90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
import streamlit as st
|
|
|
|
from open_notebook.domain.models import Model
|
|
from open_notebook.domain.notebook import text_search, vector_search
|
|
from open_notebook.graphs.ask import graph as ask_graph
|
|
from pages.stream_app.utils import convert_source_references, setup_page
|
|
|
|
setup_page("🔍 Search")
|
|
|
|
ask_tab, search_tab = st.tabs(["Ask Your Knowledge Base (beta)", "Search"])
|
|
|
|
if "search_results" not in st.session_state:
|
|
st.session_state["search_results"] = []
|
|
|
|
|
|
def results_card(item):
|
|
score = item.get("relevance", item.get("similarity", item.get("score", 0)))
|
|
with st.container(border=True):
|
|
st.markdown(
|
|
f"[{score:.2f}] **[{item['title']}](/?object_id={item['parent_id']})**"
|
|
)
|
|
with st.expander("Matches"):
|
|
for match in item["matches"]:
|
|
st.markdown(match)
|
|
|
|
|
|
with ask_tab:
|
|
st.subheader("Ask Your Knowledge Base (beta)")
|
|
st.caption(
|
|
"The LLM will answer your query based on the documents in your knowledge base. "
|
|
)
|
|
question = st.text_input("Question", "")
|
|
models = Model.get_models_by_type("language")
|
|
strategy_model: Model = st.selectbox(
|
|
"Query Strategy Model",
|
|
models,
|
|
format_func=lambda x: x.name,
|
|
help="This is the LLM that will be responsible for strategizing the search",
|
|
)
|
|
answer_model: Model = st.selectbox(
|
|
"Indivual Answer Model",
|
|
models,
|
|
format_func=lambda x: x.name,
|
|
help="This is the LLM that will be responsible for processing individual subqueries",
|
|
)
|
|
final_answer_model: Model = st.selectbox(
|
|
"Final Answer Model",
|
|
models,
|
|
format_func=lambda x: x.name,
|
|
help="This is the LLM that will be responsible for processing the final answer",
|
|
)
|
|
if st.button("Ask"):
|
|
st.write(f"Searching for {question}")
|
|
rag_results = ask_graph.invoke(
|
|
dict(
|
|
question=question,
|
|
),
|
|
config=dict(
|
|
configurable=dict(
|
|
strategy_model=strategy_model.id,
|
|
answer_model=answer_model.id,
|
|
final_answer_model=final_answer_model.id,
|
|
)
|
|
),
|
|
)
|
|
st.markdown(convert_source_references(rag_results["final_answer"]))
|
|
with st.expander("Details (for debugging)"):
|
|
st.json(rag_results)
|
|
|
|
with search_tab:
|
|
with st.container(border=True):
|
|
st.subheader("🔍 Search")
|
|
st.caption("Search your knowledge base for specific keywords or concepts")
|
|
search_term = st.text_input("Search", "")
|
|
search_type = st.radio("Search Type", ["Text Search", "Vector Search"])
|
|
search_sources = st.checkbox("Search Sources", value=True)
|
|
search_notes = st.checkbox("Search Notes", value=True)
|
|
if st.button("Search"):
|
|
if search_type == "Text Search":
|
|
st.write(f"Searching for {search_term}")
|
|
st.session_state["search_results"] = text_search(
|
|
search_term, 100, search_sources, search_notes
|
|
)
|
|
elif search_type == "Vector Search":
|
|
st.write(f"Searching for {search_term}")
|
|
st.session_state["search_results"] = vector_search(
|
|
search_term, 100, search_sources, search_notes
|
|
)
|
|
for item in st.session_state["search_results"]:
|
|
results_card(item)
|