mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-07 03:48:44 +08:00
244 lines
7.4 KiB
Go
244 lines
7.4 KiB
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 (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/cloudwego/eino/components/tool"
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
const bgptToolName = "bgpt_search"
|
|
|
|
const bgptToolDescription = "Search scientific papers via BGPT (bgpt.pro) and return structured evidence " +
|
|
"extracted from full-text studies: methods, sample sizes, results, limitations, conflicts of interest, " +
|
|
"data/code availability, study blind spots, quality scores, and falsification prompts."
|
|
|
|
// bgptEndpoint is the BGPT search URL. Exposed as a package var for tests.
|
|
var bgptEndpoint = "https://bgpt.pro/api/mcp-search"
|
|
|
|
// bgptParams is the JSON shape the model sends into InvokableRun.
|
|
type bgptParams struct {
|
|
Query string `json:"query"`
|
|
MaxResults int `json:"num_results"`
|
|
APIKey string `json:"api_key,omitempty"`
|
|
DaysBack int `json:"days_back,omitempty"`
|
|
}
|
|
|
|
// bgptResult is one paper in the result list.
|
|
type bgptResult struct {
|
|
Title string `json:"title"`
|
|
Authors string `json:"authors"`
|
|
Journal string `json:"journal"`
|
|
Year string `json:"year"`
|
|
DOI string `json:"doi"`
|
|
URL string `json:"url"`
|
|
Abstract string `json:"abstract"`
|
|
Methods string `json:"methods"`
|
|
SampleSize string `json:"sample_size"`
|
|
Results string `json:"results"`
|
|
Limitations string `json:"limitations"`
|
|
ConflictOfInterest string `json:"conflict_of_interest"`
|
|
DataAvailability string `json:"data_availability"`
|
|
BlindSpots string `json:"blind_spots"`
|
|
Falsify string `json:"falsify"`
|
|
}
|
|
|
|
// bgptEnv is what the model sees.
|
|
type bgptEnv struct {
|
|
Results []bgptResult `json:"results"`
|
|
Error string `json:"_ERROR,omitempty"`
|
|
}
|
|
|
|
// bgptResponse is the upstream BGPT API envelope.
|
|
type bgptResponse struct {
|
|
Results []map[string]interface{} `json:"results"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
// BGPTTool searches scientific papers via bgpt.pro.
|
|
type BGPTTool struct {
|
|
helper *HTTPHelper
|
|
}
|
|
|
|
// NewBGPTTool returns a BGPTTool using the default HTTPHelper.
|
|
func NewBGPTTool() *BGPTTool {
|
|
return NewBGPTToolWith(NewHTTPHelper())
|
|
}
|
|
|
|
// NewBGPTToolWith returns a BGPTTool with the given helper.
|
|
func NewBGPTToolWith(h *HTTPHelper) *BGPTTool {
|
|
if h == nil {
|
|
h = NewHTTPHelper()
|
|
}
|
|
return &BGPTTool{helper: h}
|
|
}
|
|
|
|
// Info returns the tool's metadata for the chat model.
|
|
func (b *BGPTTool) Info(_ context.Context) (*schema.ToolInfo, error) {
|
|
return &schema.ToolInfo{
|
|
Name: bgptToolName,
|
|
Desc: bgptToolDescription,
|
|
ParamsOneOf: schema.NewParamsOneOfByParams(map[string]*schema.ParameterInfo{
|
|
"query": {
|
|
Type: schema.String,
|
|
Desc: "Natural-language scientific search query.",
|
|
Required: true,
|
|
},
|
|
"num_results": {
|
|
Type: schema.Integer,
|
|
Desc: "Maximum number of results. Defaults to 10.",
|
|
Required: false,
|
|
},
|
|
"api_key": {
|
|
Type: schema.String,
|
|
Desc: "Optional BGPT API key. Leave blank for the free tier.",
|
|
Required: false,
|
|
},
|
|
"days_back": {
|
|
Type: schema.Integer,
|
|
Desc: "Optional recency filter (e.g. 365 for last year).",
|
|
Required: false,
|
|
},
|
|
}),
|
|
}, nil
|
|
}
|
|
|
|
// InvokableRun performs the BGPT search.
|
|
func (b *BGPTTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
|
|
var p bgptParams
|
|
if err := json.Unmarshal([]byte(argsJSON), &p); err != nil {
|
|
return bgptErrJSON(fmt.Errorf("bgpt: parse arguments: %w", err)),
|
|
fmt.Errorf("bgpt: parse arguments: %w", err)
|
|
}
|
|
if strings.TrimSpace(p.Query) == "" {
|
|
return bgptErrJSON(fmt.Errorf("query is required")),
|
|
fmt.Errorf("bgpt: query is required")
|
|
}
|
|
if p.MaxResults <= 0 {
|
|
p.MaxResults = 10
|
|
}
|
|
|
|
reqBody := map[string]interface{}{
|
|
"query": strings.TrimSpace(p.Query),
|
|
"num_results": p.MaxResults,
|
|
}
|
|
if p.APIKey != "" {
|
|
reqBody["api_key"] = p.APIKey
|
|
}
|
|
if p.DaysBack > 0 {
|
|
reqBody["days_back"] = p.DaysBack
|
|
}
|
|
|
|
bodyBytes, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return bgptErrJSON(err), err
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, bgptEndpoint, bytes.NewReader(bodyBytes))
|
|
if err != nil {
|
|
return bgptErrJSON(err), err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json")
|
|
|
|
client := &http.Client{Timeout: 25 * time.Second}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return bgptErrJSON(fmt.Errorf("bgpt: request failed: %w", err)),
|
|
fmt.Errorf("bgpt: request failed: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return bgptErrJSON(fmt.Errorf("bgpt: read response: %w", err)),
|
|
fmt.Errorf("bgpt: read response: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return bgptErrJSON(fmt.Errorf("bgpt: upstream returned %d: %s", resp.StatusCode, string(body))),
|
|
fmt.Errorf("bgpt: upstream returned %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
var raw bgptResponse
|
|
if err := json.Unmarshal(body, &raw); err != nil {
|
|
return bgptErrJSON(fmt.Errorf("bgpt: parse response: %w", err)),
|
|
fmt.Errorf("bgpt: parse response: %w", err)
|
|
}
|
|
if strings.TrimSpace(raw.Error) != "" {
|
|
err := fmt.Errorf("bgpt: %s", strings.TrimSpace(raw.Error))
|
|
return bgptErrJSON(err), err
|
|
}
|
|
|
|
results := make([]bgptResult, 0, len(raw.Results))
|
|
for _, r := range raw.Results {
|
|
results = append(results, bgptResult{
|
|
Title: strVal(r["title"]),
|
|
Authors: strVal(r["authors"]),
|
|
Journal: strVal(r["journal"]),
|
|
Year: strVal(r["year"]),
|
|
DOI: strVal(r["doi"]),
|
|
URL: strVal(r["url"]),
|
|
Abstract: strVal(r["abstract"]),
|
|
Methods: firstStr(r, "methods_and_experimental_techniques", "methods"),
|
|
SampleSize: firstStr(r, "sample_size_and_population_characteristics", "sample_size_and_population"),
|
|
Results: firstStr(r, "results_and_conclusions", "results"),
|
|
Limitations: firstStr(r, "paper_limitations_and_biases", "limitations"),
|
|
ConflictOfInterest: firstStr(r, "conflict_of_interest_statements", "conflict_of_interest"),
|
|
DataAvailability: firstStr(r, "data_availability_statements", "data_availability"),
|
|
BlindSpots: strVal(r["study_blindspots"]),
|
|
Falsify: strVal(r["how_to_falsify"]),
|
|
})
|
|
}
|
|
|
|
env := bgptEnv{Results: results}
|
|
out, _ := json.Marshal(env)
|
|
return string(out), nil
|
|
}
|
|
|
|
func bgptErrJSON(err error) string {
|
|
env := bgptEnv{Error: err.Error()}
|
|
b, _ := json.Marshal(env)
|
|
return string(b)
|
|
}
|
|
|
|
func strVal(v interface{}) string {
|
|
if v == nil {
|
|
return ""
|
|
}
|
|
s, _ := v.(string)
|
|
return s
|
|
}
|
|
|
|
func firstStr(m map[string]interface{}, keys ...string) string {
|
|
for _, k := range keys {
|
|
if s := strVal(m[k]); s != "" {
|
|
return s
|
|
}
|
|
}
|
|
return ""
|
|
}
|