changed pipelin

.gitattributes

@@ -34,4 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *db* filter=lfs diff=lfs merge=lfs -text
-vaccine_coverage_study.db filter=lfs diff=lfs merge=lfs -text
+vaccine_coverage_study.db filter=lfs diff=lfs merge=lfs -text
+*.db filter=lfs diff=lfs merge=lfs -text

database/vaccine_coverage_db.py

@@ -1,46 +1,3 @@
-import sqlite3
-from typing import List, Dict, Any
-
-
-class VaccineCoverageDB:
-    def __init__(self, db_path: str):
-        self.conn = sqlite3.connect(db_path)
-        self.conn.row_factory = sqlite3.Row
-
-    def get_all_items(self) -> List[Dict[str, Any]]:
-        cursor = self.conn.execute("SELECT * FROM items")
-        return [dict(row) for row in cursor.fetchall()]
-
-    def get_item_by_key(self, key: str) -> Dict[str, Any]:
-        cursor = self.conn.execute("SELECT * FROM items WHERE key = ?", (key,))
-        return dict(cursor.fetchone())
-
-    def get_attachments_for_item(self, item_key: str) -> List[Dict[str, Any]]:
-        cursor = self.conn.execute(
-            "SELECT * FROM attachments WHERE parent_key = ?", (item_key,)
-        )
-        return [dict(row) for row in cursor.fetchall()]
-
-    def get_pdf_content(self, attachment_key: str) -> bytes:
-        cursor = self.conn.execute(
-            "SELECT content FROM attachments WHERE key = ?", (attachment_key,)
-        )
-        result = cursor.fetchone()
-        return result["content"] if result else None
-
-    def save_pdf_to_file(self, attachment_key: str, output_path: str) -> bool:
-        pdf_content = self.get_pdf_content(attachment_key)
-        if pdf_content:
-            try:
-                with open(output_path, "wb") as f:
-                    f.write(pdf_content)
-                return True
-            except Exception as e:
-                print(f"Error saving PDF: {str(e)}")
-                return False
-        else:
-            print(f"No PDF content found for attachment key: {attachment_key}")
-            return False
-
-    def close(self):
-        self.conn.close()
+version https://git-lfs.github.com/spec/v1
+oid sha256:42a0645cdd38f2d7ede525768eb21a4cbe08b4d86959cb4eb2349887f2bcf70e
+size 1774

initialize_db.py

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c0cb2cf50f14d131b1e999cee44652575fd1029141514dfc2e028af1419b0d46
-size 2344
+oid sha256:08030c4783a86d9a9afb9437b102dde959405b6b2857725eec02b6d9c2699e97
+size 2346

rag/rag_pipeline.py

@@ -2,116 +2,81 @@ import json
 import os
 from typing import Dict, Any
 from llama_index.core import (
-    SimpleDirectoryReader,
     VectorStoreIndex,
     Document,
-    StorageContext,
-    load_index_from_storage,
+    SentenceWindowNodeParser,
+)
+from llama_index.core.node_parser import (
+    SentenceSplitter,
 )
-from llama_index.core.node_parser import SentenceSplitter, SemanticSplitterNodeParser
-from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.core import PromptTemplate
 
 
 class RAGPipeline:
     def __init__(
-        self, metadata_file: str, pdf_dir: str, use_semantic_splitter: bool = False
+        self,
+        study_json,
+        use_semantic_splitter=False,
     ):
-        self.metadata_file = metadata_file
-        self.pdf_dir = pdf_dir
-        self.use_semantic_splitter = use_semantic_splitter
+
+        self.study_json = study_json
         self.index = None
+        self.use_semantic_splitter = use_semantic_splitter
         self.load_documents()
         self.build_index()
 
     def load_documents(self):
-        if not os.path.exists(self.metadata_file):
-            print(f"Metadata file not found: {self.metadata_file}")
-            self.documents = []
-            return
-
-        with open(self.metadata_file, "r") as f:
-            self.metadata = json.load(f)
+        with open(self.study_json, "r") as f:
+            self.data = json.load(f)
 
         self.documents = []
-        for item_key, item_data in self.metadata.items():
-            metadata = item_data["metadata"]
-            pdf_path = item_data.get("pdf_path")
-
-            if pdf_path:
-                full_pdf_path = os.path.join(self.pdf_dir, os.path.basename(pdf_path))
-                if os.path.exists(full_pdf_path):
-                    pdf_content = (
-                        SimpleDirectoryReader(input_files=[full_pdf_path])
-                        .load_data()[0]
-                        .text
-                    )
-                else:
-                    pdf_content = "PDF file not found"
-            else:
-                pdf_content = "PDF path not available in metadata"
 
+        for index, doc_data in enumerate(self.data):
             doc_content = (
-                f"Title: {metadata['title']}\n"
-                f"Abstract: {metadata['abstract']}\n"
-                f"Authors: {metadata['authors']}\n"
-                f"Year: {metadata['year']}\n"
-                f"DOI: {metadata['doi']}\n"
-                f"Full Text: {pdf_content}"
+                f"Title: {doc_data['title']}\n"
+                f"Abstract: {doc_data['abstract']}\n"
+                f"Authors: {', '.join(doc_data['authors'])}\n"
+                f"Year: {doc_data['year']}\n"
+                f"DOI: {doc_data['doi']}\n"
+                f"Full Text: {doc_data['full_text']}"
             )
 
