ursa-ai 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

ursa/agents/__init__.py

@@ -14,6 +14,8 @@ from .lammps_agent import LammpsState as LammpsState
 from .mp_agent import MaterialsProjectAgent as MaterialsProjectAgent
 from .planning_agent import PlanningAgent as PlanningAgent
 from .planning_agent import PlanningState as PlanningState
+from .rag_agent import RAGAgent as RAGAgent
+from .rag_agent import RAGState as RAGState
 from .recall_agent import RecallAgent as RecallAgent
 from .websearch_agent import WebSearchAgent as WebSearchAgent
 from .websearch_agent import WebSearchState as WebSearchState

ursa/agents/arxiv_agent.py

@@ -1,7 +1,6 @@
 import base64
 import os
 import re
-import statistics
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from io import BytesIO
 from urllib.parse import quote
@@ -9,8 +8,6 @@ from urllib.parse import quote
 import feedparser
 import pymupdf
 import requests
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_chroma import Chroma
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
@@ -20,15 +17,13 @@ from tqdm import tqdm
 from typing_extensions import List, TypedDict
 
 from .base import BaseAgent
+from .rag_agent import RAGAgent
 
 try:
     from openai import OpenAI
 except Exception:
     pass
 
-# embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
-# embeddings = OpenAIEmbeddings()
-
 
 class PaperMetadata(TypedDict):
     arxiv_id: str
@@ -242,27 +237,6 @@ class ArxivAgent(BaseAgent):
         papers = self._fetch_papers(state["query"])
         return {**state, "papers": papers}
 
