confoundsearch 0.1.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.
- confoundsearch-0.1.0/LICENSE +21 -0
- confoundsearch-0.1.0/MANIFEST.in +11 -0
- confoundsearch-0.1.0/PKG-INFO +98 -0
- confoundsearch-0.1.0/README.md +82 -0
- confoundsearch-0.1.0/README.pypi.md +63 -0
- confoundsearch-0.1.0/confoundsearch/__init__.py +72 -0
- confoundsearch-0.1.0/confoundsearch/__main__.py +5 -0
- confoundsearch-0.1.0/confoundsearch/data/__init__.py +2 -0
- confoundsearch-0.1.0/confoundsearch/data/corpus.py +36 -0
- confoundsearch-0.1.0/confoundsearch/data/es.py +72 -0
- confoundsearch-0.1.0/confoundsearch/data/sources.py +148 -0
- confoundsearch-0.1.0/confoundsearch/dotenv.py +21 -0
- confoundsearch-0.1.0/confoundsearch/errors.py +17 -0
- confoundsearch-0.1.0/confoundsearch/http.py +26 -0
- confoundsearch-0.1.0/confoundsearch/interface/__init__.py +1 -0
- confoundsearch-0.1.0/confoundsearch/interface/cli.py +211 -0
- confoundsearch-0.1.0/confoundsearch/interface/export.py +108 -0
- confoundsearch-0.1.0/confoundsearch/interface/result.py +63 -0
- confoundsearch-0.1.0/confoundsearch/interface/server.py +269 -0
- confoundsearch-0.1.0/confoundsearch/interface/skill/SKILL.md +85 -0
- confoundsearch-0.1.0/confoundsearch/interface/skill/__init__.py +0 -0
- confoundsearch-0.1.0/confoundsearch/interface/ui/__init__.py +12 -0
- confoundsearch-0.1.0/confoundsearch/interface/ui/app.html +128 -0
- confoundsearch-0.1.0/confoundsearch/interface/ui/render.py +1234 -0
- confoundsearch-0.1.0/confoundsearch/logic/__init__.py +2 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/__init__.py +36 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/add_confound.py +21 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/add_edge.py +24 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/attach_passage.py +48 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/enumerate.py +38 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/merge.py +24 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/retrieve.py +32 -0
- confoundsearch-0.1.0/confoundsearch/logic/actions/stop.py +17 -0
- confoundsearch-0.1.0/confoundsearch/logic/audit.py +435 -0
- confoundsearch-0.1.0/confoundsearch/logic/episode.py +113 -0
- confoundsearch-0.1.0/confoundsearch/logic/graph.py +171 -0
- confoundsearch-0.1.0/confoundsearch/logic/policy.py +111 -0
- confoundsearch-0.1.0/confoundsearch/logic/reward.py +28 -0
- confoundsearch-0.1.0/confoundsearch/logic/state.py +40 -0
- confoundsearch-0.1.0/confoundsearch/logic/transition.py +55 -0
- confoundsearch-0.1.0/confoundsearch/model/__init__.py +2 -0
- confoundsearch-0.1.0/confoundsearch/model/llm.py +210 -0
- confoundsearch-0.1.0/confoundsearch/py.typed +0 -0
- confoundsearch-0.1.0/confoundsearch.egg-info/PKG-INFO +98 -0
- confoundsearch-0.1.0/confoundsearch.egg-info/SOURCES.txt +49 -0
- confoundsearch-0.1.0/confoundsearch.egg-info/dependency_links.txt +1 -0
- confoundsearch-0.1.0/confoundsearch.egg-info/entry_points.txt +2 -0
- confoundsearch-0.1.0/confoundsearch.egg-info/requires.txt +19 -0
- confoundsearch-0.1.0/confoundsearch.egg-info/top_level.txt +1 -0
- confoundsearch-0.1.0/pyproject.toml +69 -0
- confoundsearch-0.1.0/setup.cfg +4 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 The ConfoundSearch authors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
# The source distribution carries the package, its description and its licence, not the repository around it.
|
|
2
|
+
include README.pypi.md
|
|
3
|
+
include LICENSE*
|
|
4
|
+
prune tests
|
|
5
|
+
prune docs
|
|
6
|
+
prune examples
|
|
7
|
+
prune elastic
|
|
8
|
+
prune tools
|
|
9
|
+
prune badges
|
|
10
|
+
prune .github
|
|
11
|
+
global-exclude *.py[cod] __pycache__ .DS_Store .env .env.*
|
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: confoundsearch
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: ConfoundSearch: finding confounds as a Markov Decision Process
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Keywords: causal inference,confounding,markov decision process,llm,openai
|
|
7
|
+
Classifier: Programming Language :: Python :: 3
|
|
8
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
9
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
12
|
+
Classifier: Development Status :: 3 - Alpha
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: Natural Language :: English
|
|
15
|
+
Classifier: Operating System :: OS Independent
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering :: Information Analysis
|
|
17
|
+
Classifier: Typing :: Typed
|
|
18
|
+
Requires-Python: >=3.10
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
License-File: LICENSE
|
|
21
|
+
Provides-Extra: dev
|
|
22
|
+
Requires-Dist: pytest>=8; extra == "dev"
|
|
23
|
+
Requires-Dist: pytest-cov>=5; extra == "dev"
|
|
24
|
+
Requires-Dist: hypothesis>=6; extra == "dev"
|
|
25
|
+
Requires-Dist: ruff>=0.5; extra == "dev"
|
|
26
|
+
Requires-Dist: mypy>=1.10; extra == "dev"
|
|
27
|
+
Requires-Dist: pyyaml>=6; extra == "dev"
|
|
28
|
+
Requires-Dist: rust-just>=1.17; extra == "dev"
|
|
29
|
+
Requires-Dist: tomli>=2; python_version < "3.11" and extra == "dev"
|
|
30
|
+
Provides-Extra: release
|
|
31
|
+
Requires-Dist: build>=1; extra == "release"
|
|
32
|
+
Requires-Dist: twine>=5; extra == "release"
|
|
33
|
+
Requires-Dist: tomli>=2; python_version < "3.11" and extra == "release"
|
|
34
|
+
Dynamic: license-file
|
|
35
|
+
|
|
36
|
+
# ConfoundSearch
|
|
37
|
+
|
|
38
|
+
Find candidate **confounds** for a study of whether X affects Y, together with the published passages that
|
|
39
|
+
support each one and a verifiable audit trail. An AI agent searches a collection of papers and builds a
|
|
40
|
+
diagram of confounds one small step at a time: it searches, adds a confound, links it to your two variables,
|
|
41
|
+
and cites passages that support the link. It suggests; the researcher decides.
|
|
42
|
+
|
|
43
|
+
Supporting code for our *Organization Science* proposal, *Using AI to find confounds*.
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
pip install confoundsearch # Python 3.10 or newer; no dependencies
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
## From Python
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from confoundsearch import *
|
|
53
|
+
|
|
54
|
+
corpus = InMemoryCorpus(
|
|
55
|
+
[
|
|
56
|
+
Passage("p1", "Larger firms invest more in information technology.", title="Firm size and IT"),
|
|
57
|
+
Passage(
|
|
58
|
+
"p2",
|
|
59
|
+
"Firm size is associated with firm performance and profitability.",
|
|
60
|
+
title="Firm size and profits",
|
|
61
|
+
),
|
|
62
|
+
]
|
|
63
|
+
)
|
|
64
|
+
client = AuditedClient(OpenAICompatibleClient("gpt-4.1")) # reads OPENAI_API_KEY; logs every model call
|
|
65
|
+
env = Environment(corpus, LLMJudge(client))
|
|
66
|
+
policy = EpsilonGreedy(LookaheadQ(env.preview), epsilon=0.1, explore_stop=False)
|
|
67
|
+
run = run_episode(
|
|
68
|
+
env, policy, Proposer(LLMProposer(client)), "IT investment", "firm performance", max_steps=20
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
print(run.final.graph.labels, run.score) # the confounds found, and the score
|
|
72
|
+
print(run.audit_trail().to_markdown()) # the evidence behind each, and where it came from
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
## From the command line
|
|
76
|
+
|
|
77
|
+
```bash
|
|
78
|
+
confoundsearch serve --model gpt-4.1 --corpus-file passages.jsonl # the web interface, on your computer
|
|
79
|
+
confoundsearch run --json --model gpt-4.1 --corpus-file passages.jsonl --iv "IT investment" --dv "firm performance"
|
|
80
|
+
confoundsearch audit confoundsearch.audit.jsonl # check and read a run's audit trail
|
|
81
|
+
confoundsearch install-skill # add the Claude Code skill
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
## What you get
|
|
85
|
+
|
|
86
|
+
- **Any papers.** Search our hosted Elasticsearch index (set `ES_URL` and `ES_API_KEY`), Semantic Scholar
|
|
87
|
+
(`--corpus semanticscholar`; a [free API key](https://www.semanticscholar.org/product/api) in `S2_KEY` makes it
|
|
88
|
+
much faster), OpenAlex (`--corpus openalex`), or a file of your own passages (`--corpus-file`).
|
|
89
|
+
- **Any model.** OpenAI, [OpenRouter](https://openrouter.ai) (`--openrouter`, with your own `OPENROUTER_API_KEY`), or
|
|
90
|
+
any OpenAI-compatible server, including one on your own computer.
|
|
91
|
+
- **Auditable.** Every run keeps a hash-chained log of each search, each passage retrieved, the judge's verdict on
|
|
92
|
+
each citation, and each model prompt and reply, so any edit to it is detectable.
|
|
93
|
+
- **Four modular layers**, one directory each: `data`, `logic`, `model` and `interface`.
|
|
94
|
+
- **Agent-ready.** `run --json` returns the result as JSON, and a Claude Code skill is included.
|
|
95
|
+
|
|
96
|
+
Each run calls the AI model, which costs money and sends your text to its provider.
|
|
97
|
+
|
|
98
|
+
Released under the MIT licence.
|
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
# ConfoundSearch
|
|
2
|
+
|
|
3
|
+
Supporting code for our *Organization Science* proposal, *Using AI to find confounds*.
|
|
4
|
+
|
|
5
|
+

|
|
6
|
+

|
|
7
|
+
-brightgreen)
|
|
8
|
+

|
|
9
|
+
[](https://github.com/astral-sh/ruff)
|
|
10
|
+

|
|
11
|
+

|
|
12
|
+
|
|
13
|
+

|
|
14
|
+
|
|
15
|
+
*Example output on made-up demo passages, with a scripted stand-in for the AI model (no key needed): open
|
|
16
|
+
[`examples/example_output.html`](examples/example_output.html). **Reviewing this artifact?** See the [reviewer guide](docs/artifact-evaluation.md).*
|
|
17
|
+
|
|
18
|
+
## How do I run this?
|
|
19
|
+
|
|
20
|
+
You need Python 3.10+ and an API key for OpenAI, [OpenRouter](https://openrouter.ai) (`--openrouter`) or a compatible service. Run `pip install .` in the unzipped
|
|
21
|
+
repository, then use it in any of four ways. We plan to push updates to pip as we keep developing ConfoundSearch.
|
|
22
|
+
|
|
23
|
+
**1. From Python.**
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
from confoundsearch import *
|
|
27
|
+
|
|
28
|
+
client = AuditedClient(OpenAICompatibleClient("gpt-4.1")) # reads OPENAI_API_KEY; logs every model call
|
|
29
|
+
env = Environment(InMemoryCorpus.from_jsonl("examples/passages.jsonl"), LLMJudge(client))
|
|
30
|
+
policy = EpsilonGreedy(LookaheadQ(env.preview), epsilon=0.1, explore_stop=False)
|
|
31
|
+
run = run_episode(
|
|
32
|
+
env, policy, Proposer(LLMProposer(client)), "IT investment", "firm performance", max_steps=40
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
print(run.final.graph.labels, run.score) # the confounds found, and the score
|
|
36
|
+
print(run.audit_trail().to_markdown()) # the evidence behind each, and where it came from
|
|
37
|
+
save_export(run.final, "run.zip", audit=run.audit) # export for the interface, with its audit trail
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
**2. In the web interface, on your computer.**
|
|
41
|
+
`confoundsearch serve --model gpt-4.1 --api-key YOUR_KEY --corpus-file examples/passages.jsonl`,
|
|
42
|
+
then open <http://127.0.0.1:7860>, type two variables, and click **Find confounds**.
|
|
43
|
+
|
|
44
|
+
**3. In the hosted web app,** which serves the same interface with nothing to install (link withheld during blind review).
|
|
45
|
+
|
|
46
|
+
**4. From Claude Code or another agent.** `confoundsearch install-skill` adds a skill, and `confoundsearch run --json`
|
|
47
|
+
returns the results and the audit-trail path as JSON. More: [docs/claude-code.md](docs/claude-code.md).
|
|
48
|
+
|
|
49
|
+
**Papers.** The commands above use the demo file. By default the agent searches our **hosted Elasticsearch index**
|
|
50
|
+
(put `ES_URL` and `ES_API_KEY` in a `.env` file). Or add `--corpus semanticscholar` (a [free API key](https://www.semanticscholar.org/product/api),
|
|
51
|
+
set as `S2_KEY`, makes it much faster), `--corpus openalex`, or `--corpus-file your.jsonl`.
|
|
52
|
+
|
|
53
|
+
Each run calls the AI model, which costs money and sends your text to its provider. Details: [docs/running.md](docs/running.md).
|
|
54
|
+
|
|
55
|
+
## What's going on under the hood?
|
|
56
|
+
|
|
57
|
+
An AI agent builds a diagram of candidate confounds one small step at a time: it searches the papers,
|
|
58
|
+
adds a confound, links it to your two variables, and cites passages that support the link. A score
|
|
59
|
+
rewards confounds that are linked to both variables and well supported. It suggests; the researcher decides.
|
|
60
|
+
|
|
61
|
+
**Auditable.** ConfoundSearch keeps a log of the evidence it collects and processes during search: every query,
|
|
62
|
+
every passage retrieved, the judge's verdict on every citation, and every model prompt and reply, in order and
|
|
63
|
+
hash-chained so any edit is detectable. Read it with `run.audit_trail()` or `confoundsearch audit FILE`, or open
|
|
64
|
+
[`examples/example_audit.jsonl`](examples/example_audit.jsonl). More: [docs/auditability.md](docs/auditability.md).
|
|
65
|
+
|
|
66
|
+
**Decoupled.** The package never draws anything: it exports a plain data file (`save_export`), and the interface only
|
|
67
|
+
reads such files, never running the agent. So one output can be shown in the local interface, in the hosted app,
|
|
68
|
+
or as a single HTML page (`confoundsearch run --out page.html`). More: [docs/how-it-works.md](docs/how-it-works.md).
|
|
69
|
+
|
|
70
|
+
## What's in this code base?
|
|
71
|
+
|
|
72
|
+
The application is divided into four layers, as described in the technical appendix. Each is one directory:
|
|
73
|
+
|
|
74
|
+
- **Data** ([`confoundsearch/data/`](confoundsearch/data)): where papers come from (hosted Elasticsearch, Semantic Scholar, OpenAlex, or your file).
|
|
75
|
+
- **Logic** ([`confoundsearch/logic/`](confoundsearch/logic)): the Markov decision process (graph, moves, reward, policy) and the audit log.
|
|
76
|
+
- **Language model** ([`confoundsearch/model/`](confoundsearch/model)): any OpenAI-compatible model, including OpenRouter.
|
|
77
|
+
- **Interface** ([`confoundsearch/interface/`](confoundsearch/interface)): the web app, command line, JSON output and agent skill.
|
|
78
|
+
|
|
79
|
+
Also here: `examples/`, `tests/` (100% branch coverage, and they enforce the layer boundaries), and `elastic/`
|
|
80
|
+
(loads the corpus into Elasticsearch Serverless: [how](elastic/README.md)).
|
|
81
|
+
|
|
82
|
+
More: [paper-to-code map and how we know it works](docs/code-map.md) · [for programmers](docs/for-programmers.md)
|
|
@@ -0,0 +1,63 @@
|
|
|
1
|
+
# ConfoundSearch
|
|
2
|
+
|
|
3
|
+
Find candidate **confounds** for a study of whether X affects Y, together with the published passages that
|
|
4
|
+
support each one and a verifiable audit trail. An AI agent searches a collection of papers and builds a
|
|
5
|
+
diagram of confounds one small step at a time: it searches, adds a confound, links it to your two variables,
|
|
6
|
+
and cites passages that support the link. It suggests; the researcher decides.
|
|
7
|
+
|
|
8
|
+
Supporting code for our *Organization Science* proposal, *Using AI to find confounds*.
|
|
9
|
+
|
|
10
|
+
```bash
|
|
11
|
+
pip install confoundsearch # Python 3.10 or newer; no dependencies
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
## From Python
|
|
15
|
+
|
|
16
|
+
```python
|
|
17
|
+
from confoundsearch import *
|
|
18
|
+
|
|
19
|
+
corpus = InMemoryCorpus(
|
|
20
|
+
[
|
|
21
|
+
Passage("p1", "Larger firms invest more in information technology.", title="Firm size and IT"),
|
|
22
|
+
Passage(
|
|
23
|
+
"p2",
|
|
24
|
+
"Firm size is associated with firm performance and profitability.",
|
|
25
|
+
title="Firm size and profits",
|
|
26
|
+
),
|
|
27
|
+
]
|
|
28
|
+
)
|
|
29
|
+
client = AuditedClient(OpenAICompatibleClient("gpt-4.1")) # reads OPENAI_API_KEY; logs every model call
|
|
30
|
+
env = Environment(corpus, LLMJudge(client))
|
|
31
|
+
policy = EpsilonGreedy(LookaheadQ(env.preview), epsilon=0.1, explore_stop=False)
|
|
32
|
+
run = run_episode(
|
|
33
|
+
env, policy, Proposer(LLMProposer(client)), "IT investment", "firm performance", max_steps=20
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
print(run.final.graph.labels, run.score) # the confounds found, and the score
|
|
37
|
+
print(run.audit_trail().to_markdown()) # the evidence behind each, and where it came from
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
## From the command line
|
|
41
|
+
|
|
42
|
+
```bash
|
|
43
|
+
confoundsearch serve --model gpt-4.1 --corpus-file passages.jsonl # the web interface, on your computer
|
|
44
|
+
confoundsearch run --json --model gpt-4.1 --corpus-file passages.jsonl --iv "IT investment" --dv "firm performance"
|
|
45
|
+
confoundsearch audit confoundsearch.audit.jsonl # check and read a run's audit trail
|
|
46
|
+
confoundsearch install-skill # add the Claude Code skill
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
## What you get
|
|
50
|
+
|
|
51
|
+
- **Any papers.** Search our hosted Elasticsearch index (set `ES_URL` and `ES_API_KEY`), Semantic Scholar
|
|
52
|
+
(`--corpus semanticscholar`; a [free API key](https://www.semanticscholar.org/product/api) in `S2_KEY` makes it
|
|
53
|
+
much faster), OpenAlex (`--corpus openalex`), or a file of your own passages (`--corpus-file`).
|
|
54
|
+
- **Any model.** OpenAI, [OpenRouter](https://openrouter.ai) (`--openrouter`, with your own `OPENROUTER_API_KEY`), or
|
|
55
|
+
any OpenAI-compatible server, including one on your own computer.
|
|
56
|
+
- **Auditable.** Every run keeps a hash-chained log of each search, each passage retrieved, the judge's verdict on
|
|
57
|
+
each citation, and each model prompt and reply, so any edit to it is detectable.
|
|
58
|
+
- **Four modular layers**, one directory each: `data`, `logic`, `model` and `interface`.
|
|
59
|
+
- **Agent-ready.** `run --json` returns the result as JSON, and a Claude Code skill is included.
|
|
60
|
+
|
|
61
|
+
Each run calls the AI model, which costs money and sends your text to its provider.
|
|
62
|
+
|
|
63
|
+
Released under the MIT licence.
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
"""ConfoundSearch: an agent that builds a confound graph as a Markov Decision Process."""
|
|
2
|
+
|
|
3
|
+
from confoundsearch.data.corpus import Corpus, InMemoryCorpus
|
|
4
|
+
from confoundsearch.data.es import ElasticsearchCorpus
|
|
5
|
+
from confoundsearch.data.sources import OpenAlexCorpus, SemanticScholarCorpus
|
|
6
|
+
from confoundsearch.errors import CorpusError, InvalidAction, LLMError, RemoteError
|
|
7
|
+
from confoundsearch.interface.export import export_zip, save_export, to_share_data
|
|
8
|
+
from confoundsearch.logic.actions import (
|
|
9
|
+
Action,
|
|
10
|
+
AddConfound,
|
|
11
|
+
AddEdge,
|
|
12
|
+
AttachPassage,
|
|
13
|
+
Merge,
|
|
14
|
+
Retrieve,
|
|
15
|
+
Stop,
|
|
16
|
+
enumerate_actions,
|
|
17
|
+
)
|
|
18
|
+
from confoundsearch.logic.audit import AuditedClient, AuditLog, AuditTrail, verify
|
|
19
|
+
from confoundsearch.logic.episode import Step, Trajectory, run_episode
|
|
20
|
+
from confoundsearch.logic.graph import Graph, Passage, Vertex
|
|
21
|
+
from confoundsearch.logic.policy import EpsilonGreedy, LookaheadQ, Softmax
|
|
22
|
+
from confoundsearch.logic.reward import phi, reward
|
|
23
|
+
from confoundsearch.logic.state import State
|
|
24
|
+
from confoundsearch.logic.transition import Environment, Judge
|
|
25
|
+
from confoundsearch.model.llm import LLMJudge, LLMProposer, OpenAICompatibleClient, Proposer
|
|
26
|
+
|
|
27
|
+
__version__ = "0.1.0"
|
|
28
|
+
|
|
29
|
+
__all__ = [
|
|
30
|
+
"Action",
|
|
31
|
+
"AddConfound",
|
|
32
|
+
"AddEdge",
|
|
33
|
+
"AttachPassage",
|
|
34
|
+
"AuditLog",
|
|
35
|
+
"AuditTrail",
|
|
36
|
+
"AuditedClient",
|
|
37
|
+
"Corpus",
|
|
38
|
+
"CorpusError",
|
|
39
|
+
"ElasticsearchCorpus",
|
|
40
|
+
"Environment",
|
|
41
|
+
"EpsilonGreedy",
|
|
42
|
+
"Graph",
|
|
43
|
+
"InMemoryCorpus",
|
|
44
|
+
"InvalidAction",
|
|
45
|
+
"Judge",
|
|
46
|
+
"LLMError",
|
|
47
|
+
"LLMJudge",
|
|
48
|
+
"LLMProposer",
|
|
49
|
+
"LookaheadQ",
|
|
50
|
+
"Merge",
|
|
51
|
+
"OpenAICompatibleClient",
|
|
52
|
+
"OpenAlexCorpus",
|
|
53
|
+
"Passage",
|
|
54
|
+
"Proposer",
|
|
55
|
+
"RemoteError",
|
|
56
|
+
"Retrieve",
|
|
57
|
+
"SemanticScholarCorpus",
|
|
58
|
+
"Softmax",
|
|
59
|
+
"State",
|
|
60
|
+
"Step",
|
|
61
|
+
"Stop",
|
|
62
|
+
"Trajectory",
|
|
63
|
+
"Vertex",
|
|
64
|
+
"enumerate_actions",
|
|
65
|
+
"export_zip",
|
|
66
|
+
"phi",
|
|
67
|
+
"reward",
|
|
68
|
+
"run_episode",
|
|
69
|
+
"save_export",
|
|
70
|
+
"to_share_data",
|
|
71
|
+
"verify",
|
|
72
|
+
]
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
"""The corpus D, behind a one-method protocol."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
from collections.abc import Iterable
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from typing import Protocol
|
|
9
|
+
|
|
10
|
+
from confoundsearch.logic.graph import Passage, normalize
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class Corpus(Protocol):
|
|
14
|
+
def search(self, query: str, k: int) -> list[Passage]:
|
|
15
|
+
"""Return up to k passages relevant to the query."""
|
|
16
|
+
...
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class InMemoryCorpus:
|
|
20
|
+
"""Ranks passages by how many query words they contain. Deterministic."""
|
|
21
|
+
|
|
22
|
+
def __init__(self, passages: Iterable[Passage]) -> None:
|
|
23
|
+
self.passages = tuple(passages)
|
|
24
|
+
|
|
25
|
+
@classmethod
|
|
26
|
+
def from_jsonl(cls, path: str | Path) -> InMemoryCorpus:
|
|
27
|
+
"""One JSON object per line with `id` and `text`, and optionally `title` and `url`."""
|
|
28
|
+
rows = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
|
|
29
|
+
return cls(Passage(str(r["id"]), r["text"], r.get("title", ""), r.get("url", "")) for r in rows)
|
|
30
|
+
|
|
31
|
+
def search(self, query: str, k: int) -> list[Passage]:
|
|
32
|
+
words = set(normalize(query).split())
|
|
33
|
+
scored = [(len(words & set(normalize(p.text).split())), p) for p in self.passages]
|
|
34
|
+
hits = [(n, p) for n, p in scored if n > 0]
|
|
35
|
+
hits.sort(key=lambda hit: (-hit[0], hit[1].id))
|
|
36
|
+
return [p for _, p in hits[: max(k, 0)]]
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
"""An Elasticsearch-backed corpus, over the plain REST API.
|
|
2
|
+
|
|
3
|
+
Works with a local node or Elastic Cloud. Documents need a text field (default
|
|
4
|
+
`abstract`) and may have `title`, `url` and `corpusid`, as in the Semantic
|
|
5
|
+
Scholar index the paper's experiments use.
|
|
6
|
+
|
|
7
|
+
ElasticsearchCorpus("https://my-deployment.es.us-east-1.aws.elastic.cloud", api_key="...")
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
from typing import Any
|
|
14
|
+
|
|
15
|
+
from confoundsearch.errors import CorpusError, RemoteError
|
|
16
|
+
from confoundsearch.http import Transport, urllib_transport
|
|
17
|
+
from confoundsearch.logic.graph import Passage
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class ElasticsearchCorpus:
|
|
21
|
+
def __init__(
|
|
22
|
+
self,
|
|
23
|
+
url: str,
|
|
24
|
+
index: str = "papers",
|
|
25
|
+
api_key: str | None = None,
|
|
26
|
+
text_field: str = "abstract",
|
|
27
|
+
title_field: str = "title",
|
|
28
|
+
min_citations: int = 0,
|
|
29
|
+
timeout: float = 30.0,
|
|
30
|
+
transport: Transport = urllib_transport,
|
|
31
|
+
) -> None:
|
|
32
|
+
self.url = url.rstrip("/")
|
|
33
|
+
self.index = index
|
|
34
|
+
self.api_key = api_key
|
|
35
|
+
self.text_field = text_field
|
|
36
|
+
self.title_field = title_field
|
|
37
|
+
self.min_citations = min_citations
|
|
38
|
+
self.timeout = timeout
|
|
39
|
+
self.transport = transport
|
|
40
|
+
|
|
41
|
+
def _query(self, query: str, k: int) -> dict[str, Any]:
|
|
42
|
+
match = {"multi_match": {"query": query, "fields": [f"{self.title_field}^2", self.text_field]}}
|
|
43
|
+
bool_query: dict[str, Any] = {"must": match, "filter": [{"exists": {"field": self.text_field}}]}
|
|
44
|
+
if self.min_citations:
|
|
45
|
+
bool_query["filter"].append({"range": {"citationcount": {"gte": self.min_citations}}})
|
|
46
|
+
return {"query": {"bool": bool_query}, "size": k}
|
|
47
|
+
|
|
48
|
+
def search(self, query: str, k: int) -> list[Passage]:
|
|
49
|
+
if k <= 0:
|
|
50
|
+
return []
|
|
51
|
+
headers = {"Content-Type": "application/json"}
|
|
52
|
+
if self.api_key:
|
|
53
|
+
headers["Authorization"] = f"ApiKey {self.api_key}"
|
|
54
|
+
body = json.dumps(self._query(query, k)).encode()
|
|
55
|
+
try:
|
|
56
|
+
raw = self.transport(f"{self.url}/{self.index}/_search", headers, body, self.timeout)
|
|
57
|
+
except RemoteError as err:
|
|
58
|
+
raise CorpusError(str(err)) from err
|
|
59
|
+
try:
|
|
60
|
+
hits = json.loads(raw)["hits"]["hits"]
|
|
61
|
+
return [self._passage(hit) for hit in hits]
|
|
62
|
+
except (ValueError, KeyError, TypeError) as err:
|
|
63
|
+
raise CorpusError(f"unexpected response: {raw[:200]!r}") from err
|
|
64
|
+
|
|
65
|
+
def _passage(self, hit: dict[str, Any]) -> Passage:
|
|
66
|
+
source = hit["_source"]
|
|
67
|
+
return Passage(
|
|
68
|
+
id=str(source.get("corpusid", hit["_id"])),
|
|
69
|
+
text=source[self.text_field],
|
|
70
|
+
title=source.get(self.title_field) or "",
|
|
71
|
+
url=source.get("url") or "",
|
|
72
|
+
)
|
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
"""Paper sources that need no corpus of your own: Semantic Scholar's API and OpenAlex's.
|
|
2
|
+
|
|
3
|
+
Both implement the same one-method `Corpus` protocol as the file and Elasticsearch corpora, so the agent
|
|
4
|
+
does not know or care which it is searching. Only papers that have an abstract are returned, because an
|
|
5
|
+
abstract is the passage the agent reads and cites.
|
|
6
|
+
|
|
7
|
+
SemanticScholarCorpus(api_key="...") # a free key raises the rate limit; request one at S2_KEY_URL
|
|
8
|
+
OpenAlexCorpus(mailto="you@example.org") # an email (or an api_key) puts you in OpenAlex's polite pool
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
import json
|
|
14
|
+
from collections.abc import Mapping
|
|
15
|
+
from typing import Any
|
|
16
|
+
from urllib.parse import urlencode
|
|
17
|
+
|
|
18
|
+
from confoundsearch.errors import CorpusError, RemoteError
|
|
19
|
+
from confoundsearch.http import Transport, urllib_transport
|
|
20
|
+
from confoundsearch.logic.graph import Passage
|
|
21
|
+
|
|
22
|
+
MAX_PAGE = 100
|
|
23
|
+
S2_KEY_URL = "https://www.semanticscholar.org/product/api" # where to request a free Semantic Scholar API key
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _get(transport: Transport, url: str, headers: Mapping[str, str], timeout: float, hint: str) -> Any:
|
|
27
|
+
try:
|
|
28
|
+
raw = transport(url, headers, None, timeout)
|
|
29
|
+
except RemoteError as err:
|
|
30
|
+
extra = f" {hint}" if "HTTP 429" in str(err) else ""
|
|
31
|
+
raise CorpusError(f"{err}{extra}") from err
|
|
32
|
+
try:
|
|
33
|
+
return json.loads(raw)
|
|
34
|
+
except ValueError as err:
|
|
35
|
+
raise CorpusError(f"unexpected response: {raw[:200]!r}") from err
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _page_size(k: int) -> int:
|
|
39
|
+
"""Ask for extra results, since some papers have no abstract and are dropped."""
|
|
40
|
+
return min(max(2 * k, k), MAX_PAGE)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class SemanticScholarCorpus:
|
|
44
|
+
def __init__(
|
|
45
|
+
self,
|
|
46
|
+
api_key: str | None = None,
|
|
47
|
+
base_url: str = "https://api.semanticscholar.org/graph/v1",
|
|
48
|
+
timeout: float = 30.0,
|
|
49
|
+
transport: Transport = urllib_transport,
|
|
50
|
+
) -> None:
|
|
51
|
+
self.api_key = api_key
|
|
52
|
+
self.base_url = base_url.rstrip("/")
|
|
53
|
+
self.timeout = timeout
|
|
54
|
+
self.transport = transport
|
|
55
|
+
|
|
56
|
+
def search(self, query: str, k: int) -> list[Passage]:
|
|
57
|
+
if k <= 0 or not query.strip():
|
|
58
|
+
return []
|
|
59
|
+
params = urlencode(
|
|
60
|
+
{
|
|
61
|
+
"query": query,
|
|
62
|
+
"limit": _page_size(k),
|
|
63
|
+
"fields": "title,abstract,url,year,citationCount,corpusId",
|
|
64
|
+
}
|
|
65
|
+
)
|
|
66
|
+
headers = {"Accept": "application/json"}
|
|
67
|
+
if self.api_key:
|
|
68
|
+
headers["x-api-key"] = self.api_key
|
|
69
|
+
hint = (
|
|
70
|
+
f"Semantic Scholar is rate limiting you. Request a free API key at {S2_KEY_URL} "
|
|
71
|
+
"and set it as S2_KEY (or --s2-key); it raises the limit."
|
|
72
|
+
)
|
|
73
|
+
data = _get(self.transport, f"{self.base_url}/paper/search?{params}", headers, self.timeout, hint)
|
|
74
|
+
try:
|
|
75
|
+
papers = data.get("data") or [] # an empty result has no "data" at all
|
|
76
|
+
passages = [self._passage(p) for p in papers if p.get("abstract")]
|
|
77
|
+
except (AttributeError, KeyError, TypeError) as err:
|
|
78
|
+
raise CorpusError(f"unexpected response: {str(data)[:200]!r}") from err
|
|
79
|
+
return passages[:k]
|
|
80
|
+
|
|
81
|
+
@staticmethod
|
|
82
|
+
def _passage(paper: dict[str, Any]) -> Passage:
|
|
83
|
+
return Passage(
|
|
84
|
+
id=str(paper.get("corpusId") or paper["paperId"]),
|
|
85
|
+
text=paper["abstract"],
|
|
86
|
+
title=paper.get("title") or "",
|
|
87
|
+
url=paper.get("url") or "",
|
|
88
|
+
)
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def abstract_from_inverted_index(inverted: Mapping[str, list[int]]) -> str:
|
|
92
|
+
"""OpenAlex stores an abstract as {word: [positions]}. Put the words back in order."""
|
|
93
|
+
placed = sorted((position, word) for word, positions in inverted.items() for position in positions)
|
|
94
|
+
return " ".join(word for _, word in placed)
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
class OpenAlexCorpus:
|
|
98
|
+
def __init__(
|
|
99
|
+
self,
|
|
100
|
+
api_key: str | None = None,
|
|
101
|
+
mailto: str | None = None,
|
|
102
|
+
base_url: str = "https://api.openalex.org",
|
|
103
|
+
timeout: float = 30.0,
|
|
104
|
+
transport: Transport = urllib_transport,
|
|
105
|
+
) -> None:
|
|
106
|
+
self.api_key = api_key
|
|
107
|
+
self.mailto = mailto
|
|
108
|
+
self.base_url = base_url.rstrip("/")
|
|
109
|
+
self.timeout = timeout
|
|
110
|
+
self.transport = transport
|
|
111
|
+
|
|
112
|
+
def search(self, query: str, k: int) -> list[Passage]:
|
|
113
|
+
if k <= 0 or not query.strip():
|
|
114
|
+
return []
|
|
115
|
+
params = {
|
|
116
|
+
"search": query,
|
|
117
|
+
"per_page": _page_size(k),
|
|
118
|
+
"select": "id,title,abstract_inverted_index,doi,publication_year",
|
|
119
|
+
}
|
|
120
|
+
if self.api_key:
|
|
121
|
+
params["api_key"] = self.api_key
|
|
122
|
+
if self.mailto:
|
|
123
|
+
params["mailto"] = self.mailto
|
|
124
|
+
hint = (
|
|
125
|
+
"OpenAlex is rate limiting you; add an email (OPENALEX_EMAIL) or an API key (OPENALEX_API_KEY)."
|
|
126
|
+
)
|
|
127
|
+
data = _get(
|
|
128
|
+
self.transport,
|
|
129
|
+
f"{self.base_url}/works?{urlencode(params)}",
|
|
130
|
+
{"Accept": "application/json"},
|
|
131
|
+
self.timeout,
|
|
132
|
+
hint,
|
|
133
|
+
)
|
|
134
|
+
try:
|
|
135
|
+
works = data["results"]
|
|
136
|
+
passages = [self._passage(w) for w in works if w.get("abstract_inverted_index")]
|
|
137
|
+
except (KeyError, TypeError, AttributeError) as err:
|
|
138
|
+
raise CorpusError(f"unexpected response: {str(data)[:200]!r}") from err
|
|
139
|
+
return passages[:k]
|
|
140
|
+
|
|
141
|
+
@staticmethod
|
|
142
|
+
def _passage(work: dict[str, Any]) -> Passage:
|
|
143
|
+
return Passage(
|
|
144
|
+
id=work["id"].rsplit("/", 1)[-1],
|
|
145
|
+
text=abstract_from_inverted_index(work["abstract_inverted_index"]),
|
|
146
|
+
title=work.get("title") or "",
|
|
147
|
+
url=work.get("doi") or work["id"],
|
|
148
|
+
)
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
"""Read settings from a `.env` file, so API keys need not be typed into a terminal."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def load_dotenv(path: Path) -> None:
|
|
10
|
+
"""Put KEY=VALUE lines from `path` into the environment. Never overrides a variable already set."""
|
|
11
|
+
if not path.is_file():
|
|
12
|
+
return
|
|
13
|
+
for raw in path.read_text(encoding="utf-8").splitlines():
|
|
14
|
+
line = raw.strip().removeprefix("export ")
|
|
15
|
+
if not line or line.startswith("#") or "=" not in line:
|
|
16
|
+
continue
|
|
17
|
+
key, _, value = line.partition("=")
|
|
18
|
+
value = value.strip()
|
|
19
|
+
if len(value) >= 2 and value[0] == value[-1] and value[0] in "\"'":
|
|
20
|
+
value = value[1:-1]
|
|
21
|
+
os.environ.setdefault(key.strip(), value)
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""Exceptions shared across the package."""
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class InvalidAction(ValueError):
|
|
5
|
+
"""An edit or action violates a graph invariant or the MDP's rules."""
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class RemoteError(RuntimeError):
|
|
9
|
+
"""A remote HTTP service failed or returned something unusable."""
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class LLMError(RemoteError):
|
|
13
|
+
"""An OpenAI-compatible endpoint failed or returned something unusable."""
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class CorpusError(RemoteError):
|
|
17
|
+
"""The search backend failed or returned something unusable."""
|