initial modify

README.md

@@ -1,5 +1,5 @@
 ---
-title: Github Issue Search
+title: Document Search
 emoji: 🐠
 colorFrom: green
 colorTo: purple

app.py

@@ -1,5 +1,5 @@
 from time import time
-from datetime import datetime, date, timedelta
+from datetime import datetime
 from typing import Iterable
 import streamlit as st
 import torch
@@ -13,7 +13,7 @@ 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
+from model import Doc
 
 
 @st.cache_resource
@@ -108,12 +108,12 @@ def get_similay(query: str, filter: Filter):
     db = Qdrant(
         client=client, collection_name=db_collection_name, embeddings=EMBEDDINGS
     )
-    docs = db.similarity_search_with_score(
+    qdocs = db.similarity_search_with_score(
         query,
         k=20,
         filter=filter,
     )
-    return docs
+    return qdocs
 
 
 def get_retrieval_qa(filter: Filter, llm):
@@ -150,49 +150,20 @@ def _get_related_url(metadata) -> Iterable[str]:
 
 def _get_query_str_filter(
     query: str,
-    repo_name: str,
-    query_options: str,
-    start_date: date,
-    end_date: date,
-    include_comments: bool,
+    project_name: 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(
-            {
-                "key": "metadata.created_at",
-                "range": {
-                    "gte": int(datetime.fromisoformat(str(start_date)).timestamp()),
-                    "lte": int(
-                        datetime.fromisoformat(
-                            str(end_date + timedelta(days=1))
-                        ).timestamp()
-                    ),
-                },
-            }
-        )
-    if not include_comments:
-        options.append({"key": "metadata.type_", "value": "issue"})
+    options = [{"key": "metadata.project_name", "value": project_name}]
     filter = make_filter_obj(options=options)
-    if query_options == "Empty":
-        query_options = ""
-    query_str = f"{query_options}{query}"
-    return query_str, filter
+    return query, filter
 
 
 def run_qa(
     llm,
     query: str,
-    repo_name: str,
-    query_options: str,
-    start_date: date,
-    end_date: date,
-    include_comments: bool,
+    project_name: str,
 ) -> tuple[str, str]:
     now = time()
-    query_str, filter = _get_query_str_filter(
-        query, repo_name, query_options, start_date, end_date, include_comments
-    )
+    query_str, filter = _get_query_str_filter(query, project_name)
     qa = get_retrieval_qa(filter, llm)
     try:
         result = qa(query_str)
@@ -207,71 +178,29 @@ def run_qa(
 
 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
-        metadata = doc.metadata
+    project_name: str,
+) -> Iterable[tuple[Doc, float, str]]:
+    query_str, filter = _get_query_str_filter(query, project_name)
+    qdocs = get_similay(query_str, filter)
+    for qdoc, score in qdocs:
+        text = qdoc.page_content
+        metadata = qdoc.metadata
         # print(metadata)
-        issue = Issue(
-            repo_name=repo_name,
+        doc = Doc(
+            project_name=project_name,
             id=metadata.get("id"),
             title=metadata.get("title"),
-            created_at=metadata.get("created_at"),
+            ctime=metadata.get("ctime"),
             user=metadata.get("user"),
             url=metadata.get("url"),
-            labels=metadata.get("labels"),
-            type_=metadata.get("type_"),
         )
-        yield issue, score, text
+        yield doc, score, text
 
 
 with st.form("my_form"):
-    st.title("GitHub Issue Search")
+    st.title("Document Search")
     query = st.text_input(label="query")
-    repo_name = st.radio(
-        options=[
-            "cpython",
-            "pyvista",
-            "plone",
-            "volto",
-            "plone.restapi",
-            "nvda",
-            "nvdajp",
-            "cocoa",
-        ],
-        label="Repo name",
-    )
-    query_options = st.radio(
-        options=[
-            "query: ",
-            "query: passage: ",
-            "Empty",
-        ],
-        label="Query options",
-    )
-    date_min = date(2022, 1, 1)
-    date_max = date.today()
-    date_col1, date_col2 = st.columns(2)
-    start_date = date_col1.date_input(
-        label="Select a start date",
-        value=date_min,
-        format="YYYY-MM-DD",
-    )
-    end_date = date_col2.date_input(
-        label="Select a end date",
-        value=date_max,
-        format="YYYY-MM-DD",
-    )
-    include_comments = st.checkbox(label="Include Issue comments", value=True)
+    project_name = st.text_input(label="project")
 
     submit_col1, submit_col2 = st.columns(2)
     searched = submit_col1.form_submit_button("Search")
@@ -280,28 +209,19 @@ with st.form("my_form"):
         st.header("Search Results")
         st.divider()
         with st.spinner("Searching..."):
-            results = run_search(
-                query, repo_name, query_options, start_date, end_date, include_comments
-            )
-            for issue, score, text in results:
-                title = issue.title
-                url = issue.url
-                id_ = issue.id
+            results = run_search(query, project_name)
+            for doc, score, text in results:
+                title = doc.title
+                url = doc.url
+                id_ = doc.id
                 score = round(score, 3)
-                created_at = datetime.fromtimestamp(issue.created_at)
-                user = issue.user
-                labels = issue.labels
-                is_comment = issue.type_ == "comment"
+                ctime = datetime.fromtimestamp(doc.ctime)
+                user = doc.user
                 with st.container():
-                    if not is_comment:
-                        st.subheader(f"#{id_} - {title}")
-                    else:
-                        st.subheader(f"comment with {title}")
+                    st.subheader(f"#{id_} - {title}")
                     st.write(url)
                     st.write(text)
-                    st.write("score:", score, "Date:", created_at.date(), "User:", user)
-                    st.write(f"{labels=}")
-                    # st.markdown(html, unsafe_allow_html=True)
+                    st.write("score:", score, "Date:", ctime.date(), "User:", user)
                     st.divider()
     qa_searched = submit_col2.form_submit_button("QA Search by OpenAI")
     if qa_searched:
@@ -312,11 +232,7 @@ with st.form("my_form"):
             results = run_qa(
                 LLM,
                 query,
-                repo_name,
-                query_options,
-                start_date,
-                end_date,
-                include_comments,
+                project_name,
             )
             answer, html = results
             with st.container():
@@ -333,11 +249,7 @@ with st.form("my_form"):
                 results = run_qa(
                     VICUNA_LLM,
                     query,
-                    repo_name,
-                    query_options,
-                    start_date,
-                    end_date,
-                    include_comments,
+                    project_name,
                 )
                 answer, html = results
                 with st.container():

config.py

@@ -7,14 +7,14 @@ SAAS = True
 def get_db_config():
     url = os.environ["QDRANT_URL"]
     api_key = os.environ["QDRANT_API_KEY"]
-    collection_name = "gh-issue-search"
+    collection_name = "document-search"
     return url, api_key, collection_name
 
 
 def get_local_db_congin():
     url = "localhost"
     # api_key = os.environ["QDRANT_API_KEY"]
-    collection_name = "gh-issues"
+    collection_name = "document-search"
     return url, None, collection_name
 
 

gh_issue_loader.py → doc_loader.py

@@ -4,7 +4,8 @@ from typing import Iterator
 from dateutil.parser import parse
 from langchain.docstore.document import Document
 from langchain.document_loaders.base import BaseLoader
-from gh_issue_loader import Issue
+
+from model import Doc
 
 
 def date_to_int(dt_str: str) -> int:
@@ -12,49 +13,39 @@ def date_to_int(dt_str: str) -> int:
     return int(dt.timestamp())
 
 
-def get_contents(repo_name: str, filename: str) -> Iterator[tuple[Issue, str]]:
+def get_contents(project_name: str, filename: str) -> Iterator[tuple[Doc, str]]:
+    """filename for file with ndjson
+
+        {"title": <page title>, "body": <page body>, "id": <page_id>, "ctime": ..., "user": <name>, "url": "https:..."}
+        {"title": ...}
+    """
     with open(filename, "r") as f:
         obj = [json.loads(line) for line in f]
     for data in obj:
         title = data["title"]
         body = data["body"]
-        issue = Issue(
-            repo_name=repo_name,
-            id=data["number"],
+        doc = Doc(
+            project_name=project_name,
+            id=data["id"],
             title=title,
-            created_at=date_to_int(data["created_at"]),
-            user=data["user.login"],
-            url=data["html_url"],
-            labels=data["labels_"],
-            type_="issue",
+            created_at=date_to_int(data["ctime"]),
+            user=data["user"],
+            url=data["url"],
         )
         text = title
         if body:
             text += "\n\n" + body
-        yield issue, text
-        comments = data["comments_"]
-        for comment in comments:
-            issue = Issue(
-                repo_name=repo_name,
-                id=comment["id"],
-                title=data["title"],
-                created_at=date_to_int(comment["created_at"]),
-                user=comment["user.login"],
-                url=comment["html_url"],
-                labels=data["labels_"],
-                type_="comment",
-            )
-            yield issue, comment["body"]
-
-
-class GHLoader(BaseLoader):
-    def __init__(self, repo_name: str, filename: str):
-        self.repo_name = repo_name
+        yield doc, text
+
+
+class DocLoader(BaseLoader):
+    def __init__(self, project_name: str, filename: str):
+        self.project_name = project_name
         self.filename = filename
 
     def lazy_load(self) -> Iterator[Document]:
-        for issue, text in get_contents(self.repo_name, self.filename):
-            metadata = asdict(issue)
+        for doc, text in get_contents(self.project_name, self.filename):
+            metadata = asdict(doc)
             yield Document(page_content=text, metadata=metadata)
 
     def load(self) -> list[Document]:

model.py

@@ -2,12 +2,10 @@ from dataclasses import dataclass
 
 
 @dataclass(frozen=True)
-class Issue:
-    repo_name: str
+class Doc:
+    project_name: str
     id: int
     title: str
     created_at: int
     user: str
     url: str
-    labels: list[str]
-    type_: str

store.py

@@ -2,7 +2,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Qdrant
 
-from gh_issue_loader import GHLoader
+from doc_loader import DocLoader
 from config import DB_CONFIG
 
 
@@ -36,8 +36,8 @@ def store(texts):
     )
 
 
-def main(repo_name: str, path: str) -> None:
-    loader = GHLoader(repo_name, path)
+def main(project_name: str, path: str) -> None:
+    loader = DocLoader(project_name, path)
     docs = loader.load()
     texts = get_text_chunk(docs)
     store(texts)
@@ -45,8 +45,8 @@ def main(repo_name: str, path: str) -> None:
 
 if __name__ == "__main__":
     """
-    $ python store.py "REPO_NAME" "FILE_PATH"
-    $ python store.py cocoa data/cocoa-issues.json
+    $ python store.py "PROJECT_NAME" "FILE_PATH"
+    $ python store.py hoge data/hoge-docs.json
     """
     import sys
 
@@ -54,6 +54,6 @@ if __name__ == "__main__":
     if len(args) != 3:
         print("No args, you need two args for repo_name, json_file_path")
     else:
-        repo_name = args[1]
+        project_name = args[1]
         path = args[2]
-        main(repo_name, path)
+        main(project_name, path)