-    def _get_or_build_vectorstore(self, paper_text: str, arxiv_id: str):
-        os.makedirs(self.vectorstore_path, exist_ok=True)
-
-        persist_directory = os.path.join(self.vectorstore_path, arxiv_id)
-
-        if os.path.exists(persist_directory):
-            vectorstore = Chroma(
-                persist_directory=persist_directory,
-                embedding_function=self.rag_embedding,
-            )
-        else:
-            splitter = RecursiveCharacterTextSplitter(
-                chunk_size=1000, chunk_overlap=200
-            )
-            docs = splitter.create_documents([paper_text])
-            vectorstore = Chroma.from_documents(
-                docs, self.rag_embedding, persist_directory=persist_directory
-            )
-
-        return vectorstore.as_retriever(search_kwargs={"k": 5})
-
     def _summarize_node(self, state: PaperState) -> PaperState:
         prompt = ChatPromptTemplate.from_template("""
         You are a scientific assistant responsible for summarizing extracts from research papers, in the context of the following task: {context}
@@ -285,33 +259,8 @@ class ArxivAgent(BaseAgent):
 
             try:
                 cleaned_text = remove_surrogates(paper["full_text"])
-                if self.rag_embedding:
-                    retriever = self._get_or_build_vectorstore(
-                        cleaned_text, arxiv_id
-                    )
-
-                    relevant_docs_with_scores = (
-                        retriever.vectorstore.similarity_search_with_score(
-                            state["context"], k=5
-                        )
-                    )
-
-                    if relevant_docs_with_scores:
-                        score = sum([
-                            s for _, s in relevant_docs_with_scores
-                        ]) / len(relevant_docs_with_scores)
-                        relevancy_scores[i] = abs(1.0 - score)
-                    else:
-                        relevancy_scores[i] = 0.0
-
-                    retrieved_content = "\n\n".join([
-                        doc.page_content for doc, _ in relevant_docs_with_scores
-                    ])
-                else:
-                    retrieved_content = cleaned_text
-
                 summary = chain.invoke({
-                    "retrieved_content": retrieved_content,
+                    "retrieved_content": cleaned_text,
                     "context": state["context"],
                 })
 
@@ -346,15 +295,18 @@ class ArxivAgent(BaseAgent):
                 i, result = future.result()
                 summaries[i] = result
 
-        if self.rag_embedding:
-            print(f"\nMax Relevancy Score: {max(relevancy_scores)}")
-            print(f"Min Relevancy Score: {min(relevancy_scores)}")
-            print(
-                f"Median Relevancy Score: {statistics.median(relevancy_scores)}\n"
-            )
-
         return {**state, "summaries": summaries}
 
+    def _rag_node(self, state: PaperState) -> PaperState:
+        new_state = state.copy()
+        rag_agent = RAGAgent(
+            llm=self.llm,
+            embedding=self.rag_embedding,
+            database_path=self.database_path,
+        )
+        new_state["final_summary"] = rag_agent.run(context=state["context"])
+        return new_state
+
     def _aggregate_node(self, state: PaperState) -> PaperState:
         summaries = state["summaries"]
         papers = state["papers"]
@@ -404,13 +356,20 @@ class ArxivAgent(BaseAgent):
         builder.add_node("fetch_papers", self._fetch_node)
 
         if self.summarize:
-            builder.add_node("summarize_each", self._summarize_node)
-            builder.add_node("aggregate", self._aggregate_node)
-
-            builder.set_entry_point("fetch_papers")
-            builder.add_edge("fetch_papers", "summarize_each")
-            builder.add_edge("summarize_each", "aggregate")
-            builder.set_finish_point("aggregate")
+            if self.rag_embedding:
+                builder.add_node("rag_summarize", self._rag_node)
+
+                builder.set_entry_point("fetch_papers")
+                builder.add_edge("fetch_papers", "rag_summarize")
+                builder.set_finish_point("rag_summarize")
+            else:
+                builder.add_node("summarize_each", self._summarize_node)
+                builder.add_node("aggregate", self._aggregate_node)
+
+                builder.set_entry_point("fetch_papers")
+                builder.add_edge("fetch_papers", "summarize_each")
+                builder.add_edge("summarize_each", "aggregate")
+                builder.set_finish_point("aggregate")
 
         else:
             builder.set_entry_point("fetch_papers")

ursa/agents/rag_agent.py

@@ -0,0 +1,272 @@
+import os
+import re
+import statistics
+from threading import Lock
+from typing import List, Optional, TypedDict
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_chroma import Chroma
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_core.embeddings import Embeddings
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langgraph.graph import StateGraph
+
+from ursa.agents.base import BaseAgent
+
+
+class RAGState(TypedDict, total=False):
+    context: str
+    doc_texts: List[str]
+    doc_ids: List[str]
+    summary: str
+
+
+def remove_surrogates(text: str) -> str:
+    return re.sub(r"[\ud800-\udfff]", "", text)
+
+
+class RAGAgent(BaseAgent):
+    def __init__(
+        self,
+        llm="openai/o3-mini",
+        embedding: Optional[Embeddings] = None,
+        return_k: int = 10,
+        chunk_size: int = 1000,
+        chunk_overlap: int = 200,
+        database_path: str = "database",
+        summaries_path: str = "database",
+        vectorstore_path: str = "vectorstore",
+        **kwargs,
+    ):
+        super().__init__(llm, **kwargs)
+        self.retriever = None
+        self._vs_lock = Lock()
+        self.return_k = return_k
+        self.embedding = embedding
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.database_path = database_path
+        self.summaries_path = summaries_path
+        self.vectorstore_path = vectorstore_path
+        self.graph = self._build_graph()
+
+        os.makedirs(self.vectorstore_path, exist_ok=True)
+        self.vectorstore = self._open_global_vectorstore()
+
+    @property
+    def manifest_path(self) -> str:
+        return os.path.join(self.vectorstore_path, "_ingested_ids.txt")
+
+    @property
+    def manifest_exists(self) -> bool:
+        return os.path.exists(self.manifest_path)
+
+    def _open_global_vectorstore(self) -> Chroma:
+        return Chroma(
+            persist_directory=self.vectorstore_path,
+            embedding_function=self.embedding,
+        )
+
+    def _paper_exists_in_vectorstore(self, doc_id: str) -> bool:
+        try:
+            col = self.vectorstore._collection
+            res = col.get(where={"id": doc_id}, limit=1)
+            return len(res.get("ids", [])) > 0
+        except Exception:
+            if not self.manifest_exists:
+                return False
+            with open(self.manifest_path, "r") as f:
+                return any(line.strip() == doc_id for line in f)
+
+    def _mark_paper_ingested(self, arxiv_id: str) -> None:
+        with open(self.manifest_path, "a") as f:
+            f.write(f"{arxiv_id}\n")
+
+    def _ensure_doc_in_vectorstore(self, paper_text: str, doc_id: str) -> None:
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
+        )
+        docs = splitter.create_documents(
+            [paper_text], metadatas=[{"id": doc_id}]
+        )
+        with self._vs_lock:
+            if not self._paper_exists_in_vectorstore(doc_id):
+                ids = [f"{doc_id}::{i}" for i, _ in enumerate(docs)]
+                self.vectorstore.add_documents(docs, ids=ids)
+                self._mark_paper_ingested(doc_id)
+
+    def _get_global_retriever(self, k: int = 5):
+        return self.vectorstore, self.vectorstore.as_retriever(
+            search_kwargs={"k": k}
+        )
+
+    def _read_docs(self, state: RAGState) -> RAGState:
+        print("[RAG Agent] Reading Documents....")
+        papers = []
+        new_state = state.copy()
+
+        pdf_files = [
+            f
+            for f in os.listdir(self.database_path)
+            if f.lower().endswith(".pdf")
+        ]
+
+        doc_ids = [
+            pdf_filename.rsplit(".pdf", 1)[0] for pdf_filename in pdf_files
+        ]
+        pdf_files = [
+            pdf_filename
+            for pdf_filename, id in zip(pdf_files, doc_ids)
+            if not self._paper_exists_in_vectorstore(id)
+        ]
+
+        for pdf_filename in pdf_files:
+            full_text = ""
+
+            try:
+                loader = PyPDFLoader(
+                    os.path.join(self.database_path, pdf_filename)
+                )
+                pages = loader.load()
+                full_text = "\n".join([p.page_content for p in pages])
+
+            except Exception as e:
+                full_text = f"Error loading paper: {e}"
+
+            papers.append(full_text)
+
+        new_state["doc_texts"] = papers
+        new_state["doc_ids"] = doc_ids
+
+        return new_state
+
+    def _ingest_docs(self, state: RAGState) -> RAGState:
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
+        )
+
+        batch_docs, batch_ids = [], []
+        for paper, id in zip(state["doc_texts"], state["doc_ids"]):
+            cleaned_text = remove_surrogates(paper)
+            docs = splitter.create_documents(
+                [cleaned_text], metadatas=[{"id": id}]
+            )
+            ids = [f"{id}::{i}" for i, _ in enumerate(docs)]
+            batch_docs.extend(docs)
+            batch_ids.extend(ids)
+
+        if state["doc_texts"]:
+            print("[RAG Agent] Ingesting Documents Into RAG Database....")
+            with self._vs_lock:
+                self.vectorstore.add_documents(batch_docs, ids=batch_ids)
+                for id in ids:
+                    self._mark_paper_ingested(id)
+
+        return state
+
+    def _summarize_node(self, state: RAGState) -> RAGState:
+        print(
+            "[RAG Agent] Retrieving Contextually Relevant Information From Database..."
+        )
+        prompt = ChatPromptTemplate.from_template("""
+        You are a scientific assistant responsible for summarizing extracts from research papers, in the context of the following task: {context}
+
+        Summarize the retrieved scientific content below.
+        Cite sources by ID when relevant: {source_ids}
+                                                  
+        {retrieved_content}
+        """)
+        chain = prompt | self.llm | StrOutputParser()
+
+        # 2) One retrieval over the global DB with the task context
+        try:
+            results = self.vectorstore.similarity_search_with_score(
+                state["context"], k=self.return_k
+            )
+        except Exception as e:
+            print(f"RAG failed due to: {e}")
+            return {**state, "summary": ""}
+
+        source_ids_list = []
+        for doc, _ in results:
+            aid = doc.metadata.get("id")
+            if aid and aid not in source_ids_list:
+                source_ids_list.append(aid)
+        source_ids = ", ".join(source_ids_list)
+
+        # Compute a simple similarity-based quality score
+        relevancy_scores = []
+        if results:
+            distances = [score for _, score in results]
+            sims = [1.0 / (1.0 + d) for d in distances]  # map distance -> [0,1)
+            relevancy_scores = sims
+
+        retrieved_content = (
+            "\n\n".join(doc.page_content for doc, _ in results)
+            if results
+            else ""
+        )
+
+        print("[RAG Agent] Summarizing Retrieved Information From Database...")
+        # 3) One summary based on retrieved chunks
+        rag_summary = chain.invoke({
+            "retrieved_content": retrieved_content,
+            "context": state["context"],
+            "source_ids": source_ids,
+        })
+
+        # Persist a single file for the batch (optional)
+        batch_name = "RAG_summary.txt"
+        os.makedirs(self.summaries_path, exist_ok=True)
+        with open(os.path.join(self.summaries_path, batch_name), "w") as f:
+            f.write(rag_summary)
+
+        # Diagnostics
+        if relevancy_scores:
+            print(f"\nMax Relevancy Score: {max(relevancy_scores):.4f}")
+            print(f"Min Relevancy Score: {min(relevancy_scores):.4f}")
+            print(
+                f"Median Relevancy Score: {statistics.median(relevancy_scores):.4f}\n"
+            )
+        else:
+            print("\nNo RAG results retrieved (score list empty).\n")
+
+        # Return a single-element list by default (preferred)
+        return {
+            **state,
+            "summary": rag_summary,
+            "rag_metadata": {
+                "k": self.return_k,
+                "num_results": len(results),
+                "relevancy_scores": relevancy_scores,
+            },
+        }
+
+    def _build_graph(self):
+        builder = StateGraph(RAGState)
+        builder.add_node("Read Documents", self._read_docs)
+        builder.add_node("Ingest Documents", self._ingest_docs)
+        builder.add_node("Retrieve and Summarize", self._summarize_node)
+        builder.add_edge("Read Documents", "Ingest Documents")
+        builder.add_edge("Ingest Documents", "Retrieve and Summarize")
+
+        builder.set_entry_point("Read Documents")
+        builder.set_finish_point("Retrieve and Summarize")
+
+        graph = builder.compile()
+        return graph
+
+    def run(self, context: str) -> str:
+        result = self.graph.invoke({"context": context})
+
+        return result.get("summary", "No summary generated.")
+
+
+if __name__ == "__main__":
+    agent = RAGAgent(database_path="workspace/arxiv_papers_neutron_star")
+    result = agent.run(
+        context="What are the constraints on the neutron star radius and what uncertainties are there on the constraints?",
+    )
+
+    print(result)

{ursa_ai-0.4.2.dist-info → ursa_ai-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ursa-ai
-Version: 0.4.2
+Version: 0.5.0
 Summary: Agents for science at LANL
 Author-email: Mike Grosskopf <mikegros@lanl.gov>, Nathan Debardeleben <ndebard@lanl.gov>, Rahul Somasundaram <rsomasundaram@lanl.gov>, Isaac Michaud <imichaud@lanl.gov>, Avanish Mishra <avanish@lanl.gov>, Arthur Lui <alui@lanl.gov>, Russell Bent <rbent@lanl.gov>, Earl Lawrence <earl@lanl.gov>
 License-Expression: BSD-3-Clause

{ursa_ai-0.4.2.dist-info → ursa_ai-0.5.0.dist-info}/RECORD

@@ -1,5 +1,5 @@
-ursa/agents/__init__.py,sha256=HG95_hMlR0cEAF8Vbr7dF_7I2cWM-rMW1p8ULRXTzTg,1114
-ursa/agents/arxiv_agent.py,sha256=hS6RKqP0fz9KdAyFbXIxl7DuZyuhxpR-jJLfq2tCnlY,15106
+ursa/agents/__init__.py,sha256=u5ClncJ-w4nIJLQJTOC-NBv6Hu4pXxALgZJJxDt3tZw,1202
+ursa/agents/arxiv_agent.py,sha256=kiRPshNI-fR-uEFZGp6dyfQ38sIFFb6rUiiNVw3lAXc,13494
 ursa/agents/base.py,sha256=uFhRLVzqhFbTZVA7IePKbUi03ATCXuvga7rzwaHy1B0,1321
 ursa/agents/code_review_agent.py,sha256=aUDq5gT-jdl9Qs-Wewj2oz1d60xov9sN-DOYRfGNTU0,11550
 ursa/agents/execution_agent.py,sha256=okVTsZhG0S92evHtmxge3Ymq3pH0QLYY2VqzO39WG5Y,16581
@@ -7,6 +7,7 @@ ursa/agents/hypothesizer_agent.py,sha256=pUwFDWGBJAqL7CDXxWYJrQrknC3DgRe82Poc_Q_
 ursa/agents/lammps_agent.py,sha256=16eKOtAXEm-clnIZcfEaoxQONqbUJm_1dZhZhm2C2uM,14099
 ursa/agents/mp_agent.py,sha256=UyJSheMGHZpWQJL3EgYgPPqArfv6F8sndN05q4CPtyo,6015
 ursa/agents/planning_agent.py,sha256=ayyNDQifPvYtQ-JYnFk3TaXWZcd_6k8qUheJGariqG8,5574
+ursa/agents/rag_agent.py,sha256=kEL5F3lZeSKtdXXyUWlXJYYwP2ZNQ1Hz9IWl6pzQnlY,9409
 ursa/agents/recall_agent.py,sha256=bQk7ZJtiO5pj89A50OBDzAJ4G2F7ZdsMwmKnp1WWR7g,813
 ursa/agents/websearch_agent.py,sha256=rCv4AWbqe5Us4FmuypM6jptri21nKoNg044ncsu9u3E,8014
 ursa/prompt_library/code_review_prompts.py,sha256=-HuhwW9W_p2LDn44bXLntxLADHCOyl-2KIXxRHto66w,2444
@@ -20,8 +21,8 @@ ursa/tools/write_code.py,sha256=DtCsUMZegYm0mk-HMPG5Zo3Ba1gbGfnXHsv1NZTdDs8,1220
 ursa/util/diff_renderer.py,sha256=1L1q2qWWb8gLhR532-LgJn2TrqXDx0gUpPVOWD_sqeU,4086
 ursa/util/memory_logger.py,sha256=GiKYbQBpxlNRLKyqKFJyrbSbVCkXpRB7Yr5so43tUAw,6097
 ursa/util/parse.py,sha256=M0cjyQWmjatxX4WbVmDRUiirTLyW-t_Aemlrlrsc5nA,2811
-ursa_ai-0.4.2.dist-info/licenses/LICENSE,sha256=4Vr6_u2zTHIUvYjoOBg9ztDbfpV3hyCFv3mTCS87gYU,1482
-ursa_ai-0.4.2.dist-info/METADATA,sha256=QDR_81uiAgRFxGGjqPMsFmIH9gQ93-_gFHbAS4WY54s,6898
-ursa_ai-0.4.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ursa_ai-0.4.2.dist-info/top_level.txt,sha256=OjA1gRYSUAeiXGnpqPC8iOOGfcjFO1IlP848qMnYSdY,5
-ursa_ai-0.4.2.dist-info/RECORD,,
+ursa_ai-0.5.0.dist-info/licenses/LICENSE,sha256=4Vr6_u2zTHIUvYjoOBg9ztDbfpV3hyCFv3mTCS87gYU,1482
+ursa_ai-0.5.0.dist-info/METADATA,sha256=Um85Xzvs6LSGUg9X0tWitsLVwYqYgMJ0I3EhYzK8oIE,6898
+ursa_ai-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ursa_ai-0.5.0.dist-info/top_level.txt,sha256=OjA1gRYSUAeiXGnpqPC8iOOGfcjFO1IlP848qMnYSdY,5
+ursa_ai-0.5.0.dist-info/RECORD,,