From 156a11c56bbaf229111e05005d9b8f9257ccc262 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Sat, 11 Jul 2026 17:49:59 +0800 Subject: [PATCH] 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/...` image --------- Co-authored-by: Jin Hai --- internal/agent/component/fixture_stubs.go | 1 + .../agent/component/pubmed_component_test.go | 150 ++++++++ .../agent/component/universe_a_wrappers.go | 124 +++++++ internal/agent/component/verify_p1_test.go | 8 +- internal/agent/tool/pubmed.go | 332 +++++++++--------- internal/agent/tool/pubmed_test.go | 257 ++++++++------ internal/agent/tool/registry.go | 26 +- 7 files changed, 613 insertions(+), 285 deletions(-) create mode 100644 internal/agent/component/pubmed_component_test.go diff --git a/internal/agent/component/fixture_stubs.go b/internal/agent/component/fixture_stubs.go index d20ce08e89..ee1c0d8735 100644 --- a/internal/agent/component/fixture_stubs.go +++ b/internal/agent/component/fixture_stubs.go @@ -522,5 +522,6 @@ func init() { Register("DuckDuckGo", newDuckDuckGoComponent) Register("Google", newGoogleComponent) Register("GoogleScholar", newGoogleScholarComponent) + Register("PubMed", newPubMedComponent) Register("YahooFinance", newYahooFinanceComponent) } diff --git a/internal/agent/component/pubmed_component_test.go b/internal/agent/component/pubmed_component_test.go new file mode 100644 index 0000000000..fdad06faf9 --- /dev/null +++ b/internal/agent/component/pubmed_component_test.go @@ -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) + } +} diff --git a/internal/agent/component/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index 8a3f907132..f7414618f1 100644 --- a/internal/agent/component/universe_a_wrappers.go +++ b/internal/agent/component/universe_a_wrappers.go @@ -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) diff --git a/internal/agent/component/verify_p1_test.go b/internal/agent/component/verify_p1_test.go index a0fabb24a5..f8ea31b7f3 100644 --- a/internal/agent/component/verify_p1_test.go +++ b/internal/agent/component/verify_p1_test.go @@ -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) diff --git a/internal/agent/tool/pubmed.go b/internal/agent/tool/pubmed.go index bde92c4d27..5ae3cbcb97 100644 --- a/internal/agent/tool/pubmed.go +++ b/internal/agent/tool/pubmed.go @@ -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), ¶ms); 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, ", ") } diff --git a/internal/agent/tool/pubmed_test.go b/internal/agent/tool/pubmed_test.go index ad9ca05c17..efb8102b3d 100644 --- a/internal/agent/tool/pubmed_test.go +++ b/internal/agent/tool/pubmed_test.go @@ -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(` + + + 12345678 +
+ Deep learning for retrieval augmented generation + A short abstract. + + Nature Machine Intelligence + 102 + + 101-110 + + KhanFurqan + SmithJane + +
+ + + + 10.1000/example.doi + + + +`)) + 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()) } } diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index 4386cc4438..89636bdc84 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -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