ursa-ai 0.4.2__tar.gz → 0.6.0rc1__tar.gz

{ursa_ai-0.4.2 → ursa_ai-0.6.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ursa-ai
-Version: 0.4.2
+Version: 0.6.0rc1
 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
@@ -38,8 +38,7 @@ Requires-Dist: langchain-anthropic<0.4,>=0.3.19
 Requires-Dist: langgraph-checkpoint-sqlite<3.0,>=2.0.10
 Requires-Dist: langchain-ollama<0.4,>=0.3.6
 Requires-Dist: ddgs>=9.5.5
-Requires-Dist: atomman>=1.5.2
-Requires-Dist: trafilatura>=1.6.1
+Requires-Dist: typer>=0.16.1
 Dynamic: license-file
 
 # URSA - The Universal Research and Scientific Agent
@@ -81,7 +80,68 @@ Documentation for combining agents:
 - [ArXiv -> Execution for Materials](docs/combining_arxiv_and_execution.md)
 - [ArXiv -> Execution for Neutron Star Properties](docs/combining_arxiv_and_execution_neutronStar.md)
 
-# Sandboxing
+
+## Command line usage
+
+You can install `ursa` as a command line app with `pip install`; or with `uv` via
+
+```bash
+uv tool install ursa-ai
+```
+
+To use the command line app, run
+
+```
+ursa run
+```
+
+This will start a REPL in your terminal.
+
+```
+  __  ________________ _
+ / / / / ___/ ___/ __ `/
+/ /_/ / /  (__  ) /_/ /
+\__,_/_/  /____/\__,_/
+
+For help, type: ? or help. Exit with Ctrl+d.
+ursa>
+```
+
+Within the REPL, you can get help by typing `?` or `help`. 
+
+You can chat with an LLM by simply typing into the terminal.
+
+```
+ursa> How are you?
+Thanks for asking! I’m doing well. How are you today? What can I help you with?
+```
+
+You can run various agents by typing the name of the agent. For example,
+
+```
+ursa> plan
+Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
+```
+
+If you run subsequent agents, the last output will be appended to the prompt for the next agent.
+
+So, to run the Planning Agent followed by the Execution Agent:
+```
+ursa> plan
+Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
+
+...
+
+ursa> execute
+Enter your prompt for Execution Agent: Execute the plan.
+```
+
+You can get a list of available command line options via
+```
+ursa run --help
+```
+
+## Sandboxing
 The Execution Agent is allowed to run system commands and write/run code. Being able to execute arbitrary system commands or write
 and execute code has the potential to cause problems like:
 - Damage code or data on the computer
@@ -98,6 +158,65 @@ Some suggestions for sandboxing the agent:
 
 You have a duty for ensuring that you use URSA responsibly.
 
+## Container image
+
+To enable limited sandboxing insofar as containerization does this, you can run
+the following commands:
+
+### Docker
+
+```shell
+# Build a local container using the Docker runtime
+docker buildx build --progress=plain -t ursa .
+
+# Run included example
+docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY ursa \
+    bash -c "uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
+
+# Run script from host system
+mkdir -p scripts
+echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
+docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY \
+    --mount type=bind,src=$PWD/scripts,dst=/mnt/workspace \
+    ursa \
+    bash -c "uv run /mnt/workspace/my_script.py"
+```
+
+### Charliecloud
+
+[Charliecloud](https://charliecloud.io/) is a rootless alternative to docker
+that is sometimes preferred on HPC. The following commands replicate the
+behaviors above for docker.
+
+```shell
+# Build a local container using the Docker runtime
+ch-image build -t ursa
+
+# Convert image to sqfs, for use on another system
+ch-convert ursa ursa.sqfs
+
+# Run included example (if wanted, replace ursa with /path/to/ursa.sqfs)
+ch-run -W ursa \
+    --unset-env="*" \
+    --set-env \
+    --set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
+    --cd /app \
+    -- bash -c \
+    "uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
+
+# Run script from host system (if wanted, replace ursa with /path/to/ursa.sqfs)
+mkdir -p scripts
+echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
+ch-run -W ursa \
+    --unset-env="*" \
+    --set-env \
+    --set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
+    --bind ${PWD}/scripts:/mnt/workspace \
+    --cd /app \
+    -- bash -c \
+    "uv run python /mnt/workspace/integer_sum.py"
+```
+
 ## Development Dependencies
 
 * [`uv`](https://docs.astral.sh/uv/)

{ursa_ai-0.4.2 → ursa_ai-0.6.0rc1}/README.md

@@ -37,7 +37,68 @@ Documentation for combining agents:
 - [ArXiv -> Execution for Materials](docs/combining_arxiv_and_execution.md)
 - [ArXiv -> Execution for Neutron Star Properties](docs/combining_arxiv_and_execution_neutronStar.md)
 
-# Sandboxing
+
+## Command line usage
+
+You can install `ursa` as a command line app with `pip install`; or with `uv` via
+
+```bash
+uv tool install ursa-ai
+```
+
+To use the command line app, run
+
+```
+ursa run
+```
+
+This will start a REPL in your terminal.
+
+```
+  __  ________________ _
+ / / / / ___/ ___/ __ `/
+/ /_/ / /  (__  ) /_/ /
+\__,_/_/  /____/\__,_/
+
+For help, type: ? or help. Exit with Ctrl+d.
+ursa>
+```
+
+Within the REPL, you can get help by typing `?` or `help`. 
+
+You can chat with an LLM by simply typing into the terminal.
+
+```
+ursa> How are you?
+Thanks for asking! I’m doing well. How are you today? What can I help you with?
+```
+
+You can run various agents by typing the name of the agent. For example,
+
+```
+ursa> plan
+Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
+```
+
+If you run subsequent agents, the last output will be appended to the prompt for the next agent.
+
+So, to run the Planning Agent followed by the Execution Agent:
+```
+ursa> plan
+Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
+
+...
+
+ursa> execute
+Enter your prompt for Execution Agent: Execute the plan.
+```
+
+You can get a list of available command line options via
+```
+ursa run --help
+```
+
+## Sandboxing
 The Execution Agent is allowed to run system commands and write/run code. Being able to execute arbitrary system commands or write
 and execute code has the potential to cause problems like:
 - Damage code or data on the computer
@@ -54,6 +115,65 @@ Some suggestions for sandboxing the agent:
 
 You have a duty for ensuring that you use URSA responsibly.
 
+## Container image
+
+To enable limited sandboxing insofar as containerization does this, you can run
+the following commands:
+
+### Docker
+
+```shell
+# Build a local container using the Docker runtime
+docker buildx build --progress=plain -t ursa .
+
+# Run included example
+docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY ursa \
+    bash -c "uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
+
+# Run script from host system
+mkdir -p scripts
+echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
+docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY \
+    --mount type=bind,src=$PWD/scripts,dst=/mnt/workspace \
+    ursa \
+    bash -c "uv run /mnt/workspace/my_script.py"
+```
+
+### Charliecloud
+
+[Charliecloud](https://charliecloud.io/) is a rootless alternative to docker
+that is sometimes preferred on HPC. The following commands replicate the
+behaviors above for docker.
+
+```shell
+# Build a local container using the Docker runtime
+ch-image build -t ursa
+
+# Convert image to sqfs, for use on another system
+ch-convert ursa ursa.sqfs
+
+# Run included example (if wanted, replace ursa with /path/to/ursa.sqfs)
+ch-run -W ursa \
+    --unset-env="*" \
+    --set-env \
+    --set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
+    --cd /app \
+    -- bash -c \
+    "uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
+
+# Run script from host system (if wanted, replace ursa with /path/to/ursa.sqfs)
+mkdir -p scripts
+echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
+ch-run -W ursa \
+    --unset-env="*" \
+    --set-env \
+    --set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
+    --bind ${PWD}/scripts:/mnt/workspace \
+    --cd /app \
+    -- bash -c \
+    "uv run python /mnt/workspace/integer_sum.py"
+```
+
 ## Development Dependencies
 
 * [`uv`](https://docs.astral.sh/uv/)

{ursa_ai-0.4.2 → ursa_ai-0.6.0rc1}/pyproject.toml

@@ -38,8 +38,7 @@ dependencies = [
     "langgraph-checkpoint-sqlite>=2.0.10,<3.0",
     "langchain-ollama>=0.3.6,<0.4",
     "ddgs>=9.5.5",
-    "atomman>=1.5.2",
-    "trafilatura>=1.6.1",
+    "typer>=0.16.1",
 ]
 classifiers = [
     "Operating System :: OS Independent",
@@ -50,6 +49,9 @@ classifiers = [
     "Programming Language :: Python :: 3.14",
 ]
 
+[project.scripts]
+ursa = "ursa.cli:main"
+
 [project.urls]
 Homepage = "https://github.com/lanl/ursa"
 Documentation = "https://github.com/lanl/ursa/tree/main/docs"
@@ -81,5 +83,13 @@ dev = [
     "langgraph-checkpoint-sqlite>=2.0.10",
     "notebook>=7.3.3",
     "pre-commit>=4.3.0",
+    "pytest>=8.4.2",
     "scikit-optimize>=0.10.2",
 ]
+lammps = [
+    "atomman>=1.5.2",
+    "trafilatura>=1.6.1",
+]
+opt = [
+    "ortools>=9.14,<9.15",
+]

{ursa_ai-0.4.2 → ursa_ai-0.6.0rc1}/src/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_ai-0.4.2 → ursa_ai-0.6.0rc1}/src/ursa/agents/arxiv_agent.py

@@ -1,17 +1,16 @@
 import base64
 import os
 import re
-import statistics
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from io import BytesIO
+from typing import Any, Mapping
 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.language_models import BaseChatModel
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 from langgraph.graph import StateGraph
@@ -19,16 +18,14 @@ from PIL import Image
 from tqdm import tqdm
 from typing_extensions import List, TypedDict
 
-from .base import BaseAgent
+from ursa.agents.base import BaseAgent
+from ursa.agents.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
@@ -125,7 +122,7 @@ def remove_surrogates(text: str) -> str:
 class ArxivAgent(BaseAgent):
     def __init__(
         self,
-        llm="openai/o3-mini",
+        llm: str | BaseChatModel = "openai/o3-mini",
         summarize: bool = True,
         process_images=True,
         max_results: int = 3,
@@ -146,7 +143,7 @@ class ArxivAgent(BaseAgent):
         self.download_papers = download_papers
         self.rag_embedding = rag_embedding
 
-        self.graph = self._build_graph()
+        self._action = self._build_graph()
 
         os.makedirs(self.database_path, exist_ok=True)
 
@@ -242,27 +239,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,35 +261,13 @@ 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,
-                    "context": state["context"],
-                })
+                summary = chain.invoke(
+                    {
+                        "retrieved_content": cleaned_text,
+                        "context": state["context"],
+                    },
+                    config=self.build_config(tags=["arxiv", "summarize_each"]),
+                )
 
             except Exception as e:
                 summary = f"Error summarizing paper: {e}"
@@ -346,15 +300,20 @@ 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.invoke(context=state["context"])[
+            "summary"
+        ]
+        return new_state
+
     def _aggregate_node(self, state: PaperState) -> PaperState:
         summaries = state["summaries"]
         papers = state["papers"]
@@ -389,10 +348,13 @@ class ArxivAgent(BaseAgent):
 
         chain = prompt | self.llm | StrOutputParser()
 
-        final_summary = chain.invoke({
-            "Summaries": combined,
-            "context": state["context"],
-        })
+        final_summary = chain.invoke(
+            {
+                "Summaries": combined,
+                "context": state["context"],
+            },
+            config=self.build_config(tags=["arxiv", "aggregate"]),
+        )
 
         with open(self.summaries_path + "/final_summary.txt", "w") as f:
             f.write(final_summary)
@@ -400,42 +362,69 @@ class ArxivAgent(BaseAgent):
         return {**state, "final_summary": final_summary}
 
     def _build_graph(self):
-        builder = StateGraph(PaperState)
-        builder.add_node("fetch_papers", self._fetch_node)
+        graph = StateGraph(PaperState)
 
+        self.add_node(graph, self._fetch_node)
         if self.summarize:
-            builder.add_node("summarize_each", self._summarize_node)
-            builder.add_node("aggregate", self._aggregate_node)
+            if self.rag_embedding:
+                self.add_node(graph, self._rag_node)
+                graph.set_entry_point("_fetch_node")
+                graph.add_edge("_fetch_node", "_rag_node")
+                graph.set_finish_point("_rag_node")
+            else:
+                self.add_node(graph, self._summarize_node)
+                self.add_node(graph, self._aggregate_node)
+
+                graph.set_entry_point("_fetch_node")
+                graph.add_edge("_fetch_node", "_summarize_node")
+                graph.add_edge("_summarize_node", "_aggregate_node")
+                graph.set_finish_point("_aggregate_node")
+        else:
+            graph.set_entry_point("_fetch_node")
+            graph.set_finish_point("_fetch_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")
+        return graph.compile(checkpointer=self.checkpointer)
 
-        else:
-            builder.set_entry_point("fetch_papers")
-            builder.set_finish_point("fetch_papers")
+    def _invoke(
+        self,
+        inputs: Mapping[str, Any],
+        *,
+        summarize: bool | None = None,
+        recursion_limit: int = 1000,
+        **_,
+    ) -> str:
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
 
-        graph = builder.compile()
-        return graph
+        # this seems dumb, but it's b/c sometimes we had referred to the value as
+        # 'query' other times as 'arxiv_search_query' so trying to keep it compatible
+        # aliasing: accept arxiv_search_query -> query
+        if "query" not in inputs:
+            if "arxiv_search_query" in inputs:
+                # make a shallow copy and rename the key
+                inputs = dict(inputs)
+                inputs["query"] = inputs.pop("arxiv_search_query")
+            else:
+                raise KeyError(
+                    "Missing 'query' in inputs (alias 'arxiv_search_query' also accepted)."
+                )
 
-    def run(self, arxiv_search_query: str, context: str) -> str:
-        result = self.graph.invoke({
-            "query": arxiv_search_query,
-            "context": context,
-        })
+        result = self._action.invoke(inputs, config)
 
-        if self.summarize:
-            return result.get("final_summary", "No summary generated.")
-        else:
-            return "\n\nFinished Fetching papers!"
+        use_summary = self.summarize if summarize is None else summarize
 
+        return (
+            result.get("final_summary", "No summary generated.")
+            if use_summary
+            else "\n\nFinished Fetching papers!"
+        )
 
-if __name__ == "__main__":
-    agent = ArxivAgent()
-    result = agent.run(
-        arxiv_search_query="Experimental Constraints on neutron star radius",
-        context="What are the constraints on the neutron star radius and what uncertainties are there on the constraints?",
-    )
 
-    print(result)
+# NOTE: Run test in `tests/agents/test_arxiv_agent/test_arxiv_agent.py` via:
+#
+# pytest -s tests/agents/test_arxiv_agent
+#
+# OR
+#
+# uv run pytest -s tests/agents/test_arxiv_agent