support RAG refs #5

app.py

@@ -1,3 +1,4 @@
+from time import time
 from datetime import datetime, date, timedelta
 from typing import Iterable
 import streamlit as st
@@ -6,6 +7,9 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Qdrant
 from qdrant_client import QdrantClient
 from qdrant_client.http.models import Filter, FieldCondition, MatchValue, Range
+from langchain.chains import RetrievalQA
+from openai.error import InvalidRequestError
+from langchain.chat_models import ChatOpenAI
 from config import DB_CONFIG
 from model import Issue
 
@@ -23,7 +27,14 @@ def load_embeddings():
     return embeddings
 
 
+@st.cache_resource
+def llm_model(model="gpt-3.5-turbo", temperature=0.2):
+    llm = ChatOpenAI(model=model, temperature=temperature)
+    return llm
+
+
 EMBEDDINGS = load_embeddings()
+LLM = llm_model()
 
 
 def make_filter_obj(options: list[dict[str]]):
@@ -67,14 +78,46 @@ def get_similay(query: str, filter: Filter):
     return docs
 
 
-def main(
+def get_retrieval_qa(filter: Filter):
+    db_url, db_api_key, db_collection_name = DB_CONFIG
+    client = QdrantClient(url=db_url, api_key=db_api_key)
+    db = Qdrant(
+        client=client, collection_name=db_collection_name, embeddings=EMBEDDINGS
+    )
+    retriever = db.as_retriever(
+        search_kwargs={
+            "filter": filter,
+        }
+    )
+    result = RetrievalQA.from_chain_type(
+        llm=LLM,
+        chain_type="stuff",
+        retriever=retriever,
+        return_source_documents=True,
+    )
+    return result
+
+
+def _get_related_url(metadata) -> Iterable[str]:
+    urls = set()
+    for m in metadata:
+        url = m["url"]
+        if url in urls:
+            continue
+        urls.add(url)
+        created_at = datetime.fromtimestamp(m["created_at"])
+        # print(m)
+        yield f'<p>URL: <a href="{url}">{url}</a> (created: {created_at:%Y-%m-%d})</p>'
+
+
+def _get_query_str_filter(
     query: str,
     repo_name: str,
     query_options: str,
     start_date: date,
     end_date: date,
     include_comments: bool,
-) -> Iterable[tuple[Issue, float, str]]:
+) -> tuple[str, Filter]:
     options = [{"key": "metadata.repo_name", "value": repo_name}]
     if start_date is not None and end_date is not None:
         options.append(
@@ -96,6 +139,44 @@ def main(
     if query_options == "Empty":
         query_options = ""
     query_str = f"{query_options}{query}"
+    return query_str, filter
+
+
+def run_qa(
+    query: str,
+    repo_name: str,
+    query_options: str,
+    start_date: date,
+    end_date: date,
+    include_comments: bool,
+) -> tuple[str, str]:
+    now = time()
+    query_str, filter = _get_query_str_filter(
+        query, repo_name, query_options, start_date, end_date, include_comments
+    )
+    qa = get_retrieval_qa(filter)
+    try:
+        result = qa(query_str)
+    except InvalidRequestError as e:
+        return "回答が見つかりませんでした。別な質問をしてみてください", str(e)
+    else:
+        metadata = [s.metadata for s in result["source_documents"]]
+        sec_html = f"<p>実行時間: {(time() - now):.2f}秒</p>"
+        html = "<div>" + sec_html + "\n".join(_get_related_url(metadata)) + "</div>"
+    return result["result"], html
+
+
+def run_search(
+    query: str,
+    repo_name: str,
+    query_options: str,
+    start_date: date,
+    end_date: date,
+    include_comments: bool,
+) -> Iterable[tuple[Issue, float, str]]:
+    query_str, filter = _get_query_str_filter(
+        query, repo_name, query_options, start_date, end_date, include_comments
+    )
     docs = get_similay(query_str, filter)
     for doc, score in docs:
         text = doc.page_content
@@ -153,13 +234,14 @@ with st.form("my_form"):
     )
     include_comments = st.checkbox(label="Include Issue comments", value=True)
 
-    submitted = st.form_submit_button("Submit")
-    if submitted:
+    submit_col1, submit_col2 = st.columns(2)
+    searched = submit_col1.form_submit_button("Search")
+    if searched:
         st.divider()
         st.header("Search Results")
         st.divider()
         with st.spinner("Searching..."):
-            results = main(
+            results = run_search(
                 query, repo_name, query_options, start_date, end_date, include_comments
             )
             for issue, score, text in results:
@@ -182,3 +264,17 @@ with st.form("my_form"):
                     st.write(f"{labels=}")
                     # st.markdown(html, unsafe_allow_html=True)
                     st.divider()
+    qa_searched = submit_col2.form_submit_button("QA Search by OpenAI")
+    if qa_searched:
+        st.divider()
+        st.header("QA Search Results by OpenAI GPT-3")
+        st.divider()
+        with st.spinner("QA Searching..."):
+            results = run_qa(
+                query, repo_name, query_options, start_date, end_date, include_comments
+            )
+            answer, html = results
+            with st.container():
+                st.write(answer)
+                st.markdown(html, unsafe_allow_html=True)
+                st.divider()

requirements.txt

@@ -8,3 +8,4 @@ bitsandbytes
 sentence_transformers
 streamlit
 python-dateutil
+openai