#163: Fixed issues found by linter

EinStack · Mar 13, 2024 · 8894ab0 · 8894ab0
1 parent 3c11462
commit 8894ab0
Show file tree

Hide file tree

Showing 15 changed files with 226 additions and 183 deletions.
diff --git a/pkg/providers/anthropic/chat_stream.go b/pkg/providers/anthropic/chat_stream.go
@@ -11,11 +11,6 @@ func (c *Client) SupportChatStream() bool {
 	return false
 }
 
-func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) <-chan *clients.ChatStreamResult {
-	streamResultC := make(chan *clients.ChatStreamResult)
-
-	streamResultC <- clients.NewChatStreamResult(nil, clients.ErrChatStreamNotImplemented)
-	close(streamResultC)
-
-	return streamResultC
+func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) (clients.ChatStream, error) {
+	return nil, clients.ErrChatStreamNotImplemented
 }
diff --git a/pkg/providers/azureopenai/chat_stream.go b/pkg/providers/azureopenai/chat_stream.go
@@ -11,11 +11,6 @@ func (c *Client) SupportChatStream() bool {
 	return false
 }
 
-func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) <-chan *clients.ChatStreamResult {
-	streamResultC := make(chan *clients.ChatStreamResult)
-
-	streamResultC <- clients.NewChatStreamResult(nil, clients.ErrChatStreamNotImplemented)
-	close(streamResultC)
-
-	return streamResultC
+func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) (clients.ChatStream, error) {
+	return nil, clients.ErrChatStreamNotImplemented
 }
diff --git a/pkg/providers/bedrock/chat_stream.go b/pkg/providers/bedrock/chat_stream.go
@@ -11,11 +11,6 @@ func (c *Client) SupportChatStream() bool {
 	return false
 }
 
-func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) <-chan *clients.ChatStreamResult {
-	streamResultC := make(chan *clients.ChatStreamResult)
-
-	streamResultC <- clients.NewChatStreamResult(nil, clients.ErrChatStreamNotImplemented)
-	close(streamResultC)
-
-	return streamResultC
+func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) (clients.ChatStream, error) {
+	return nil, clients.ErrChatStreamNotImplemented
 }
diff --git a/pkg/providers/clients/stream.go b/pkg/providers/clients/stream.go
@@ -4,6 +4,12 @@ import (
 	"glide/pkg/api/schemas"
 )
 
+type ChatStream interface {
+	Open() error
+	Recv() (*schemas.ChatStreamChunk, error)
+	Close() error
+}
+
 type ChatStreamResult struct {
 	chunk *schemas.ChatStreamChunk
 	err   error

diff --git a/pkg/providers/cohere/chat_stream.go b/pkg/providers/cohere/chat_stream.go
@@ -11,11 +11,6 @@ func (c *Client) SupportChatStream() bool {
 	return false
 }
 
-func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) <-chan *clients.ChatStreamResult {
-	streamResultC := make(chan *clients.ChatStreamResult)
-
-	streamResultC <- clients.NewChatStreamResult(nil, clients.ErrChatStreamNotImplemented)
-	close(streamResultC)
-
-	return streamResultC
+func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) (clients.ChatStream, error) {
+	return nil, clients.ErrChatStreamNotImplemented
 }
diff --git a/pkg/providers/lang.go b/pkg/providers/lang.go
@@ -2,6 +2,7 @@ package providers
 
 import (
 	"context"
+	"io"
 	"time"
 
 	"glide/pkg/routers/health"
@@ -18,12 +19,14 @@ type LangProvider interface {
 	SupportChatStream() bool
 
 	Chat(ctx context.Context, req *schemas.ChatRequest) (*schemas.ChatResponse, error)
-	ChatStream(ctx context.Context, req *schemas.ChatRequest) <-chan *clients.ChatStreamResult
+	ChatStream(ctx context.Context, req *schemas.ChatRequest) (clients.ChatStream, error)
 }
 
 type LangModel interface {
-	LangProvider
 	Model
+	Provider() string
+	Chat(ctx context.Context, req *schemas.ChatRequest) (*schemas.ChatResponse, error)
+	ChatStream(ctx context.Context, req *schemas.ChatRequest) (<-chan *clients.ChatStreamResult, error)
 }
 
 // LanguageModel wraps provider client and expend it with health & latency tracking
@@ -99,35 +102,65 @@ func (m *LanguageModel) Chat(ctx context.Context, request *schemas.ChatRequest)
 	return resp, err
 }
 
