ragflow/internal/agent/tool/arxiv_test.go

//
//  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 tool

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"
)

func TestArxiv_BuildURL(t *testing.T) {
	t.Parallel()

	got := buildArxivURL("transformer", 3, "lastUpdatedDate")
	u, err := url.Parse(got)
	if err != nil {
		t.Fatalf("url.Parse(%q): %v", got, err)
	}
	if u.Host != "export.arxiv.org" {
		t.Errorf("host = %q, want export.arxiv.org", u.Host)
	}
	if u.Path != "/api/query" {
		t.Errorf("path = %q, want /api/query", u.Path)
	}
	q := u.Query()
	if q.Get("search_query") != "all:transformer" {
		t.Errorf("search_query = %q, want all:transformer", q.Get("search_query"))
	}
	if q.Get("max_results") != "3" {
		t.Errorf("max_results = %q, want 3", q.Get("max_results"))
	}
	if q.Get("sortBy") != "lastUpdatedDate" {
		t.Errorf("sortBy = %q, want lastUpdatedDate", q.Get("sortBy"))
	}
}

func TestArxiv_ParseAtomEntry(t *testing.T) {
	t.Parallel()

	const canned = `<?xml version="1.0" encoding="UTF-8"?>
<feed xmlns="http://www.w3.org/2005/Atom" xmlns:arxiv="http://arxiv.org/schemas/atom">
  <title>arXiv Query: all:rag</title>
  <entry>
    <id>http://arxiv.org/abs/2501.12345v1</id>
    <title>Retrieval-Augmented Generation for
      Open-Domain Question Answering</title>
    <summary>  We present a method
      combining retrieval and generation.  </summary>
    <author><name>Alice Liddell</name></author>
    <author><name>Bob Builder</name></author>
    <link href="http://arxiv.org/abs/2501.12345v1" rel="alternate" type="text/html"/>
    <link href="http://arxiv.org/pdf/2501.12345v1" rel="related" type="application/pdf"/>
  </entry>
  <entry>
    <id>http://arxiv.org/abs/2409.99999v2</id>
    <title>Single-Author Paper</title>
    <summary>Brief.</summary>
    <author><name>Carol Danvers</name></author>
    <!-- no pdf link: pickArxivPDF must fall back to abs id derived pdf -->
  </entry>
</feed>`

	results, err := parseArxivAtom([]byte(canned))
	if err != nil {
		t.Fatalf("parseArxivAtom: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("len(results) = %d, want 2", len(results))
	}

	// First entry: explicit pdf link
	r0 := results[0]
	if !strings.Contains(r0.Title, "Retrieval-Augmented Generation") {
		t.Errorf("r0.Title = %q, want to contain 'Retrieval-Augmented Generation'", r0.Title)
	}
	if !strings.HasPrefix(r0.Title, "Retrieval-") {
		// whitespace normalized
		if strings.Contains(r0.Title, "\n") {
			t.Errorf("r0.Title contains raw newline: %q", r0.Title)
		}
	}
	if len(r0.Authors) != 2 {
		t.Errorf("r0.Authors len = %d, want 2", len(r0.Authors))
	}
	if r0.Authors[0] != "Alice Liddell" || r0.Authors[1] != "Bob Builder" {
		t.Errorf("r0.Authors = %v, want [Alice Liddell Bob Builder]", r0.Authors)
	}
	if r0.PDFURL != "http://arxiv.org/pdf/2501.12345v1" {
		t.Errorf("r0.PDFURL = %q, want http://arxiv.org/pdf/2501.12345v1", r0.PDFURL)
	}
	if r0.EntryID != "http://arxiv.org/abs/2501.12345v1" {
		t.Errorf("r0.EntryID = %q, want http://arxiv.org/abs/2501.12345v1", r0.EntryID)
	}
	if !strings.HasPrefix(r0.Summary, "We present") {
		t.Errorf("r0.Summary = %q, want to start with 'We present' (whitespace normalized)", r0.Summary)
	}

	// Second entry: fallback pdf derived from abs id
	r1 := results[1]
	if r1.PDFURL != "http://arxiv.org/pdf/2409.99999v2" {
		t.Errorf("r1.PDFURL = %q, want http://arxiv.org/pdf/2409.99999v2 (derived)", r1.PDFURL)
	}
	if len(r1.Authors) != 1 || r1.Authors[0] != "Carol Danvers" {
		t.Errorf("r1.Authors = %v, want [Carol Danvers]", r1.Authors)
	}
}

func TestArxiv_Info(t *testing.T) {
	t.Parallel()

	tool := NewArxivTool()
	info, err := tool.Info(context.Background())
	if err != nil {
		t.Fatalf("Info: %v", err)
	}
	if info.Name != "arxiv_search" {
		t.Errorf("Name = %q, want arxiv_search", info.Name)
	}
	if !strings.Contains(info.Desc, "arXiv") {
		t.Errorf("Desc = %q, want to mention arXiv", info.Desc)
	}
	params, err := json.Marshal(info.ParamsOneOf)
	if err != nil {
		t.Fatalf("marshal params: %v", err)
	}
	if strings.Contains(string(params), "top_n") || strings.Contains(string(params), "sort_by") || strings.Contains(string(params), "max_results") {
		t.Errorf("schema must only expose query: %s", params)
	}
}

func TestArxiv_EmptyQuery(t *testing.T) {
	t.Parallel()

	tool := NewArxivTool()
	out, err := tool.InvokableRun(context.Background(), `{"query":""}`)
	if err != nil {
		t.Fatalf("InvokableRun(empty): %v", err)
	}
	var envelope arxivEnvelope
	if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 {
		t.Fatalf("empty result = %s / %v", out, err)
	}
}

func TestArxiv_FullRoundtrip(t *testing.T) {
	t.Parallel()

	const canned = `<?xml version="1.0" encoding="UTF-8"?>
<feed xmlns="http://www.w3.org/2005/Atom">
  <title>arXiv Query</title>
  <entry>
    <id>http://arxiv.org/abs/2501.12345v1</id>
    <title>Test Paper</title>
    <summary>Summary.</summary>
    <author><name>Author One</name></author>
    <link href="http://arxiv.org/pdf/2501.12345v1" rel="related" type="application/pdf"/>
  </entry>
</feed>`

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/atom+xml")
		_, _ = w.Write([]byte(canned))
	}))
	defer srv.Close()

	// rewriteHostTransport points the hard-coded export.arxiv.org at the
	// test server.
	helper := NewHTTPHelper().WithClient(&http.Client{
		Transport: rewriteHostTransport(srv.URL),
	})
	tool := NewArxivToolWith(helper)
	out, err := tool.InvokableRun(context.Background(), `{"query":"rag"}`)
	if err != nil {
		t.Fatalf("InvokableRun: %v", err)
	}

	var env arxivEnvelope
	if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
		t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
	}
	if env.Error != "" {
		t.Errorf("Error = %q, want empty", env.Error)
	}
	if len(env.Results) != 1 {
		t.Fatalf("Results len = %d, want 1", len(env.Results))
	}
	if env.Results[0].Title != "Test Paper" {
		t.Errorf("Title = %q, want Test Paper", env.Results[0].Title)
	}
	if env.Results[0].PDFURL != "http://arxiv.org/pdf/2501.12345v1" {
		t.Errorf("PDFURL = %q, want http://arxiv.org/pdf/2501.12345v1", env.Results[0].PDFURL)
	}
}

