add more metadata fields and metadata search (#185)

added authors, source and research area fields in chunks metadata

Author: AsyaOrlova

ChemCoScientist/paper_analysis/chroma_db_operations.py

@@ -3,6 +3,7 @@
 import os
 import uuid
 from pathlib import Path
+from typing import Literal
 
 import chromadb
 import numpy as np
@@ -18,6 +19,7 @@
 import requests
 
 from ChemCoScientist.chemical_utils.openchemie_functions import extract_molecules_from_figure, extract_reactions_from_figure
+from ChemCoScientist.paper_analysis.constants import ResearchArea
 from ChemCoScientist.paper_analysis.prompts import summarisation_prompt
 from ChemCoScientist.paper_analysis.settings import allowed_providers
 from ChemCoScientist.paper_analysis.settings import settings as default_settings
@@ -53,10 +55,18 @@ class ExpandedSummary(BaseModel):
         description="Title of the paper. If the title is not explicitly specified, use the default value - 'NO TITLE'"
     )
     publication_year: int = Field(
-        description=(
-            "Year of publication of the paper. If the publication year is not explicitly specified, use the default"
+        description="Year of publication of the paper. If the publication year is not explicitly specified, use the default"
             " value - 9999."
-        )
+    )
+    paper_authors: str = Field(
+        description="Authors of the paper. If the authors are not explicitly specified, use the default value - 'NO AUTHORS'"
+    )
+    publication_source: str = Field(
+        description="Source where the paper was published. If the source is not explicitly specified, use the default value - 'NO SOURCE'"
+    )
+    research_area: ResearchArea = Field(
+        description="Area or field of chemistry the paper is about. Must be one of the predefined values."
+        " If the area has no match in the predefined list or is hard to determine, use the default value - 'OTHER'"
     )
 
 
