faiss server

ai-chat-service/services/faiss/faiss.go

@@ -0,0 +1,89 @@
+package faiss
+
+import (
+	"ai-chat-service/pkg/config"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+)
+
+type Client interface {
+	Insert(ctx context.Context, embedding []float32) (string, error)
+	Search(ctx context.Context, embedding []float32, k int) (*SearchResponse, error)
+}
+
+type client struct {
+	baseURL    string
+	httpClient *http.Client
+}
+
+type SearchResponse struct {
+	IDs              []int64   `json:"ids"`
+	Distances        []float32 `json:"distances"`
+	SimilarityScores []float32 `json:"similarity_scores"`
+}
+
+type insertRequest struct {
+	Embedding []float32 `json:"embedding"`
+}
+
+type insertResponse struct {
+	ID int64 `json:"id"`
+}
+
+type searchRequest struct {
+	Embedding []float32 `json:"embedding"`
+	K         int       `json:"k"`
+}
+
+func NewClient(cnf *config.Config) Client {
+	return &client{
+		baseURL:    strings.TrimRight(cnf.Faiss.BaseUrl, "/"),
+		httpClient: &http.Client{Timeout: time.Duration(cnf.Faiss.Timeout) * time.Second},
+	}
+}
+
+func (c *client) Insert(ctx context.Context, embedding []float32) (string, error) {
+	reqBody := &insertRequest{Embedding: embedding}
+	result := &insertResponse{}
+	if err := c.postJSON(ctx, "/insert", reqBody, result); err != nil {
+		return "", err
+	}
+	return fmt.Sprintf("%d", result.ID), nil
+}
+
+func (c *client) Search(ctx context.Context, embedding []float32, k int) (*SearchResponse, error) {
+	reqBody := &searchRequest{Embedding: embedding, K: k}
+	result := &SearchResponse{}
+	if err := c.postJSON(ctx, "/search", reqBody, result); err != nil {
+		return nil, err
+	}
+	return result, nil
+}
+
+func (c *client) postJSON(ctx context.Context, path string, requestData any, responseData any) error {
+	body, err := json.Marshal(requestData)
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+path, bytes.NewReader(body))
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return fmt.Errorf("faiss request failed: status=%d", resp.StatusCode)
+	}
+	return json.NewDecoder(resp.Body).Decode(responseData)
+}