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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ursa-ai
-Version: 0.5.0
+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.5.0 → 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.5.0 → 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.5.0 → ursa_ai-0.6.0rc1}/src/ursa/agents/arxiv_agent.py

@@ -3,12 +3,14 @@ import os
 import re
 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_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
@@ -16,8 +18,8 @@ from PIL import Image
 from tqdm import tqdm
 from typing_extensions import List, TypedDict
 
-from .base import BaseAgent
-from .rag_agent import RAGAgent
+from ursa.agents.base import BaseAgent
+from ursa.agents.rag_agent import RAGAgent
 
 try:
     from openai import OpenAI
@@ -120,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,
@@ -141,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)
 
@@ -259,10 +261,13 @@ class ArxivAgent(BaseAgent):
 
             try:
                 cleaned_text = remove_surrogates(paper["full_text"])
-                summary = chain.invoke({
-                    "retrieved_content": cleaned_text,
-                    "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}"
@@ -304,7 +309,9 @@ class ArxivAgent(BaseAgent):
             embedding=self.rag_embedding,
             database_path=self.database_path,
         )
-        new_state["final_summary"] = rag_agent.run(context=state["context"])
+        new_state["final_summary"] = rag_agent.invoke(context=state["context"])[
+            "summary"
+        ]
         return new_state
 
     def _aggregate_node(self, state: PaperState) -> PaperState:
@@ -341,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)
@@ -352,49 +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:
             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")
+                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:
-                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")
+                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:
-            builder.set_entry_point("fetch_papers")
-            builder.set_finish_point("fetch_papers")
+            graph.set_entry_point("_fetch_node")
+            graph.set_finish_point("_fetch_node")
 
-        graph = builder.compile()
-        return graph
+        return graph.compile(checkpointer=self.checkpointer)
 
-    def run(self, arxiv_search_query: str, context: str) -> str:
-        result = self.graph.invoke({
-            "query": arxiv_search_query,
-            "context": context,
-        })
+    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"]
+        )
 
-        if self.summarize:
-            return result.get("final_summary", "No summary generated.")
-        else:
-            return "\n\nFinished Fetching papers!"
+        # 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)."
+                )
 
+        result = self._action.invoke(inputs, config)
+
+        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