initial
This commit is contained in:
Regis David Souza Mesquita
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