@@ -455,7 +465,7 @@ def search_for_papers(self,
         return res
 
     def retrieve_context(
-            self, query: str, relevant_papers: dict = None
+            self, query: str, relevant_papers: dict = None, meta_filter: dict = None
     ) -> tuple[list, dict, dict]:
         """
         Retrieves relevant information from text and images associated with scientific papers based on a user query.
@@ -468,14 +478,15 @@ def retrieve_context(
             query (str): The search query used to identify relevant information.
             relevant_papers (list, optional): A list of pre-identified relevant papers. Defaults to None, in which case
                 a search for relevant papers is initiated.
+            meta_filter (dict, optional): A dictionary of metadata filters to apply during the search.
 
         Returns:
             tuple[list, dict]: A tuple containing the retrieved text and image context.
                 - text_context (list): A list of text chunks deemed most relevant to the query.
                 - image_context (dict): A dictionary containing image data associated with the query.
         """
         if not relevant_papers:
-            relevant_papers = self.search_for_papers(query)
+            relevant_papers = self.search_for_papers(query, meta_filter=meta_filter)
 
         raw_text_context = self.client.query_chromadb(
             self.txt_collection,
@@ -558,7 +569,24 @@ def search_with_reranker(
 
         return scored_docs[:top_k]
 
-    def add_paper_summary_to_db(self, paper_name: str, parsed_paper: str, llm) -> None:
+    def _generate_expanded_summary(self, parsed_paper: str, llm) -> ExpandedSummary:
+        """
+        Generates an expanded summary of a paper using a language model.
+
+        This method takes parsed paper content and uses an LLM to extract and structure
+        key information including summary, title, authors, publication year, and source.
+
+        Args:
+            parsed_paper (str): The text content of the parsed paper.
+            llm: The language model used to generate the summary.
+
+        Returns:
+            ExpandedSummary: An object containing the paper's structured summary information.
+        """
+        expanded_summary: ExpandedSummary = llm.invoke([HumanMessage(content=summarisation_prompt + parsed_paper)])
+        return expanded_summary
+
+    def add_paper_summary_to_db(self, paper_name: str, parsed_paper: str, expanded_summary: ExpandedSummary) -> None:
         """
         Adds a paper summary to the document collection for efficient information retrieval.
 
@@ -574,13 +602,15 @@ def add_paper_summary_to_db(self, paper_name: str, parsed_paper: str, llm) -> No
         Returns:
             None
         """
-        expanded_summary: ExpandedSummary = llm.invoke([HumanMessage(content=summarisation_prompt + parsed_paper)])
         doc = Document(
             page_content=expanded_summary.paper_summary,
             metadata={
                 "source": paper_name,
                 "paper_title": expanded_summary.paper_title,
-                "publication_year": expanded_summary.publication_year
+                "publication_year": expanded_summary.publication_year,
+                "paper_authors": expanded_summary.paper_authors,
+                "publication_source": expanded_summary.publication_source,
+                "research_area": expanded_summary.research_area
             }
         )
         embedding = self.get_embeddings([doc.page_content])
@@ -771,13 +801,15 @@ def process_single_document(folder_path: Path, s3_service: S3BucketService, s3_p
         else:
             parsed_paper, mapping = clean_up_html(folder_path, paper_name, text)
         print(f"Finished post-processing paper: {paper_name}")
-        documents = html_chunking(parsed_paper, paper_name)
 
         llm = create_llm_connector(SUMMARY_LLM_URL, extra_body={"provider": {"only": allowed_providers}})
         struct_llm = llm.with_structured_output(schema=ExpandedSummary)
+        paper_summary = process_local_store._generate_expanded_summary(parsed_paper, struct_llm)
+        
+        documents = html_chunking(parsed_paper, paper_name, paper_summary)
 
         print(f"Starting loading paper: {paper_name}")
-        process_local_store.add_paper_summary_to_db(str(paper_name_to_load), parsed_paper, struct_llm)
+        process_local_store.add_paper_summary_to_db(str(paper_name_to_load), parsed_paper, paper_summary)
         process_local_store.store_text_chunks_in_chromadb(documents)
         process_local_store.store_images_in_chromadb_txt_format(str(folder_path), str(paper_name_to_load), mapping)
         print(f"Finished loading paper: {paper_name}")

ChemCoScientist/paper_analysis/constants.py

@@ -0,0 +1,15 @@
+from typing import Literal, TypeAlias
+
+
+ResearchArea: TypeAlias = Literal[
+	"Polymer Chemistry",
+	"Organic Chemistry",
+	"Nanomaterials",
+	"Molecular Dynamics",
+	"Membrane Chemistry",
+	"Electrochemistry",
+	"DFT",
+	"Biological Macromolecules",
+	"Biological Chemistry",
+	"Analytical Chemistry",
+]

ChemCoScientist/paper_analysis/prompts.py

@@ -55,7 +55,9 @@
     " list separately all tables with its names, all images with its names, and all main"
     " substances that are in the article. Keywords/terms, as well as lists of tables, images, and"
     " substances are also part of the summary.\n"
-    " Also try to determine the title of the article and the year of its publication.\n\n"
+    " Also try to determine the title of the article, the year of its publication, authors,"
+    " its publication source (journal name, venue name, preprint server name etc.)"
+    " and the research area of the paper.\n\n"
     "Article in Markdown markup:\n"
 )
 
@@ -107,3 +109,30 @@
     "Do not include molecules where the required property is not mentioned."
     "Output only the CSV data (no explanations, no markdown, no additional text)."
 )
+
+extract_query_filters_prompt = (
+    "You are an assistant that extracts metadata filters from user questions about scientific papers. "
+    "Your task is to analyze the USER QUESTION and identify any mentions of:"
+    "\n1. Author names (e.g., 'What did Smith say', 'According to John Doe', 'research by Dr. Jane')"
+    "\n2. Publication year or year range (e.g., 'papers from 2020', 'research since 2018', 'recent studies')"
+    "\n3. Publication source/journal (e.g., 'papers in Nature', 'from ACS Catalysis', 'published in Science')"
+    "\n4. Research area (e.g., 'polymer chemistry papers', 'nanomaterials research', 'DFT studies')"
+    "\n\nFor year filters:"
+    "\n- 'recent' or 'latest' should translate to publication_year_min = current_year - 2"
+    "\n- 'since YEAR' should translate to publication_year_min = YEAR"
+    "\n- 'in YEAR' or 'from YEAR' should translate to publication_year_exact = YEAR"
+    "\n- 'between YEAR1 and YEAR2' should translate to publication_year_min = YEAR1, publication_year_max = YEAR2"
+    "\n\nFor research areas, use only these values if detected:"
+    "\n'Polymer Chemistry', 'Organic Chemistry', 'Nanomaterials', 'Molecular Dynamics', 'Membrane Chemistry', "
+    "'Electrochemistry', 'DFT', 'Biological Macromolecules', 'Biological Chemistry', 'Analytical Chemistry'"
+    "\n\nIf no specific filter is mentioned for a field, leave it as null."
+    "\n\nExamples:"
+    "\nQ: 'What did Sam Smith say about catalysis?'"
+    "\nA: {\"paper_authors\": \"Sam Smith\", \"publication_year_min\": null, \"publication_year_max\": null, \"publication_year_exact\": null, \"publication_source\": null, \"research_area\": null}"
+    "\n\nQ: 'What are recent advances in polymer chemistry?'"
+    "\nA: {\"paper_authors\": null, \"publication_year_min\": 2024, \"publication_year_max\": null, \"publication_year_exact\": null, \"publication_source\": null, \"research_area\": \"Polymer Chemistry\"}"
+    "\n\nQ: 'Show me DFT studies from Nature Chemistry published in 2023'"
+    "\nA: {\"paper_authors\": null, \"publication_year_min\": null, \"publication_year_max\": null, \"publication_year_exact\": 2023, \"publication_source\": \"Nature Chemistry\", \"research_area\": \"DFT\"}"
+    "\n\nQ: 'What synthesis methods are used for nanoparticles?'"
+    "\nA: {\"paper_authors\": null, \"publication_year_min\": null, \"publication_year_max\": null, \"publication_year_exact\": null, \"publication_source\": null, \"research_area\": null}"
+)

ChemCoScientist/paper_analysis/question_processing.py

@@ -12,7 +12,8 @@
 from io import BytesIO
 
 from ChemCoScientist.paper_analysis.chroma_db_operations import ChromaDBPaperStore
-from ChemCoScientist.paper_analysis.prompts import sys_prompt, explore_my_papers_prompt
+from ChemCoScientist.paper_analysis.constants import ResearchArea
+from ChemCoScientist.paper_analysis.prompts import sys_prompt, explore_my_papers_prompt, extract_query_filters_prompt
 from ChemCoScientist.paper_analysis.settings import allowed_providers
 from CoScientist.paper_parser.utils import convert_to_base64, prompt_func, load_image_as_binary
 from ChemCoScientist.chemical_utils.openchemie_functions import *
@@ -22,6 +23,111 @@
 
 VISION_LLM_URL = os.environ["VISION_LLM_URL"]
 
+
+class QueryFilters(BaseModel):
+    """Metadata filters extracted from user question."""
+    paper_authors: str | None = Field(
+        description="Author name(s) mentioned in the question",
+        default=None
+    )
+    publication_year_min: int | None = Field(
+        description="Minimum publication year for filtering",
+        default=None
+    )
+    publication_year_max: int | None = Field(
+        description="Maximum publication year for filtering",
+        default=None
+    )
+    publication_year_exact: int | None = Field(
+        description="Exact publication year when specified",
+        default=None
+    )
+    publication_source: str | None = Field(
+        description="Journal or publication source name",
+        default=None
+    )
+    research_area: ResearchArea | None = Field(
+        description="Research area/field of chemistry",
+        default=None
+    )
+
+
+def extract_metadata_filters(question: str) -> QueryFilters:
+    """
+    Uses LLM to extract metadata filters from user question.
+    
+    Args:
+        question: The user's question string
+        
+    Returns:
+        QueryFilters: Structured filters including authors, year, source, and research area
+    """
+    llm = create_llm_connector(
+        VISION_LLM_URL,
+        extra_body={"provider": {"only": allowed_providers}},
+        temperature=0.0
+    )
+    
+    struct_llm = llm.with_structured_output(schema=QueryFilters)
+    
+    prompt = extract_query_filters_prompt + f"\n\nUSER QUESTION: {question}"
+    
+    from langchain_core.messages import HumanMessage
+    filters: QueryFilters = struct_llm.invoke([HumanMessage(content=prompt)])
+    
+    return filters
+
+
+def build_chroma_where_filter(filters: QueryFilters) -> dict | None:
+    """
+    Converts QueryFilters to ChromaDB where clause format.
+    
+    Args:
+        filters: QueryFilters object with extracted metadata
+        
+    Returns:
+        dict: ChromaDB where clause ready for collection.query(), or None if no filters
+        
+    Example output:
+        {"paper_authors": {"$eq": "Smith"}}
+        {
+            "$and": [
+                {"paper_authors": {"$eq": "Smith"}},
+                {"publication_year": {"$gte": 2020}}
+            ]
+        }
+    """
+    conditions = []
+    
+    if filters.paper_authors:
+        conditions.append({"paper_authors": {"$eq": filters.paper_authors}})
+    
+    if filters.publication_year_exact:
+        conditions.append({"publication_year": {"$eq": filters.publication_year_exact}})
+    elif filters.publication_year_min or filters.publication_year_max:
+        year_condition = {}
+        if filters.publication_year_min:
+            year_condition["$gte"] = filters.publication_year_min
+        if filters.publication_year_max:
+            year_condition["$lte"] = filters.publication_year_max
+        if year_condition:
+            conditions.append({"publication_year": year_condition})
+    
+    if filters.publication_source:
+        conditions.append({"publication_source": {"$eq": filters.publication_source}})
+    
+    if filters.research_area and filters.research_area != "OTHER":
+        conditions.append({"research_area": {"$eq": filters.research_area}})
+    
+    if not conditions:
+        return None
+    
+    if len(conditions) == 1:
+        return conditions[0]
+    
+    return {"$and": conditions}
+
+
 def query_llm(
     model_url: str, question: str, txt_context: str, img_paths: list[str]
 ) -> tuple:
@@ -159,7 +265,10 @@ def process_question(question: str, store: ChromaDBPaperStore) -> dict:
                 'image_context' - the set of image paths identified as relevant to the question;
                 'metadata' - Additional metadata returned by the LLM query.
     """
-    txt_data, img_data = store.retrieve_context(question)
+    meta_filter = extract_metadata_filters(question)
+    meta_filter_chroma = build_chroma_where_filter(meta_filter)
+    
+    txt_data, img_data = store.retrieve_context(question, meta_filter=meta_filter_chroma)
     txt_context = ""
     relevant_txt_context = []
     img_paths = []
@@ -276,7 +385,7 @@ def process_question(question: str, store: ChromaDBPaperStore) -> dict:
     #######################################################
 
     paper_store = ChromaDBPaperStore()
-    question = 'What is the title of an article?'
+    question = 'What are papers since 2023 about analytical chemistry are focused on?'
       # question = 'What components are involved in the synthesis of BASHY dyes, and what are the uses of these dyes?'
     # question = 'What IC50 values do weakly active and highly active Bruton\'s tyrosine kinase inhibitors have?'
     # question = 'How does the synthesis of Glionitrin A/B happen?'

CoScientist/paper_parser/parse_and_split.py

@@ -178,7 +178,7 @@ def clean_up_html(
     return soup.prettify(), image_url_mapping
 
 
-def html_chunking(html_string: str, paper_name: str) -> list:
+def html_chunking(html_string: str, paper_name: str, paper_summary) -> list:
     """
     Chunks an HTML string into semantic passages for efficient information retrieval.
     
@@ -216,8 +216,15 @@ def custom_table_extractor(table_tag):
     documents = splitter.split_text(html_string)
     for doc in documents:
         doc.page_content = "passage: " + doc.page_content  # Maybe delete "passage: " addition
-        doc.metadata["imgs_in_chunk"] = str(extract_img_url(doc.page_content, paper_name))
-        doc.metadata["source"] = paper_name + ".pdf"
+        doc.metadata.update({
+            "imgs_in_chunk": str(extract_img_url(doc.page_content, paper_name)),
+            "source": f"{paper_name}.pdf",
+            "paper_title": paper_summary.paper_title,
+            "publication_year": paper_summary.publication_year,
+            "paper_authors": paper_summary.paper_authors,
+            "publication_source": paper_summary.publication_source,
+            "research_area": paper_summary.research_area,
+        })
 
     return documents
 

CoScientist/paper_parser/s3_connection.py

@@ -92,23 +92,22 @@ def list_objects(self, prefix: str) -> list[str]:
         Lists all objects in the S3 bucket with the given prefix.
 
         Args:
-            prefix: The prefix to filter objects in the S3 bucket
+            prefix: The prefix to filter objects in the S3 bucket. If empty, lists all objects.
 
         Returns:
             A list of object keys (file paths) that match the prefix
         """
         client = self.create_s3_client()
 
-        response = client.list_objects_v2(Bucket=self.bucket_name, Prefix=prefix)
-        storage_content: list[str] = []
+        paginator = client.get_paginator("list_objects_v2")
+        page_iterator = paginator.paginate(Bucket=self.bucket_name, Prefix=prefix)
 
-        try:
-            contents = response["Contents"]
-        except KeyError:
-            return storage_content
+        storage_content: list[str] = []
 
-        for item in contents:
-            storage_content.append(item["Key"])
+        for page in page_iterator:
+            contents = page.get("Contents", [])
+            for item in contents:
+                storage_content.append(item["Key"])
 
         return storage_content
 
@@ -248,5 +247,4 @@ def clean_up_by_prefix(self, prefix_to_delete: str):
     # for bucket in buckets["Buckets"]:
     #     print(bucket["Name"], bucket["CreationDate"])
     # objects = s3_service.list_objects(prefix="")
-    # print(objects)
-    
+    # print(objects)