func TestArxiv_ComponentReferencesAndDefaults(t *testing.T) {
	t.Parallel()

	built, err := BuildByName("arxiv", map[string]any{
		"top_n":   float64(7),
		"sort_by": "relevance",
		"outputs": map[string]any{"json": map[string]any{}},
	})
	if err != nil {
		t.Fatalf("BuildByName: %v", err)
	}
	arxiv := built.(*ArxivTool)
	if arxiv.defaults.TopN != 7 || arxiv.defaults.SortBy != "relevance" {
		t.Fatalf("defaults = %+v", arxiv.defaults)
	}
	spec := arxiv.ComponentSpec()
	if query, ok := spec.InputForm["query"].(map[string]any); !ok || query["type"] != "line" {
		t.Fatalf("query input form = %#v", spec.InputForm["query"])
	}
	envelope := map[string]any{"results": []any{map[string]any{
		"title": "Paper", "summary": "Paper summary.", "pdf_url": "https://arxiv.org/pdf/1", "entry_id": "kept",
	}}}
	chunks, docAggs := arxiv.BuildReferences(context.Background(), envelope)
	if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["content"] != "Paper summary." {
		t.Fatalf("references = %#v / %#v", chunks, docAggs)
	}
	outputs := arxiv.BuildComponentOutputs(envelope)
	results := outputs["json"].([]any)
	if results[0].(map[string]any)["entry_id"] != "kept" {
		t.Fatalf("json output = %#v", results)
	}
	if !strings.Contains(outputs["formalized_content"].(string), "Paper summary.") {
		t.Fatalf("formalized_content = %q", outputs["formalized_content"])
	}
	if _, exists := envelope["chunks"]; exists {
		t.Fatalf("output conversion mutated envelope: %#v", envelope)
	}
}