diff --git a/backend/src/search/openai.go b/backend/src/search/openai.go index 13f2b1524..73a88c2eb 100644 --- a/backend/src/search/openai.go +++ b/backend/src/search/openai.go @@ -24,17 +24,25 @@ func NewOpenAIClient(settings config.OpenAISettings) *OpenAIClient { return &OpenAIClient{Settings: settings} } +type CreateEmbeddingRequestBody struct { + Input string `json:"input"` + Model string `json:"model"` +} + +type Embedding struct { + Embedding []float32 `json:"embedding"` +} + type CreateEmbeddingResponseBody struct { - Data []struct { - Embedding []float32 `json:"embedding"` - } `json:"data"` + Data []Embedding `json:"data"` } func (c *OpenAIClient) CreateEmbedding(payload string) ([]float32, *errors.Error) { - embeddingBody, err := json.Marshal(map[string]interface{}{ - "input": payload, - "model": "text-embedding-ada-002", - }) + embeddingBody, err := json.Marshal( + CreateEmbeddingRequestBody{ + Input: payload, + Model: "text-embedding-ada-002", + }) if err != nil { return nil, &errors.FailedToCreateEmbedding } diff --git a/backend/src/search/pinecone.go b/backend/src/search/pinecone.go index 293f367a5..ca145f69a 100644 --- a/backend/src/search/pinecone.go +++ b/backend/src/search/pinecone.go @@ -55,7 +55,7 @@ func (c *PineconeClient) Upsert(item Searchable) *errors.Error { return &errors.FailedToUpsertToPinecone } - upsertBody, _ := json.Marshal( + upsertBody, err := json.Marshal( PineconeUpsertRequestBody{ Vectors: []Vector{ { @@ -65,6 +65,9 @@ func (c *PineconeClient) Upsert(item Searchable) *errors.Error { }, Namespace: item.Namespace(), }) + if err != nil { + return &errors.FailedToUpsertToPinecone + } req, err := http.NewRequest(fiber.MethodPost, fmt.Sprintf("%s/vectors/upsert", c.Settings.IndexHost.Expose()), @@ -141,13 +144,15 @@ type PineconeSearchRequestBody struct { Namespace string `json:"namespace"` } +type Match struct { + Id string `json:"id"` + Score float32 `json:"score"` + Values []float32 `json:"values"` +} + type PineconeSearchResponseBody struct { - Matches []struct { - Id string `json:"id"` - Score float32 `json:"score"` - Values []float32 `json:"values"` - } `json:"matches"` - Namespace string `json:"namespace"` + Matches []Match `json:"matches"` + Namespace string `json:"namespace"` } func (c *PineconeClient) Search(item Searchable, topK int) ([]string, *errors.Error) { @@ -184,10 +189,6 @@ func (c *PineconeClient) Search(item Searchable, topK int) ([]string, *errors.Er defer resp.Body.Close() - if err != nil { - return []string{}, &errors.FailedToSearchToPinecone - } - if resp.StatusCode != fiber.StatusOK { return []string{}, &errors.FailedToSearchToPinecone } diff --git a/backend/tests/search_test.go b/backend/tests/search_test.go index 589f16222..8c3c789e9 100644 --- a/backend/tests/search_test.go +++ b/backend/tests/search_test.go @@ -30,7 +30,7 @@ type mockConfig struct { } func newMockConfig() *mockConfig { - pineconeIndexHost := "indexHost" + pineconeIndexHost := "https://indexHost.com" pineconeApiKey := "pinecone" pineconeIndexHostSecret, err := m.NewSecret(pineconeIndexHost) @@ -79,14 +79,14 @@ func TestPineconeUpsertWorks(t *testing.T) { MatchHeader("accept", "application/json"). MatchHeader("content-type", "application/json"). MatchType("json"). - JSON(map[string]interface{}{ - "vectors": []interface{}{ - map[string]interface{}{ - "id": mockSearchId, - "values": mockValues, + JSON(search.PineconeUpsertRequestBody{ + Vectors: []search.Vector{ + { + ID: mockSearchId, + Values: mockValues, }, }, - "namespace": mockNamespace, + Namespace: mockNamespace, }). Reply(200) @@ -95,15 +95,15 @@ func TestPineconeUpsertWorks(t *testing.T) { MatchHeader("Authorization", "Bearer "+mockConfig.OpenAI.APIKey.Expose()). MatchHeader("content-type", "application/json"). MatchType("json"). - JSON(map[string]interface{}{ - "input": mockSearchString, - "model": "text-embedding-ada-002", + JSON(search.CreateEmbeddingRequestBody{ + Input: mockSearchString, + Model: "text-embedding-ada-002", }). Reply(200). - JSON(map[string]interface{}{ - "data": []map[string]interface{}{ + JSON(search.CreateEmbeddingResponseBody{ + Data: []search.Embedding{ { - "embedding": []float32{1.0, 1.0, 1.0, 1.0}, + Embedding: mockValues, }, }, }) @@ -129,12 +129,10 @@ func TestPineconeDeleteWorks(t *testing.T) { MatchHeader("accept", "application/json"). MatchHeader("content-type", "application/json"). MatchType("json"). - JSON(map[string]interface{}{ - "deleteAll": false, - "ids": []string{ - mockSearchId, - }, - "namespace": mockNamespace, + JSON(search.PineconeDeleteRequestBody{ + Namespace: mockNamespace, + IDs: []string{mockSearchId}, + DeleteAll: false, }). Reply(200) @@ -162,23 +160,23 @@ func TestPineconeSearchWorks(t *testing.T) { MatchHeader("accept", "application/json"). MatchHeader("content-type", "application/json"). MatchType("json"). - JSON(map[string]interface{}{ - "includeValues": false, - "includeMetadata": false, - "topK": topK, - "vector": mockValues, - "namespace": mockNamespace, + JSON(search.PineconeSearchRequestBody{ + IncludeValues: false, + IncludeMetadata: false, + TopK: topK, + Vector: mockValues, + Namespace: mockNamespace, }). Reply(200). - JSON(map[string]interface{}{ - "matches": []map[string]interface{}{ + JSON(search.PineconeSearchResponseBody{ + Matches: []search.Match{ { - "id": mockSearchId, - "score": 1.0, - "values": []float32{1.0, 1.0, 1.0, 1.0}, + Id: mockSearchId, + Score: 1.0, + Values: mockValues, }, }, - "namespace": mockNamespace, + Namespace: mockNamespace, }) gock.New("https://api.openai.com"). @@ -186,15 +184,15 @@ func TestPineconeSearchWorks(t *testing.T) { MatchHeader("Authorization", "Bearer "+mockConfig.OpenAI.APIKey.Expose()). MatchHeader("content-type", "application/json"). MatchType("json"). - JSON(map[string]interface{}{ - "input": mockSearchString, - "model": "text-embedding-ada-002", + JSON(search.CreateEmbeddingRequestBody{ + Input: mockSearchString, + Model: "text-embedding-ada-002", }). Reply(200). - JSON(map[string]interface{}{ - "data": []map[string]interface{}{ + JSON(search.CreateEmbeddingResponseBody{ + Data: []search.Embedding{ { - "embedding": []float32{1.0, 1.0, 1.0, 1.0}, + Embedding: []float32{1.0, 1.0, 1.0, 1.0}, }, }, }) @@ -220,15 +218,15 @@ func TestOpenAIEmbeddingWorks(t *testing.T) { MatchHeader("Authorization", "Bearer "+mockConfig.OpenAI.APIKey.Expose()). MatchHeader("content-type", "application/json"). MatchType("json"). - JSON(map[string]interface{}{ - "input": testString, - "model": "text-embedding-ada-002", + JSON(search.CreateEmbeddingRequestBody{ + Input: testString, + Model: "text-embedding-ada-002", }). Reply(200). - JSON(map[string]interface{}{ - "data": []map[string]interface{}{ + JSON(search.CreateEmbeddingResponseBody{ + Data: []search.Embedding{ { - "embedding": []float32{1.0, 1.0, 1.0, 1.0}, + Embedding: []float32{1.0, 1.0, 1.0, 1.0}, }, }, }) diff --git a/config/.example_backend_env b/config/.example_backend_env index 25c1a3bb7..977a98269 100644 --- a/config/.example_backend_env +++ b/config/.example_backend_env @@ -1,3 +1,3 @@ -SAC_PINECONE_INDEX_HOST="SAC_PINECONE_INDEX_HOST" +SAC_PINECONE_INDEX_HOST="https://SAC_PINECONE_INDEX_HOST.com" SAC_PINECONE_API_KEY="SAC_PINECONE_API_KEY" SAC_OPENAI_API_KEY="SAC_OPENAI_API_KEY" \ No newline at end of file