Add support for authorization headers, closes #26

api.go

@@ -4,11 +4,13 @@ package txtai
 import (
     "fmt"
     "github.com/go-resty/resty/v2"
+    "os"
 )
 
 // Base API definition
 type API struct {
     url string
+    token string
 }
 
 // Index result
@@ -17,6 +19,28 @@ type IndexResult struct {
     Score float64 `json:"score"`
 }
 
+// Creates a new API instance
+func NewAPI(params ...string) API {
+    // Parse url and token
+    var url, token string
+    if len(params) > 0 {
+        url = params[0]
+    }
+    if len(params) > 1 {
+        token = params[1]
+    }
+
+    // Fallback to env variables
+    if url == "" {
+        url = os.Getenv("TXTAI_API_URL")
+    }
+    if token == "" {
+        token = os.Getenv("TXTAI_API_TOKEN")
+    }
+
+    return API{url, token}
+}
+
 // Creates a new Resty request and returns it.
 func (api *API) Request(method string, result interface{}) *resty.Request {
     request := resty.New().R()
@@ -34,6 +58,8 @@ func (api *API) Get(method string, params map[string]string, result interface{})
         request.SetQueryParams(params)
     }
 
+    // Set headers and run request
+    request = api.Headers(request)
     _, err := request.Get(api.url + "/" + method)
 
     if err != nil {
@@ -48,9 +74,20 @@ func (api *API) Post(method string, params interface{}, result interface{}) {
         request.SetBody(params)
     }
 
+    // Set headers and run request
+    request = api.Headers(request)
     _, err := request.Post(api.url + "/" + method)
 
     if err != nil {
         fmt.Println(err)
     }
 }
+
+// Adds headers
+func (api *API) Headers(request *resty.Request) *resty.Request {
+    if api.token != "" {
+        request.SetHeader("Authorization", "Bearer " + api.token)
+    }
+
+    return request
+}

embeddings.go

@@ -5,7 +5,6 @@ import "strconv"
 
 // Embeddings definition
 type EmbeddingsAPI struct {
-    url string
     api API
 }
 
@@ -22,8 +21,8 @@ type SearchResult struct {
 }
 
 // Creates an Embeddings instance.
-func Embeddings(url string) EmbeddingsAPI {
-    return EmbeddingsAPI{url, API{url}}
+func Embeddings(params ...string) EmbeddingsAPI {
+    return EmbeddingsAPI{NewAPI(params...)}
 }
 
 // Finds documents in the embeddings model most similar to the input query.

extractor.go

@@ -3,7 +3,6 @@ package txtai
 
 // Extractor definition
 type ExtractorAPI struct {
-    url string
     api API
 }
 
@@ -22,8 +21,8 @@ type Answer struct {
 }
 
 // Creates an Extractor instance.
-func Extractor(url string) ExtractorAPI {
-    return ExtractorAPI{url, API{url}}
+func Extractor(params ...string) ExtractorAPI {
+    return ExtractorAPI{NewAPI(params...)}
 }
 
 // Extracts answers to input questions.

labels.go

@@ -3,13 +3,12 @@ package txtai
 
 // Labels definition
 type LabelsAPI struct {
-    url string
     api API
 }
 
 // Creates a Labels instance.
-func Labels(url string) LabelsAPI {
-    return LabelsAPI{url, API{url}}
+func Labels(params ...string) LabelsAPI {
+    return LabelsAPI{NewAPI(params...)}
 }
 
 // Applies a zero shot classifier to a text using a list of labels.

segmentation.go

@@ -3,13 +3,12 @@ package txtai
 
 // Segmentation definition
 type SegmentationAPI struct {
-    url string
     api API
 }
 
 // Creates a Segmentation instance.
-func Segmentation(url string) SegmentationAPI {
-    return SegmentationAPI{url, API{url}}
+func Segmentation(params ...string) SegmentationAPI {
+    return SegmentationAPI{NewAPI(params...)}
 }
 
 // Segments text into semantic units.

similarity.go

@@ -3,13 +3,12 @@ package txtai
 
 // Similarity definition
 type SimilarityAPI struct {
-    url string
     api API
 }
 
 // Creates a Similarity instance.
-func Similarity(url string) SimilarityAPI {
-    return SimilarityAPI{url, API{url}}
+func Similarity(params ...string) SimilarityAPI {
+    return SimilarityAPI{NewAPI(params...)}
 }
 
 // Computes the similarity between query and list of text.

summary.go

@@ -7,13 +7,12 @@ import (
 
 // Summary definition
 type SummaryAPI struct {
-    url string
     api API
 }
 
 // Creates a Summary instance.
-func Summary(url string) SummaryAPI {
-    return SummaryAPI{url, API{url}}
+func Summary(params ...string) SummaryAPI {
+    return SummaryAPI{NewAPI(params...)}
 }
 
 // Runs a summarization model against a block of text.

textractor.go

@@ -3,13 +3,12 @@ package txtai
 
 // Textractor definition
 type TextractorAPI struct {
-    url string
     api API
 }
 
 // Creates a Textractor instance.
-func Textractor(url string) TextractorAPI {
-    return TextractorAPI{url, API{url}}
+func Textractor(params ...string) TextractorAPI {
+    return TextractorAPI{NewAPI(params...)}
 }
 
 // Extracts text from a file at path.

transcription.go

@@ -3,13 +3,12 @@ package txtai
 
 // Transcription definition
 type TranscriptionAPI struct {
-    url string
     api API
 }
 
 // Creates a Transcription instance.
-func Transcription(url string) TranscriptionAPI {
-    return TranscriptionAPI{url, API{url}}
+func Transcription(params ...string) TranscriptionAPI {
+    return TranscriptionAPI{NewAPI(params...)}
 }
 
 // Transcribes audio files to text.

translation.go

@@ -3,13 +3,12 @@ package txtai
 
 // Translation definition
 type TranslationAPI struct {
-    url string
     api API
 }
 
 // Creates a Translation instance.
-func Translation(url string) TranslationAPI {
-    return TranslationAPI{url, API{url}}
+func Translation(params ...string) TranslationAPI {
+    return TranslationAPI{NewAPI(params...)}
 }
 
 // Translates text from source language into target language.

workflow.go

@@ -3,13 +3,12 @@ package txtai
 
 // Workflow definition
 type WorkflowAPI struct {
-    url string
     api API
 }
 
 // Creates a Workflow instance.
-func Workflow(url string) WorkflowAPI {
-    return WorkflowAPI{url, API{url}}
+func Workflow(params ...string) WorkflowAPI {
+    return WorkflowAPI{NewAPI(params...)}
 }
 
 // Segments text into semantic units.