@fbraza/pi-cite 0.1.0

package/src/literature-search.ts

@@ -0,0 +1,345 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type, type Static } from "typebox";
+import { searchPubmed } from "./pubmed.ts";
+import {
+  compactPapersForDisplay,
+  renderLiteratureSearchResult,
+  type LiteratureSearchDisplayEvent,
+  type LiteratureSearchDisplaySearch,
+} from "./rendering.ts";
+import { searchSemanticScholar } from "./semantic-scholar.ts";
+import { formatPaperText, normalizeDoi, unique } from "./shared.ts";
+import { emitProgress, textResult, type TextToolUpdate } from "./tool-output.ts";
+import type { PaperRecord } from "./types.ts";
+
+export const LITERATURE_SEARCH_PARAMS = Type.Object({
+  pubmed_query: Type.String({
+    description:
+      "PubMed-ready query using PubMed syntax such as MeSH [mh], title/abstract [tiab], publication type [pt], substance [nm], and Boolean logic.",
+  }),
+  semantic_scholar_query: Type.Optional(
+    Type.String({
+      description:
+        "Optional natural-language Semantic Scholar query for supplementary search. If omitted and Semantic Scholar is configured, a simplified query is derived from pubmed_query.",
+    }),
+  ),
+  max_results: Type.Optional(
+    Type.Number({ description: "Maximum results per provider (default 20)" }),
+  ),
+  date_from: Type.Optional(
+    Type.String({ description: "PubMed publication start date as YYYY/MM/DD" }),
+  ),
+  date_to: Type.Optional(
+    Type.String({ description: "PubMed publication end date as YYYY/MM/DD" }),
+  ),
+  publication_types: Type.Optional(
+    Type.Array(Type.String({ description: "PubMed publication type" })),
+  ),
+  fetch_abstracts: Type.Optional(
+    Type.Boolean({ description: "Whether PubMed should fetch abstracts (default true)" }),
+  ),
+});
+
+export type LiteratureSearchParams = Static<typeof LITERATURE_SEARCH_PARAMS>;
+
+type ProviderExecution =
+  | { searched: true; count: number; query: string; total?: number }
+  | { searched: false; reason: string };
+
+export type LiteratureSearchResult = {
+  count: number;
+  papers: PaperRecord[];
+  providers: {
+    pubmed: ProviderExecution;
+    semantic_scholar: ProviderExecution;
+  };
+  searches: LiteratureSearchDisplaySearch[];
+  events: LiteratureSearchDisplayEvent[];
+};
+
+function firstYear(value?: string): number | undefined {
+  const match = value?.match(/^(\d{4})/);
+  return match?.[1] ? Number(match[1]) : undefined;
+}
+
+export function simplifyPubmedQueryForSemanticScholar(query: string): string {
+  const simplified = query
+    .replace(/\[[^\]]+\]/g, " ")
+    .replace(/\b(?:AND|OR|NOT)\b/gi, " ")
+    .replace(/[()"']/g, " ")
+    .replace(/\s+/g, " ")
+    .trim();
+  return simplified || query.trim();
+}
+
+function sourceList(paper: PaperRecord): string[] {
+  return unique([
+    ...(paper.sources ?? []),
+    ...(paper.source ? paper.source.split(";") : []),
+  ].map((source) => source.trim()).filter(Boolean));
+}
+
+function normalizedTitle(title: string): string {
+  return title
+    .toLowerCase()
+    .replace(/[^a-z0-9]+/g, " ")
+    .replace(/\s+/g, " ")
+    .trim();
+}
+
+function dedupeKeys(paper: PaperRecord): string[] {
+  const doi = normalizeDoi(paper.doi)?.toLowerCase();
+  const keys = [
+    doi ? `doi:${doi}` : undefined,
+    paper.pmid ? `pmid:${paper.pmid}` : undefined,
+    paper.s2_id ? `s2:${paper.s2_id}` : undefined,
+  ];
+  const title = normalizedTitle(paper.title);
+  if (title && paper.year) keys.push(`title-year:${title}:${paper.year}`);
+  return unique(keys);
+}
+
+function mergePapers(existing: PaperRecord, incoming: PaperRecord): PaperRecord {
+  const sources = unique([...sourceList(existing), ...sourceList(incoming)]);
+  return {
+    ...incoming,
+    ...existing,
+    doi: normalizeDoi(existing.doi) ?? normalizeDoi(incoming.doi),
+    pmid: existing.pmid ?? incoming.pmid,
+    s2_id: existing.s2_id ?? incoming.s2_id,
+    title: existing.title !== "Untitled" ? existing.title : incoming.title,
+    abstract: existing.abstract ?? incoming.abstract,
+    authors: unique([...(existing.authors ?? []), ...(incoming.authors ?? [])]),
+    journal: existing.journal ?? incoming.journal,
+    year: existing.year ?? incoming.year,
+    publication_types: unique([
+      ...(existing.publication_types ?? []),
+      ...(incoming.publication_types ?? []),
+    ]),
+    mesh_terms: unique([...(existing.mesh_terms ?? []), ...(incoming.mesh_terms ?? [])]),
+    citation_count: existing.citation_count ?? incoming.citation_count,
+    tldr: existing.tldr ?? incoming.tldr,
+    open_access_pdf: existing.open_access_pdf ?? incoming.open_access_pdf,
+    external_ids: { ...(incoming.external_ids ?? {}), ...(existing.external_ids ?? {}) },
+    source: sources.join(";"),
+    sources,
+  };
+}
+
+export function dedupeLiteraturePapers(papers: PaperRecord[]): PaperRecord[] {
+  const merged: PaperRecord[] = [];
+  const keyToIndex = new Map<string, number>();
+
+  for (const paper of papers) {
+    const keys = dedupeKeys(paper);
+    const existingIndex = keys
+      .map((key) => keyToIndex.get(key))
+      .find((index) => index !== undefined);
+
+    if (existingIndex === undefined) {
+      const index = merged.length;
+      const sources = sourceList(paper);
+      merged.push({ ...paper, source: sources.join(";"), sources });
+      for (const key of keys) keyToIndex.set(key, index);
+      continue;
+    }
+
+    merged[existingIndex] = mergePapers(merged[existingIndex], paper);
+    for (const key of dedupeKeys(merged[existingIndex])) keyToIndex.set(key, existingIndex);
+  }
+
+  return merged;
+}
+
+export async function searchLiterature(
+  params: LiteratureSearchParams,
+  signal?: AbortSignal,
+  onUpdate?: TextToolUpdate,
+): Promise<LiteratureSearchResult> {
+  const maxResults = Math.min(
+    200,
+    Math.max(1, Math.floor(params.max_results ?? 20)),
+  );
+
+  const events: LiteratureSearchDisplayEvent[] = [{ phase: "start" }];
+  const searches: LiteratureSearchDisplaySearch[] = [];
+  const emitEvent = (text: string) => {
+    emitProgress(onUpdate, text, { events: [...events] });
+  };
+
+  emitEvent("Starting literature search...");
+
+  events.push({
+    phase: "query_start",
+    provider: "pubmed",
+    query_index: 1,
+    query: params.pubmed_query,
+  });
+  emitEvent(`Searching PubMed q1: ${params.pubmed_query}`);
+
+  const pubmed = await searchPubmed(
+    {
+      query: params.pubmed_query,
+      max_results: maxResults,
+      date_from: params.date_from,
+      date_to: params.date_to,
+      publication_types: params.publication_types,
+      fetch_abstracts: params.fetch_abstracts,
+    },
+    signal,
+    undefined,
+  );
+
+  const pubmedDisplayPapers = compactPapersForDisplay(pubmed.papers);
+  searches.push({
+    provider: "pubmed",
+    query_index: 1,
+    query: pubmed.query ?? params.pubmed_query,
+    count: pubmed.count,
+    papers: pubmedDisplayPapers,
+  });
+  events.push({
+    phase: "query_results",
+    provider: "pubmed",
+    query_index: 1,
+    query: pubmed.query ?? params.pubmed_query,
+    count: pubmed.count,
+    papers: pubmedDisplayPapers,
+  });
+  emitEvent(`PubMed q1 found ${pubmed.count} candidate papers.`);
+
+  const semanticScholarApiKey = process.env.SEMANTIC_SCHOLAR_API_KEY?.trim();
+  let semanticScholar: ProviderExecution = {
+    searched: false,
+    reason: "SEMANTIC_SCHOLAR_API_KEY not configured",
+  };
+  let semanticScholarPapers: PaperRecord[] = [];
+
+  if (semanticScholarApiKey) {
+    const semanticScholarQuery =
+      params.semantic_scholar_query?.trim() ||
+      simplifyPubmedQueryForSemanticScholar(params.pubmed_query);
+
+    events.push({
+      phase: "query_start",
+      provider: "semantic_scholar",
+      query_index: 1,
+      query: semanticScholarQuery,
+    });
+    emitEvent(`Searching Semantic Scholar q1: ${semanticScholarQuery}`);
+
+    try {
+      const semanticScholarResult = await searchSemanticScholar(
+        {
+          query: semanticScholarQuery,
+          max_results: Math.min(100, maxResults),
+          year_from: firstYear(params.date_from),
+          year_to: firstYear(params.date_to),
+        },
+        signal,
+        undefined,
+      );
+      semanticScholarPapers = semanticScholarResult.papers;
+      const semanticScholarDisplayPapers = compactPapersForDisplay(
+        semanticScholarResult.papers,
+      );
+      searches.push({
+        provider: "semantic_scholar",
+        query_index: 1,
+        query: semanticScholarQuery,
+        count: semanticScholarResult.count,
+        papers: semanticScholarDisplayPapers,
+      });
+      events.push({
+        phase: "query_results",
+        provider: "semantic_scholar",
+        query_index: 1,
+        query: semanticScholarQuery,
+        count: semanticScholarResult.count,
+        papers: semanticScholarDisplayPapers,
+      });
+      emitEvent(
+        `Semantic Scholar q1 found ${semanticScholarResult.count} candidate papers.`,
+      );
+      semanticScholar = {
+        searched: true,
+        count: semanticScholarResult.count,
+        query: semanticScholarQuery,
+      };
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      events.push({
+        phase: "query_error",
+        provider: "semantic_scholar",
+        query_index: 1,
+        query: semanticScholarQuery,
+        error: message,
+      });
+      semanticScholar = {
+        searched: false,
+        reason: `Semantic Scholar search failed: ${message}`,
+      };
+      emitEvent(`Semantic Scholar q1 failed: ${message}`);
+    }
+  }
+
+  events.push({ phase: "dedupe" });
+  emitEvent("Deduplicating literature results...");
+
+  const papers = dedupeLiteraturePapers([
+    ...pubmed.papers,
+    ...semanticScholarPapers,
+  ]);
+  events.push({
+    phase: "complete",
+    count: papers.length,
+    papers: compactPapersForDisplay(papers),
+  });
+  emitEvent(`Literature search complete: ${papers.length} merged papers.`);
+
+  return {
+    count: papers.length,
+    papers,
+    providers: {
+      pubmed: {
+        searched: true,
+        count: pubmed.count,
+        query: pubmed.query ?? params.pubmed_query,
+        total: pubmed.total,
+      },
+      semantic_scholar: semanticScholar,
+    },
+    searches,
+    events,
+  };
+}
+
+export function createLiteratureSearchTool() {
+  return {
+    name: "literature_search",
+    label: "Literature Search",
+    description:
+      "Run the literature workflow search: PubMed is always searched first with a PubMed-ready query; Semantic Scholar is searched as supplementary metadata when SEMANTIC_SCHOLAR_API_KEY is configured.",
+    parameters: LITERATURE_SEARCH_PARAMS,
+    async execute(
+      _toolCallId: string,
+      params: LiteratureSearchParams,
+      signal?: AbortSignal,
+      onUpdate?: TextToolUpdate,
+    ) {
+      const result = await searchLiterature(params, signal, onUpdate);
+      return textResult(formatPaperText(result.papers), result);
+    },
+    renderResult(
+      result: Parameters<typeof renderLiteratureSearchResult>[0],
+      options: Parameters<typeof renderLiteratureSearchResult>[1],
+      theme: Parameters<typeof renderLiteratureSearchResult>[2],
+    ) {
+      return renderLiteratureSearchResult(result, options, theme);
+    },
+  };
+}
+
+export function registerLiteratureSearchTool(pi: ExtensionAPI): void {
+  pi.registerTool(createLiteratureSearchTool());
+}

package/src/pubmed.ts

@@ -0,0 +1,286 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type, type Static } from "typebox";
+import { renderProviderSearchResult } from "./rendering.ts";
+import { emitProgress, textResult, type TextToolUpdate } from "./tool-output.ts";
+import {
+  fetchJson,
+  fetchText,
+  formatPaperText,
+  normalizeDoi,
+  pickAll,
+  pickOne,
+  sleep,
+  unique,
+  xmlDecode,
+} from "./shared.ts";
+import type { PaperRecord } from "./types.ts";
+
+export const PUBMED_SEARCH_PARAMS = Type.Object({
+  query: Type.String({
+    description:
+      "PubMed query string (supports field tags like [tiab], [mh], [pt])",
+  }),
+  max_results: Type.Optional(
+    Type.Number({
+      description: "Maximum results to return (default 20, max 200)",
+    }),
+  ),
+  date_from: Type.Optional(
+    Type.String({ description: "Start date as YYYY/MM/DD" }),
+  ),
+  date_to: Type.Optional(
+    Type.String({ description: "End date as YYYY/MM/DD" }),
+  ),
+  publication_types: Type.Optional(
+    Type.Array(Type.String({ description: "PubMed publication type" })),
+  ),
+  fetch_abstracts: Type.Optional(
+    Type.Boolean({ description: "Whether to fetch abstracts (default true)" }),
+  ),
+  sort: Type.Optional(
+    Type.Union([
+      Type.Literal("relevance"),
+      Type.Literal("pub_date"),
+      Type.Literal("first_author"),
+    ]),
+  ),
+  api_key: Type.Optional(
+    Type.String({
+      description: "Environment variable name containing an NCBI API key (defaults to NCBI_API_KEY when omitted)",
+    }),
+  ),
+});
+
+export type PubmedSearchParams = Static<typeof PUBMED_SEARCH_PARAMS>;
+
+export const DEFAULT_NCBI_API_KEY_ENV = "NCBI_API_KEY";
+
+export function getNcbiApiKey(envVarName?: string): string | undefined {
+  const keyEnv = envVarName?.trim() || DEFAULT_NCBI_API_KEY_ENV;
+  const apiKey = process.env[keyEnv]?.trim();
+  return apiKey || undefined;
+}
+
+export function addNcbiApiKeyParam(url: URL, envVarName?: string): boolean {
+  const apiKey = getNcbiApiKey(envVarName);
+  if (!apiKey) return false;
+  url.searchParams.set("api_key", apiKey);
+  return true;
+}
+
+export function normalizePubmedQuery(
+  query: string,
+  publicationTypes?: string[],
+  dateFrom?: string,
+  dateTo?: string,
+): string {
+  const fragments = [query.trim()].filter(Boolean);
+  if (publicationTypes && publicationTypes.length > 0) {
+    fragments.push(
+      `(${publicationTypes.map((item) => `\"${item}\"[Publication Type]`).join(" OR ")})`,
+    );
+  }
+  if (dateFrom || dateTo) {
+    const start = dateFrom ?? "1000/01/01";
+    const end = dateTo ?? "3000/12/31";
+    fragments.push(`(${start}:${end}[Date - Publication])`);
+  }
+  return fragments.join(" AND ");
+}
+
+export function parsePubmedArticle(articleXml: string): PaperRecord {
+  const pmid = pickOne(/<PMID[^>]*>(.*?)<\/PMID>/i, articleXml);
+  const title =
+    pickOne(/<ArticleTitle>([\s\S]*?)<\/ArticleTitle>/i, articleXml) ??
+    "Untitled";
+  const abstractSections = pickAll(
+    /<AbstractText[^>]*>([\s\S]*?)<\/AbstractText>/gi,
+    articleXml,
+  );
+  const abstract = abstractSections.join(" ").trim() || undefined;
+  const journal =
+    pickOne(/<Title>([\s\S]*?)<\/Title>/i, articleXml) ??
+    pickOne(/<ISOAbbreviation>(.*?)<\/ISOAbbreviation>/i, articleXml);
+  const yearText =
+    pickOne(/<PubDate>[\s\S]*?<Year>(\d{4})<\/Year>/i, articleXml) ??
+    pickOne(/<ArticleDate[^>]*>[\s\S]*?<Year>(\d{4})<\/Year>/i, articleXml) ??
+    pickOne(
+      /<PubMedPubDate[^>]*PubStatus=\"pubmed\">[\s\S]*?<Year>(\d{4})<\/Year>/i,
+      articleXml,
+    );
+  const doi =
+    normalizeDoi(
+      pickOne(
+        /<ELocationID[^>]*EIdType=\"doi\"[^>]*>(.*?)<\/ELocationID>/i,
+        articleXml,
+      ),
+    ) ??
+    normalizeDoi(
+      pickOne(
+        /<ArticleId[^>]*IdType=\"doi\"[^>]*>(.*?)<\/ArticleId>/i,
+        articleXml,
+      ),
+    );
+  const publicationTypes = unique(
+    pickAll(
+      /<PublicationType[^>]*>([\s\S]*?)<\/PublicationType>/gi,
+      articleXml,
+    ),
+  );
+  const meshTerms = unique(
+    pickAll(/<DescriptorName[^>]*>([\s\S]*?)<\/DescriptorName>/gi, articleXml),
+  );
+  const authors = unique(
+    Array.from(
+      articleXml.matchAll(
+        /<Author[\s\S]*?<LastName>(.*?)<\/LastName>[\s\S]*?(?:<ForeName>(.*?)<\/ForeName>|<Initials>(.*?)<\/Initials>)/gi,
+      ),
+    ).map((match) => {
+      const last = xmlDecode(match[1] ?? "");
+      const fore = xmlDecode(match[2] ?? match[3] ?? "");
+      return [fore, last].filter(Boolean).join(" ").trim();
+    }),
+  );
+  const collectiveAuthors = pickAll(
+    /<CollectiveName>([\s\S]*?)<\/CollectiveName>/gi,
+    articleXml,
+  );
+  return {
+    pmid,
+    doi,
+    title,
+    abstract,
+    authors: unique([...authors, ...collectiveAuthors]),
+    journal,
+    year: yearText ? Number(yearText) : undefined,
+    publication_types: publicationTypes,
+    mesh_terms: meshTerms,
+    source: "pubmed",
+  };
+}
+
+export function parsePubmedArticles(xml: string): PaperRecord[] {
+  const chunks = xml.match(/<PubmedArticle>[\s\S]*?<\/PubmedArticle>/gi) ?? [];
+  return chunks.map(parsePubmedArticle);
+}
+
+export async function lookupPubmedIdentifiers(
+  pmid: string,
+  signal?: AbortSignal,
+): Promise<{ doi?: string; title?: string }> {
+  const url = new URL(
+    "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi",
+  );
+  url.searchParams.set("db", "pubmed");
+  url.searchParams.set("id", pmid);
+  url.searchParams.set("retmode", "xml");
+  addNcbiApiKeyParam(url);
+  const xml = await fetchText(url.toString(), signal);
+  const article = parsePubmedArticles(xml)[0];
+  return { doi: article?.doi, title: article?.title };
+}
+
+export type PubmedSearchResult = {
+  count: number;
+  papers: PaperRecord[];
+  query?: string;
+  total?: number;
+};
+
+export async function searchPubmed(
+  params: PubmedSearchParams,
+  signal?: AbortSignal,
+  onUpdate?: TextToolUpdate,
+): Promise<PubmedSearchResult> {
+  const maxResults = Math.min(
+    200,
+    Math.max(1, Math.floor(params.max_results ?? 20)),
+  );
+  const query = normalizePubmedQuery(
+    params.query,
+    params.publication_types,
+    params.date_from,
+    params.date_to,
+  );
+  const esearchUrl = new URL(
+    "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi",
+  );
+  esearchUrl.searchParams.set("db", "pubmed");
+  esearchUrl.searchParams.set("retmode", "json");
+  esearchUrl.searchParams.set("retmax", String(maxResults));
+  esearchUrl.searchParams.set("sort", params.sort ?? "relevance");
+  esearchUrl.searchParams.set("term", query);
+  const hasApiKey = addNcbiApiKeyParam(esearchUrl, params.api_key);
+  emitProgress(onUpdate, `Searching PubMed for: ${params.query}`);
+  const esearch = await fetchJson<{
+    esearchresult?: { idlist?: string[]; count?: string };
+  }>(esearchUrl.toString(), signal);
+  const ids = esearch.esearchresult?.idlist ?? [];
+  if (ids.length === 0) {
+    return { count: 0, papers: [] };
+  }
+  if (params.fetch_abstracts === false) {
+    const papers = ids.map((pmid) => ({
+      pmid,
+      title: "PubMed record",
+      source: "pubmed",
+    }));
+    return { count: papers.length, papers, query };
+  }
+  const rateLimitMs = hasApiKey ? 120 : 350;
+  const batchSize = 50;
+  const papers: PaperRecord[] = [];
+  for (let start = 0; start < ids.length; start += batchSize) {
+    const batch = ids.slice(start, start + batchSize);
+    emitProgress(
+      onUpdate,
+      `Searching PubMed... found ${ids.length} PMIDs, fetching abstracts ${start + 1}-${Math.min(start + batch.length, ids.length)}...`,
+    );
+    const efetchUrl = new URL(
+      "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi",
+    );
+    efetchUrl.searchParams.set("db", "pubmed");
+    efetchUrl.searchParams.set("retmode", "xml");
+    efetchUrl.searchParams.set("id", batch.join(","));
+    addNcbiApiKeyParam(efetchUrl, params.api_key);
+    const xml = await fetchText(efetchUrl.toString(), signal);
+    papers.push(...parsePubmedArticles(xml));
+    if (start + batchSize < ids.length) await sleep(rateLimitMs, signal);
+  }
+  return {
+    count: papers.length,
+    papers,
+    query,
+    total: Number(esearch.esearchresult?.count ?? papers.length),
+  };
+}
+
+export function createPubmedSearchTool() {
+  return {
+    name: "pubmed_search",
+    label: "PubMed Search",
+    description:
+      "Search PubMed using typed parameters and return metadata with abstracts when available.",
+    parameters: PUBMED_SEARCH_PARAMS,
+    async execute(
+      _toolCallId: string,
+      params: PubmedSearchParams,
+      signal?: AbortSignal,
+      onUpdate?: TextToolUpdate,
+    ) {
+      const result = await searchPubmed(params, signal, onUpdate);
+      return textResult(formatPaperText(result.papers), result);
+    },
+    renderResult(
+      result: Parameters<typeof renderProviderSearchResult>[1],
+      options: Parameters<typeof renderProviderSearchResult>[2],
+      theme: Parameters<typeof renderProviderSearchResult>[3],
+    ) {
+      return renderProviderSearchResult("pubmed", result, options, theme);
+    },
+  };
+}
+
+export function registerPubmedSearchTool(pi: ExtensionAPI): void {
+  pi.registerTool(createPubmedSearchTool());
+}