+            metadata = {
+                "title": doc_data.get("title"),
+                "abstract": doc_data.get("abstract"),
+                "authors": doc_data.get("authors", []),
+                "year": doc_data.get("year"),
+                "doi": doc_data.get("doi"),
+            }
+
             self.documents.append(
-                Document(text=doc_content, id_=item_key, metadata=metadata)
+                Document(
+                    text=doc_content,
+                    id_=f"doc_{index}",
+                    metadata=metadata,
+                )
             )
 
     def build_index(self):
-        if self.use_semantic_splitter:
-            embed_model = OpenAIEmbedding()
-            splitter = SemanticSplitterNodeParser(
-                buffer_size=1,
-                breakpoint_percentile_threshold=95,
-                embed_model=embed_model,
-            )
-        else:
-            splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=20)
-
-        nodes = splitter.get_nodes_from_documents(self.documents)
-        self.index = VectorStoreIndex(nodes)
+        sentence_splitter = SentenceSplitter(chunk_size=128, chunk_overlap=13)
 
-    def query(self, question: str, prompt_type: str = "default") -> Dict[str, Any]:
-        prompt_template = self._get_prompt_template(prompt_type)
+        def _split(text: str) -> List[str]:
+            return sentence_splitter.split_text(text)
 
-        query_engine = self.index.as_query_engine(
-            text_qa_template=prompt_template, similarity_top_k=5
+        node_parser = SentenceWindowNodeParser.from_defaults(
+            sentence_splitter=_split,
+            window_size=3,
+            window_metadata_key="window",
+            original_text_metadata_key="original_text",
         )
-        response = query_engine.query(question)
 
-        return response
+        nodes = node_parser.get_nodes_from_documents(self.documents)
 
-    def _get_prompt_template(self, prompt_type: str) -> PromptTemplate:
-        if prompt_type == "highlight":
-            return PromptTemplate(
-                "Context information is below.\n"
-                "---------------------\n"
-                "{context_str}\n"
-                "---------------------\n"
-                "Given this information, please answer the question: {query_str}\n"
-                "Include all relevant information from the provided context. "
-                "Highlight key information by enclosing it in **asterisks**. "
-                "When quoting specific information, please use square brackets to indicate the source, e.g. [1], [2], etc."
-            )
-        elif prompt_type == "evidence_based":
-            return PromptTemplate(
-                "Context information is below.\n"
-                "---------------------\n"
-                "{context_str}\n"
-                "---------------------\n"
-                "Given this information, please answer the question: {query_str}\n"
-                "Provide an answer to the question using evidence from the context above. "
-                "Cite sources using square brackets."
-            )
-        else:
-            return PromptTemplate(
+        self.index = VectorStoreIndex(nodes)
+
+    def query(self, question, prompt_template=None):
+        if prompt_template is None:
+            prompt_template = PromptTemplate(
                 "Context information is below.\n"
                 "---------------------\n"
                 "{context_str}\n"
@@ -122,3 +87,11 @@ class RAGPipeline:
                 "If the information is not available in the context, please state that clearly. "
                 "When quoting specific information, please use square brackets to indicate the source, e.g. [1], [2], etc."
             )
+
+        query_engine = self.index.as_query_engine(
+            text_qa_template=prompt_template,
+            similarity_top_k=5,
+        )
+        response = query_engine.query(question)
+
+        return response

utils/helpers.py

@@ -1,6 +1,7 @@
 from typing import Dict, Any
 from llama_index.core import Response
 
+
 def process_response(response: Response) -> Dict[str, Any]:
     source_nodes = response.source_nodes
     sources = {}
@@ -18,12 +19,13 @@ def process_response(response: Response) -> Dict[str, Any]:
 
     return {"markdown": markdown_text, "raw": raw_text, "sources": sources}
 
+
 def format_source(metadata: Dict[str, Any]) -> str:
-    authors = metadata.get('authors', 'Unknown Author')
-    year = metadata.get('year', 'n.d.')
-    title = metadata.get('title', 'Untitled')
+    authors = metadata.get("authors", "Unknown Author")
+    year = metadata.get("year", "n.d.")
+    title = metadata.get("title", "Untitled")
 
-    author_list = authors.split(',')
+    author_list = authors.split(",")
     if len(author_list) > 2:
         formatted_authors = f"{author_list[0].strip()} et al."
     elif len(author_list) == 2:
@@ -31,10 +33,10 @@ def format_source(metadata: Dict[str, Any]) -> str:
     else:
         formatted_authors = author_list[0].strip()
 
-    year = 'n.d.' if year is None or year == 'None' else str(year)
+    year = "n.d." if year is None or year == "None" else str(year)
 
     max_title_length = 250
     if len(title) > max_title_length:
-        title = title[:max_title_length] + '...'
+        title = title[:max_title_length] + "..."
 
     return f"{formatted_authors} ({year}). {title}"

vaccine_coverage_study.db

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:121fd525453b27b5008a3714840c929402ec01b74aea4d21bdd87be1a60bc008
-size 41222144