diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index d3afa298e..db32dd990 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 } diff --git a/pkg/tools/stock.go b/pkg/tools/stock.go new file mode 100644 index 000000000..9806bffc5 --- /dev/null +++ b/pkg/tools/stock.go @@ -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) +} diff --git a/pkg/tools/stock_test.go b/pkg/tools/stock_test.go new file mode 100644 index 000000000..65b63b474 --- /dev/null +++ b/pkg/tools/stock_test.go @@ -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) + } +}