mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-12 06:35:44 +08:00
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:
@@ -522,5 +522,6 @@ func init() {
|
||||
Register("DuckDuckGo", newDuckDuckGoComponent)
|
||||
Register("Google", newGoogleComponent)
|
||||
Register("GoogleScholar", newGoogleScholarComponent)
|
||||
Register("PubMed", newPubMedComponent)
|
||||
Register("YahooFinance", newYahooFinanceComponent)
|
||||
}
|
||||
|
||||
150
internal/agent/component/pubmed_component_test.go
Normal file
150
internal/agent/component/pubmed_component_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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, ", ")
|
||||
}
|
||||
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user