fix(go-agent): Added PubMed component support (#16817)

## Summary

- Merge upstream main and retain PubMed component support.
- Preserve newly registered tool components and update registry
verification.

## Tests

- `bash build.sh --test ./internal/agent/component/...`
- `bash build.sh --test ./internal/agent/tool/...`

<img width="1817" height="972" alt="image"
src="https://github.com/user-attachments/assets/9fcb9448-9e26-41b9-940c-a9bfde9835e9"
/>

---------

Co-authored-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Hz_
2026-07-11 17:49:59 +08:00
committed by GitHub
parent 8a3699fa87
commit 156a11c56b
7 changed files with 613 additions and 285 deletions

View File

@@ -522,5 +522,6 @@ func init() {
Register("DuckDuckGo", newDuckDuckGoComponent)
Register("Google", newGoogleComponent)
Register("GoogleScholar", newGoogleScholarComponent)
Register("PubMed", newPubMedComponent)
Register("YahooFinance", newYahooFinanceComponent)
}

View File

@@ -0,0 +1,150 @@
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package component
import (
"context"
"encoding/json"
"errors"
"strings"
"testing"
einotool "github.com/cloudwego/eino/components/tool"
)
type fakePubMedInvoker struct {
args map[string]any
err error
out string
}
func (f *fakePubMedInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) {
if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil {
return "", err
}
if f.out != "" || f.err != nil {
return f.out, f.err
}
return `{"results":[{"title":"Deep learning for retrieval augmented generation","url":"https://pubmed.ncbi.nlm.nih.gov/12345678","content":"Title: Deep learning for retrieval augmented generation\nAuthors: Furqan Khan, Jane Smith\nJournal: Nature Machine Intelligence\nVolume: 10\nIssue: 2\nPages: 101-110\nDOI: 10.1000/example.doi\nAbstract: A short abstract."}]}`, nil
}
func TestPubMed_RegisteredFactory(t *testing.T) {
t.Parallel()
c, err := New("PubMed", map[string]any{
"top_n": 8,
"email": "node@example.com",
"outputs": map[string]any{"formalized_content": map[string]any{}},
"setups": map[string]any{"query": "configured query"},
})
if err != nil {
t.Fatalf("New(PubMed) errored: %v", err)
}
if got := c.Name(); got != "PubMed" {
t.Fatalf("Name() = %q, want PubMed", got)
}
formGetter, ok := c.(interface{ GetInputForm() map[string]any })
if !ok {
t.Fatal("PubMed component does not expose GetInputForm")
}
form := formGetter.GetInputForm()
if len(form) != 1 {
t.Fatalf("GetInputForm size = %d, want 1", len(form))
}
query, ok := form["query"].(map[string]any)
if !ok {
t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"])
}
if query["type"] != "line" {
t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"])
}
if _, ok := c.Outputs()["formalized_content"]; !ok {
t.Fatal("Outputs() missing formalized_content")
}
if _, ok := c.Outputs()["json"]; !ok {
t.Fatal("Outputs() missing json")
}
}
func TestPubMed_InvokeOnlyPassesQuery(t *testing.T) {
t.Parallel()
fake := &fakePubMedInvoker{}
c := newPubMedComponentWithInvoker(fake)
out, err := c.Invoke(context.Background(), map[string]any{
"query": " retrieval augmented generation ",
"top_n": float64(8),
"email": "ignored@example.com",
"unused": true,
})
if err != nil {
t.Fatalf("Invoke errored: %v", err)
}
if got := fake.args["query"]; got != "retrieval augmented generation" {
t.Fatalf("query arg = %v, want trimmed query", got)
}
if len(fake.args) != 1 {
t.Fatalf("runtime args = %#v, want only query", fake.args)
}
formalized, _ := out["formalized_content"].(string)
for _, want := range []string{"ID: 0", "Title: Deep learning for retrieval augmented generation", "URL: https://pubmed.ncbi.nlm.nih.gov/12345678", "Content:", "Abstract: A short abstract."} {
if !strings.Contains(formalized, want) {
t.Fatalf("formalized_content missing %q: %s", want, formalized)
}
}
results, ok := out["json"].([]any)
if !ok || len(results) != 1 {
t.Fatalf("json output = %#v, want one result", out["json"])
}
}
func TestPubMed_InvokeEmptyQueryReturnsEmptyPayload(t *testing.T) {
t.Parallel()
c := newPubMedComponentWithInvoker(&fakePubMedInvoker{})
out, err := c.Invoke(context.Background(), map[string]any{"query": " "})
if err != nil {
t.Fatalf("Invoke errored: %v", err)
}
if got := out["formalized_content"]; got != "" {
t.Fatalf("formalized_content = %v, want empty string", got)
}
results, ok := out["json"].([]any)
if !ok || len(results) != 0 {
t.Fatalf("json output = %#v, want empty []any", out["json"])
}
}
func TestPubMed_InvokeSurfacesToolErrorEnvelope(t *testing.T) {
t.Parallel()
fake := &fakePubMedInvoker{
out: `{"results":[],"_ERROR":"upstream down"}`,
err: errors.New("boom"),
}
c := newPubMedComponentWithInvoker(fake)
out, err := c.Invoke(context.Background(), map[string]any{"query": "pubmed"})
if err != nil {
t.Fatalf("Invoke errored: %v", err)
}
if got := out["_ERROR"]; got != "upstream down" {
t.Fatalf("_ERROR = %v, want upstream down", got)
}
if got := out["formalized_content"]; got != "" {
t.Fatalf("formalized_content = %v, want empty string", got)
}
}

View File

@@ -1511,6 +1511,10 @@ type googleScholarInvoker interface {
InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error)
}
type pubmedInvoker interface {
InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error)
}
func newGoogleScholarComponent(params map[string]any) (Component, error) {
cloned := make(map[string]any, len(params))
for k, v := range params {
@@ -1650,6 +1654,124 @@ func renderGoogleScholarResults(results []any) string {
return strings.Join(blocks, "\n\n")
}
// pubMedComponent delegates to the PubMed tool. Its node parameters are
// consumed at construction time, leaving query as the sole runtime input.
type pubMedComponent struct {
inner pubmedInvoker
}
func newPubMedComponent(params map[string]any) (Component, error) {
toolParams := make(map[string]any, 2)
for _, key := range []string{"top_n", "email"} {
if value, ok := params[key]; ok {
toolParams[key] = value
}
}
inner, err := agenttool.BuildByName("pubmed", toolParams)
if err != nil {
return nil, err
}
invoker, ok := inner.(pubmedInvoker)
if !ok {
return nil, fmt.Errorf("PubMed: tool does not implement InvokableRun")
}
return newPubMedComponentWithInvoker(invoker), nil
}
func newPubMedComponentWithInvoker(inner pubmedInvoker) Component {
return &pubMedComponent{inner: inner}
}
func (c *pubMedComponent) Name() string { return "PubMed" }
func (c *pubMedComponent) Inputs() map[string]string {
return map[string]string{
"query": "PubMed search query.",
}
}
func (c *pubMedComponent) Outputs() map[string]string {
return map[string]string{
"formalized_content": "Rendered PubMed references for downstream LLM prompts.",
"json": "Raw PubMed result list.",
}
}
func (c *pubMedComponent) GetInputForm() map[string]any {
return map[string]any{
"query": map[string]any{
"name": "Query",
"type": "line",
},
}
}
func (c *pubMedComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
query := strings.TrimSpace(stringParam(inputs["query"]))
if query == "" {
return map[string]any{"formalized_content": "", "json": []any{}}, nil
}
argsJSON, err := json.Marshal(map[string]any{"query": query})
if err != nil {
return nil, fmt.Errorf("canvas: PubMed: encode query: %w", err)
}
out, err := c.inner.InvokableRun(ctx, string(argsJSON))
decoded := parseToolEnvelope(out)
results := anySlice(decoded["results"])
if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" {
return map[string]any{
"formalized_content": "",
"json": results,
"_ERROR": existing,
}, nil
}
if err != nil {
if len(decoded) > 0 {
return map[string]any{
"formalized_content": "",
"json": results,
"_ERROR": decoded["_ERROR"],
}, nil
}
return nil, fmt.Errorf("canvas: PubMed: %w", err)
}
return map[string]any{
"formalized_content": renderPubMedResults(results),
"json": results,
}, nil
}
func (c *pubMedComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) {
return nil, nil
}
func renderPubMedResults(results []any) string {
if len(results) == 0 {
return ""
}
blocks := make([]string, 0, len(results))
for i, item := range results {
result, ok := item.(map[string]any)
if !ok {
continue
}
content := strings.TrimSpace(stringParam(result["content"]))
if content == "" {
continue
}
lines := []string{fmt.Sprintf("ID: %d", i)}
if title := strings.TrimSpace(stringParam(result["title"])); title != "" {
lines = append(lines, "Title: "+title)
}
if link := strings.TrimSpace(stringParam(result["url"])); link != "" {
lines = append(lines, "URL: "+link)
}
lines = append(lines, "Content:", content)
blocks = append(blocks, strings.Join(lines, "\n"))
}
return strings.Join(blocks, "\n\n")
}
// Compile-time interface checks.
var (
_ Component = (*retrievalComponent)(nil)
@@ -1661,6 +1783,7 @@ var (
_ Component = (*codeExecComponent)(nil)
_ Component = (*wikipediaComponent)(nil)
_ Component = (*googleScholarComponent)(nil)
_ Component = (*pubMedComponent)(nil)
_ Component = (*yahooFinanceComponent)(nil)
)
@@ -1673,3 +1796,4 @@ var _ einotool.InvokableTool = (*agenttool.TavilyExtractTool)(nil)
var _ einotool.InvokableTool = (*agenttool.DuckDuckGoTool)(nil)
var _ einotool.InvokableTool = (*agenttool.YahooFinanceTool)(nil)
var _ einotool.InvokableTool = (*agenttool.GoogleScholarTool)(nil)
var _ einotool.InvokableTool = (*agenttool.PubMedTool)(nil)

View File

@@ -11,10 +11,10 @@ import (
// read from plan §2.11.10 — P0 (8) + P1 (5) + P2 (4) + P3 (2) + P4 (3) = 22
// at plan completion, plus v1 fixture wrappers/stubs (including Retrieval,
// TavilySearch, TavilyExtract, ExeSQL, CodeExec, Google, BGPT, YahooFinance,
// Wikipedia, GoogleScholar, DuckDuckGo, Generate, Answer, Iteration, and IterationItem)
// registered by fixture_stubs.go to keep the dsl-examples and canvas tool
// surface compiling. The test allows counts between 12 (P0+P1 minus the
// removed ExitLoop) and 55 to roll forward as subsequent batches land.
// Wikipedia, GoogleScholar, PubMed, DuckDuckGo, Generate, Answer, Iteration,
// and IterationItem) registered by fixture_stubs.go to keep the dsl-examples
// and canvas tool surface compiling. The test allows counts between 12 (P0+P1
// minus the removed ExitLoop) and 55 to roll forward as subsequent batches land.
//
// Note: ExitLoop is intentionally NOT in the registry anymore. The
// canvas engine (internal/agent/canvas/canvas.go's legacyNoOpNames)

View File

@@ -19,8 +19,8 @@ package tool
import (
"context"
"encoding/json"
"encoding/xml"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
@@ -30,67 +30,71 @@ import (
"github.com/cloudwego/eino/schema"
)
const pubmedToolName = "pubmed"
const (
pubmedToolName = "pubmed"
pubmedToolDescription = "Search PubMed for life sciences and biomedical references."
defaultPubMedTopN = 12
defaultPubMedEmail = "A.N.Other@example.com"
)
const pubmedToolDescription = "Search PubMed via NCBI E-utilities. Returns {pmid, title, authors, journal, year}."
// pubmedParams is the JSON shape the model sends into InvokableRun.
// pubmedParams mirrors Python PubMedParam. Info() still only exposes query to
// the LLM; top_n and email are canvas-side params merged with constructor
// defaults at runtime.
type pubmedParams struct {
Query string `json:"query"`
MaxResults int `json:"max_results"`
Query string `json:"query"`
TopN int `json:"top_n"`
Email string `json:"email"`
}
// pubmedResult is one row in the returned record list.
// pubmedResult is one PubMed reference rendered with the same fields and
// fallbacks as Python PubMed._format_pubmed_content.
type pubmedResult struct {
PMID string `json:"pmid"`
Title string `json:"title"`
Authors string `json:"authors"`
Journal string `json:"journal"`
Year string `json:"year"`
URL string `json:"url"`
Content string `json:"content"`
}
// pubmedEnvelope is what the model sees.
type pubmedEnvelope struct {
Results []pubmedResult `json:"results"`
Error string `json:"_ERROR,omitempty"`
}
// pubmedUserAgent is the User-Agent that NCBI requires for all
// E-utilities requests. Without it, requests are silently dropped
// or rate-limited to a single IP.
const pubmedUserAgent = "ragflow/1.0"
// pubmedESearchEndpoint is the E-utilities esearch URL. Exposed as a
// package var so tests can substitute a httptest.Server URL.
var pubmedESearchEndpoint = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
// pubmedESummaryEndpoint is the E-utilities esummary URL. Exposed
// as a package var so tests can substitute a httptest.Server URL.
var pubmedESummaryEndpoint = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
var pubmedEFetchEndpoint = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
// PubMedTool is the PubMed
// search tool. It uses NCBI
// E-utilities: esearch returns a list of PMIDs, then esummary fetches
// the full records for those PMIDs.
// PubMedTool queries NCBI E-utilities with esearch followed by efetch XML,
// which is the same retrieval path as the Python PubMed component.
type PubMedTool struct {
helper *HTTPHelper
helper *HTTPHelper
defaults pubmedParams
}
// NewPubMedTool returns a PubMedTool using the default HTTPHelper.
func NewPubMedTool() *PubMedTool {
return NewPubMedToolWith(NewHTTPHelper())
}
// NewPubMedToolWith returns a PubMedTool that uses the provided
// HTTPHelper. Useful for tests.
func NewPubMedToolWith(h *HTTPHelper) *PubMedTool {
return NewPubMedToolWithDefaults(h, pubmedParams{})
}
func NewPubMedToolWithDefaults(h *HTTPHelper, defaults pubmedParams) *PubMedTool {
if h == nil {
h = NewHTTPHelper()
}
return &PubMedTool{helper: h}
if defaults.TopN == 0 {
defaults.TopN = defaultPubMedTopN
}
if strings.TrimSpace(defaults.Email) == "" {
defaults.Email = defaultPubMedEmail
}
return &PubMedTool{helper: h, defaults: defaults}
}
// Info returns the tool's metadata for the chat model.
// Info exposes only query to the LLM. Node parameters belong to canvas setup,
// not model-emitted runtime arguments.
func (p *PubMedTool) Info(_ context.Context) (*schema.ToolInfo, error) {
return &schema.ToolInfo{
Name: pubmedToolName,
@@ -98,208 +102,196 @@ func (p *PubMedTool) Info(_ context.Context) (*schema.ToolInfo, error) {
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
"query": {
Type: schema.String,
Desc: "PubMed search query (full PubMed query syntax supported).",
Desc: "The search keywords to execute with PubMed. The keywords should be the most important words or terms, including synonyms, from the original request.",
Required: true,
},
"max_results": {
Type: schema.Integer,
Desc: "Maximum number of records to return. Defaults to 5 (max 100 per request).",
Required: false,
},
}),
}, nil
}
// buildPubMedESearchURL composes the esearch URL. Centralized for
// testability.
func buildPubMedESearchURL(query string, maxResults int) string {
if maxResults <= 0 {
maxResults = 5
}
if maxResults > 100 {
maxResults = 100
}
func buildPubMedESearchURL(query string, topN int, email string) string {
q := url.Values{}
q.Set("db", "pubmed")
q.Set("term", query)
q.Set("retmax", strconv.Itoa(maxResults))
q.Set("retmax", strconv.Itoa(topN))
q.Set("retmode", "json")
q.Set("email", email)
return pubmedESearchEndpoint + "?" + q.Encode()
}
// buildPubMedESummaryURL composes the esummary URL for a list of
// PMIDs. Centralized for testability.
func buildPubMedESummaryURL(pmids []string) string {
func buildPubMedEFetchURL(pmids []string, email string) string {
q := url.Values{}
q.Set("db", "pubmed")
q.Set("id", strings.Join(pmids, ","))
q.Set("retmode", "json")
return pubmedESummaryEndpoint + "?" + q.Encode()
q.Set("retmode", "xml")
q.Set("email", email)
return pubmedEFetchEndpoint + "?" + q.Encode()
}
// pubmedESearchResponse is the upstream esearch envelope.
type pubmedESearchResponse struct {
ESearchResult struct {
IDList []string `json:"idlist"`
} `json:"esearchresult"`
}
// pubmedESummaryAuthor is one author in the esummary author list.
type pubmedESummaryAuthor struct {
Name string `json:"name"`
type pubmedXMLAuthor struct {
LastName string `xml:"LastName"`
ForeName string `xml:"ForeName"`
}
// pubmedESummaryArticle is one article in the esummary result map.
type pubmedESummaryArticle struct {
Title string `json:"title"`
Authors []pubmedESummaryAuthor `json:"authors"`
FullJournalName string `json:"fulljournalname"`
PubDate string `json:"pubdate"`
type pubmedXMLArticleID struct {
Type string `xml:"IdType,attr"`
Value string `xml:",chardata"`
}
// pubmedESummaryResponse is the upstream esummary envelope. The
// `result` value mixes per-PMID article objects with a string-list
// `uids` key, so we decode it as a map of RawMessage and then walk
// the entries to extract the article objects (skipping the `uids`
// array). See decodePubMedESummary.
type pubmedESummaryResponse struct {
Result map[string]json.RawMessage `json:"result"`
type pubmedXMLArticle struct {
PMID string `xml:"MedlineCitation>PMID"`
Article struct {
Title string `xml:"ArticleTitle"`
Abstract struct {
Text []string `xml:"AbstractText"`
} `xml:"Abstract"`
Journal struct {
Title string `xml:"Title"`
Issue struct {
Volume string `xml:"Volume"`
Number string `xml:"Issue"`
} `xml:"JournalIssue"`
} `xml:"Journal"`
Pages struct {
MedlinePgn string `xml:"MedlinePgn"`
} `xml:"Pagination"`
Authors []pubmedXMLAuthor `xml:"AuthorList>Author"`
} `xml:"MedlineCitation>Article"`
ArticleIDs []pubmedXMLArticleID `xml:"PubmedData>ArticleIdList>ArticleId"`
}
// mustReadAll is a tiny helper to read the entire response body. We
// keep it private because pubmed.go owns the only two-step HTTP dance.
func mustReadAll(r io.Reader) []byte {
b, err := io.ReadAll(r)
if err != nil {
return nil
}
return b
type pubmedXMLResponse struct {
Articles []pubmedXMLArticle `xml:"PubmedArticle"`
}
// decodePubMedESummary parses the raw esummary response and returns
// a PMID → article map. The upstream response is a flat object whose
// keys are PMIDs (article values) plus a `uids` key (string list).
// A single JSON decode into map[string]Article would fail on `uids`
// because Go's encoding/json cannot store arrays in a struct-typed
// map; the RawMessage indirection sidesteps that.
func decodePubMedESummary(body []byte) (map[string]pubmedESummaryArticle, error) {
var raw pubmedESummaryResponse
if err := json.Unmarshal(body, &raw); err != nil {
return nil, err
}
out := make(map[string]pubmedESummaryArticle, len(raw.Result))
for k, v := range raw.Result {
// The `uids` key is a JSON array of strings; skip it.
if len(v) > 0 && v[0] == '[' {
continue
}
var article pubmedESummaryArticle
if err := json.Unmarshal(v, &article); err != nil {
continue
}
out[k] = article
}
return out, nil
}
// InvokableRun performs the two-step PubMed lookup.
func (p *PubMedTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
var params pubmedParams
if err := json.Unmarshal([]byte(argsJSON), &params); err != nil {
return pubmedErrJSON(fmt.Errorf("pubmed: parse arguments: %w", err)),
fmt.Errorf("pubmed: parse arguments: %w", err)
err = fmt.Errorf("pubmed: parse arguments: %w", err)
return pubmedErrJSON(err), err
}
if strings.TrimSpace(params.Query) == "" {
return pubmedErrJSON(fmt.Errorf("query is required")),
fmt.Errorf("pubmed: query is required")
}
if params.MaxResults <= 0 {
params.MaxResults = 5
params = mergePubMedDefaults(p.defaults, params)
params.Query = strings.TrimSpace(params.Query)
if params.Query == "" {
err := fmt.Errorf("pubmed: query is required")
return pubmedErrJSON(err), err
}
headers := map[string]string{
"User-Agent": pubmedUserAgent,
"Accept": "application/json",
}
// Step 1: esearch → list of PMIDs
searchURL := buildPubMedESearchURL(params.Query, params.MaxResults)
searchResp, err := p.helper.Do(ctx, http.MethodGet, searchURL, "", "", headers)
searchResp, err := p.helper.Do(ctx, http.MethodGet, buildPubMedESearchURL(params.Query, params.TopN, params.Email), "", "", headers)
if err != nil {
return pubmedErrJSON(err), err
}
var searchBody pubmedESearchResponse
if decErr := json.NewDecoder(searchResp.Body).Decode(&searchBody); decErr != nil {
_ = searchResp.Body.Close()
return pubmedErrJSON(fmt.Errorf("pubmed: decode esearch: %w", decErr)),
fmt.Errorf("pubmed: decode esearch: %w", decErr)
}
_ = searchResp.Body.Close()
if searchResp.StatusCode < 200 || searchResp.StatusCode >= 300 {
return pubmedErrJSON(fmt.Errorf("pubmed: esearch returned %d", searchResp.StatusCode)),
fmt.Errorf("pubmed: esearch returned %d", searchResp.StatusCode)
defer searchResp.Body.Close()
if searchResp.StatusCode < http.StatusOK || searchResp.StatusCode >= http.StatusMultipleChoices {
err := fmt.Errorf("pubmed: esearch returned %d", searchResp.StatusCode)
return pubmedErrJSON(err), err
}
pmids := searchBody.ESearchResult.IDList
if len(pmids) == 0 {
var searchBody pubmedESearchResponse
if err := json.NewDecoder(searchResp.Body).Decode(&searchBody); err != nil {
err = fmt.Errorf("pubmed: decode esearch: %w", err)
return pubmedErrJSON(err), err
}
if len(searchBody.ESearchResult.IDList) == 0 {
return pubmedJSON(pubmedEnvelope{Results: []pubmedResult{}}), nil
}
// Step 2: esummary → full records
summaryURL := buildPubMedESummaryURL(pmids)
summaryResp, err := p.helper.Do(ctx, http.MethodGet, summaryURL, "", "", headers)
headers["Accept"] = "application/xml"
fetchResp, err := p.helper.Do(ctx, http.MethodGet, buildPubMedEFetchURL(searchBody.ESearchResult.IDList, params.Email), "", "", headers)
if err != nil {
return pubmedErrJSON(err), err
}
defer summaryResp.Body.Close()
if summaryResp.StatusCode < 200 || summaryResp.StatusCode >= 300 {
return pubmedErrJSON(fmt.Errorf("pubmed: esummary returned %d", summaryResp.StatusCode)),
fmt.Errorf("pubmed: esummary returned %d", summaryResp.StatusCode)
defer fetchResp.Body.Close()
if fetchResp.StatusCode < http.StatusOK || fetchResp.StatusCode >= http.StatusMultipleChoices {
err := fmt.Errorf("pubmed: efetch returned %d", fetchResp.StatusCode)
return pubmedErrJSON(err), err
}
articles, err := decodePubMedESummary(mustReadAll(summaryResp.Body))
if err != nil {
return pubmedErrJSON(fmt.Errorf("pubmed: parse esummary: %w", err)),
fmt.Errorf("pubmed: parse esummary: %w", err)
var fetched pubmedXMLResponse
if err := xml.NewDecoder(fetchResp.Body).Decode(&fetched); err != nil {
err = fmt.Errorf("pubmed: decode efetch: %w", err)
return pubmedErrJSON(err), err
}
results := make([]pubmedResult, 0, len(pmids))
for _, pmid := range pmids {
article, ok := articles[pmid]
if !ok {
continue
}
results = append(results, pubmedResult{
PMID: pmid,
Title: strings.TrimSpace(article.Title),
Authors: joinAuthorNames(article.Authors),
Journal: article.FullJournalName,
Year: firstFourDigitYear(article.PubDate),
})
results := make([]pubmedResult, 0, len(fetched.Articles))
for _, article := range fetched.Articles {
results = append(results, formatPubMedResult(article))
}
return pubmedJSON(pubmedEnvelope{Results: results}), nil
}
// joinAuthorNames joins the first N authors with ", " and adds
// "et al." for any beyond N. We use N=3 to mirror the convention
// common in academic citation styles.
func joinAuthorNames(authors []pubmedESummaryAuthor) string {
const cap = 3
if len(authors) == 0 {
return ""
}
if len(authors) <= cap {
names := make([]string, len(authors))
for i, a := range authors {
names[i] = a.Name
func formatPubMedResult(article pubmedXMLArticle) pubmedResult {
title := fallbackPubMedField(article.Article.Title, "No title")
abstract := fallbackPubMedField(strings.Join(article.Article.Abstract.Text, " "), "No abstract available")
journal := fallbackPubMedField(article.Article.Journal.Title, "Unknown Journal")
volume := fallbackPubMedField(article.Article.Journal.Issue.Volume, "-")
issue := fallbackPubMedField(article.Article.Journal.Issue.Number, "-")
pages := fallbackPubMedField(article.Article.Pages.MedlinePgn, "-")
authors := formatPubMedAuthors(article.Article.Authors)
doi := "-"
for _, id := range article.ArticleIDs {
if id.Type == "doi" && strings.TrimSpace(id.Value) != "" {
doi = strings.TrimSpace(id.Value)
break
}
return strings.Join(names, ", ")
}
names := make([]string, 0, cap+1)
for i := range cap {
names = append(names, authors[i].Name)
content := strings.Join([]string{
"Title: " + title,
"Authors: " + authors,
"Journal: " + journal,
"Volume: " + volume,
"Issue: " + issue,
"Pages: " + pages,
"DOI: " + doi,
"Abstract: " + abstract,
}, "\n")
return pubmedResult{
Title: title,
URL: "https://pubmed.ncbi.nlm.nih.gov/" + strings.TrimSpace(article.PMID),
Content: content,
}
}
func mergePubMedDefaults(defaults, params pubmedParams) pubmedParams {
if params.Query == "" {
params.Query = defaults.Query
}
if params.TopN == 0 {
params.TopN = defaults.TopN
}
if strings.TrimSpace(params.Email) == "" {
params.Email = defaults.Email
}
return params
}
func fallbackPubMedField(value, fallback string) string {
if value = strings.TrimSpace(value); value != "" {
return value
}
return fallback
}
func formatPubMedAuthors(authors []pubmedXMLAuthor) string {
names := make([]string, 0, len(authors))
for _, author := range authors {
name := strings.TrimSpace(strings.TrimSpace(author.ForeName) + " " + strings.TrimSpace(author.LastName))
if name != "" {
names = append(names, name)
}
}
if len(names) == 0 {
return "Unknown Authors"
}
names = append(names, "et al.")
return strings.Join(names, ", ")
}

View File

@@ -32,14 +32,11 @@ func TestPubMed_BuildURL(t *testing.T) {
t.Run("esearch", func(t *testing.T) {
t.Parallel()
got := buildPubMedESearchURL("covid vaccine", 7)
got := buildPubMedESearchURL("covid vaccine", 7, "user@example.com")
u, err := url.Parse(got)
if err != nil {
t.Fatalf("url.Parse(%q): %v", got, err)
}
if u.Host != "eutils.ncbi.nlm.nih.gov" {
t.Errorf("host = %q, want eutils.ncbi.nlm.nih.gov", u.Host)
}
q := u.Query()
if q.Get("db") != "pubmed" {
t.Errorf("db = %q, want pubmed", q.Get("db"))
@@ -53,166 +50,148 @@ func TestPubMed_BuildURL(t *testing.T) {
if q.Get("retmode") != "json" {
t.Errorf("retmode = %q, want json", q.Get("retmode"))
}
if q.Get("email") != "user@example.com" {
t.Errorf("email = %q, want user@example.com", q.Get("email"))
}
})
t.Run("esummary", func(t *testing.T) {
t.Run("efetch", func(t *testing.T) {
t.Parallel()
got := buildPubMedESummaryURL([]string{"12345", "67890"})
got := buildPubMedEFetchURL([]string{"12345", "67890"}, "user@example.com")
u, err := url.Parse(got)
if err != nil {
t.Fatalf("url.Parse(%q): %v", got, err)
}
q := u.Query()
if q.Get("db") != "pubmed" {
t.Errorf("db = %q, want pubmed", q.Get("db"))
}
if q.Get("id") != "12345,67890" {
t.Errorf("id = %q, want 12345,67890", q.Get("id"))
}
if q.Get("retmode") != "json" {
t.Errorf("retmode = %q, want json", q.Get("retmode"))
if q.Get("retmode") != "xml" {
t.Errorf("retmode = %q, want xml", q.Get("retmode"))
}
})
t.Run("esearch clamp high", func(t *testing.T) {
t.Parallel()
got := buildPubMedESearchURL("x", 999)
u, _ := url.Parse(got)
if u.Query().Get("retmax") != "100" {
t.Errorf("retmax = %q, want 100 (clamped)", u.Query().Get("retmax"))
if q.Get("email") != "user@example.com" {
t.Errorf("email = %q, want user@example.com", q.Get("email"))
}
})
}
func TestPubMed_ParseESummary(t *testing.T) {
func TestPubMed_InvokableRunParsesXML(t *testing.T) {
t.Parallel()
var esearchHits, esummaryHits int32
var esearchHits, efetchHits int32
var lastUA string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
lastUA = r.Header.Get("User-Agent")
w.Header().Set("Content-Type", "application/json")
if strings.Contains(r.URL.Path, "esearch") {
switch {
case strings.Contains(r.URL.Path, "esearch"):
atomic.AddInt32(&esearchHits, 1)
_, _ = w.Write([]byte(`{
"esearchresult": {
"idlist": ["11111", "22222"]
}
}`))
return
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"esearchresult":{"idlist":["12345678"]}}`))
case strings.Contains(r.URL.Path, "efetch"):
atomic.AddInt32(&efetchHits, 1)
w.Header().Set("Content-Type", "application/xml")
_, _ = w.Write([]byte(`<PubmedArticleSet>
<PubmedArticle>
<MedlineCitation>
<PMID>12345678</PMID>
<Article>
<ArticleTitle>Deep learning for retrieval augmented generation</ArticleTitle>
<Abstract><AbstractText>A short abstract.</AbstractText></Abstract>
<Journal>
<Title>Nature Machine Intelligence</Title>
<JournalIssue><Volume>10</Volume><Issue>2</Issue></JournalIssue>
</Journal>
<Pagination><MedlinePgn>101-110</MedlinePgn></Pagination>
<AuthorList>
<Author><LastName>Khan</LastName><ForeName>Furqan</ForeName></Author>
<Author><LastName>Smith</LastName><ForeName>Jane</ForeName></Author>
</AuthorList>
</Article>
</MedlineCitation>
<PubmedData>
<ArticleIdList>
<ArticleId IdType="doi">10.1000/example.doi</ArticleId>
</ArticleIdList>
</PubmedData>
</PubmedArticle>
</PubmedArticleSet>`))
default:
http.NotFound(w, r)
}
if strings.Contains(r.URL.Path, "esummary") {
atomic.AddInt32(&esummaryHits, 1)
_, _ = w.Write([]byte(`{
"result": {
"uids": ["11111","22222"],
"11111": {
"title": "Cochrane review of masks",
"authors": [{"name":"Smith J"},{"name":"Doe A"}],
"fulljournalname": "Cochrane Database Syst Rev",
"pubdate": "2020 Nov 1"
},
"22222": {
"title": "Vaccine efficacy meta-analysis",
"authors": [{"name":"Alice"},{"name":"Bob"},{"name":"Carol"},{"name":"Dave"}],
"fulljournalname": "Lancet",
"pubdate": "2021 Mar-Apr"
}
}
}`))
return
}
http.NotFound(w, r)
}))
defer srv.Close()
helper := NewHTTPHelper().WithClient(&http.Client{
Transport: rewriteHostTransport(srv.URL),
})
tool := NewPubMedToolWith(helper)
out, err := tool.InvokableRun(context.Background(),
`{"query":"covid","max_results":5}`)
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)})
tool := NewPubMedToolWithDefaults(helper, pubmedParams{TopN: 3, Email: "tester@example.com"})
out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow"}`)
if err != nil {
t.Fatalf("InvokableRun: %v", err)
}
if esearchHits != 1 {
t.Errorf("esearch calls = %d, want 1", esearchHits)
}
if esummaryHits != 1 {
t.Errorf("esummary calls = %d, want 1", esummaryHits)
if esearchHits != 1 || efetchHits != 1 {
t.Fatalf("calls = esearch:%d efetch:%d, want 1/1", esearchHits, efetchHits)
}
if !strings.Contains(lastUA, "ragflow") {
t.Errorf("User-Agent = %q, want to contain ragflow", lastUA)
t.Fatalf("User-Agent = %q, want ragflow marker", lastUA)
}
var env pubmedEnvelope
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
if err := json.Unmarshal([]byte(out), &env); err != nil {
t.Fatalf("output is not valid JSON: %v (raw=%s)", err, out)
}
if env.Error != "" {
t.Errorf("Error = %q, want empty", env.Error)
t.Fatalf("Error = %q, want empty", env.Error)
}
if len(env.Results) != 2 {
t.Fatalf("Results len = %d, want 2", len(env.Results))
if len(env.Results) != 1 {
t.Fatalf("Results len = %d, want 1", len(env.Results))
}
if env.Results[0].PMID != "11111" {
t.Errorf("Results[0].PMID = %q, want 11111", env.Results[0].PMID)
result := env.Results[0]
if result.Title != "Deep learning for retrieval augmented generation" {
t.Fatalf("Title = %q", result.Title)
}
if env.Results[0].Title != "Cochrane review of masks" {
t.Errorf("Results[0].Title = %q, want Cochrane review of masks", env.Results[0].Title)
if result.URL != "https://pubmed.ncbi.nlm.nih.gov/12345678" {
t.Fatalf("URL = %q", result.URL)
}
if env.Results[0].Authors != "Smith J, Doe A" {
t.Errorf("Results[0].Authors = %q, want Smith J, Doe A", env.Results[0].Authors)
}
if env.Results[0].Journal != "Cochrane Database Syst Rev" {
t.Errorf("Results[0].Journal = %q, want Cochrane Database Syst Rev", env.Results[0].Journal)
}
if env.Results[0].Year != "2020" {
t.Errorf("Results[0].Year = %q, want 2020", env.Results[0].Year)
}
// 4 authors → first 3 + "et al."
if !strings.HasSuffix(env.Results[1].Authors, "et al.") {
t.Errorf("Results[1].Authors = %q, want to end with et al.", env.Results[1].Authors)
for _, want := range []string{
"Title: Deep learning for retrieval augmented generation",
"Authors: Furqan Khan, Jane Smith",
"Journal: Nature Machine Intelligence",
"Volume: 10",
"Issue: 2",
"Pages: 101-110",
"DOI: 10.1000/example.doi",
"Abstract: A short abstract.",
} {
if !strings.Contains(result.Content, want) {
t.Fatalf("Content missing %q: %s", want, result.Content)
}
}
}
func TestPubMed_EmptyResults(t *testing.T) {
func TestPubMed_InvokableRunEmptyResults(t *testing.T) {
t.Parallel()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
if strings.Contains(r.URL.Path, "esearch") {
_, _ = w.Write([]byte(`{"esearchresult":{"idlist":[]}}`))
return
}
http.NotFound(w, r)
_, _ = w.Write([]byte(`{"esearchresult":{"idlist":[]}}`))
}))
defer srv.Close()
helper := NewHTTPHelper().WithClient(&http.Client{
Transport: rewriteHostTransport(srv.URL),
})
helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)})
tool := NewPubMedToolWith(helper)
out, err := tool.InvokableRun(context.Background(),
`{"query":"noresultsfound-zzz-9999","max_results":5}`)
out, err := tool.InvokableRun(context.Background(), `{"query":"missing"}`)
if err != nil {
t.Fatalf("InvokableRun: %v", err)
}
var env pubmedEnvelope
if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
if err := json.Unmarshal([]byte(out), &env); err != nil {
t.Fatalf("output is not valid JSON: %v (raw=%s)", err, out)
}
if env.Error != "" {
t.Errorf("Error = %q, want empty", env.Error)
}
if len(env.Results) != 0 {
t.Errorf("Results len = %d, want 0", len(env.Results))
if len(env.Results) != 0 || env.Error != "" {
t.Fatalf("env = %+v, want empty results and no error", env)
}
}
func TestPubMed_RequiresQuery(t *testing.T) {
func TestPubMed_InvokableRunRequiresQuery(t *testing.T) {
t.Parallel()
tool := NewPubMedTool()
@@ -221,11 +200,11 @@ func TestPubMed_RequiresQuery(t *testing.T) {
t.Fatal("expected error for empty query")
}
if !strings.Contains(err.Error(), "query") {
t.Errorf("err = %v, want to mention query", err)
t.Fatalf("err = %q, want query validation", err.Error())
}
}
func TestPubMed_Info(t *testing.T) {
func TestPubMed_InfoOnlyExposesQuery(t *testing.T) {
t.Parallel()
tool := NewPubMedTool()
@@ -234,9 +213,67 @@ func TestPubMed_Info(t *testing.T) {
t.Fatalf("Info: %v", err)
}
if info.Name != "pubmed" {
t.Errorf("Name = %q, want pubmed", info.Name)
t.Fatalf("Name = %q, want pubmed", info.Name)
}
if !strings.Contains(info.Desc, "PubMed") {
t.Errorf("Desc = %q, want to mention PubMed", info.Desc)
schema, err := info.ParamsOneOf.ToJSONSchema()
if err != nil {
t.Fatalf("ToJSONSchema: %v", err)
}
raw, err := json.Marshal(schema)
if err != nil {
t.Fatalf("marshal params schema: %v", err)
}
params := string(raw)
if !strings.Contains(params, `"query"`) {
t.Fatalf("schema missing query: %s", params)
}
if strings.Contains(params, `"top_n"`) || strings.Contains(params, `"email"`) {
t.Fatalf("schema leaked node params: %s", params)
}
if !strings.Contains(params, `"required":["query"]`) {
t.Fatalf("schema does not require query: %s", params)
}
}
func TestPubMed_BuildByNameAcceptsNodeParams(t *testing.T) {
t.Parallel()
built, err := BuildByName("pubmed", map[string]any{"top_n": 8, "email": "node@example.com"})
if err != nil {
t.Fatalf("BuildByName: %v", err)
}
tool, ok := built.(*PubMedTool)
if !ok {
t.Fatalf("built type = %T, want *PubMedTool", built)
}
if tool.defaults.TopN != 8 {
t.Fatalf("defaults.TopN = %d, want 8", tool.defaults.TopN)
}
if tool.defaults.Email != "node@example.com" {
t.Fatalf("defaults.Email = %q, want node@example.com", tool.defaults.Email)
}
}
func TestPubMed_MergeDefaults(t *testing.T) {
t.Parallel()
got := mergePubMedDefaults(
pubmedParams{Query: "configured query", TopN: 8, Email: "node@example.com"},
pubmedParams{Query: "runtime query"},
)
if got.Query != "runtime query" || got.TopN != 8 || got.Email != "node@example.com" {
t.Fatalf("merged params = %+v, want runtime query with node defaults", got)
}
}
func TestPubMed_BuildByNameRejectsInvalidTopN(t *testing.T) {
t.Parallel()
_, err := BuildByName("pubmed", map[string]any{"top_n": 0})
if err == nil {
t.Fatal("expected top_n validation error")
}
if !strings.Contains(err.Error(), "positive integer") {
t.Fatalf("err = %q, want positive integer validation", err.Error())
}
}

View File

@@ -45,7 +45,7 @@ var registry = map[string]Factory{
"google_scholar_search": buildGoogleScholarTool,
"jin10": noConfig("jin10", func() einotool.BaseTool { return NewJin10Tool() }),
"keenable": buildKeenableTool,
"pubmed": noConfig("pubmed", func() einotool.BaseTool { return NewPubMedTool() }),
"pubmed": buildPubMedTool,
"qweather": noConfig("qweather", func() einotool.BaseTool { return NewQWeatherTool() }),
"retrieval": noConfig("retrieval", func() einotool.BaseTool { return NewRetrievalTool() }),
"search_my_dataset": noConfig("search_my_dataset", func() einotool.BaseTool { return NewRetrievalTool() }),
@@ -202,6 +202,30 @@ func buildGoogleScholarTool(params map[string]any) (einotool.BaseTool, error) {
return NewGoogleScholarToolWithDefaults(nil, defaults), nil
}
func buildPubMedTool(params map[string]any) (einotool.BaseTool, error) {
defaults := pubmedParams{}
for key := range params {
switch key {
case "top_n", "email":
default:
return nil, fmt.Errorf("agent tool: tool %q does not accept node-level param %s", "pubmed", key)
}
}
if topN, ok := intParam(params, "top_n"); ok {
if topN <= 0 {
return nil, fmt.Errorf("agent tool: tool %q requires positive integer node-level param top_n", "pubmed")
}
defaults.TopN = topN
}
if email, ok := stringParam(params, "email"); ok {
if strings.TrimSpace(email) == "" {
return nil, fmt.Errorf("agent tool: tool %q requires non-empty string node-level param email", "pubmed")
}
defaults.Email = email
}
return NewPubMedToolWithDefaults(nil, defaults), nil
}
func buildKeenableTool(params map[string]any) (einotool.BaseTool, error) {
if len(params) == 0 {
return NewKeenableTool(), nil