-func (m *LanguageModel) ChatStream(ctx context.Context, req *schemas.ChatRequest) <-chan *clients.ChatStreamResult {
+func (m *LanguageModel) ChatStream(ctx context.Context, req *schemas.ChatRequest) (<-chan *clients.ChatStreamResult, error) {
+	stream, err := m.client.ChatStream(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
 	streamResultC := make(chan *clients.ChatStreamResult)
-	resultC := m.client.ChatStream(ctx, req)
 
 	go func() {
 		defer close(streamResultC)
 
-		var chunkLatency *time.Duration
+		startedAt := time.Now()
+		err = stream.Open()
+		chunkLatency := time.Since(startedAt)
 
-		for chunkResult := range resultC {
-			if chunkResult.Error() == nil {
-				streamResultC <- chunkResult
+		// the first chunk latency
+		m.chatStreamLatency.Add(float64(chunkLatency))
 
-				chunkLatency = chunkResult.Chunk().Latency
+		if err != nil {
+			streamResultC <- clients.NewChatStreamResult(nil, err)
+
+			m.healthTracker.TrackErr(err)
+
+			return
+		}
 
-				if chunkLatency != nil {
-					m.chatStreamLatency.Add(float64(*chunkLatency))
+		defer stream.Close()
+
+		for {
+			startedAt = time.Now()
+			chunk, err := stream.Recv()
+			chunkLatency = time.Since(startedAt)
+
+			if err != nil {
+				if err == io.EOF {
+					// end of the stream
+					return
 				}
 
-				continue
+				streamResultC <- clients.NewChatStreamResult(nil, err)
+
+				m.healthTracker.TrackErr(err)
+
+				return
 			}
 
-			m.healthTracker.TrackErr(chunkResult.Error())
+			streamResultC <- clients.NewChatStreamResult(chunk, nil)
 
-			streamResultC <- chunkResult
+			if chunkLatency > 1*time.Millisecond {
+				// All events are read in a bigger chunks of bytes, so one chunk may contain more than one event.
+				//  Each byte chunk is then parsed, so there is no easy way to precisely guess latency per chunk,
+				//  So we assume that if we spent more than 1ms waiting for a chunk it's likely
+				//  we were trying to read from the connection (otherwise, it would take nanoseconds)
+				m.chatStreamLatency.Add(float64(chunkLatency))
+			}
 		}
 	}()
 
-	return streamResultC
+	return streamResultC, nil
 }
 
 func (m *LanguageModel) Provider() string {

diff --git a/pkg/providers/octoml/chat_stream.go b/pkg/providers/octoml/chat_stream.go
@@ -11,11 +11,6 @@ func (c *Client) SupportChatStream() bool {
 	return false
 }
 
-func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) <-chan *clients.ChatStreamResult {
-	streamResultC := make(chan *clients.ChatStreamResult)
-
-	streamResultC <- clients.NewChatStreamResult(nil, clients.ErrChatStreamNotImplemented)
-	close(streamResultC)
-
-	return streamResultC
+func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) (clients.ChatStream, error) {
+	return nil, clients.ErrChatStreamNotImplemented
 }
diff --git a/pkg/providers/ollama/chat_stream.go b/pkg/providers/ollama/chat_stream.go
@@ -11,11 +11,6 @@ func (c *Client) SupportChatStream() bool {
 	return false
 }
 
-func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) <-chan *clients.ChatStreamResult {
-	streamResultC := make(chan *clients.ChatStreamResult)
-
-	streamResultC <- clients.NewChatStreamResult(nil, clients.ErrChatStreamNotImplemented)
-	close(streamResultC)
-
-	return streamResultC
+func (c *Client) ChatStream(_ context.Context, _ *schemas.ChatRequest) (clients.ChatStream, error) {
+	return nil, clients.ErrChatStreamNotImplemented
 }
diff --git a/pkg/providers/openai/chat.go b/pkg/providers/openai/chat.go
@@ -7,9 +7,6 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"time"
-
-	"glide/pkg/providers/clients"
 
 	"glide/pkg/api/schemas"
 	"go.uber.org/zap"
@@ -106,7 +103,7 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 	defer resp.Body.Close()
 
 	if resp.StatusCode != http.StatusOK {
-		return nil, c.handleChatReqErrs(resp)
+		return nil, c.errMapper.Map(resp)
 	}
 
 	// Read the response body into a byte slice
@@ -161,39 +158,3 @@ func (c *Client) doChatRequest(ctx context.Context, payload *ChatRequest) (*sche
 
 	return &response, nil
 }
-
-func (c *Client) handleChatReqErrs(resp *http.Response) error {
-	bodyBytes, err := io.ReadAll(resp.Body)
-	if err != nil {
-		c.tel.Logger.Error(
-			"Failed to unmarshal chat response error",
-			zap.String("provider", c.Provider()),
-			zap.Error(err),
-			zap.ByteString("rawResponse", bodyBytes),
-		)
-	}
-
-	c.tel.Logger.Error(
-		"Chat request failed",
-		zap.String("provider", c.Provider()),
-		zap.Int("statusCode", resp.StatusCode),
-		zap.String("response", string(bodyBytes)),
-		zap.Any("headers", resp.Header),
-	)
-
-	if resp.StatusCode == http.StatusTooManyRequests {
-		// Read the value of the "Retry-After" header to get the cooldown delay
-		retryAfter := resp.Header.Get("Retry-After")
-
-		// Parse the value to get the duration
-		cooldownDelay, err := time.ParseDuration(retryAfter)
-		if err != nil {
-			return fmt.Errorf("failed to parse cooldown delay from headers: %w", err)
-		}
-
-		return clients.NewRateLimitError(&cooldownDelay)
-	}
-
-	// Server & client errors result in the same error to keep gateway resilient
-	return clients.ErrProviderUnavailable
-}