An open API service indexing awesome lists of open source software.

https://github.com/atopx/teiclient

go client for text-embedding-inference (https://github.com/huggingface/text-embeddings-inference)
https://github.com/atopx/teiclient

rerank tei text-embedding text-embeddings-inference

Last synced: 12 months ago
JSON representation

go client for text-embedding-inference (https://github.com/huggingface/text-embeddings-inference)

Awesome Lists containing this project

README

          

# text-embeddings-inference go client

[![Go Reference](https://pkg.go.dev/badge/github.com/atopx/teiclient)](https://pkg.go.dev/github.com/atopx/teiclient)

An elegant, idiomatic Go client for the Hugging Face Text Embeddings Inference (TEI) service, supporting both HTTP REST and gRPC transports. Simplify your workflow by choosing the interface that best fits your use case.

---

## Repository Structure

```
.
├── proto/ # Auto-generated protobuf code
│ ├── tei.proto # gRPC + HTTP OpenAPI definitions
│ ├── tei.pb.go # Go structs + HTTP client interfaces
│ └── tei_grpc.pb.go # Go gRPC client stubs
├── teiapi/ # HTTP REST client implementation
│ └── api.go # HTTP client wrapper
└── teigrpc/ # gRPC client implementation
└── grpc.go # gRPC client wrapper
```

---

## Features

* **Dual Transport**: Use HTTP REST or gRPC with the same protobuf definitions.
* **Unified API**: Access Info, Embed, Predict, Rerank, Tokenize, and Decode endpoints.
* **Streaming Support**: Bidirectional streaming for large or batched operations.
* **Context-aware**: Pass `context.Context` to manage timeouts and cancellations.
* **Lightweight**: Zero external dependencies beyond `grpc-go` and the standard library.

---

## Installation

```bash
go get -u github.com/atopx/teiclient
```

---

## HTTP REST Client Quick Start
> [OpenAPI Specification](https://huggingface.github.io/text-embeddings-inference/#/ )

```go
package main

import (
"context"
"fmt"
"time"

"github.com/atopx/teiclient/teiapi"
"github.com/atopx/teiclient/proto"
)

func main() {
// 1. Create HTTP client
client := teiapi.New("http://localhost:8080", 30*time.Second)

// 2. Info
infoResp, err := client.Info(context.Background(), &proto.InfoRequest{})
if err != nil {
panic(err)
}
fmt.Println("Model ID:", infoResp.ModelId)

// 3. Embed
embedReq := &proto.EmbedRequest{Inputs: "Hello, HTTP!", Normalize: true}
embedResp, err := client.Embed(context.Background(), embedReq)
if err != nil {
panic(err)
}
fmt.Printf("Embedding length: %d\n", len(embedResp.Embeddings))
}
```

---

## gRPC Client Quick Start
> [proto definition](./proto/tei.proto)

```go
package main

import (
"context"
"fmt"
"time"

"github.com/atopx/teiclient/teigrpc"
"github.com/atopx/teiclient/proto"
)

func main() {
// 1. Create gRPC client
client, err := teigrpc.New("localhost:50051")
if err != nil {
panic(err)
}
defer client.Close()

// 2. Rerank
rerankReq := &proto.RerankRequest{
Query: "search term",
Texts: []string{"A", "B"},
ReturnText: true,
}
rerankResp, err := client.Rerank(context.Background(), rerankReq)
if err != nil {
panic(err)
}
for _, r := range rerankResp.Ranks {
fmt.Printf("#%d: %s (%.2f)\n", r.Index, r.Text, r.Score)
}
}
```

---

## Common API Reference

All methods correspond to `tei.proto` definitions. Key interfaces:

* HTTP client in `teiapi/api.go` uses methods:

* `Info(ctx, *InfoRequest) (*InfoResponse, error)`
* `Embed(ctx, *EmbedRequest) (*EmbedResponse, error)`
* `Predict(ctx, *PredictRequest) (*PredictResponse, error)`
* `Rerank(ctx, *RerankRequest) (*RerankResponse, error)`
* `Tokenize(ctx, *EncodeRequest) (*EncodeResponse, error)`
* `Decode(ctx, *DecodeRequest) (*DecodeResponse, error)`

* gRPC client in `teigrpc/grpc.go` exposes:

* `Info`, `Embed`, `EmbedStream`, `Predict`, `PredictStream`, `Rerank`, `RerankStream`, `Tokenize`, `Decode`, etc.

Refer to the GoDoc for full signatures and streaming variants.

---

## Contributing

Contributions, issues, and feature requests are welcome! Please open a GitHub issue or submit a pull request.

---

## License

MIT. See [LICENSE](LICENSE).