#196: Fixed tests after major refactoring

pkg/providers/azureopenai/chat_stream.go

@@ -123,8 +123,8 @@ func (s *ChatStream) Recv() (*schemas.ChatStreamChunk, error) {
 
 		// TODO: use objectpool here
 		return &schemas.ChatStreamChunk{
-			Provider:  providerName,
 			Cached:    false,
+			Provider:  providerName,
 			ModelName: completionChunk.ModelName,
 			ModelResponse: schemas.ModelChunkResponse{
 				Metadata: &schemas.Metadata{

pkg/providers/cohere/chat_stream.go

@@ -5,11 +5,12 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"glide/pkg/providers/clients"
-	"glide/pkg/telemetry"
 	"io"
 	"net/http"
 
+	"glide/pkg/providers/clients"
+	"glide/pkg/telemetry"
+
 	"go.uber.org/zap"
 
 	"glide/pkg/api/schemas"
@@ -29,9 +30,7 @@ var (
 type ChatStream struct {
 	client             *http.Client
 	req                *http.Request
-	reqID              string
 	modelName          string
-	reqMetadata        *schemas.Metadata
 	resp               *http.Response
 	generationID       string
 	streamFinished     bool
@@ -46,7 +45,6 @@ func NewChatStream(
 	client *http.Client,
 	req *http.Request,
 	modelName string,
-	reqMetadata *schemas.Metadata,
 	errMapper *ErrorMapper,
 	finishReasonMapper *FinishReasonMapper,
 ) *ChatStream {
@@ -55,7 +53,6 @@ func NewChatStream(
 		client:             client,
 		req:                req,
 		modelName:          modelName,
-		reqMetadata:        reqMetadata,
 		errMapper:          errMapper,
 		streamFinished:     false,
 		finishReasonMapper: finishReasonMapper,
@@ -133,13 +130,13 @@ func (s *ChatStream) Recv() (*schemas.ChatStreamChunk, error) {
 
 			// TODO: use objectpool here
 			return &schemas.ChatStreamChunk{
-				Provider:  providerName,
 				Cached:    false,
+				Provider:  providerName,
 				ModelName: s.modelName,
 				ModelResponse: schemas.ModelChunkResponse{
 					Metadata: &schemas.Metadata{
-						"generationId": s.generationID,
-						"responseId":   responseChunk.Response.ResponseID,
+						"generation_id": s.generationID,
+						"response_id":   responseChunk.Response.ResponseID,
 					},
 					Message: schemas.ChatMessage{
 						Role:    "model",
@@ -152,12 +149,12 @@ func (s *ChatStream) Recv() (*schemas.ChatStreamChunk, error) {
 
 		// TODO: use objectpool here
 		return &schemas.ChatStreamChunk{
-			Provider:  providerName,
 			Cached:    false,
+			Provider:  providerName,
 			ModelName: s.modelName,
 			ModelResponse: schemas.ModelChunkResponse{
 				Metadata: &schemas.Metadata{
-					"generationId": s.generationID,
+					"generation_id": s.generationID,
 				},
 				Message: schemas.ChatMessage{
 					Role:    "model",
@@ -192,7 +189,6 @@ func (c *Client) ChatStream(ctx context.Context, req *schemas.ChatStreamRequest)
 		c.httpClient,
 		httpRequest,
 		c.chatRequestTemplate.Model,
-		req.Metadata,
 		c.errMapper,
 		c.finishReasonMapper,
 	), nil

pkg/providers/lang.go

@@ -87,8 +87,8 @@ func (m LanguageModel) ChatStreamLatency() *latency.MovingAverage {
 
 func (m *LanguageModel) Chat(ctx context.Context, request *schemas.ChatRequest) (*schemas.ChatResponse, error) {
 	startedAt := time.Now()
-	resp, err := m.client.Chat(ctx, request)
 
+	resp, err := m.client.Chat(ctx, request)
 	if err != nil {
 		m.healthTracker.TrackErr(err)
 

pkg/providers/openai/chat.go

@@ -84,9 +84,8 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 	req.Header.Set("Authorization", fmt.Sprintf("Bearer %v", string(c.config.APIKey)))
 
 	// TODO: this could leak information from messages which may not be a desired thing to have
-	c.tel.Logger.Debug(
+	c.logger.Debug(
 		"Chat Request",
-		zap.String("provider", c.Provider()),
 		zap.String("chatURL", c.chatURL),
 		zap.Any("payload", payload),
 	)
@@ -105,9 +104,9 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 	// Read the response body into a byte slice
 	bodyBytes, err := io.ReadAll(resp.Body)
 	if err != nil {
-		c.tel.Logger.Error(
+		c.logger.Error(
 			"Failed to read chat response",
-			zap.String("provider", c.Provider()), zap.Error(err),
+			zap.Error(err),
 			zap.ByteString("rawResponse", bodyBytes),
 		)
 
@@ -119,9 +118,8 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 
 	err = json.Unmarshal(bodyBytes, &chatCompletion)
 	if err != nil {
-		c.tel.Logger.Error(
+		c.logger.Error(
 			"Failed to unmarshal chat response",
-			zap.String("provider", c.Provider()),
 			zap.ByteString("rawResponse", bodyBytes),
 			zap.Error(err),
 		)

pkg/providers/openai/chat_stream.go

@@ -10,8 +10,6 @@ import (
 
 	"github.com/r3labs/sse/v2"
 	"glide/pkg/providers/clients"
-	"glide/pkg/telemetry"
-
 	"go.uber.org/zap"
 
 	"glide/pkg/api/schemas"
@@ -21,30 +19,28 @@ var StreamDoneMarker = []byte("[DONE]")
 
 // ChatStream represents OpenAI chat stream for a specific request
 type ChatStream struct {
-	tel                *telemetry.Telemetry
 	client             *http.Client
 	req                *http.Request
-	reqID              schemas.StreamRequestID
-	reqMetadata        *schemas.Metadata
 	resp               *http.Response
 	reader             *sse.EventStreamReader
 	finishReasonMapper *FinishReasonMapper
 	errMapper          *ErrorMapper
+	logger             *zap.Logger
 }
 
 func NewChatStream(
-	tel *telemetry.Telemetry,
 	client *http.Client,
 	req *http.Request,
 	finishReasonMapper *FinishReasonMapper,
 	errMapper *ErrorMapper,
+	logger *zap.Logger,
 ) *ChatStream {
 	return &ChatStream{
-		tel:                tel,
 		client:             client,
 		req:                req,
 		finishReasonMapper: finishReasonMapper,
 		errMapper:          errMapper,
+		logger:             logger,
 	}
 }
 
@@ -70,9 +66,8 @@ func (s *ChatStream) Recv() (*schemas.ChatStreamChunk, error) {
 	for {
 		rawEvent, err := s.reader.ReadEvent()
 		if err != nil {
-			s.tel.L().Warn(
+			s.logger.Warn(
 				"Chat stream is unexpectedly disconnected",
-				zap.String("provider", providerName),
 				zap.Error(err),
 			)
 
@@ -82,9 +77,8 @@ func (s *ChatStream) Recv() (*schemas.ChatStreamChunk, error) {
 			return nil, clients.ErrProviderUnavailable
 		}
 
-		s.tel.L().Debug(
+		s.logger.Debug(
 			"Raw chat stream chunk",
-			zap.String("provider", providerName),
 			zap.ByteString("rawChunk", rawEvent),
 		)
 
@@ -99,9 +93,8 @@ func (s *ChatStream) Recv() (*schemas.ChatStreamChunk, error) {
 		}
 
 		if !event.HasContent() {
-			s.tel.L().Debug(
+			s.logger.Debug(
 				"Received an empty message in chat stream, skipping it",
-				zap.String("provider", providerName),
 				zap.Any("msg", event),
 			)
 
@@ -117,8 +110,8 @@ func (s *ChatStream) Recv() (*schemas.ChatStreamChunk, error) {
 
 		// TODO: use objectpool here
 		return &schemas.ChatStreamChunk{
-			Provider:  providerName,
 			Cached:    false,
+			Provider:  providerName,
 			ModelName: completionChunk.ModelName,
 			ModelResponse: schemas.ModelChunkResponse{
 				Metadata: &schemas.Metadata{
@@ -156,11 +149,11 @@ func (c *Client) ChatStream(ctx context.Context, req *schemas.ChatStreamRequest)
 	}
 
 	return NewChatStream(
-		c.tel,
 		c.httpClient,
 		httpRequest,
 		c.finishReasonMapper,
 		c.errMapper,
+		c.logger,
 	), nil
 }
 
@@ -202,7 +195,7 @@ func (c *Client) makeStreamReq(ctx context.Context, req *schemas.ChatStreamReque
 	request.Header.Set("Connection", "keep-alive")
 
 	// TODO: this could leak information from messages which may not be a desired thing to have
-	c.tel.L().Debug(
+	c.logger.Debug(
 		"Stream chat request",
 		zap.String("chatURL", c.chatURL),
 		zap.Any("payload", chatRequest),

pkg/providers/openai/client.go

@@ -5,6 +5,8 @@ import (
 	"net/http"
 	"net/url"
 
+	"go.uber.org/zap"
+
 	"glide/pkg/providers/clients"
 	"glide/pkg/telemetry"
 )
@@ -28,6 +30,7 @@ type Client struct {
 	config              *Config
 	httpClient          *http.Client
 	tel                 *telemetry.Telemetry
+	logger              *zap.Logger
 }
 
 // NewClient creates a new OpenAI client for the OpenAI API.
@@ -37,6 +40,10 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel *
 		return nil, err
 	}
 
+	logger := tel.L().With(
+		zap.String("provider", providerName),
+	)
+
 	c := &Client{
 		baseURL:             providerConfig.BaseURL,
 		chatURL:             chatURL,
@@ -52,7 +59,8 @@ func NewClient(providerConfig *Config, clientConfig *clients.ClientConfig, tel *
 				MaxIdleConnsPerHost: 2,
 			},
 		},
-		tel: tel,
+		tel:    tel,
+		logger: logger,
 	}
 
 	return c, nil

pkg/providers/testing/lang.go

@@ -30,7 +30,6 @@ func (m *RespMock) Resp() *schemas.ChatResponse {
 
 func (m *RespMock) RespChunk() *schemas.ChatStreamChunk {
 	return &schemas.ChatStreamChunk{
-		ID: "rsp0001",
 		ModelResponse: schemas.ModelChunkResponse{
 			Message: schemas.ChatMessage{
 				Content: m.Msg,