ursa-ai 0.5.0__tar.gz → 0.6.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/PKG-INFO +123 -4
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/README.md +121 -1
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/pyproject.toml +18 -2
- ursa_ai-0.6.0/src/ursa/__init__.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/arxiv_agent.py +77 -47
- ursa_ai-0.6.0/src/ursa/agents/base.py +408 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/code_review_agent.py +3 -1
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/execution_agent.py +92 -48
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/hypothesizer_agent.py +39 -42
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/lammps_agent.py +51 -29
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/mp_agent.py +45 -20
- ursa_ai-0.6.0/src/ursa/agents/optimization_agent.py +405 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/planning_agent.py +63 -28
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/rag_agent.py +75 -44
- ursa_ai-0.6.0/src/ursa/agents/recall_agent.py +53 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/websearch_agent.py +44 -54
- ursa_ai-0.6.0/src/ursa/cli/__init__.py +127 -0
- ursa_ai-0.6.0/src/ursa/cli/hitl.py +426 -0
- ursa_ai-0.6.0/src/ursa/observability/pricing.py +319 -0
- ursa_ai-0.6.0/src/ursa/observability/timing.py +1441 -0
- ursa_ai-0.6.0/src/ursa/prompt_library/__init__.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/prompt_library/execution_prompts.py +7 -0
- ursa_ai-0.6.0/src/ursa/prompt_library/optimization_prompts.py +131 -0
- ursa_ai-0.6.0/src/ursa/tools/__init__.py +0 -0
- ursa_ai-0.6.0/src/ursa/tools/feasibility_checker.py +114 -0
- ursa_ai-0.6.0/src/ursa/tools/feasibility_tools.py +1075 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/tools/write_code.py +1 -1
- ursa_ai-0.6.0/src/ursa/util/__init__.py +0 -0
- ursa_ai-0.6.0/src/ursa/util/helperFunctions.py +142 -0
- ursa_ai-0.6.0/src/ursa/util/optimization_schema.py +78 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/util/parse.py +1 -1
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa_ai.egg-info/PKG-INFO +123 -4
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa_ai.egg-info/SOURCES.txt +15 -0
- ursa_ai-0.6.0/src/ursa_ai.egg-info/entry_points.txt +2 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa_ai.egg-info/requires.txt +1 -2
- ursa_ai-0.5.0/src/ursa/agents/base.py +0 -41
- ursa_ai-0.5.0/src/ursa/agents/recall_agent.py +0 -23
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/LICENSE +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/setup.cfg +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/agents/__init__.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/prompt_library/code_review_prompts.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/prompt_library/hypothesizer_prompts.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/prompt_library/literature_prompts.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/prompt_library/planning_prompts.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/prompt_library/websearch_prompts.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/tools/run_command.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/util/diff_renderer.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa/util/memory_logger.py +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa_ai.egg-info/dependency_links.txt +0 -0
- {ursa_ai-0.5.0 → ursa_ai-0.6.0}/src/ursa_ai.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: ursa-ai
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.6.0
|
|
4
4
|
Summary: Agents for science at LANL
|
|
5
5
|
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>
|
|
6
6
|
License-Expression: BSD-3-Clause
|
|
@@ -38,8 +38,7 @@ Requires-Dist: langchain-anthropic<0.4,>=0.3.19
|
|
|
38
38
|
Requires-Dist: langgraph-checkpoint-sqlite<3.0,>=2.0.10
|
|
39
39
|
Requires-Dist: langchain-ollama<0.4,>=0.3.6
|
|
40
40
|
Requires-Dist: ddgs>=9.5.5
|
|
41
|
-
Requires-Dist:
|
|
42
|
-
Requires-Dist: trafilatura>=1.6.1
|
|
41
|
+
Requires-Dist: typer>=0.16.1
|
|
43
42
|
Dynamic: license-file
|
|
44
43
|
|
|
45
44
|
# URSA - The Universal Research and Scientific Agent
|
|
@@ -81,7 +80,68 @@ Documentation for combining agents:
|
|
|
81
80
|
- [ArXiv -> Execution for Materials](docs/combining_arxiv_and_execution.md)
|
|
82
81
|
- [ArXiv -> Execution for Neutron Star Properties](docs/combining_arxiv_and_execution_neutronStar.md)
|
|
83
82
|
|
|
84
|
-
|
|
83
|
+
|
|
84
|
+
## Command line usage
|
|
85
|
+
|
|
86
|
+
You can install `ursa` as a command line app with `pip install`; or with `uv` via
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
uv tool install ursa-ai
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
To use the command line app, run
|
|
93
|
+
|
|
94
|
+
```
|
|
95
|
+
ursa run
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
This will start a REPL in your terminal.
|
|
99
|
+
|
|
100
|
+
```
|
|
101
|
+
__ ________________ _
|
|
102
|
+
/ / / / ___/ ___/ __ `/
|
|
103
|
+
/ /_/ / / (__ ) /_/ /
|
|
104
|
+
\__,_/_/ /____/\__,_/
|
|
105
|
+
|
|
106
|
+
For help, type: ? or help. Exit with Ctrl+d.
|
|
107
|
+
ursa>
|
|
108
|
+
```
|
|
109
|
+
|
|
110
|
+
Within the REPL, you can get help by typing `?` or `help`.
|
|
111
|
+
|
|
112
|
+
You can chat with an LLM by simply typing into the terminal.
|
|
113
|
+
|
|
114
|
+
```
|
|
115
|
+
ursa> How are you?
|
|
116
|
+
Thanks for asking! I’m doing well. How are you today? What can I help you with?
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
You can run various agents by typing the name of the agent. For example,
|
|
120
|
+
|
|
121
|
+
```
|
|
122
|
+
ursa> plan
|
|
123
|
+
Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
If you run subsequent agents, the last output will be appended to the prompt for the next agent.
|
|
127
|
+
|
|
128
|
+
So, to run the Planning Agent followed by the Execution Agent:
|
|
129
|
+
```
|
|
130
|
+
ursa> plan
|
|
131
|
+
Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
|
|
132
|
+
|
|
133
|
+
...
|
|
134
|
+
|
|
135
|
+
ursa> execute
|
|
136
|
+
Enter your prompt for Execution Agent: Execute the plan.
|
|
137
|
+
```
|
|
138
|
+
|
|
139
|
+
You can get a list of available command line options via
|
|
140
|
+
```
|
|
141
|
+
ursa run --help
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
## Sandboxing
|
|
85
145
|
The Execution Agent is allowed to run system commands and write/run code. Being able to execute arbitrary system commands or write
|
|
86
146
|
and execute code has the potential to cause problems like:
|
|
87
147
|
- Damage code or data on the computer
|
|
@@ -98,6 +158,65 @@ Some suggestions for sandboxing the agent:
|
|
|
98
158
|
|
|
99
159
|
You have a duty for ensuring that you use URSA responsibly.
|
|
100
160
|
|
|
161
|
+
## Container image
|
|
162
|
+
|
|
163
|
+
To enable limited sandboxing insofar as containerization does this, you can run
|
|
164
|
+
the following commands:
|
|
165
|
+
|
|
166
|
+
### Docker
|
|
167
|
+
|
|
168
|
+
```shell
|
|
169
|
+
# Build a local container using the Docker runtime
|
|
170
|
+
docker buildx build --progress=plain -t ursa .
|
|
171
|
+
|
|
172
|
+
# Run included example
|
|
173
|
+
docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY ursa \
|
|
174
|
+
bash -c "uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
|
|
175
|
+
|
|
176
|
+
# Run script from host system
|
|
177
|
+
mkdir -p scripts
|
|
178
|
+
echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
|
|
179
|
+
docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY \
|
|
180
|
+
--mount type=bind,src=$PWD/scripts,dst=/mnt/workspace \
|
|
181
|
+
ursa \
|
|
182
|
+
bash -c "uv run /mnt/workspace/my_script.py"
|
|
183
|
+
```
|
|
184
|
+
|
|
185
|
+
### Charliecloud
|
|
186
|
+
|
|
187
|
+
[Charliecloud](https://charliecloud.io/) is a rootless alternative to docker
|
|
188
|
+
that is sometimes preferred on HPC. The following commands replicate the
|
|
189
|
+
behaviors above for docker.
|
|
190
|
+
|
|
191
|
+
```shell
|
|
192
|
+
# Build a local container using the Docker runtime
|
|
193
|
+
ch-image build -t ursa
|
|
194
|
+
|
|
195
|
+
# Convert image to sqfs, for use on another system
|
|
196
|
+
ch-convert ursa ursa.sqfs
|
|
197
|
+
|
|
198
|
+
# Run included example (if wanted, replace ursa with /path/to/ursa.sqfs)
|
|
199
|
+
ch-run -W ursa \
|
|
200
|
+
--unset-env="*" \
|
|
201
|
+
--set-env \
|
|
202
|
+
--set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
|
|
203
|
+
--cd /app \
|
|
204
|
+
-- bash -c \
|
|
205
|
+
"uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
|
|
206
|
+
|
|
207
|
+
# Run script from host system (if wanted, replace ursa with /path/to/ursa.sqfs)
|
|
208
|
+
mkdir -p scripts
|
|
209
|
+
echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
|
|
210
|
+
ch-run -W ursa \
|
|
211
|
+
--unset-env="*" \
|
|
212
|
+
--set-env \
|
|
213
|
+
--set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
|
|
214
|
+
--bind ${PWD}/scripts:/mnt/workspace \
|
|
215
|
+
--cd /app \
|
|
216
|
+
-- bash -c \
|
|
217
|
+
"uv run python /mnt/workspace/integer_sum.py"
|
|
218
|
+
```
|
|
219
|
+
|
|
101
220
|
## Development Dependencies
|
|
102
221
|
|
|
103
222
|
* [`uv`](https://docs.astral.sh/uv/)
|
|
@@ -37,7 +37,68 @@ Documentation for combining agents:
|
|
|
37
37
|
- [ArXiv -> Execution for Materials](docs/combining_arxiv_and_execution.md)
|
|
38
38
|
- [ArXiv -> Execution for Neutron Star Properties](docs/combining_arxiv_and_execution_neutronStar.md)
|
|
39
39
|
|
|
40
|
-
|
|
40
|
+
|
|
41
|
+
## Command line usage
|
|
42
|
+
|
|
43
|
+
You can install `ursa` as a command line app with `pip install`; or with `uv` via
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
uv tool install ursa-ai
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
To use the command line app, run
|
|
50
|
+
|
|
51
|
+
```
|
|
52
|
+
ursa run
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
This will start a REPL in your terminal.
|
|
56
|
+
|
|
57
|
+
```
|
|
58
|
+
__ ________________ _
|
|
59
|
+
/ / / / ___/ ___/ __ `/
|
|
60
|
+
/ /_/ / / (__ ) /_/ /
|
|
61
|
+
\__,_/_/ /____/\__,_/
|
|
62
|
+
|
|
63
|
+
For help, type: ? or help. Exit with Ctrl+d.
|
|
64
|
+
ursa>
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
Within the REPL, you can get help by typing `?` or `help`.
|
|
68
|
+
|
|
69
|
+
You can chat with an LLM by simply typing into the terminal.
|
|
70
|
+
|
|
71
|
+
```
|
|
72
|
+
ursa> How are you?
|
|
73
|
+
Thanks for asking! I’m doing well. How are you today? What can I help you with?
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
You can run various agents by typing the name of the agent. For example,
|
|
77
|
+
|
|
78
|
+
```
|
|
79
|
+
ursa> plan
|
|
80
|
+
Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
If you run subsequent agents, the last output will be appended to the prompt for the next agent.
|
|
84
|
+
|
|
85
|
+
So, to run the Planning Agent followed by the Execution Agent:
|
|
86
|
+
```
|
|
87
|
+
ursa> plan
|
|
88
|
+
Enter your prompt for Planning Agent: Write a python script to do linear regression using only numpy.
|
|
89
|
+
|
|
90
|
+
...
|
|
91
|
+
|
|
92
|
+
ursa> execute
|
|
93
|
+
Enter your prompt for Execution Agent: Execute the plan.
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
You can get a list of available command line options via
|
|
97
|
+
```
|
|
98
|
+
ursa run --help
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
## Sandboxing
|
|
41
102
|
The Execution Agent is allowed to run system commands and write/run code. Being able to execute arbitrary system commands or write
|
|
42
103
|
and execute code has the potential to cause problems like:
|
|
43
104
|
- Damage code or data on the computer
|
|
@@ -54,6 +115,65 @@ Some suggestions for sandboxing the agent:
|
|
|
54
115
|
|
|
55
116
|
You have a duty for ensuring that you use URSA responsibly.
|
|
56
117
|
|
|
118
|
+
## Container image
|
|
119
|
+
|
|
120
|
+
To enable limited sandboxing insofar as containerization does this, you can run
|
|
121
|
+
the following commands:
|
|
122
|
+
|
|
123
|
+
### Docker
|
|
124
|
+
|
|
125
|
+
```shell
|
|
126
|
+
# Build a local container using the Docker runtime
|
|
127
|
+
docker buildx build --progress=plain -t ursa .
|
|
128
|
+
|
|
129
|
+
# Run included example
|
|
130
|
+
docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY ursa \
|
|
131
|
+
bash -c "uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
|
|
132
|
+
|
|
133
|
+
# Run script from host system
|
|
134
|
+
mkdir -p scripts
|
|
135
|
+
echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
|
|
136
|
+
docker run -e "OPENAI_API_KEY"=$OPENAI_API_KEY \
|
|
137
|
+
--mount type=bind,src=$PWD/scripts,dst=/mnt/workspace \
|
|
138
|
+
ursa \
|
|
139
|
+
bash -c "uv run /mnt/workspace/my_script.py"
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
### Charliecloud
|
|
143
|
+
|
|
144
|
+
[Charliecloud](https://charliecloud.io/) is a rootless alternative to docker
|
|
145
|
+
that is sometimes preferred on HPC. The following commands replicate the
|
|
146
|
+
behaviors above for docker.
|
|
147
|
+
|
|
148
|
+
```shell
|
|
149
|
+
# Build a local container using the Docker runtime
|
|
150
|
+
ch-image build -t ursa
|
|
151
|
+
|
|
152
|
+
# Convert image to sqfs, for use on another system
|
|
153
|
+
ch-convert ursa ursa.sqfs
|
|
154
|
+
|
|
155
|
+
# Run included example (if wanted, replace ursa with /path/to/ursa.sqfs)
|
|
156
|
+
ch-run -W ursa \
|
|
157
|
+
--unset-env="*" \
|
|
158
|
+
--set-env \
|
|
159
|
+
--set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
|
|
160
|
+
--cd /app \
|
|
161
|
+
-- bash -c \
|
|
162
|
+
"uv run python examples/single_agent_examples/execution_agnet/integer_sum.py"
|
|
163
|
+
|
|
164
|
+
# Run script from host system (if wanted, replace ursa with /path/to/ursa.sqfs)
|
|
165
|
+
mkdir -p scripts
|
|
166
|
+
echo "import ursa; print('Hello from ursa')" > scripts/my_script.py
|
|
167
|
+
ch-run -W ursa \
|
|
168
|
+
--unset-env="*" \
|
|
169
|
+
--set-env \
|
|
170
|
+
--set-env="OPENAI_API_KEY"=$OPENAI_API_KEY \
|
|
171
|
+
--bind ${PWD}/scripts:/mnt/workspace \
|
|
172
|
+
--cd /app \
|
|
173
|
+
-- bash -c \
|
|
174
|
+
"uv run python /mnt/workspace/integer_sum.py"
|
|
175
|
+
```
|
|
176
|
+
|
|
57
177
|
## Development Dependencies
|
|
58
178
|
|
|
59
179
|
* [`uv`](https://docs.astral.sh/uv/)
|
|
@@ -38,8 +38,7 @@ dependencies = [
|
|
|
38
38
|
"langgraph-checkpoint-sqlite>=2.0.10,<3.0",
|
|
39
39
|
"langchain-ollama>=0.3.6,<0.4",
|
|
40
40
|
"ddgs>=9.5.5",
|
|
41
|
-
"
|
|
42
|
-
"trafilatura>=1.6.1",
|
|
41
|
+
"typer>=0.16.1",
|
|
43
42
|
]
|
|
44
43
|
classifiers = [
|
|
45
44
|
"Operating System :: OS Independent",
|
|
@@ -50,6 +49,9 @@ classifiers = [
|
|
|
50
49
|
"Programming Language :: Python :: 3.14",
|
|
51
50
|
]
|
|
52
51
|
|
|
52
|
+
[project.scripts]
|
|
53
|
+
ursa = "ursa.cli:main"
|
|
54
|
+
|
|
53
55
|
[project.urls]
|
|
54
56
|
Homepage = "https://github.com/lanl/ursa"
|
|
55
57
|
Documentation = "https://github.com/lanl/ursa/tree/main/docs"
|
|
@@ -81,5 +83,19 @@ dev = [
|
|
|
81
83
|
"langgraph-checkpoint-sqlite>=2.0.10",
|
|
82
84
|
"notebook>=7.3.3",
|
|
83
85
|
"pre-commit>=4.3.0",
|
|
86
|
+
"pytest>=8.4.2",
|
|
84
87
|
"scikit-optimize>=0.10.2",
|
|
85
88
|
]
|
|
89
|
+
docs = [
|
|
90
|
+
"mkdocs>=1.6.1",
|
|
91
|
+
"mkdocs-autorefs>=1.4.3",
|
|
92
|
+
"mkdocs-material>=9.6.21",
|
|
93
|
+
"mkdocstrings-python>=1.18.2",
|
|
94
|
+
]
|
|
95
|
+
lammps = [
|
|
96
|
+
"atomman>=1.5.2",
|
|
97
|
+
"trafilatura>=1.6.1",
|
|
98
|
+
]
|
|
99
|
+
opt = [
|
|
100
|
+
"ortools>=9.14,<9.15",
|
|
101
|
+
]
|
|
File without changes
|
|
@@ -3,12 +3,14 @@ import os
|
|
|
3
3
|
import re
|
|
4
4
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
5
5
|
from io import BytesIO
|
|
6
|
+
from typing import Any, Mapping
|
|
6
7
|
from urllib.parse import quote
|
|
7
8
|
|
|
8
9
|
import feedparser
|
|
9
10
|
import pymupdf
|
|
10
11
|
import requests
|
|
11
12
|
from langchain_community.document_loaders import PyPDFLoader
|
|
13
|
+
from langchain_core.language_models import BaseChatModel
|
|
12
14
|
from langchain_core.output_parsers import StrOutputParser
|
|
13
15
|
from langchain_core.prompts import ChatPromptTemplate
|
|
14
16
|
from langgraph.graph import StateGraph
|
|
@@ -16,8 +18,8 @@ from PIL import Image
|
|
|
16
18
|
from tqdm import tqdm
|
|
17
19
|
from typing_extensions import List, TypedDict
|
|
18
20
|
|
|
19
|
-
from .base import BaseAgent
|
|
20
|
-
from .rag_agent import RAGAgent
|
|
21
|
+
from ursa.agents.base import BaseAgent
|
|
22
|
+
from ursa.agents.rag_agent import RAGAgent
|
|
21
23
|
|
|
22
24
|
try:
|
|
23
25
|
from openai import OpenAI
|
|
@@ -120,7 +122,7 @@ def remove_surrogates(text: str) -> str:
|
|
|
120
122
|
class ArxivAgent(BaseAgent):
|
|
121
123
|
def __init__(
|
|
122
124
|
self,
|
|
123
|
-
llm="openai/o3-mini",
|
|
125
|
+
llm: str | BaseChatModel = "openai/o3-mini",
|
|
124
126
|
summarize: bool = True,
|
|
125
127
|
process_images=True,
|
|
126
128
|
max_results: int = 3,
|
|
@@ -141,7 +143,7 @@ class ArxivAgent(BaseAgent):
|
|
|
141
143
|
self.download_papers = download_papers
|
|
142
144
|
self.rag_embedding = rag_embedding
|
|
143
145
|
|
|
144
|
-
self.
|
|
146
|
+
self._action = self._build_graph()
|
|
145
147
|
|
|
146
148
|
os.makedirs(self.database_path, exist_ok=True)
|
|
147
149
|
|
|
@@ -259,10 +261,13 @@ class ArxivAgent(BaseAgent):
|
|
|
259
261
|
|
|
260
262
|
try:
|
|
261
263
|
cleaned_text = remove_surrogates(paper["full_text"])
|
|
262
|
-
summary = chain.invoke(
|
|
263
|
-
|
|
264
|
-
|
|
265
|
-
|
|
264
|
+
summary = chain.invoke(
|
|
265
|
+
{
|
|
266
|
+
"retrieved_content": cleaned_text,
|
|
267
|
+
"context": state["context"],
|
|
268
|
+
},
|
|
269
|
+
config=self.build_config(tags=["arxiv", "summarize_each"]),
|
|
270
|
+
)
|
|
266
271
|
|
|
267
272
|
except Exception as e:
|
|
268
273
|
summary = f"Error summarizing paper: {e}"
|
|
@@ -304,7 +309,9 @@ class ArxivAgent(BaseAgent):
|
|
|
304
309
|
embedding=self.rag_embedding,
|
|
305
310
|
database_path=self.database_path,
|
|
306
311
|
)
|
|
307
|
-
new_state["final_summary"] = rag_agent.
|
|
312
|
+
new_state["final_summary"] = rag_agent.invoke(context=state["context"])[
|
|
313
|
+
"summary"
|
|
314
|
+
]
|
|
308
315
|
return new_state
|
|
309
316
|
|
|
310
317
|
def _aggregate_node(self, state: PaperState) -> PaperState:
|
|
@@ -341,10 +348,13 @@ class ArxivAgent(BaseAgent):
|
|
|
341
348
|
|
|
342
349
|
chain = prompt | self.llm | StrOutputParser()
|
|
343
350
|
|
|
344
|
-
final_summary = chain.invoke(
|
|
345
|
-
|
|
346
|
-
|
|
347
|
-
|
|
351
|
+
final_summary = chain.invoke(
|
|
352
|
+
{
|
|
353
|
+
"Summaries": combined,
|
|
354
|
+
"context": state["context"],
|
|
355
|
+
},
|
|
356
|
+
config=self.build_config(tags=["arxiv", "aggregate"]),
|
|
357
|
+
)
|
|
348
358
|
|
|
349
359
|
with open(self.summaries_path + "/final_summary.txt", "w") as f:
|
|
350
360
|
f.write(final_summary)
|
|
@@ -352,49 +362,69 @@ class ArxivAgent(BaseAgent):
|
|
|
352
362
|
return {**state, "final_summary": final_summary}
|
|
353
363
|
|
|
354
364
|
def _build_graph(self):
|
|
355
|
-
|
|
356
|
-
builder.add_node("fetch_papers", self._fetch_node)
|
|
365
|
+
graph = StateGraph(PaperState)
|
|
357
366
|
|
|
367
|
+
self.add_node(graph, self._fetch_node)
|
|
358
368
|
if self.summarize:
|
|
359
369
|
if self.rag_embedding:
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
363
|
-
|
|
364
|
-
builder.set_finish_point("rag_summarize")
|
|
370
|
+
self.add_node(graph, self._rag_node)
|
|
371
|
+
graph.set_entry_point("_fetch_node")
|
|
372
|
+
graph.add_edge("_fetch_node", "_rag_node")
|
|
373
|
+
graph.set_finish_point("_rag_node")
|
|
365
374
|
else:
|
|
366
|
-
|
|
367
|
-
|
|
368
|
-
|
|
369
|
-
builder.set_entry_point("fetch_papers")
|
|
370
|
-
builder.add_edge("fetch_papers", "summarize_each")
|
|
371
|
-
builder.add_edge("summarize_each", "aggregate")
|
|
372
|
-
builder.set_finish_point("aggregate")
|
|
375
|
+
self.add_node(graph, self._summarize_node)
|
|
376
|
+
self.add_node(graph, self._aggregate_node)
|
|
373
377
|
|
|
378
|
+
graph.set_entry_point("_fetch_node")
|
|
379
|
+
graph.add_edge("_fetch_node", "_summarize_node")
|
|
380
|
+
graph.add_edge("_summarize_node", "_aggregate_node")
|
|
381
|
+
graph.set_finish_point("_aggregate_node")
|
|
374
382
|
else:
|
|
375
|
-
|
|
376
|
-
|
|
383
|
+
graph.set_entry_point("_fetch_node")
|
|
384
|
+
graph.set_finish_point("_fetch_node")
|
|
377
385
|
|
|
378
|
-
graph
|
|
379
|
-
return graph
|
|
386
|
+
return graph.compile(checkpointer=self.checkpointer)
|
|
380
387
|
|
|
381
|
-
def
|
|
382
|
-
|
|
383
|
-
|
|
384
|
-
|
|
385
|
-
|
|
388
|
+
def _invoke(
|
|
389
|
+
self,
|
|
390
|
+
inputs: Mapping[str, Any],
|
|
391
|
+
*,
|
|
392
|
+
summarize: bool | None = None,
|
|
393
|
+
recursion_limit: int = 1000,
|
|
394
|
+
**_,
|
|
395
|
+
) -> str:
|
|
396
|
+
config = self.build_config(
|
|
397
|
+
recursion_limit=recursion_limit, tags=["graph"]
|
|
398
|
+
)
|
|
386
399
|
|
|
387
|
-
|
|
388
|
-
|
|
389
|
-
|
|
390
|
-
|
|
400
|
+
# this seems dumb, but it's b/c sometimes we had referred to the value as
|
|
401
|
+
# 'query' other times as 'arxiv_search_query' so trying to keep it compatible
|
|
402
|
+
# aliasing: accept arxiv_search_query -> query
|
|
403
|
+
if "query" not in inputs:
|
|
404
|
+
if "arxiv_search_query" in inputs:
|
|
405
|
+
# make a shallow copy and rename the key
|
|
406
|
+
inputs = dict(inputs)
|
|
407
|
+
inputs["query"] = inputs.pop("arxiv_search_query")
|
|
408
|
+
else:
|
|
409
|
+
raise KeyError(
|
|
410
|
+
"Missing 'query' in inputs (alias 'arxiv_search_query' also accepted)."
|
|
411
|
+
)
|
|
391
412
|
|
|
413
|
+
result = self._action.invoke(inputs, config)
|
|
414
|
+
|
|
415
|
+
use_summary = self.summarize if summarize is None else summarize
|
|
416
|
+
|
|
417
|
+
return (
|
|
418
|
+
result.get("final_summary", "No summary generated.")
|
|
419
|
+
if use_summary
|
|
420
|
+
else "\n\nFinished Fetching papers!"
|
|
421
|
+
)
|
|
392
422
|
|
|
393
|
-
if __name__ == "__main__":
|
|
394
|
-
agent = ArxivAgent()
|
|
395
|
-
result = agent.run(
|
|
396
|
-
arxiv_search_query="Experimental Constraints on neutron star radius",
|
|
397
|
-
context="What are the constraints on the neutron star radius and what uncertainties are there on the constraints?",
|
|
398
|
-
)
|
|
399
423
|
|
|
400
|
-
|
|
424
|
+
# NOTE: Run test in `tests/agents/test_arxiv_agent/test_arxiv_agent.py` via:
|
|
425
|
+
#
|
|
426
|
+
# pytest -s tests/agents/test_arxiv_agent
|
|
427
|
+
#
|
|
428
|
+
# OR
|
|
429
|
+
#
|
|
430
|
+
# uv run pytest -s tests/agents/test_arxiv_agent
|