@getmegabrain/cli 0.1.8 → 0.1.10

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.
@@ -19,8 +19,10 @@ export function pythonTool(bridge) {
19
19
  const SYSTEM_PROMPT = `You are the MegaBrain Science research agent. You work in a live
20
20
  notebook, sharing one Python kernel with the user — variables and imports persist across
21
21
  cells and everything runs locally on the user's machine. Use run_python to write and run
22
- code and iterate on its output. Prefer small, verifiable steps. When you have the answer,
23
- stop calling tools and give a concise final summary of what you did and found.`;
22
+ code and iterate on its output, and search_literature to find and cite real papers (arXiv,
23
+ OpenAlex, PubMed) when the task calls for surveying prior work. Prefer small, verifiable
24
+ steps. When you have the answer, stop calling tools and give a concise final summary of
25
+ what you did and found.`;
24
26
  const MAX_TURNS = 16;
25
27
  export class AgentSession {
26
28
  model;
@@ -0,0 +1,116 @@
1
+ function firstMatch(block, re) {
2
+ const m = re.exec(block);
3
+ return m ? m[1].replace(/\s+/g, ' ').trim() : '';
4
+ }
5
+ /** Parse an arXiv Atom feed into papers (tolerant, entry-by-entry). */
6
+ export function parseArxivAtom(xml) {
7
+ const papers = [];
8
+ const entryRe = /<entry>([\s\S]*?)<\/entry>/g;
9
+ let entry = entryRe.exec(xml);
10
+ while (entry !== null) {
11
+ const block = entry[1];
12
+ const authors = [];
13
+ const nameRe = /<name>([\s\S]*?)<\/name>/g;
14
+ let name = nameRe.exec(block);
15
+ while (name !== null) {
16
+ authors.push(name[1].trim());
17
+ name = nameRe.exec(block);
18
+ }
19
+ papers.push({
20
+ title: firstMatch(block, /<title>([\s\S]*?)<\/title>/),
21
+ id: firstMatch(block, /<id>([\s\S]*?)<\/id>/),
22
+ published: firstMatch(block, /<published>([\s\S]*?)<\/published>/),
23
+ summary: firstMatch(block, /<summary>([\s\S]*?)<\/summary>/),
24
+ authors,
25
+ });
26
+ entry = entryRe.exec(xml);
27
+ }
28
+ return papers;
29
+ }
30
+ export function formatPapers(papers) {
31
+ if (papers.length === 0)
32
+ return 'No papers found.';
33
+ return papers
34
+ .map((p, i) => {
35
+ const year = p.published.slice(0, 4);
36
+ const who = p.authors.slice(0, 3).join(', ') + (p.authors.length > 3 ? ', et al.' : '');
37
+ return `${i + 1}. ${p.title} (${year}) — ${who} [${p.id}]`;
38
+ })
39
+ .join('\n');
40
+ }
41
+ export class ArxivClient {
42
+ endpoint;
43
+ constructor(endpoint = 'https://export.arxiv.org/api/query') {
44
+ this.endpoint = endpoint;
45
+ }
46
+ async search(query, maxResults) {
47
+ const url = `${this.endpoint}?search_query=${encodeURIComponent(`all:${query}`)}&max_results=${maxResults}`;
48
+ const response = await fetch(url);
49
+ if (!response.ok) {
50
+ throw new Error(`arXiv request failed (${response.status})`);
51
+ }
52
+ return parseArxivAtom(await response.text());
53
+ }
54
+ }
55
+ /** Dedupe papers by normalized title, preserving order. */
56
+ export function dedupePapers(papers) {
57
+ const seen = new Set();
58
+ const out = [];
59
+ for (const paper of papers) {
60
+ const key = paper.title.toLowerCase().replace(/\s+/g, ' ').trim();
61
+ if (key && !seen.has(key)) {
62
+ seen.add(key);
63
+ out.push(paper);
64
+ }
65
+ }
66
+ return out;
67
+ }
68
+ /**
69
+ * Queries several sources concurrently and merges the results. Resilient: a source that
70
+ * throws is dropped, not fatal — so one database being down doesn't break search.
71
+ */
72
+ export class MultiSourceClient {
73
+ sources;
74
+ constructor(sources) {
75
+ this.sources = sources;
76
+ }
77
+ async search(query, maxResults) {
78
+ const settled = await Promise.allSettled(this.sources.map(source => source.search(query, maxResults)));
79
+ const merged = settled.flatMap(result => (result.status === 'fulfilled' ? result.value : []));
80
+ return dedupePapers(merged).slice(0, maxResults);
81
+ }
82
+ }
83
+ function strOf(value) {
84
+ return typeof value === 'string' ? value : '';
85
+ }
86
+ /** The search_literature tool. */
87
+ export function literatureTool(client) {
88
+ return {
89
+ definition: {
90
+ name: 'search_literature',
91
+ description: 'Search the scientific literature (arXiv + OpenAlex + PubMed) for papers relevant to a query. Returns titles, authors, years, and ids (DOI/PMID/arXiv) you can cite.',
92
+ input_schema: {
93
+ type: 'object',
94
+ properties: {
95
+ query: { type: 'string', description: 'Search terms.' },
96
+ maxResults: { type: 'number', description: 'How many papers to return (default 5).' },
97
+ },
98
+ required: ['query'],
99
+ },
100
+ },
101
+ describe(input) {
102
+ return `search_literature: ${strOf(input.query)}`;
103
+ },
104
+ async run(input) {
105
+ const query = strOf(input.query);
106
+ const maxResults = Number(input.maxResults ?? 5) || 5;
107
+ try {
108
+ const papers = await client.search(query, maxResults);
109
+ return { stdout: formatPapers(papers), stderr: '' };
110
+ }
111
+ catch (error) {
112
+ return { stdout: '', stderr: error instanceof Error ? error.message : String(error) };
113
+ }
114
+ },
115
+ };
116
+ }
@@ -0,0 +1,27 @@
1
+ /** Parse an OpenAlex /works response into papers. */
2
+ export function parseOpenAlex(json) {
3
+ const results = json?.results;
4
+ if (!Array.isArray(results))
5
+ return [];
6
+ return results.map(work => ({
7
+ title: work.title ?? work.display_name ?? '(untitled)',
8
+ authors: (work.authorships ?? []).map(a => a.author?.display_name ?? '').filter(Boolean),
9
+ id: work.doi ?? work.id ?? '',
10
+ published: work.publication_year ? String(work.publication_year) : '',
11
+ summary: '',
12
+ }));
13
+ }
14
+ export class OpenAlexClient {
15
+ endpoint;
16
+ constructor(endpoint = 'https://api.openalex.org/works') {
17
+ this.endpoint = endpoint;
18
+ }
19
+ async search(query, maxResults) {
20
+ const url = `${this.endpoint}?search=${encodeURIComponent(query)}&per_page=${maxResults}`;
21
+ const response = await fetch(url);
22
+ if (!response.ok) {
23
+ throw new Error(`OpenAlex request failed (${response.status})`);
24
+ }
25
+ return parseOpenAlex(await response.json());
26
+ }
27
+ }
@@ -0,0 +1,62 @@
1
+ /**
2
+ * PubMed connector — the marquee biomedical database (NCBI E-utilities, no key). Two
3
+ * steps: esearch returns PMIDs for a query, efetch returns the article records as XML.
4
+ * Implements the same LiteratureClient interface, so it joins the multi-source search.
5
+ */
6
+ const ESEARCH = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi';
7
+ const EFETCH = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi';
8
+ /** Parse the esearch JSON into a list of PMIDs. */
9
+ export function parseEsearchIds(json) {
10
+ const ids = json?.esearchresult?.idlist;
11
+ return Array.isArray(ids) ? ids.filter(id => typeof id === 'string') : [];
12
+ }
13
+ function firstMatch(block, re) {
14
+ const m = re.exec(block);
15
+ return m ? m[1].replace(/\s+/g, ' ').trim() : '';
16
+ }
17
+ /** Parse an efetch PubmedArticleSet into papers (tolerant, article-by-article). */
18
+ export function parsePubmedXml(xml) {
19
+ const papers = [];
20
+ const articleRe = /<PubmedArticle>([\s\S]*?)<\/PubmedArticle>/g;
21
+ let article = articleRe.exec(xml);
22
+ while (article !== null) {
23
+ const block = article[1];
24
+ const pmid = firstMatch(block, /<PMID[^>]*>(\d+)<\/PMID>/);
25
+ const authors = [];
26
+ const authorRe = /<Author[^>]*>([\s\S]*?)<\/Author>/g;
27
+ let author = authorRe.exec(block);
28
+ while (author !== null) {
29
+ const last = firstMatch(author[1], /<LastName>([\s\S]*?)<\/LastName>/);
30
+ const fore = firstMatch(author[1], /<ForeName>([\s\S]*?)<\/ForeName>/);
31
+ const name = [fore, last].filter(Boolean).join(' ');
32
+ if (name)
33
+ authors.push(name);
34
+ author = authorRe.exec(block);
35
+ }
36
+ papers.push({
37
+ title: firstMatch(block, /<ArticleTitle>([\s\S]*?)<\/ArticleTitle>/),
38
+ id: pmid ? `PMID:${pmid}` : '',
39
+ published: firstMatch(block, /<PubDate>[\s\S]*?<Year>(\d{4})<\/Year>/),
40
+ summary: '',
41
+ authors,
42
+ });
43
+ article = articleRe.exec(xml);
44
+ }
45
+ return papers;
46
+ }
47
+ export class PubMedClient {
48
+ async search(query, maxResults) {
49
+ const searchUrl = `${ESEARCH}?db=pubmed&retmode=json&retmax=${maxResults}&term=${encodeURIComponent(query)}`;
50
+ const searchRes = await fetch(searchUrl);
51
+ if (!searchRes.ok)
52
+ throw new Error(`PubMed esearch failed (${searchRes.status})`);
53
+ const ids = parseEsearchIds(await searchRes.json());
54
+ if (ids.length === 0)
55
+ return [];
56
+ const fetchUrl = `${EFETCH}?db=pubmed&retmode=xml&id=${ids.join(',')}`;
57
+ const fetchRes = await fetch(fetchUrl);
58
+ if (!fetchRes.ok)
59
+ throw new Error(`PubMed efetch failed (${fetchRes.status})`);
60
+ return parsePubmedXml(await fetchRes.text());
61
+ }
62
+ }
@@ -13,9 +13,14 @@ cell are visible in the next.
13
13
  import contextlib
14
14
  import io
15
15
  import json
16
+ import os
16
17
  import sys
17
18
  import traceback
18
19
 
20
+ # Render matplotlib to a non-interactive backend so plt.show() never blocks; we capture
21
+ # figures ourselves after each cell (see emit_figures). Set before any user import.
22
+ os.environ.setdefault("MPLBACKEND", "Agg")
23
+
19
24
  # Persistent namespace shared across all cells in this session.
20
25
  NAMESPACE: dict = {"__name__": "__mb_science__"}
21
26
 
@@ -25,6 +30,30 @@ def emit(frame: dict) -> None:
25
30
  sys.stdout.flush()
26
31
 
27
32
 
33
+ def emit_figures(cell_id: str) -> int:
34
+ """Emit any open matplotlib figures as base64 PNG image frames, then close them.
35
+
36
+ This is what makes plots appear inline in the transcript (a science-notebook staple).
37
+ No-op if matplotlib hasn't been used in this session."""
38
+ pyplot = sys.modules.get("matplotlib.pyplot")
39
+ if pyplot is None:
40
+ return 0
41
+ import base64
42
+
43
+ count = 0
44
+ try:
45
+ for num in pyplot.get_fignums():
46
+ buf = io.BytesIO()
47
+ pyplot.figure(num).savefig(buf, format="png", bbox_inches="tight", dpi=110)
48
+ data = base64.b64encode(buf.getvalue()).decode("ascii")
49
+ emit({"type": "image", "cellId": cell_id, "mime": "image/png", "data": data})
50
+ count += 1
51
+ pyplot.close("all")
52
+ except Exception: # noqa: BLE001 - figure capture is best-effort
53
+ pass
54
+ return count
55
+
56
+
28
57
  def run_cell(cell_id: str, code: str) -> None:
29
58
  emit({"type": "status", "cellId": cell_id, "state": "busy"})
30
59
  stdout, stderr = io.StringIO(), io.StringIO()
@@ -41,6 +70,11 @@ def run_cell(cell_id: str, code: str) -> None:
41
70
  emit({"type": "stream", "cellId": cell_id, "kind": "stdout", "data": out})
42
71
  if err:
43
72
  emit({"type": "stream", "cellId": cell_id, "kind": "stderr", "data": err})
73
+ figures = emit_figures(cell_id)
74
+ if figures:
75
+ # A note in stdout so the (text-only) model knows a figure was produced.
76
+ emit({"type": "stream", "cellId": cell_id, "kind": "stdout",
77
+ "data": f"[{figures} figure(s) displayed]\n"})
44
78
  emit({"type": "status", "cellId": cell_id, "state": "idle"})
45
79
 
46
80
 
@@ -307,6 +307,8 @@ export function sessionPage(projectId, sessionId) {
307
307
  white-space:pre-wrap; overflow-x:auto; border-top:1px solid var(--border); }
308
308
  .step .out { color:var(--muted); background:var(--bg); }
309
309
  .step .out.err { color:#c0392b; }
310
+ img.figure { max-width:calc(100% - 24px); border-radius:8px; margin:10px 12px; display:block;
311
+ background:#fff; border:1px solid var(--border); }
310
312
  /* Composer */
311
313
  .composer { flex-shrink:0; border-top:1px solid var(--border); padding:14px 20px; }
312
314
  .composer-inner { max-width:760px; margin:0 auto; display:flex; gap:10px; align-items:flex-end;
@@ -9,6 +9,9 @@ import { homePage, loginPage, nonceInvalidPage, nonceLandingPage, projectPage, s
9
9
  import { FEATURED_CONNECTORS, NETWORK_CATEGORIES, WorkspaceStore } from './store.js';
10
10
  import { KernelManager } from './kernel.js';
11
11
  import { AgentSession, pythonTool } from './agent/agent.js';
12
+ import { ArxivClient, literatureTool, MultiSourceClient } from './agent/literature.js';
13
+ import { OpenAlexClient } from './agent/openalex.js';
14
+ import { PubMedClient } from './agent/pubmed.js';
12
15
  import { GatewayModel } from './agent/model.js';
13
16
  /**
14
17
  * The `megabrain science` local daemon's HTTP server. Serves the sign-in hand-off on
@@ -264,7 +267,10 @@ export async function startScienceServer(host = '127.0.0.1') {
264
267
  const project = session ? store.getProject(session.projectId) : null;
265
268
  store.touchSession(sid, session && !session.summary ? { summary: prompt } : undefined);
266
269
  fanout(sid, 'agent', { type: 'user', text: prompt });
267
- const agent = new AgentSession(new GatewayModel({ apiKey: creds.token }), [pythonTool(kernels.get(sid))], evt => fanout(sid, 'agent', evt), project?.agentContext ?? '', store.getMessages(sid));
270
+ const agent = new AgentSession(new GatewayModel({ apiKey: creds.token }), [
271
+ pythonTool(kernels.get(sid)),
272
+ literatureTool(new MultiSourceClient([new ArxivClient(), new OpenAlexClient(), new PubMedClient()])),
273
+ ], evt => fanout(sid, 'agent', evt), project?.agentContext ?? '', store.getMessages(sid));
268
274
  // Persist the full history when the turn ends, so the session remembers it and the
269
275
  // transcript survives a refresh.
270
276
  void agent.run(prompt).finally(() => store.setMessages(sid, agent.getMessages()));
@@ -64,6 +64,18 @@
64
64
  if (stderr) out.classList.add('err');
65
65
  scroll();
66
66
  }
67
+ // Inline a rendered figure (matplotlib PNG) under the current step, or on its own if none.
68
+ function addFigure(mime, data) {
69
+ var img = el('img', 'figure');
70
+ img.src = 'data:' + mime + ';base64,' + data;
71
+ if (lastCard) lastCard.appendChild(img);
72
+ else {
73
+ var row = el('div', 'msg agent');
74
+ row.appendChild(img);
75
+ transcript.appendChild(row);
76
+ }
77
+ scroll();
78
+ }
67
79
 
68
80
  function setRunning(on) {
69
81
  sendBtn.disabled = on;
@@ -112,6 +124,7 @@
112
124
  if (m.channel === 'kernel') {
113
125
  var f = m.payload;
114
126
  if (f.type === 'status') kstate.textContent = f.state === 'busy' ? 'running' : f.state;
127
+ else if (f.type === 'image') addFigure(f.mime, f.data);
115
128
  return;
116
129
  }
117
130
  var p = m.payload;
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@getmegabrain/cli",
3
- "version": "0.1.8",
3
+ "version": "0.1.10",
4
4
  "description": "MegaBrain CLI — sign in once, then code from your terminal via OpenCode wired to the MegaBrain Gateway, or run `megabrain science` for the local research workbench.",
5
5
  "license": "MIT",
6
6
  "type": "module",