feat: integrate Yahoo finance stock price tool
This commit is contained in:
parent
0d16525fab
commit
02e07ba9c9
3 changed files with 283 additions and 0 deletions
|
|
@ -104,6 +104,9 @@ func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msg
|
|||
})
|
||||
registry.Register(messageTool)
|
||||
|
||||
// Stock price tool
|
||||
registry.Register(tools.NewStockTool())
|
||||
|
||||
return registry
|
||||
}
|
||||
|
||||
|
|
|
|||
103
pkg/tools/stock.go
Normal file
103
pkg/tools/stock.go
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
// StockTool allows you to retrieve the current price of a stock.
|
||||
type StockTool struct {
|
||||
baseURL string
|
||||
}
|
||||
|
||||
func NewStockTool() *StockTool {
|
||||
return &StockTool{
|
||||
baseURL: "https://query1.finance.yahoo.com", // Default URL
|
||||
}
|
||||
}
|
||||
|
||||
func (t *StockTool) Name() string {
|
||||
return "get_stock_price"
|
||||
}
|
||||
|
||||
func (t *StockTool) Description() string {
|
||||
return "Get the current stock price and currency from Yahoo Finance using the ticker symbol (e.g., AAPL, TSLA, MSFT)."
|
||||
}
|
||||
|
||||
func (t *StockTool) Parameters() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"ticker": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "The stock symbol (ticker) of the company, e.g., AAPL for Apple.",
|
||||
},
|
||||
},
|
||||
"required": []string{"ticker"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *StockTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
|
||||
ticker, ok := args["ticker"].(string)
|
||||
if !ok || ticker == "" {
|
||||
return ErrorResult("The parameter 'ticker' is mandatory")
|
||||
}
|
||||
|
||||
url := fmt.Sprintf("%s/v8/finance/chart/%s?interval=1d&range=1d", t.baseURL, ticker)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
|
||||
if err != nil {
|
||||
return ErrorResult(fmt.Sprintf("Error creating HTTP request: %v", err))
|
||||
}
|
||||
|
||||
// Yahoo Finance blocks requests without User-Agent
|
||||
req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36")
|
||||
|
||||
client := &http.Client{Timeout: 10 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return ErrorResult(fmt.Sprintf("Error connecting to Yahoo Finance: %v", err))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return ErrorResult(fmt.Sprintf("Yahoo Finance returned an error status: %d", resp.StatusCode))
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Chart struct {
|
||||
Result []struct {
|
||||
Meta struct {
|
||||
Currency string `json:"currency"`
|
||||
Symbol string `json:"symbol"`
|
||||
RegularMarketPrice float64 `json:"regularMarketPrice"`
|
||||
} `json:"meta"`
|
||||
} `json:"result"`
|
||||
Error *struct {
|
||||
Description string `json:"description"`
|
||||
} `json:"error"`
|
||||
} `json:"chart"`
|
||||
}
|
||||
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return ErrorResult(fmt.Sprintf("Error in reading JSON data: %v", err))
|
||||
}
|
||||
|
||||
// API-side error checking
|
||||
if result.Chart.Error != nil {
|
||||
return ErrorResult(fmt.Sprintf("Yahoo API Error: %s", result.Chart.Error.Description))
|
||||
}
|
||||
if len(result.Chart.Result) == 0 {
|
||||
return ErrorResult(fmt.Sprintf("No data found for ticker: %s", ticker))
|
||||
}
|
||||
|
||||
// Extract data
|
||||
meta := result.Chart.Result[0].Meta
|
||||
|
||||
output := fmt.Sprintf("Yahoo Finance data for %s: Current price %.2f %s.", meta.Symbol, meta.RegularMarketPrice, meta.Currency)
|
||||
|
||||
return UserResult(output)
|
||||
}
|
||||
177
pkg/tools/stock_test.go
Normal file
177
pkg/tools/stock_test.go
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestStockTool_NameAndDescription(t *testing.T) {
|
||||
tool := NewStockTool()
|
||||
|
||||
if tool.Name() != "get_stock_price" {
|
||||
t.Errorf("Expected name 'get_stock_price', got '%s'", tool.Name())
|
||||
}
|
||||
|
||||
if tool.Description() == "" {
|
||||
t.Error("Expected description to not be empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStockTool_Parameters(t *testing.T) {
|
||||
tool := NewStockTool()
|
||||
params := tool.Parameters()
|
||||
|
||||
if params["type"] != "object" {
|
||||
t.Errorf("Expected type 'object', got '%v'", params["type"])
|
||||
}
|
||||
|
||||
req, ok := params["required"].([]string)
|
||||
if !ok || len(req) == 0 || req[0] != "ticker" {
|
||||
t.Error("Expected 'ticker' to be in required parameters")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStockTool_MissingTicker(t *testing.T) {
|
||||
tool := NewStockTool()
|
||||
ctx := context.Background()
|
||||
|
||||
// We pass empty args
|
||||
args := map[string]interface{}{}
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
if !result.IsError {
|
||||
t.Error("Expected error when 'ticker' is missing")
|
||||
}
|
||||
|
||||
if !strings.Contains(result.ForLLM, "mandatory") {
|
||||
t.Errorf("Expected error message about missing ticker, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStockTool_Success(t *testing.T) {
|
||||
mockResponse := `{
|
||||
"chart": {
|
||||
"result": [
|
||||
{
|
||||
"meta": {
|
||||
"currency": "USD",
|
||||
"symbol": "AAPL",
|
||||
"regularMarketPrice": 150.50
|
||||
}
|
||||
}
|
||||
],
|
||||
"error": null
|
||||
}
|
||||
}`
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !strings.Contains(r.URL.Path, "AAPL") {
|
||||
t.Errorf("Expected request to contain 'AAPL', got %s", r.URL.Path)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockResponse))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := &StockTool{baseURL: server.URL}
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"ticker": "AAPL",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
if result.IsError {
|
||||
t.Errorf("Expected success, got error: %s", result.ForLLM)
|
||||
}
|
||||
|
||||
expectedOutput := "Yahoo Finance data for AAPL: Current price 150.50 USD."
|
||||
if !strings.Contains(result.ForUser, expectedOutput) {
|
||||
t.Errorf("Expected ForUser to contain '%s', got: %s", expectedOutput, result.ForUser)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStockTool_HttpError(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := &StockTool{baseURL: server.URL}
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"ticker": "INVALID",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
if !result.IsError {
|
||||
t.Error("Expected error for 404 response")
|
||||
}
|
||||
|
||||
if !strings.Contains(result.ForLLM, "404") {
|
||||
t.Errorf("Expected error to contain '404', got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStockTool_ApiError(t *testing.T) {
|
||||
mockResponse := `{
|
||||
"chart": {
|
||||
"result": null,
|
||||
"error": {
|
||||
"description": "Not Found"
|
||||
}
|
||||
}
|
||||
}`
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockResponse))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := &StockTool{baseURL: server.URL}
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"ticker": "UNKNOWN_TICKER",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
if !result.IsError {
|
||||
t.Error("Expected error when API returns an error description")
|
||||
}
|
||||
|
||||
if !strings.Contains(result.ForLLM, "Not Found") {
|
||||
t.Errorf("Expected error message to contain 'Not Found', got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStockTool_InvalidJson(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(`{ "broken json": `))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
tool := &StockTool{baseURL: server.URL}
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"ticker": "AAPL",
|
||||
}
|
||||
|
||||
result := tool.Execute(ctx, args)
|
||||
|
||||
if !result.IsError {
|
||||
t.Error("Expected error when JSON is invalid")
|
||||
}
|
||||
|
||||
if !strings.Contains(result.ForLLM, "Error in reading JSON data") {
|
||||
t.Errorf("Expected JSON parse error, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue