initial
This commit is contained in:
commit
4c0797a45d
11 changed files with 809 additions and 0 deletions
76
adapters/serper.go
Normal file
76
adapters/serper.go
Normal file
|
|
@ -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
|
||||
}
|
||||
67
adapters/serper_test.go
Normal file
67
adapters/serper_test.go
Normal file
|
|
@ -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
|
||||
}
|
||||
130
adapters/tavily.go
Normal file
130
adapters/tavily.go
Normal file
|
|
@ -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
|
||||
}
|
||||
116
adapters/tavily_test.go
Normal file
116
adapters/tavily_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
36
adapters/types.go
Normal file
36
adapters/types.go
Normal file
|
|
@ -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"`
|
||||
}
|
||||
92
cmd/search/main.go
Normal file
92
cmd/search/main.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
3
go.mod
Normal file
3
go.mod
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
module search-cli
|
||||
|
||||
go 1.21
|
||||
32
internal/formatter/formatter.go
Normal file
32
internal/formatter/formatter.go
Normal file
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
}
|
||||
115
main.go
Normal file
115
main.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
86
main_test.go
Normal file
86
main_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
56
search.py
Normal file
56
search.py
Normal file
|
|
@ -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()
|
||||
Loading…
Add table
Reference in a new issue