commit 4c0797a45dce497362bedd0ab08dc547382c4860 Author: Regis David Souza Mesquita Date: Sat Dec 7 00:54:31 2024 +0000 initial diff --git a/adapters/serper.go b/adapters/serper.go new file mode 100644 index 0000000..31f961c --- /dev/null +++ b/adapters/serper.go @@ -0,0 +1,76 @@ +package adapters + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" +) + +type SerperProvider struct { + apiKey string +} + +type serperRequest struct { + Query string `json:"q"` +} + +type serperResponse struct { + Organic []struct { + Link string `json:"link"` + Description string `json:"snippet"` + } `json:"organic"` +} + +func NewSerperProvider() (*SerperProvider, error) { + apiKey := os.Getenv("SERPER_API_KEY") + if apiKey == "" { + return nil, fmt.Errorf("SERPER_API_KEY environment variable is not set") + } + return &SerperProvider{apiKey: apiKey}, nil +} + +func (s *SerperProvider) Search(query string, options map[string]string) (*SearchResponse, error) { + payload := serperRequest{Query: query} + jsonData, err := json.Marshal(payload) + if err != nil { + return nil, err + } + + req, err := http.NewRequest("POST", "https://google.serper.dev/search", bytes.NewBuffer(jsonData)) + if err != nil { + return nil, err + } + + req.Header.Set("X-API-KEY", s.apiKey) + req.Header.Set("Content-Type", "application/json") + + client := &http.Client{} + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API error: %s - %s", resp.Status, string(body)) + } + + var serperResp serperResponse + if err := json.NewDecoder(resp.Body).Decode(&serperResp); err != nil { + return nil, err + } + + results := make([]SearchResult, 0, len(serperResp.Organic)) + for _, result := range serperResp.Organic { + results = append(results, SearchResult{ + URL: result.Link, + Content: result.Description, + }) + } + + return &SearchResponse{Results: results}, nil +} diff --git a/adapters/serper_test.go b/adapters/serper_test.go new file mode 100644 index 0000000..0b1dabe --- /dev/null +++ b/adapters/serper_test.go @@ -0,0 +1,67 @@ +package adapters + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "testing" +) + +func TestSerperProvider_Search(t *testing.T) { + // Mock server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + if r.Header.Get("X-API-KEY") != "test-key" { + t.Errorf("Expected API key header, got %s", r.Header.Get("X-API-KEY")) + } + + // Return mock response + response := map[string]interface{}{ + "organic": []map[string]interface{}{ + { + "link": "https://example.com", + "snippet": "Example content", + }, + }, + } + json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + // Set test environment + os.Setenv("SERPER_API_KEY", "test-key") + defer os.Unsetenv("SERPER_API_KEY") + + provider, err := NewSerperProvider() + if err != nil { + t.Fatalf("Failed to create provider: %v", err) + } + + // Override API URL for testing + originalURL := "https://google.serper.dev/search" + http.DefaultClient = server.Client() + + results, err := provider.Search("test query", nil) + if err != nil { + t.Fatalf("Search failed: %v", err) + } + + if len(results.Results) != 1 { + t.Errorf("Expected 1 result, got %d", len(results.Results)) + } + + if results.Results[0].URL != "https://example.com" { + t.Errorf("Expected URL https://example.com, got %s", results.Results[0].URL) + } + + if results.Results[0].Content != "Example content" { + t.Errorf("Expected content 'Example content', got %s", results.Results[0].Content) + } + + // Restore original URL + _ = originalURL +} diff --git a/adapters/tavily.go b/adapters/tavily.go new file mode 100644 index 0000000..8d8c01b --- /dev/null +++ b/adapters/tavily.go @@ -0,0 +1,130 @@ +package adapters + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" +) + +type TavilyProvider struct { + apiKey string +} + +func NewTavilyProvider() (*TavilyProvider, error) { + apiKey := os.Getenv("TAVILY_API_KEY") + if apiKey == "" { + return nil, fmt.Errorf("TAVILY_API_KEY environment variable is not set") + } + return &TavilyProvider{apiKey: apiKey}, nil +} + +func (t *TavilyProvider) Extract(urls []string) (*ExtractResponse, error) { + params := struct { + URLs []string `json:"urls"` + ApiKey string `json:"api_key"` + }{ + URLs: urls, + ApiKey: t.apiKey, + } + + jsonData, err := json.Marshal(params) + if err != nil { + return nil, err + } + + req, err := http.NewRequest("POST", "https://api.tavily.com/extract", bytes.NewBuffer(jsonData)) + if err != nil { + return nil, err + } + + req.Header.Set("Content-Type", "application/json") + + client := &http.Client{} + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API error: %s - %s", resp.Status, string(body)) + } + + var result ExtractResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, err + } + + return &result, nil +} + +func (t *TavilyProvider) Search(query string, options map[string]string) (*SearchResponse, error) { + params := struct { + Query string `json:"query"` + ApiKey string `json:"api_key"` + SearchDepth string `json:"search_depth"` + IncludeDomains []string `json:"include_domains"` + ExcludeDomains []string `json:"exclude_domains"` + }{ + Query: query, + ApiKey: t.apiKey, + SearchDepth: "basic", + IncludeDomains: []string{}, + ExcludeDomains: []string{}, + } + + // Apply options if provided + if depth, ok := options["depth"]; ok { + params.SearchDepth = depth + } + + jsonData, err := json.Marshal(params) + if err != nil { + return nil, err + } + + req, err := http.NewRequest("POST", "https://api.tavily.com/search", bytes.NewBuffer(jsonData)) + if err != nil { + return nil, err + } + + req.Header.Set("Content-Type", "application/json") + + client := &http.Client{} + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API error: %s - %s", resp.Status, string(body)) + } + + var result struct { + Results []struct { + URL string `json:"url"` + Content string `json:"content"` + } `json:"results"` + } + + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return nil, err + } + + // Convert to common SearchResponse format + searchResults := make([]SearchResult, 0, len(result.Results)) + for _, r := range result.Results { + searchResults = append(searchResults, SearchResult{ + URL: r.URL, + Content: r.Content, + }) + } + + return &SearchResponse{Results: searchResults}, nil +} diff --git a/adapters/tavily_test.go b/adapters/tavily_test.go new file mode 100644 index 0000000..eb7730f --- /dev/null +++ b/adapters/tavily_test.go @@ -0,0 +1,116 @@ +package adapters + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "testing" +) + +func TestTavilyProvider_Search(t *testing.T) { + // Mock server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + + // Verify request body contains API key + var reqBody map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil { + t.Fatalf("Failed to decode request body: %v", err) + } + if reqBody["api_key"] != "test-key" { + t.Errorf("Expected API key in body, got %v", reqBody["api_key"]) + } + + // Return mock response + response := map[string]interface{}{ + "results": []map[string]interface{}{ + { + "url": "https://example.com", + "content": "Example content", + }, + }, + } + json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + // Set test environment + os.Setenv("TAVILY_API_KEY", "test-key") + defer os.Unsetenv("TAVILY_API_KEY") + + provider, err := NewTavilyProvider() + if err != nil { + t.Fatalf("Failed to create provider: %v", err) + } + + // Override API URL for testing + http.DefaultClient = server.Client() + + results, err := provider.Search("test query", map[string]string{"depth": "advanced"}) + if err != nil { + t.Fatalf("Search failed: %v", err) + } + + if len(results.Results) != 1 { + t.Errorf("Expected 1 result, got %d", len(results.Results)) + } + + if results.Results[0].URL != "https://example.com" { + t.Errorf("Expected URL https://example.com, got %s", results.Results[0].URL) + } + + if results.Results[0].Content != "Example content" { + t.Errorf("Expected content 'Example content', got %s", results.Results[0].Content) + } +} + +func TestTavilyProvider_Extract(t *testing.T) { + // Mock server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Verify request + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + + // Return mock response + response := ExtractResponse{ + Results: []ExtractResult{ + { + URL: "https://example.com", + RawContent: "Example content", + }, + }, + } + json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + // Set test environment + os.Setenv("TAVILY_API_KEY", "test-key") + defer os.Unsetenv("TAVILY_API_KEY") + + provider, err := NewTavilyProvider() + if err != nil { + t.Fatalf("Failed to create provider: %v", err) + } + + // Override API URL for testing + http.DefaultClient = server.Client() + + results, err := provider.Extract([]string{"https://example.com"}) + if err != nil { + t.Fatalf("Extract failed: %v", err) + } + + if len(results.Results) != 1 { + t.Errorf("Expected 1 result, got %d", len(results.Results)) + } + + if results.Results[0].URL != "https://example.com" { + t.Errorf("Expected URL https://example.com, got %s", results.Results[0].URL) + } +} diff --git a/adapters/types.go b/adapters/types.go new file mode 100644 index 0000000..d357dca --- /dev/null +++ b/adapters/types.go @@ -0,0 +1,36 @@ +package adapters + +// Common interfaces +type SearchProvider interface { + Search(query string, options map[string]string) (*SearchResponse, error) +} + +type ExtractProvider interface { + Extract(urls []string) (*ExtractResponse, error) +} + +// Response types +type SearchResult struct { + URL string `json:"url"` + Content string `json:"content"` +} + +type SearchResponse struct { + Results []SearchResult `json:"results"` +} + +type ExtractResult struct { + URL string `json:"url"` + RawContent string `json:"raw_content"` +} + +type FailedResult struct { + URL string `json:"url"` + Error string `json:"error"` +} + +type ExtractResponse struct { + Results []ExtractResult `json:"results"` + FailedResults []FailedResult `json:"failed_results"` + ResponseTime float64 `json:"response_time"` +} diff --git a/cmd/search/main.go b/cmd/search/main.go new file mode 100644 index 0000000..cf50b13 --- /dev/null +++ b/cmd/search/main.go @@ -0,0 +1,92 @@ +package main + +import ( + "encoding/json" + "flag" + "fmt" + "os" + "strings" + + "search-cli/internal/adapters" + "search-cli/internal/formatter" +) + +func main() { + jsonOutput := flag.Bool("json", false, "Output in JSON format") + tavilyMode := flag.Bool("tavily", false, "Use Tavily as provider") + serperMode := flag.Bool("serper", false, "Use Serper as provider (default)") + extractMode := flag.Bool("extract", false, "Extract content from URLs (Tavily only)") + depth := flag.String("depth", "basic", "Search depth (basic or advanced) - only for Tavily") + flag.Parse() + + if flag.NArg() < 1 { + fmt.Fprintln(os.Stderr, "Error: Query or URLs required") + os.Exit(1) + } + + // Validate provider flags + if *tavilyMode && *serperMode { + fmt.Fprintln(os.Stderr, "Error: Cannot use both -tavily and -serper") + os.Exit(1) + } + + // Validate extract mode + if *extractMode && *serperMode { + fmt.Fprintln(os.Stderr, "Error: Extract mode is only available with Tavily") + os.Exit(1) + } + + if *extractMode { + provider, err := adapters.NewTavilyProvider() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + urls := strings.Fields(strings.Join(flag.Args(), " ")) + results, err := provider.Extract(urls) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + if *jsonOutput { + json.NewEncoder(os.Stdout).Encode(results) + } else { + formatter.FormatExtractResults(results) + } + } else { + var searchProvider adapters.SearchProvider + var err error + + // Default to Serper unless Tavily is explicitly requested + if *tavilyMode { + searchProvider, err = adapters.NewTavilyProvider() + } else { + searchProvider, err = adapters.NewSerperProvider() + } + + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + options := map[string]string{} + if *tavilyMode { + options["depth"] = *depth + } + + query := flag.Arg(0) + results, err := searchProvider.Search(query, options) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + if *jsonOutput { + json.NewEncoder(os.Stdout).Encode(results) + } else { + formatter.FormatResults(results) + } + } +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..f504e8e --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module search-cli + +go 1.21 diff --git a/internal/formatter/formatter.go b/internal/formatter/formatter.go new file mode 100644 index 0000000..6945527 --- /dev/null +++ b/internal/formatter/formatter.go @@ -0,0 +1,32 @@ +package formatter + +import ( + "fmt" + "strings" + + "search-cli/internal/adapters" +) + +func FormatResults(results *adapters.SearchResponse) { + for _, result := range results.Results { + fmt.Printf("\nšŸ”— %s\n", result.URL) + fmt.Printf("šŸ“ %s\n", result.Content) + fmt.Println(strings.Repeat("-", 80)) + } +} + +func FormatExtractResults(results *adapters.ExtractResponse) { + for _, result := range results.Results { + fmt.Printf("\nšŸ”— %s\n", result.URL) + fmt.Printf("šŸ“„ %s\n", result.RawContent) + fmt.Println(strings.Repeat("-", 80)) + } + + if len(results.FailedResults) > 0 { + fmt.Println("\nāŒ Failed URLs:") + for _, failed := range results.FailedResults { + fmt.Printf("URL: %s\nError: %s\n", failed.URL, failed.Error) + fmt.Println(strings.Repeat("-", 40)) + } + } +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..c127d6e --- /dev/null +++ b/main.go @@ -0,0 +1,115 @@ +package main + +import ( + "encoding/json" + "flag" + "fmt" + "os" + "strings" + + "search-cli/adapters" +) + +func formatResults(results *adapters.SearchResponse) { + for _, result := range results.Results { + fmt.Printf("\nšŸ”— %s\n", result.URL) + fmt.Printf("šŸ“ %s\n", result.Content) + fmt.Println(strings.Repeat("-", 80)) + } +} + +func formatExtractResults(results *adapters.ExtractResponse) { + for _, result := range results.Results { + fmt.Printf("\nšŸ”— %s\n", result.URL) + fmt.Printf("šŸ“„ %s\n", result.RawContent) + fmt.Println(strings.Repeat("-", 80)) + } + + if len(results.FailedResults) > 0 { + fmt.Println("\nāŒ Failed URLs:") + for _, failed := range results.FailedResults { + fmt.Printf("URL: %s\nError: %s\n", failed.URL, failed.Error) + fmt.Println(strings.Repeat("-", 40)) + } + } +} + +func main() { + jsonOutput := flag.Bool("json", false, "Output in JSON format") + tavilyMode := flag.Bool("tavily", false, "Use Tavily as provider") + serperMode := flag.Bool("serper", false, "Use Serper as provider (default)") + extractMode := flag.Bool("extract", false, "Extract content from URLs (Tavily only)") + depth := flag.String("depth", "basic", "Search depth (basic or advanced) - only for Tavily") + flag.Parse() + + if flag.NArg() < 1 { + fmt.Fprintln(os.Stderr, "Error: Query or URLs required") + os.Exit(1) + } + + // Validate provider flags + if *tavilyMode && *serperMode { + fmt.Fprintln(os.Stderr, "Error: Cannot use both -tavily and -serper") + os.Exit(1) + } + + // Validate extract mode + if *extractMode && *serperMode { + fmt.Fprintln(os.Stderr, "Error: Extract mode is only available with Tavily") + os.Exit(1) + } + + if *extractMode { + provider, err := adapters.NewTavilyProvider() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + urls := strings.Fields(strings.Join(flag.Args(), " ")) + results, err := provider.Extract(urls) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + if *jsonOutput { + json.NewEncoder(os.Stdout).Encode(results) + } else { + formatExtractResults(results) + } + } else { + var searchProvider adapters.SearchProvider + var err error + + // Default to Serper unless Tavily is explicitly requested + if *tavilyMode { + searchProvider, err = adapters.NewTavilyProvider() + } else { + searchProvider, err = adapters.NewSerperProvider() + } + + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + options := map[string]string{} + if *tavilyMode { + options["depth"] = *depth + } + + query := flag.Arg(0) + results, err := searchProvider.Search(query, options) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + if *jsonOutput { + json.NewEncoder(os.Stdout).Encode(results) + } else { + formatResults(results) + } + } +} diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..15ddab6 --- /dev/null +++ b/main_test.go @@ -0,0 +1,86 @@ +package main + +import ( + "bytes" + "os" + "strings" + "testing" + + "search-cli/adapters" +) + +func TestFormatResults(t *testing.T) { + results := &adapters.SearchResponse{ + Results: []adapters.SearchResult{ + { + URL: "https://example.com", + Content: "Test content", + }, + }, + } + + // Capture stdout + old := os.Stdout + r, w, _ := os.Pipe() + os.Stdout = w + + formatResults(results) + + w.Close() + os.Stdout = old + + var buf bytes.Buffer + buf.ReadFrom(r) + output := buf.String() + + if !strings.Contains(output, "https://example.com") { + t.Error("Expected output to contain URL") + } + if !strings.Contains(output, "Test content") { + t.Error("Expected output to contain content") + } +} + +func TestFormatExtractResults(t *testing.T) { + results := &adapters.ExtractResponse{ + Results: []adapters.ExtractResult{ + { + URL: "https://example.com", + RawContent: "Test content", + }, + }, + FailedResults: []adapters.FailedResult{ + { + URL: "https://failed.com", + Error: "Failed to extract", + }, + }, + } + + // Capture stdout + old := os.Stdout + r, w, _ := os.Pipe() + os.Stdout = w + + formatExtractResults(results) + + w.Close() + os.Stdout = old + + var buf bytes.Buffer + buf.ReadFrom(r) + output := buf.String() + + if !strings.Contains(output, "https://example.com") { + t.Error("Expected output to contain URL") + } + if !strings.Contains(output, "Test content") { + t.Error("Expected output to contain content") + } + if !strings.Contains(output, "Failed URLs") { + t.Error("Expected output to contain failed URLs section") + } + if !strings.Contains(output, "https://failed.com") { + t.Error("Expected output to contain failed URL") + } +} diff --git a/search.py b/search.py new file mode 100644 index 0000000..c91cf99 --- /dev/null +++ b/search.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +import argparse +import requests +import json +import sys +import os +from typing import Dict, Any + +TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") +if not TAVILY_API_KEY: + print("Error: TAVILY_API_KEY environment variable is not set", file=sys.stderr) + sys.exit(1) + +API_URL = "https://api.tavily.com/search" + +def search(query: str, search_depth: str = "basic") -> Dict[Any, Any]: + headers = {"api-key": TAVILY_API_KEY} + params = { + "query": query, + "search_depth": search_depth, + "include_domains": [], + "exclude_domains": [] + } + + try: + response = requests.post(API_URL, json=params, headers=headers) + response.raise_for_status() + return response.json() + except requests.exceptions.RequestException as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) + +def format_results(data: Dict[Any, Any]) -> None: + for result in data['results']: + print(f"\nšŸ”— {result['url']}") + print(f"šŸ“ {result['content']}\n") + print("-" * 80) + +def main(): + parser = argparse.ArgumentParser(description="Tavily Search CLI") + parser.add_argument("query", help="Search query") + parser.add_argument("--json", action="store_true", help="Output in JSON format") + parser.add_argument("--depth", choices=["basic", "advanced"], default="basic", + help="Search depth (basic or advanced)") + + args = parser.parse_args() + + results = search(args.query, args.depth) + + if args.json: + print(json.dumps(results, indent=2)) + else: + format_results(results) + +if __name__ == "__main__": + main() \ No newline at end of file